@@ -1,3 +1,12 @@
1+{-# LANGUAGE BangPatterns #-}
2+{-# LANGUAGE GADTs #-}
3+{-# LANGUAGE KindSignatures #-}
4+{-# LANGUAGE MagicHash #-}
5+{-# LANGUAGE PolyKinds #-}
6+{-# LANGUAGE TypeFamilies #-}
7+
8+{-# OPTIONS_GHC -ddump-simpl -dsuppress-idinfo -dsuppress-coercions -dsuppress-type-applications -dsuppress-uniques -dsuppress-module-prefixes #-}
9+
110module Data.TypeRep.Map
211 ( -- * Map type
312 TypeRepMap (..)
@@ -15,4 +24,147 @@ module Data.TypeRep.Map
1524
1625import Prelude hiding (lookup)
1726
18−import Data.TypeRep.CacheMap
27+import Control.Arrow ((&&&))
28+import Data.Kind (Type)
29+import Data.Proxy (Proxy (..))
30+import Data.Typeable (Typeable, typeRep, typeRepFingerprint)
31+import Data.Word (Word64)
32+import GHC.Base (Any, Int (..), Int#, (*#), (+#), (<#))
33+import GHC.Exts (inline, sortWith)
34+import GHC.Fingerprint (Fingerprint (..))
35+import GHC.Prim (eqWord#, ltWord#)
36+import GHC.Word (Word64 (..))
37+import Unsafe.Coerce (unsafeCoerce)
38+
39+import qualified Data.Vector as V
40+import qualified Data.Vector.Unboxed as Unboxed
41+
42+data TypeRepMap (f :: k -> Type) = TypeRepMap
43+ { fingerprintAs :: Unboxed.Vector Word64
44+ , fingerprintBs :: Unboxed.Vector Word64
45+ , anys :: V.Vector Any
46+ }
47+
48+fromAny :: Any -> f a
49+fromAny = unsafeCoerce
50+
51+-- | Empty structure.
52+empty :: TypeRepMap f
53+empty = TypeRepMap mempty mempty mempty
54+
55+-- | Inserts the value with its type as a key.
56+insert :: forall a f . Typeable a => f a -> TypeRepMap f -> TypeRepMap f
57+insert x = fromListPairs . addX . toPairList
58+ where
59+ toPairList :: TypeRepMap f -> [(Fingerprint, Any)]
60+ toPairList (TypeRepMap as bs ans) = zip (zipWith Fingerprint
61+ (Unboxed.toList as)
62+ (Unboxed.toList bs))
63+ (V.toList ans)
64+
65+ pairX :: (Fingerprint, Any)
66+ pairX@(fpX, _) = (calcFp x, unsafeCoerce x)
67+
68+ addX :: [(Fingerprint, Any)] -> [(Fingerprint, Any)]
69+ addX l = pairX : filter ((/= fpX) . fst) l
70+{-# INLINE insert #-}
71+
72+-- | Looks up the value at the type.
73+-- >>> let x = lookup $ insert (11 :: Int) empty
74+-- >>> x :: Maybe Int
75+-- Just 11
76+-- >>> x :: Maybe ()
77+-- Nothing
78+lookup :: forall a f . Typeable a => TypeRepMap f -> Maybe (f a)
79+lookup tVect = fromAny . (anys tVect V.!)
80+ <$> cachedBinarySearch (typeRepFingerprint $ typeRep $ Proxy @a)
81+ (fingerprintAs tVect)
82+ (fingerprintBs tVect)
83+-- {-# INLINE lookup #-}
84+
85+-- | Returns the size of the 'TypeRepMap'.
86+size :: TypeRepMap f -> Int
87+size = Unboxed.length . fingerprintAs
88+
89+-- | Binary searched based on this article
90+-- http://bannalia.blogspot.com/2015/06/cache-friendly-binary-search.html
91+-- with modification for our two-vector search case.
92+cachedBinarySearch :: Fingerprint -> Unboxed.Vector Word64 -> Unboxed.Vector Word64 -> Maybe Int
93+cachedBinarySearch (Fingerprint (W64# a) (W64# b)) fpAs fpBs = inline (go 0#)
94+ where
95+ go :: Int# -> Maybe Int
96+ go i = case i <# len of
97+ 0# -> Nothing
98+ _ -> let !(W64# valA) = Unboxed.unsafeIndex fpAs (I# i) in case a `ltWord#` valA of
99+ 0# -> case a `eqWord#` valA of
100+ 0# -> go (2# *# i +# 2#)
101+ _ -> let !(W64# valB) = Unboxed.unsafeIndex fpBs (I# i) in case b `eqWord#` valB of
102+ 0# -> case b `ltWord#` valB of
103+ 0# -> go (2# *# i +# 2#)
104+ _ -> go (2# *# i +# 1#)
105+ _ -> Just (I# i)
106+ _ -> go (2# *# i +# 1#)
107+
108+ len :: Int#
109+ len = let !(I# l) = Unboxed.length fpAs in l
110+-- {-# INLINE cachedBinarySearch #-}
111+
112+----------------------------------------------------------------------------
113+-- Functions for testing and benchmarking
114+----------------------------------------------------------------------------
115+
116+data TF f where
117+ TF :: Typeable a => f a -> TF f
118+
119+fromF :: Typeable a => f a -> Proxy a
120+fromF _ = Proxy
121+
122+calcFp :: Typeable a => f a -> Fingerprint
123+calcFp = typeRepFingerprint . typeRep . fromF
124+
125+fromListPairs :: [(Fingerprint, Any)] -> TypeRepMap f
126+fromListPairs kvs = TypeRepMap (Unboxed.fromList fpAs) (Unboxed.fromList fpBs) (V.fromList ans)
127+ where
128+ (fpAs, fpBs) = unzip $ map (\(Fingerprint a b) -> (a, b)) fps
129+ (fps, ans) = unzip $ breadthFirst $ fromListToTree $ sortWith fst kvs
130+
131+fromList :: forall f . [TF f] -> TypeRepMap f
132+fromList = fromListPairs . map (fp &&& an)
133+ where
134+ fp :: TF f -> Fingerprint
135+ fp (TF x) = calcFp x
136+
137+ an :: TF f -> Any
138+ an (TF x) = unsafeCoerce x
139+
140+----------------------------------------------------------------------------
141+-- Tree
142+----------------------------------------------------------------------------
143+
144+data Tree a = Leaf | Node a (Tree a) (Tree a)
145+ deriving (Show)
146+
147+fromListToTree :: [a] -> Tree a
148+fromListToTree [] = Leaf
149+fromListToTree xs =
150+ let len = length xs
151+ in case splitAt (len `div` 2) xs of
152+ ([], []) -> Leaf
153+ (ls, []) -> Node (last ls) (fromListToTree $ init ls) Leaf
154+ ([], (r:rs)) -> Node r Leaf $ fromListToTree rs
155+ (ls, r:rs) -> Node r (fromListToTree ls) (fromListToTree rs)
156+
157+breadthFirst :: Tree a -> [a]
158+breadthFirst tree = bf [tree]
159+ where
160+ bf :: [Tree a] -> [a]
161+ bf [] = []
162+ bf xs = mapValue xs ++ bf (concatMap children xs)
163+
164+ mapValue :: [Tree a] -> [a]
165+ mapValue [] = []
166+ mapValue (Leaf : xs) = mapValue xs
167+ mapValue (Node a _ _ : xs) = a : mapValue xs
168+
169+ children Leaf = []
170+ children (Node _ l r) = [l, r]