git.iko.soy/typerep-map
on develop-bytearray help sign in
typerep-map/src/Data/TypeRep/OptimalVector.hs
119 lines 3.5 kB history raw
1{-# OPTIONS_GHC -fno-warn-orphans #-}
3{-# LANGUAGE BangPatterns #-}
4{-# LANGUAGE ExplicitForAll #-}
5{-# LANGUAGE GADTs #-}
6{-# LANGUAGE KindSignatures #-}
7{-# LANGUAGE MagicHash #-}
8{-# LANGUAGE MultiParamTypeClasses #-}
9{-# LANGUAGE PolyKinds #-}
10{-# LANGUAGE ScopedTypeVariables #-}
11{-# LANGUAGE TypeFamilies #-}
12{-# LANGUAGE RecordWildCards #-}
15module Data.TypeRep.OptimalVector
16 ( TypeRepVector (..)
17 , TF (..)
18 , empty
19 , insert
20 , lookup
21 , size
22 , fromList
23 ) where
25import Prelude hiding (lookup)
27import Control.Arrow ((&&&))
28import Data.Proxy (Proxy (..))
29import Data.Typeable (Typeable, typeRep, typeRepFingerprint)
30import GHC.Base hiding (empty)
31import GHC.Exts (sortWith)
32import GHC.Fingerprint (Fingerprint (..))
33import Unsafe.Coerce (unsafeCoerce)
35import qualified Data.Vector as V
36import qualified Data.Vector.Primitive as P
38import Data.Primitive.ByteArray
39import GHC.Word
41data TypeRepVector (f :: k -> *) = TypeRepVect
42 { fingerprintAs :: ByteArray#
43 , fingerprintBs :: ByteArray#
44 , anys :: V.Vector Any
45 }
47-- | Empty structure.
48empty :: TypeRepVector f
49empty = fromList []
51-- | Inserts the value with its type as a key.
52insert :: forall a f . Typeable a => a -> TypeRepVector f -> TypeRepVector f
53insert _ = id
55-- | Looks up the value at the type.
56-- >>> let x = lookup $ insert (11 :: Int) empty
57-- >>> x :: Maybe Int
58-- Just 11
59-- >>> x :: Maybe ()
60-- Nothing
61lookup :: forall a f . Typeable a => TypeRepVector f -> Maybe (f a)
62lookup TypeRepVect{..} =
63 fromAny (binarySearch anys fp fingerprintAs fingerprintBs)
64 where
65 fp = typeRepFingerprint (typeRep (Proxy :: Proxy a))
66 fromAny = unsafeCoerce :: Maybe Any -> Maybe (f a)
68-- | Returns the size of the 'TypeRepVect'.
69size :: TypeRepVector f -> Int
70size = V.length . anys
72data TF f where
73 TF :: Typeable a => f a -> TF f
75fromF :: Typeable a => f a -> Proxy a
76fromF _ = Proxy
78fromList :: forall f . [TF f] -> TypeRepVector f
79fromList tfs = TypeRepVect fpAs' fpBs' (V.fromList ans)
80 where
81 !(P.Vector _ _ (ByteArray fpAs')) = P.fromList fpAs
82 !(P.Vector _ _ (ByteArray fpBs')) = P.fromList fpBs
83 (fpAs, fpBs) = unzip $ fmap (\(Fingerprint a b) -> (a, b)) fps
84 (fps, ans) = unzip $ sortWith fst $ map (fp &&& an) tfs
86 fp :: TF f -> Fingerprint
87 fp (TF x) = typeRepFingerprint $ typeRep $ fromF x
89 an :: TF f -> Any
90 an (TF x) = unsafeCoerce x
92-- | Returns the index is found.
93binarySearch :: forall a. V.Vector a -> Fingerprint -> ByteArray# -> ByteArray# -> Maybe a
94binarySearch !anys !(Fingerprint (W64# a) (W64# b)) !fpAs !fpBs =
95 let
96 !(I# len) = V.length anys
97 checkfpBs :: Int# -> Maybe a
98 checkfpBs !i =
99 case i <# len of
100 0# -> Nothing
101 _ ->
102 case eqWord# a (indexWord64Array# fpAs i) of
103 0# -> Nothing
104 _ ->
105 case eqWord# b (indexWord64Array# fpBs i) of
106 0# -> checkfpBs (i +# 1#)
107 _ -> Just (V.unsafeIndex anys (I# i))
108 in
109 inline (checkfpBs (binSearchHelp (-1#) len))
110 where
111 binSearchHelp :: Int# -> Int# -> Int#
112 binSearchHelp !l !r = case l <# (r -# 1#) of
113 0# -> r
114 _ ->
115 let m = uncheckedIShiftRA# (l +# r) 1# in
116 case ltWord# (indexWord64Array# fpAs m) a of
117 0# -> binSearchHelp l m
118 _ -> binSearchHelp m r