@@ -39,6 +39,8 @@ import Data.IntMap.Strict (IntMap)
3939import Data.Kind (Type)
4040import Data.List (nubBy)
4141import Data.Maybe (fromJust)
42+import Data.Primitive.Array (Array, indexArray, mapArray')
43+import Data.Primitive.PrimArray (PrimArray, indexPrimArray, sizeofPrimArray)
4244import Data.Proxy (Proxy (..))
4345import Data.Typeable (Typeable, typeRep, typeRepFingerprint)
4446import GHC.Base (Any, Int (..), Int#, (*#), (+#), (<#))
@@ -49,14 +51,13 @@ import GHC.Word (Word64 (..))
4951import Unsafe.Coerce (unsafeCoerce)
5052
5153import qualified Data.IntMap.Strict as IM
52−import qualified Data.Vector as V
53−import qualified Data.Vector.Unboxed as Unboxed
54+import qualified GHC.Exts as GHC (fromList, toList)
5455
5556-- | Map-like data structure that keeps types as keys.
5657data TypeRepMap (f :: k -> Type) = TypeRepMap
57− { fingerprintAs :: {-# UNPACK #-} !(Unboxed.Vector Word64)
58− , fingerprintBs :: {-# UNPACK #-} !(Unboxed.Vector Word64)
59− , anys :: {-# UNPACK #-} !(V.Vector Any)
58+ { fingerprintAs :: {-# UNPACK #-} !(PrimArray Word64)
59+ , fingerprintBs :: {-# UNPACK #-} !(PrimArray Word64)
60+ , anys :: {-# UNPACK #-} !(Array Any)
6061 }
6162
6263-- | Shows only 'Fingerprint's.
@@ -66,8 +67,8 @@ instance Show (TypeRepMap f) where
6667-- | Returns the list of 'Fingerprint's from 'TypeRepMap'.
6768toFps :: TypeRepMap f -> [Fingerprint]
6869toFps TypeRepMap{..} = zipWith Fingerprint
69− (Unboxed.toList fingerprintAs)
70− (Unboxed.toList fingerprintBs)
70+ (GHC.toList fingerprintAs)
71+ (GHC.toList fingerprintBs)
7172
7273
7374-- | Empty structure.
@@ -121,7 +122,7 @@ Just [True]
121122Just "a"
122123-}
123124hoist :: (forall x. f x -> g x) -> TypeRepMap f -> TypeRepMap g
124−hoist f (TypeRepMap as bs ans) = TypeRepMap as bs $ V.map (toAny . f . fromAny) ans
125+hoist f (TypeRepMap as bs ans) = TypeRepMap as bs $ mapArray' (toAny . f . fromAny) ans
125126
126127{- | Returns 'True' if there exist value of given type.
127128
@@ -145,7 +146,7 @@ Just (Identity 11)
145146Nothing
146147-}
147148lookup :: forall a f . Typeable a => TypeRepMap f -> Maybe (f a)
148−lookup tVect = fromAny . (anys tVect V.!)
149+lookup tVect = fromAny . (anys tVect `indexArray`)
149150 <$> cachedBinarySearch (typeFp @a)
150151 (fingerprintAs tVect)
151152 (fingerprintBs tVect)
@@ -153,22 +154,22 @@ lookup tVect = fromAny . (anys tVect V.!)
153154
154155-- | Returns the size of the 'TypeRepMap'.
155156size :: TypeRepMap f -> Int
156−size = Unboxed.length . fingerprintAs
157+size = sizeofPrimArray . fingerprintAs
157158{-# INLINE size #-}
158159
159160-- | Binary searched based on this article
160161-- http://bannalia.blogspot.com/2015/06/cache-friendly-binary-search.html
161162-- with modification for our two-vector search case.
162−cachedBinarySearch :: Fingerprint -> Unboxed.Vector Word64 -> Unboxed.Vector Word64 -> Maybe Int
163+cachedBinarySearch :: Fingerprint -> PrimArray Word64 -> PrimArray Word64 -> Maybe Int
163164cachedBinarySearch (Fingerprint (W64# a) (W64# b)) fpAs fpBs = inline (go 0#)
164165 where
165166 go :: Int# -> Maybe Int
166167 go i = case i <# len of
167168 0# -> Nothing
168− _ -> let !(W64# valA) = Unboxed.unsafeIndex fpAs (I# i) in case a `ltWord#` valA of
169+ _ -> let !(W64# valA) = indexPrimArray fpAs (I# i) in case a `ltWord#` valA of
169170 0# -> case a `eqWord#` valA of
170171 0# -> go (2# *# i +# 2#)
171− _ -> let !(W64# valB) = Unboxed.unsafeIndex fpBs (I# i) in case b `eqWord#` valB of
172+ _ -> let !(W64# valB) = indexPrimArray fpBs (I# i) in case b `eqWord#` valB of
172173 0# -> case b `ltWord#` valB of
173174 0# -> go (2# *# i +# 2#)
174175 _ -> go (2# *# i +# 1#)
@@ -176,7 +177,7 @@ cachedBinarySearch (Fingerprint (W64# a) (W64# b)) fpAs fpBs = inline (go 0#)
176177 _ -> go (2# *# i +# 1#)
177178
178179 len :: Int#
179− len = let !(I# l) = Unboxed.length fpAs in l
180+ len = let !(I# l) = sizeofPrimArray fpAs in l
180181{-# INLINE cachedBinarySearch #-}
181182
182183----------------------------------------------------------------------------
@@ -194,7 +195,7 @@ typeFp = typeRepFingerprint $ typeRep $ Proxy @a
194195{-# INLINE typeFp #-}
195196
196197toPairList :: TypeRepMap f -> [(Fingerprint, Any)]
197−toPairList trMap = zip (toFps trMap) (V.toList $ anys trMap)
198+toPairList trMap = zip (toFps trMap) (GHC.toList $ anys trMap)
198199
199200deleteByFst :: Eq a => a -> [(a, b)] -> [(a, b)]
200201deleteByFst x = filter ((/= x) . fst)
@@ -236,7 +237,7 @@ calcFp :: Typeable a => f a -> Fingerprint
236237calcFp = typeRepFingerprint . typeRep . fromF
237238
238239fromListPairs :: [(Fingerprint, Any)] -> TypeRepMap f
239−fromListPairs kvs = TypeRepMap (Unboxed.fromList fpAs) (Unboxed.fromList fpBs) (V.fromList ans)
240+fromListPairs kvs = TypeRepMap (GHC.fromList fpAs) (GHC.fromList fpBs) (GHC.fromList ans)
240241 where
241242 (fpAs, fpBs) = unzip $ map (\(Fingerprint a b) -> (a, b)) fps
242243 (fps, ans) = unzip $ fromSortedList $ sortWith fst $ nubByFst kvs