@@ -1,3 +1,16 @@
1+{-# LANGUAGE AllowAmbiguousTypes #-}
2+{-# LANGUAGE BangPatterns #-}
3+{-# LANGUAGE DataKinds #-}
4+{-# LANGUAGE GADTs #-}
5+{-# LANGUAGE KindSignatures #-}
6+{-# LANGUAGE MagicHash #-}
7+{-# LANGUAGE PolyKinds #-}
8+{-# LANGUAGE Rank2Types #-}
9+{-# LANGUAGE TypeFamilies #-}
10+{-# LANGUAGE TypeInType #-}
11+
12+-- {-# OPTIONS_GHC -ddump-simpl -dsuppress-idinfo -dsuppress-coercions -dsuppress-type-applications -dsuppress-uniques -dsuppress-module-prefixes #-}
13+
114-- | This module introduces the 'TypeRepMap' data structure and functions to work with it.
215
316module Data.TypeRep.Map
@@ -23,8 +36,252 @@ module Data.TypeRep.Map
2336 -- * Helpful testing functions
2437 , TF (..)
2538 , fromList
39+ , toFingerprints
2640 ) where
2741
2842import Prelude hiding (lookup)
2943
30−import Data.TypeRep.CacheMap
44+import Control.Arrow ((&&&))
45+import Data.Function (on)
46+import Data.IntMap.Strict (IntMap)
47+import Data.Kind (Type)
48+import Data.List (nubBy)
49+import Data.Maybe (fromJust)
50+import Data.Primitive.Array (Array, indexArray, mapArray')
51+import Data.Primitive.PrimArray (PrimArray, indexPrimArray, sizeofPrimArray)
52+import Data.Proxy (Proxy (..))
53+import Data.Typeable (Typeable, typeRep, typeRepFingerprint)
54+import GHC.Base (Any, Int (..), Int#, (*#), (+#), (<#))
55+import GHC.Exts (inline, sortWith)
56+import GHC.Fingerprint (Fingerprint (..))
57+import GHC.Prim (eqWord#, ltWord#)
58+import GHC.Word (Word64 (..))
59+import Unsafe.Coerce (unsafeCoerce)
60+
61+import qualified Data.IntMap.Strict as IM
62+import qualified Data.Map.Strict as Map
63+import qualified GHC.Exts as GHC (fromList, toList)
64+
65+-- | Map-like data structure that keeps types as keys.
66+data TypeRepMap (f :: k -> Type) = TypeRepMap
67+ { fingerprintAs :: {-# UNPACK #-} !(PrimArray Word64)
68+ , fingerprintBs :: {-# UNPACK #-} !(PrimArray Word64)
69+ , anys :: {-# UNPACK #-} !(Array Any)
70+ }
71+
72+-- | Shows only 'Fingerprint's.
73+instance Show (TypeRepMap f) where
74+ show = show . toFingerprints
75+
76+-- | Returns the list of 'Fingerprint's from 'TypeRepMap'.
77+toFingerprints :: TypeRepMap f -> [Fingerprint]
78+toFingerprints TypeRepMap{..} =
79+ zipWith Fingerprint (GHC.toList fingerprintAs) (GHC.toList fingerprintBs)
80+
81+
82+-- | Empty structure.
83+empty :: TypeRepMap f
84+empty = TypeRepMap mempty mempty mempty
85+
86+-- | Construct a 'TypeRepMap' with a single element.
87+one :: forall a f . Typeable a => f a -> TypeRepMap f
88+one x = insert x empty
89+{-# INLINE one #-}
90+
91+-- | Inserts the value with its type as a key.
92+insert :: forall a f . Typeable a => f a -> TypeRepMap f -> TypeRepMap f
93+insert x = fromListPairs . addX . toPairList
94+ where
95+ pairX :: (Fingerprint, Any)
96+ pairX@(fpX, _) = (calcFp x, toAny x)
97+
98+ addX :: [(Fingerprint, Any)] -> [(Fingerprint, Any)]
99+ addX l = pairX : deleteByFst fpX l
100+{-# INLINE insert #-}
101+
102+-- Helper type to workaround some issues with types
103+type KindOf (a :: k) = k
104+
105+{- | Deletes value from list.
106+
107+>>> let trMap = delete @Bool $ insert (Identity True) $ one (Identity 'a')
108+>>> size trMap
109+1
110+>>> member @Bool trMap
111+False
112+>>> member @Char trMap
113+True
114+-}
115+delete :: forall a (f :: KindOf a -> Type) . Typeable a => TypeRepMap f -> TypeRepMap f
116+delete = fromListPairs . deleteByFst (typeFp @a) . toPairList
117+{-# INLINE delete #-}
118+
119+{- | Map over elements.
120+
121+>>> let trMap = insert (Identity True) $ one (Identity 'a')
122+>>> lookup @Bool trMap
123+Just (Identity True)
124+>>> lookup @Char trMap
125+Just (Identity 'a')
126+>>> let trMap2 = hoist ((:[]) . runIdentity) trMap
127+>>> lookup @Bool trMap2
128+Just [True]
129+>>> lookup @Char trMap2
130+Just "a"
131+-}
132+hoist :: (forall x. f x -> g x) -> TypeRepMap f -> TypeRepMap g
133+hoist f (TypeRepMap as bs ans) = TypeRepMap as bs $ mapArray' (toAny . f . fromAny) ans
134+{-# INLINE hoist #-}
135+
136+-- | The union with a combining function.
137+unionWith :: (forall x. f x -> f x -> f x) -> TypeRepMap f -> TypeRepMap f -> TypeRepMap f
138+unionWith f m1 m2 = fromListPairs
139+ $ Map.toList
140+ $ Map.unionWith combine
141+ (Map.fromList $ toPairList m1)
142+ (Map.fromList $ toPairList m2)
143+ where
144+ combine :: Any -> Any -> Any
145+ combine a b = toAny $ f (fromAny a) (fromAny b)
146+{-# INLINE unionWith #-}
147+
148+-- | The (left-biased) union of two maps. It prefers the first map when
149+-- duplicate keys are encountered, i.e. @'union' == 'unionWith' const@.
150+union :: TypeRepMap f -> TypeRepMap f -> TypeRepMap f
151+union = unionWith const
152+{-# INLINE union #-}
153+
154+{- | Returns 'True' if there exist value of given type.
155+
156+>>> member @Char $ one (Identity 'a')
157+True
158+>>> member @Bool $ one (Identity 'a')
159+False
160+-}
161+member :: forall a (f :: KindOf a -> Type) . Typeable a => TypeRepMap f -> Bool
162+member trMap = case lookup @a trMap of
163+ Nothing -> False
164+ Just _ -> True
165+{-# INLINE member #-}
166+
167+{- | Looks up the value at the type.
168+
169+>>> let x = lookup $ insert (Identity (11 :: Int)) empty
170+>>> x :: Maybe (Identity Int)
171+Just (Identity 11)
172+>>> x :: Maybe (Identity ())
173+Nothing
174+-}
175+lookup :: forall a f . Typeable a => TypeRepMap f -> Maybe (f a)
176+lookup tVect = fromAny . (anys tVect `indexArray`)
177+ <$> cachedBinarySearch (typeFp @a)
178+ (fingerprintAs tVect)
179+ (fingerprintBs tVect)
180+{-# INLINE lookup #-}
181+
182+-- | Returns the size of the 'TypeRepMap'.
183+size :: TypeRepMap f -> Int
184+size = sizeofPrimArray . fingerprintAs
185+{-# INLINE size #-}
186+
187+-- | Binary searched based on this article
188+-- http://bannalia.blogspot.com/2015/06/cache-friendly-binary-search.html
189+-- with modification for our two-vector search case.
190+cachedBinarySearch :: Fingerprint -> PrimArray Word64 -> PrimArray Word64 -> Maybe Int
191+cachedBinarySearch (Fingerprint (W64# a) (W64# b)) fpAs fpBs = inline (go 0#)
192+ where
193+ go :: Int# -> Maybe Int
194+ go i = case i <# len of
195+ 0# -> Nothing
196+ _ -> let !(W64# valA) = indexPrimArray fpAs (I# i) in case a `ltWord#` valA of
197+ 0# -> case a `eqWord#` valA of
198+ 0# -> go (2# *# i +# 2#)
199+ _ -> let !(W64# valB) = indexPrimArray fpBs (I# i) in case b `eqWord#` valB of
200+ 0# -> case b `ltWord#` valB of
201+ 0# -> go (2# *# i +# 2#)
202+ _ -> go (2# *# i +# 1#)
203+ _ -> Just (I# i)
204+ _ -> go (2# *# i +# 1#)
205+
206+ len :: Int#
207+ len = let !(I# l) = sizeofPrimArray fpAs in l
208+{-# INLINE cachedBinarySearch #-}
209+
210+----------------------------------------------------------------------------
211+-- Internal functions
212+----------------------------------------------------------------------------
213+
214+toAny :: f a -> Any
215+toAny = unsafeCoerce
216+
217+fromAny :: Any -> f a
218+fromAny = unsafeCoerce
219+
220+typeFp :: forall a . Typeable a => Fingerprint
221+typeFp = typeRepFingerprint $ typeRep $ Proxy @a
222+{-# INLINE typeFp #-}
223+
224+toPairList :: TypeRepMap f -> [(Fingerprint, Any)]
225+toPairList trMap = zip (toFingerprints trMap) (GHC.toList $ anys trMap)
226+
227+deleteByFst :: Eq a => a -> [(a, b)] -> [(a, b)]
228+deleteByFst x = filter ((/= x) . fst)
229+
230+nubByFst :: (Eq a) => [(a, b)] -> [(a, b)]
231+nubByFst = nubBy ((==) `on` fst)
232+
233+----------------------------------------------------------------------------
234+-- Functions for testing and benchmarking
235+----------------------------------------------------------------------------
236+
237+-- | Existential wrapper around 'Typeable' indexed by @f@ type parameter.
238+-- Useful for 'TypeRepMap' structure creation form list of 'TF's.
239+data TF f where
240+ TF :: Typeable a => f a -> TF f
241+
242+instance Show (TF f) where
243+ show (TF tf) = show $ calcFp tf
244+
245+{- | Creates 'TypeRepMap' from a list of 'TF's.
246+
247+>>> size $ fromList [TF $ Identity True, TF $ Identity 'a']
248+2
249+
250+-}
251+fromList :: forall f . [TF f] -> TypeRepMap f
252+fromList = fromListPairs . map (fp &&& an)
253+ where
254+ fp :: TF f -> Fingerprint
255+ fp (TF x) = calcFp x
256+
257+ an :: TF f -> Any
258+ an (TF x) = toAny x
259+
260+fromF :: Typeable a => f a -> Proxy a
261+fromF _ = Proxy
262+
263+calcFp :: Typeable a => f a -> Fingerprint
264+calcFp = typeRepFingerprint . typeRep . fromF
265+
266+fromListPairs :: [(Fingerprint, Any)] -> TypeRepMap f
267+fromListPairs kvs = TypeRepMap (GHC.fromList fpAs) (GHC.fromList fpBs) (GHC.fromList ans)
268+ where
269+ (fpAs, fpBs) = unzip $ map (\(Fingerprint a b) -> (a, b)) fps
270+ (fps, ans) = unzip $ fromSortedList $ sortWith fst $ nubByFst kvs
271+
272+----------------------------------------------------------------------------
273+-- Tree-like conversion
274+----------------------------------------------------------------------------
275+
276+fromSortedList :: forall a . [a] -> [a]
277+fromSortedList l = IM.elems $ fst $ go 0 0 mempty (IM.fromList $ zip [0..] l)
278+ where
279+ -- state monad could be used here, but it's another dependency
280+ go :: Int -> Int -> IntMap a -> IntMap a -> (IntMap a, Int)
281+ go i first result vector =
282+ if i >= IM.size vector
283+ then (result, first)
284+ else do
285+ let (newResult, newFirst) = go (2 * i + 1) first result vector
286+ let withCur = IM.insert i (fromJust $ IM.lookup newFirst vector) newResult
287+ go (2 * i + 2) (newFirst + 1) withCur vector