git.iko.soy/typerep-map
on 92b6021 help sign in
typerep-map/src/Data/TypeRepMap/Internal.hs
403 lines 14 kB history raw
1{-# LANGUAGE AllowAmbiguousTypes #-}
2{-# LANGUAGE BangPatterns #-}
3{-# LANGUAGE DataKinds #-}
4{-# LANGUAGE GADTs #-}
5{-# LANGUAGE InstanceSigs #-}
6{-# LANGUAGE KindSignatures #-}
7{-# LANGUAGE MagicHash #-}
8{-# LANGUAGE PolyKinds #-}
9{-# LANGUAGE Rank2Types #-}
10{-# LANGUAGE TypeFamilies #-}
11{-# LANGUAGE TypeInType #-}
12{-# LANGUAGE ViewPatterns #-}
14-- {-# OPTIONS_GHC -ddump-simpl -dsuppress-idinfo -dsuppress-coercions -dsuppress-type-applications -dsuppress-uniques -dsuppress-module-prefixes #-}
16-- | Internal API for 'TypeRepMap' and operations on it. The functions here do
17-- not have any stability guarantees and can change between minor versions.
18--
19-- If you need to use this module for purposes other than tests,
20-- create an issue.
21--
22module Data.TypeRepMap.Internal where
24import Prelude hiding (lookup)
26import Control.Monad.ST (runST)
27import Control.Monad.Zip (mzip)
28import Data.Function (on)
29import Data.IntMap.Strict (IntMap)
30import Data.Kind (Type)
31import Data.List (intercalate, nubBy)
32import Data.Maybe (fromJust)
33import Data.Primitive.Array (Array, freezeArray, indexArray, mapArray', readArray, sizeofArray,
34 thawArray, writeArray)
35import Data.Primitive.PrimArray (PrimArray, indexPrimArray, sizeofPrimArray)
36import Data.Semigroup (Semigroup (..))
37import GHC.Base (Any, Int (..), Int#, (*#), (+#), (<#))
38import GHC.Exts (IsList (..), inline, sortWith)
39import GHC.Fingerprint (Fingerprint (..))
40import GHC.Prim (eqWord#, ltWord#)
41import GHC.Word (Word64 (..))
42import Type.Reflection (SomeTypeRep (..), TypeRep, Typeable, typeRep, withTypeable)
43import Type.Reflection.Unsafe (typeRepFingerprint)
44import Unsafe.Coerce (unsafeCoerce)
46import qualified Data.IntMap.Strict as IM
47import qualified Data.Map.Strict as Map
48import qualified GHC.Exts as GHC (fromList, toList)
50{- |
52'TypeRepMap' is a heterogeneous data structure similar in its essence to
53'Data.Map.Map' with types as keys, where each value has the type of its key. In
54addition to that, each value is wrapped in an interpretation @f@.
56Here is an example of using 'Prelude.Maybe' as an interpretation, with a
57comparison to 'Data.Map.Map':
60 'Data.Map.Map' 'Prelude.String' ('Prelude.Maybe' 'Prelude.String') 'TypeRepMap' 'Prelude.Maybe'
61--------------------------- ---------------------
62 \"Int\" -> Just \"5\" 'Prelude.Int' -> Just 5
63 \"Bool\" -> Just \"True\" 'Prelude.Bool' -> Just 'Prelude.True'
64 \"Char\" -> Nothing 'Prelude.Char' -> Nothing
67The runtime representation of 'TypeRepMap' is an array, not a tree. This makes
68'lookup' significantly more efficient.
70-}
71data TypeRepMap (f :: k -> Type) =
72 TypeRepMap
73 { fingerprintAs :: {-# UNPACK #-} !(PrimArray Word64) -- ^ first components of key fingerprints
74 , fingerprintBs :: {-# UNPACK #-} !(PrimArray Word64) -- ^ second components of key fingerprints
75 , trAnys :: {-# UNPACK #-} !(Array Any) -- ^ values stored in the map
76 , trKeys :: {-# UNPACK #-} !(Array Any) -- ^ typerep keys
77 }
78 -- ^ an unsafe constructor for 'TypeRepMap'
80-- | Shows only keys.
81instance Show (TypeRepMap f) where
82 show TypeRepMap{..} = "TypeRepMap [" ++ showKeys ++ "]"
83 where
84 showKeys :: String
85 showKeys = intercalate ", " $ toList $ mapArray' (show . anyToTypeRep) trKeys
87-- | Uses 'union' to combine 'TypeRepMap's.
88instance Semigroup (TypeRepMap f) where
89 (<>) :: TypeRepMap f -> TypeRepMap f -> TypeRepMap f
90 (<>) = union
91 {-# INLINE (<>) #-}
93instance Monoid (TypeRepMap f) where
94 mempty = TypeRepMap mempty mempty mempty mempty
95 mappend = (<>)
96 {-# INLINE mempty #-}
97 {-# INLINE mappend #-}
99-- | Returns the list of 'Fingerprint's from 'TypeRepMap'.
100toFingerprints :: TypeRepMap f -> [Fingerprint]
101toFingerprints TypeRepMap{..} =
102 zipWith Fingerprint (GHC.toList fingerprintAs) (GHC.toList fingerprintBs)
104{- |
106A 'TypeRepMap' with no values stored in it.
108prop> size empty == 0
109prop> member @a empty == False
111-}
112empty :: TypeRepMap f
113empty = mempty
114{-# INLINE empty #-}
116{- |
118Construct a 'TypeRepMap' with a single element.
120prop> size (one x) == 1
121prop> member @a (one (x :: f a)) == True
123-}
124one :: forall a f . Typeable a => f a -> TypeRepMap f
125one x = insert x empty
126{-# INLINE one #-}
128{- |
130Insert a value into a 'TypeRepMap'.
132prop> size (insert v tm) >= size tm
133prop> member @a (insert (x :: f a) tm) == True
135-}
136insert :: forall a f . Typeable a => f a -> TypeRepMap f -> TypeRepMap f
137insert x = fromTriples . addX . toTriples
138 where
139 tripleX :: (Fingerprint, Any, Any)
140 tripleX@(fpX, _, _) = (calcFp @a, toAny x, unsafeCoerce $ typeRep @a)
142 addX :: [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
143 addX l = tripleX : deleteByFst fpX l
144{-# INLINE insert #-}
146-- Extract the kind of a type. We use it to work around lack of syntax for
147-- inferred type variables (which are not subject to type applications).
148type KindOf (a :: k) = k
150{- | Delete a value from a 'TypeRepMap'.
152prop> size (delete @a tm) <= size tm
153prop> member @a (delete @a tm) == False
155>>> tm = delete @Bool $ insert (Just True) $ one (Just 'a')
156>>> size tm
158>>> member @Bool tm
159False
160>>> member @Char tm
161True
162-}
163delete :: forall a (f :: KindOf a -> Type) . Typeable a => TypeRepMap f -> TypeRepMap f
164delete = fromTriples . deleteByFst (typeFp @a) . toTriples
165{-# INLINE delete #-}
167{- |
168Update a value at a specific key with the result of the provided function. When
169the key is not a member of the map, the original map is returned.
171>>> trmap = fromList @(TypeRepMap Identity) [WrapTypeable $ Identity "a"]
172>>> lookup @String $ adjust (fmap (++ "ww")) trmap
173Just (Identity "aww")
174-}
175adjust :: forall a f . Typeable a => (f a -> f a) -> TypeRepMap f -> TypeRepMap f
176adjust fun tr = case cachedBinarySearch (typeFp @a) (fingerprintAs tr) (fingerprintBs tr) of
177 Nothing -> tr
178 Just i -> tr {trAnys = changeAnyArr i (trAnys tr)}
179 where
180 changeAnyArr :: Int -> Array Any -> Array Any
181 changeAnyArr i trAs = runST $ do
182 let n = sizeofArray trAs
183 mutArr <- thawArray trAs 0 n
184 a <- toAny . fun . fromAny <$> readArray mutArr i
185 writeArray mutArr i a
186 freezeArray mutArr 0 n
187{-# INLINE adjust #-}
189{- | Map over the elements of a 'TypeRepMap'.
191>>> tm = insert (Identity True) $ one (Identity 'a')
192>>> lookup @Bool tm
193Just (Identity True)
194>>> lookup @Char tm
195Just (Identity 'a')
196>>> tm2 = hoist ((:[]) . runIdentity) tm
197>>> lookup @Bool tm2
198Just [True]
199>>> lookup @Char tm2
200Just "a"
201-}
202hoist :: (forall x. f x -> g x) -> TypeRepMap f -> TypeRepMap g
203hoist f (TypeRepMap as bs ans ks) = TypeRepMap as bs (mapArray' (toAny . f . fromAny) ans) ks
204{-# INLINE hoist #-}
206hoistA :: (Applicative t) => (forall x. f x -> t (g x)) -> TypeRepMap f -> t (TypeRepMap g)
207hoistA f (TypeRepMap as bs (toList -> ans) ks) = (\l -> TypeRepMap as bs (fromList $ map toAny l) ks)
208 <$> traverse (f . fromAny) ans
209{-# INLINE hoistA #-}
211hoistWithKey :: forall f g. (forall x. Typeable x => f x -> g x) -> TypeRepMap f -> TypeRepMap g
212hoistWithKey f (TypeRepMap as bs ans ks) = TypeRepMap as bs newAns ks
213 where
214 newAns = mapArray' mapAns (mzip ans ks)
215 mapAns (a, k) = toAny $ withTr (unsafeCoerce k) $ fromAny a
217 withTr :: forall x. TypeRep x -> f x -> g x
218 withTr t = withTypeable t f
219{-# INLINE hoistWithKey #-}
221-- | The union of two 'TypeRepMap's using a combining function.
222unionWith :: (forall x. f x -> f x -> f x) -> TypeRepMap f -> TypeRepMap f -> TypeRepMap f
223unionWith f m1 m2 = fromTriples
224 $ toTripleList
225 $ Map.unionWith combine
226 (fromTripleList $ toTriples m1)
227 (fromTripleList $ toTriples m2)
228 where
229 combine :: (Any, Any) -> (Any, Any) -> (Any, Any)
230 combine (av, ak) (bv, _) = (toAny $ f (fromAny av) (fromAny bv), ak)
232 fromTripleList :: Ord a => [(a, b, c)] -> Map.Map a (b, c)
233 fromTripleList = Map.fromList . map (\(a, b, c) -> (a, (b, c)))
235 toTripleList :: Map.Map a (b, c) -> [(a, b, c)]
236 toTripleList = map (\(a, (b, c)) -> (a, b, c)) . Map.toList
237{-# INLINE unionWith #-}
239-- | The (left-biased) union of two 'TypeRepMap's. It prefers the first map when
240-- duplicate keys are encountered, i.e. @'union' == 'unionWith' const@.
241union :: TypeRepMap f -> TypeRepMap f -> TypeRepMap f
242union = unionWith const
243{-# INLINE union #-}
245{- | Check if a value of the given type is present in a 'TypeRepMap'.
247>>> member @Char $ one (Identity 'a')
248True
249>>> member @Bool $ one (Identity 'a')
250False
251-}
252member :: forall a (f :: KindOf a -> Type) . Typeable a => TypeRepMap f -> Bool
253member tm = case lookup @a tm of
254 Nothing -> False
255 Just _ -> True
256{-# INLINE member #-}
258{- | Lookup a value of the given type in a 'TypeRepMap'.
260>>> x = lookup $ insert (Identity (11 :: Int)) empty
261>>> x :: Maybe (Identity Int)
262Just (Identity 11)
263>>> x :: Maybe (Identity ())
264Nothing
265-}
266lookup :: forall a f . Typeable a => TypeRepMap f -> Maybe (f a)
267lookup tVect = fromAny . (trAnys tVect `indexArray`)
268 <$> cachedBinarySearch (typeFp @a)
269 (fingerprintAs tVect)
270 (fingerprintBs tVect)
271{-# INLINE lookup #-}
273-- | Get the amount of elements in a 'TypeRepMap'.
274size :: TypeRepMap f -> Int
275size = sizeofPrimArray . fingerprintAs
276{-# INLINE size #-}
278-- | Return the list of 'SomeTypeRep' from the keys.
279keys :: TypeRepMap f -> [SomeTypeRep]
280keys TypeRepMap{..} = SomeTypeRep . anyToTypeRep <$> toList trKeys
281{-# INLINE keys #-}
283-- | Binary searched based on this article
284-- http://bannalia.blogspot.com/2015/06/cache-friendly-binary-search.html
285-- with modification for our two-vector search case.
286cachedBinarySearch :: Fingerprint -> PrimArray Word64 -> PrimArray Word64 -> Maybe Int
287cachedBinarySearch (Fingerprint (W64# a) (W64# b)) fpAs fpBs = inline (go 0#)
288 where
289 go :: Int# -> Maybe Int
290 go i = case i <# len of
291 0# -> Nothing
292 _ -> let !(W64# valA) = indexPrimArray fpAs (I# i) in case a `ltWord#` valA of
293 0# -> case a `eqWord#` valA of
294 0# -> go (2# *# i +# 2#)
295 _ -> let !(W64# valB) = indexPrimArray fpBs (I# i) in case b `eqWord#` valB of
296 0# -> case b `ltWord#` valB of
297 0# -> go (2# *# i +# 2#)
298 _ -> go (2# *# i +# 1#)
299 _ -> Just (I# i)
300 _ -> go (2# *# i +# 1#)
302 len :: Int#
303 len = let !(I# l) = sizeofPrimArray fpAs in l
304{-# INLINE cachedBinarySearch #-}
306----------------------------------------------------------------------------
307-- Internal functions
308----------------------------------------------------------------------------
310toAny :: f a -> Any
311toAny = unsafeCoerce
313fromAny :: Any -> f a
314fromAny = unsafeCoerce
316anyToTypeRep :: Any -> TypeRep f
317anyToTypeRep = unsafeCoerce
319typeFp :: forall a . Typeable a => Fingerprint
320typeFp = typeRepFingerprint $ typeRep @a
321{-# INLINE typeFp #-}
323toTriples :: TypeRepMap f -> [(Fingerprint, Any, Any)]
324toTriples tm = zip3 (toFingerprints tm) (GHC.toList $ trAnys tm) (GHC.toList $ trKeys tm)
326deleteByFst :: Eq a => a -> [(a, b, c)] -> [(a, b, c)]
327deleteByFst x = filter ((/= x) . fst3)
329nubByFst :: (Eq a) => [(a, b, c)] -> [(a, b, c)]
330nubByFst = nubBy ((==) `on` fst3)
332fst3 :: (a, b, c) -> a
333fst3 (a, _, _) = a
335----------------------------------------------------------------------------
336-- Functions for testing and benchmarking
337----------------------------------------------------------------------------
339-- | Existential wrapper around 'Typeable' indexed by @f@ type parameter.
340-- Useful for 'TypeRepMap' structure creation form list of 'WrapTypeable's.
341data WrapTypeable f where
342 WrapTypeable :: Typeable a => f a -> WrapTypeable f
344instance Show (WrapTypeable f) where
345 show (WrapTypeable (_ :: f a)) = show $ calcFp @a
347{- |
349prop> fromList . toList == 'id'
351Creates 'TypeRepMap' from a list of 'WrapTypeable's.
353>>> show $ fromList [WrapTypeable $ Identity True, WrapTypeable $ Identity 'a']
354TypeRepMap [Bool, Char]
357-}
358instance IsList (TypeRepMap f) where
359 type Item (TypeRepMap f) = WrapTypeable f
361 fromList :: [WrapTypeable f] -> TypeRepMap f
362 fromList = fromTriples . map (\x -> (fp x, an x, k x))
363 where
364 fp :: WrapTypeable f -> Fingerprint
365 fp (WrapTypeable (_ :: f a)) = calcFp @a
367 an :: WrapTypeable f -> Any
368 an (WrapTypeable x) = toAny x
370 k :: WrapTypeable f -> Any
371 k (WrapTypeable (_ :: f a)) = unsafeCoerce $ typeRep @a
373 toList :: TypeRepMap f -> [WrapTypeable f]
374 toList = map toWrapTypeable . toTriples
375 where
376 toWrapTypeable :: (Fingerprint, Any, Any) -> WrapTypeable f
377 toWrapTypeable (_, an, k) = withTypeable (unsafeCoerce k) $ fromAny an
379calcFp :: forall a . Typeable a => Fingerprint
380calcFp = typeRepFingerprint $ typeRep @a
382fromTriples :: [(Fingerprint, Any, Any)] -> TypeRepMap f
383fromTriples kvs = TypeRepMap (GHC.fromList fpAs) (GHC.fromList fpBs) (GHC.fromList ans) (GHC.fromList ks)
384 where
385 (fpAs, fpBs) = unzip $ map (\(Fingerprint a b) -> (a, b)) fps
386 (fps, ans, ks) = unzip3 $ fromSortedList $ sortWith fst3 $ nubByFst kvs
388----------------------------------------------------------------------------
389-- Tree-like conversion
390----------------------------------------------------------------------------
392fromSortedList :: forall a . [a] -> [a]
393fromSortedList l = IM.elems $ fst $ go 0 0 mempty (IM.fromList $ zip [0..] l)
394 where
395 -- state monad could be used here, but it's another dependency
396 go :: Int -> Int -> IntMap a -> IntMap a -> (IntMap a, Int)
397 go i first result vector =
398 if i >= IM.size vector
399 then (result, first)
400 else do
401 let (newResult, newFirst) = go (2 * i + 1) first result vector
402 let withCur = IM.insert i (fromJust $ IM.lookup newFirst vector) newResult
403 go (2 * i + 2) (newFirst + 1) withCur vector