@@ -9,6 +9,7 @@
99{-# LANGUAGE PolyKinds #-}
1010{-# LANGUAGE ScopedTypeVariables #-}
1111{-# LANGUAGE TypeFamilies #-}
12+{-# LANGUAGE RecordWildCards #-}
1213
1314
1415module Data.TypeRep.OptimalVector
@@ -26,32 +27,30 @@ import Prelude hiding (lookup)
2627import Control.Arrow ((&&&))
2728import Data.Proxy (Proxy (..))
2829import Data.Typeable (Typeable, typeRep, typeRepFingerprint)
29−import Data.Word (Word64)
3030import GHC.Base hiding (empty)
3131import GHC.Exts (sortWith)
3232import GHC.Fingerprint (Fingerprint (..))
3333import Unsafe.Coerce (unsafeCoerce)
34−import Debug.Trace
3534
3635import qualified Data.Vector as V
37−import qualified Data.Vector.Unboxed as Unboxed
36+import qualified Data.Vector.Primitive as P
3837
39−data TypeRepVector f = TypeRepVect
40− { fingerprintAs :: Unboxed.Vector Word64
41− , fingerprintBs :: Unboxed.Vector Word64
38+import Data.Primitive.ByteArray
39+import GHC.Word
40+
41+data TypeRepVector (f :: k -> *) = TypeRepVect
42+ { fingerprintAs :: ByteArray#
43+ , fingerprintBs :: ByteArray#
4244 , anys :: V.Vector Any
4345 }
4446
45−fromAny :: Any -> f a
46−fromAny = unsafeCoerce
47−
4847-- | Empty structure.
4948empty :: TypeRepVector f
50−empty = TypeRepVect mempty mempty mempty
49+empty = fromList []
5150
5251-- | Inserts the value with its type as a key.
5352insert :: forall a f . Typeable a => a -> TypeRepVector f -> TypeRepVector f
54−insert = undefined
53+insert _ = id
5554
5655-- | Looks up the value at the type.
5756-- >>> let x = lookup $ insert (11 :: Int) empty
@@ -60,12 +59,15 @@ insert = undefined
6059-- >>> x :: Maybe ()
6160-- Nothing
6261lookup :: forall a f . Typeable a => TypeRepVector f -> Maybe (f a)
63−lookup tVect = fromAny . (anys tVect V.!)
64− <$> binarySearch (typeRepFingerprint (typeRep (Proxy :: Proxy a))) (fingerprintAs tVect) (fingerprintBs tVect)
62+lookup 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)
6567
6668-- | Returns the size of the 'TypeRepVect'.
6769size :: TypeRepVector f -> Int
68−size = Unboxed.length . fingerprintAs
70+size = V.length . anys
6971
7072data TF f where
7173 TF :: Typeable a => f a -> TF f
@@ -74,8 +76,10 @@ fromF :: Typeable a => f a -> Proxy a
7476fromF _ = Proxy
7577
7678fromList :: forall f . [TF f] -> TypeRepVector f
77−fromList tfs = TypeRepVect (Unboxed.fromList fpAs) (Unboxed.fromList fpBs) (V.fromList ans)
79+fromList tfs = TypeRepVect fpAs' fpBs' (V.fromList ans)
7880 where
81+ !(P.Vector _ _ (ByteArray fpAs')) = P.fromList fpAs
82+ !(P.Vector _ _ (ByteArray fpBs')) = P.fromList fpBs
7983 (fpAs, fpBs) = unzip $ fmap (\(Fingerprint a b) -> (a, b)) fps
8084 (fps, ans) = unzip $ sortWith fst $ map (fp &&& an) tfs
8185
@@ -86,29 +90,30 @@ fromList tfs = TypeRepVect (Unboxed.fromList fpAs) (Unboxed.fromList fpBs) (V.fr
8690 an (TF x) = unsafeCoerce x
8791
8892-- | Returns the index is found.
89−binarySearch :: Fingerprint -> Unboxed.Vector Word64 -> Unboxed.Vector Word64 -> Maybe Int
90−binarySearch (Fingerprint a b) fpAs fpBs =
93+binarySearch :: forall a. V.Vector a -> Fingerprint -> ByteArray# -> ByteArray# -> Maybe a
94+binarySearch !anys !(Fingerprint (W64# a) (W64# b)) !fpAs !fpBs =
9195 let
92− !(I# len) = Unboxed.length fpAs
93− checkfpBs :: Int# -> Maybe Int
94− checkfpBs i =
96+ !(I# len) = V.length anys
97+ checkfpBs :: Int# -> Maybe a
98+ checkfpBs !i =
9599 case i <# len of
96100 0# -> Nothing
97101 _ ->
98− if a /= Unboxed.unsafeIndex fpAs (I# i)
99− then Nothing
100− else if b == Unboxed.unsafeIndex fpBs (I# i)
101− then Just (I# i)
102− else checkfpBs (i +# 1#)
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))
103108 in
104109 inline (checkfpBs (binSearchHelp (-1#) len))
105110 where
106111 binSearchHelp :: Int# -> Int# -> Int#
107− binSearchHelp l r = case l <# (r -# 1#) of
112+ binSearchHelp !l !r = case l <# (r -# 1#) of
108113 0# -> r
109114 _ ->
110115 let m = uncheckedIShiftRA# (l +# r) 1# in
111− if Unboxed.unsafeIndex fpAs (I# m) < a
112− then binSearchHelp m r
113− else binSearchHelp l m
116+ case ltWord# (indexWord64Array# fpAs m) a of
117+ 0# -> binSearchHelp l m
118+ _ -> binSearchHelp m r
114119