@@ -1,38 +1,58 @@
1−{-# LANGUAGE ExplicitForAll #-}
1+{-# LANGUAGE BangPatterns #-}
2+{-# LANGUAGE GADTs #-}
23{-# LANGUAGE KindSignatures #-}
4+{-# LANGUAGE MagicHash #-}
35{-# LANGUAGE PolyKinds #-}
46{-# LANGUAGE ScopedTypeVariables #-}
7+{-# LANGUAGE TypeApplications #-}
8+{-# LANGUAGE TypeFamilies #-}
59
610module Data.TypeRep.Map
7− ( TypeRepMap (..)
11+ ( -- * Map type
12+ TypeRepMap (..)
13+
14+ -- 'TypeRepMap' interface
815 , empty
916 , insert
10− , keys
1117 , lookup
1218 , size
19+
20+ -- * Helpful testing functions
21+ , TF (..)
22+ , fromList
1323 ) where
1424
1525import Prelude hiding (lookup)
1626
27+import Control.Arrow ((&&&))
28+import Data.Kind (Type)
1729import Data.Proxy (Proxy (..))
18−import Data.Typeable (TypeRep, Typeable, typeRep)
19−import GHC.Base (Any)
30+import Data.Typeable (Typeable, typeRep, typeRepFingerprint)
31+import Data.Word (Word64)
32+import GHC.Base (Any, Int (..), Int#, uncheckedIShiftRA#, (+#), (-#), (<#))
33+import GHC.Exts (inline, sortWith)
34+import GHC.Fingerprint (Fingerprint (..))
2035import Unsafe.Coerce (unsafeCoerce)
2136
22−import qualified Data.Map.Lazy as LMap
37+import qualified Data.Vector as V
38+import qualified Data.Vector.Unboxed as Unboxed
2339
24−-- | Map-like data structure with types served as the keys.
25−newtype TypeRepMap (f :: k -> *) = TypeRepMap
26− { unMap :: LMap.Map TypeRep Any
40+data TypeRepMap (f :: k -> Type) = TypeRepMap
41+ { fingerprintAs :: Unboxed.Vector Word64
42+ , fingerprintBs :: Unboxed.Vector Word64
43+ , anys :: V.Vector Any
2744 }
2845
46+fromAny :: Any -> f a
47+fromAny = unsafeCoerce
48+
2949-- | Empty structure.
3050empty :: TypeRepMap f
31−empty = TypeRepMap mempty
51+empty = TypeRepMap mempty mempty mempty
3252
3353-- | Inserts the value with its type as a key.
34−insert :: forall a f . Typeable a => f a -> TypeRepMap f -> TypeRepMap f
35−insert val = TypeRepMap . LMap.insert (typeRep (Proxy :: Proxy a)) (unsafeCoerce val) . unMap
54+insert :: forall a f . Typeable a => a -> TypeRepMap f -> TypeRepMap f
55+insert = undefined
3656
3757-- | Looks up the value at the type.
3858-- >>> let x = lookup $ insert (11 :: Int) empty
@@ -41,10 +61,60 @@ insert val = TypeRepMap . LMap.insert (typeRep (Proxy :: Proxy a)) (unsafeCoerce
4161-- >>> x :: Maybe ()
4262-- Nothing
4363lookup :: forall a f . Typeable a => TypeRepMap f -> Maybe (f a)
44−lookup = fmap unsafeCoerce . LMap.lookup (typeRep (Proxy :: Proxy a)) . unMap
64+lookup tVect = fromAny . (anys tVect V.!)
65+ <$> binarySearch (typeRepFingerprint $ typeRep $ Proxy @a)
66+ (fingerprintAs tVect)
67+ (fingerprintBs tVect)
4568
69+-- | Returns the size of the 'TypeRepMap'.
4670size :: TypeRepMap f -> Int
47−size = LMap.size . unMap
71+size = Unboxed.length . fingerprintAs
72+
73+-- | Returns the index is found.
74+binarySearch :: Fingerprint -> Unboxed.Vector Word64 -> Unboxed.Vector Word64 -> Maybe Int
75+binarySearch (Fingerprint a b) fpAs fpBs =
76+ let
77+ !(I# len) = Unboxed.length fpAs
78+ checkfpBs :: Int# -> Maybe Int
79+ checkfpBs i =
80+ case i <# len of
81+ 0# -> Nothing
82+ _ ->
83+ if a /= Unboxed.unsafeIndex fpAs (I# i)
84+ then Nothing
85+ else if b == Unboxed.unsafeIndex fpBs (I# i)
86+ then Just (I# i)
87+ else checkfpBs (i +# 1#)
88+ in
89+ inline (checkfpBs (binSearchHelp (-1#) len))
90+ where
91+ binSearchHelp :: Int# -> Int# -> Int#
92+ binSearchHelp l r = case l <# (r -# 1#) of
93+ 0# -> r
94+ _ ->
95+ let m = uncheckedIShiftRA# (l +# r) 1# in
96+ if Unboxed.unsafeIndex fpAs (I# m) < a
97+ then binSearchHelp m r
98+ else binSearchHelp l m
99+
100+----------------------------------------------------------------------------
101+-- Functions for testing and benchmarking
102+----------------------------------------------------------------------------
103+
104+data TF f where
105+ TF :: Typeable a => f a -> TF f
106+
107+fromF :: Typeable a => f a -> Proxy a
108+fromF _ = Proxy
109+
110+fromList :: forall f . [TF f] -> TypeRepMap f
111+fromList tfs = TypeRepMap (Unboxed.fromList fpAs) (Unboxed.fromList fpBs) (V.fromList ans)
112+ where
113+ (fpAs, fpBs) = unzip $ fmap (\(Fingerprint a b) -> (a, b)) fps
114+ (fps, ans) = unzip $ sortWith fst $ map (fp &&& an) tfs
115+
116+ fp :: TF f -> Fingerprint
117+ fp (TF x) = typeRepFingerprint $ typeRep $ fromF x
48118
49−keys :: TypeRepMap f -> [TypeRep]
50−keys = LMap.keys . unMap
119+ an :: TF f -> Any
120+ an (TF x) = unsafeCoerce x