@@ -1,47 +1,69 @@
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 #-}
1+{-# OPTIONS_GHC -Wno-missing-export-lists #-}
2+
3+{-# LANGUAGE AllowAmbiguousTypes #-}
4+{-# LANGUAGE BangPatterns #-}
5+{-# LANGUAGE CPP #-}
6+{-# LANGUAGE DataKinds #-}
7+{-# LANGUAGE GADTs #-}
8+{-# LANGUAGE InstanceSigs #-}
9+{-# LANGUAGE MagicHash #-}
10+{-# LANGUAGE PolyKinds #-}
11+{-# LANGUAGE Rank2Types #-}
12+{-# LANGUAGE TypeFamilies #-}
13+{-# LANGUAGE TypeInType #-}
14+{-# LANGUAGE ViewPatterns #-}
15+
16+#if __GLASGOW_HASKELL__ >= 806
17+{-# LANGUAGE QuantifiedConstraints #-}
18+#endif
1219
1320-- {-# OPTIONS_GHC -ddump-simpl -dsuppress-idinfo -dsuppress-coercions -dsuppress-type-applications -dsuppress-uniques -dsuppress-module-prefixes #-}
1421
15−-- | Internal API for 'TypeRepMap' and operations on it. The functions here do
16−-- not have any stability guarantees and can change between minor versions.
17−--
18−-- If you need to use this module for purposes other than tests,
19−-- create an issue.
20−--
22+{- |
23+Copyright: (c) 2017-2020 Kowainik
24+SPDX-License-Identifier: MPL-2.0
25+Maintainer: Kowainik <xrom.xkov@gmail.com>
26+
27+Internal API for 'TypeRepMap' and operations on it. The functions here do
28+not have any stability guarantees and can change between minor versions.
29+
30+If you need to use this module for purposes other than tests,
31+create an issue.
32+-}
33+
34+#include "MachDeps.h"
35+
2136module Data.TypeRepMap.Internal where
2237
2338import Prelude hiding (lookup)
2439
25−import Control.Arrow ((&&&))
40+import Control.DeepSeq
41+import Control.Monad.ST (ST, runST)
42+import Control.Monad.Zip (mzip)
2643import Data.Function (on)
27−import Data.IntMap.Strict (IntMap)
2844import Data.Kind (Type)
29−import Data.List (nubBy)
30−import Data.Maybe (fromJust)
31−import Data.Primitive.Array (Array, indexArray, mapArray')
32−import Data.Primitive.PrimArray (PrimArray, indexPrimArray, sizeofPrimArray)
33−import Data.Proxy (Proxy (..))
34−import Data.Semigroup (Semigroup (..))
35−import Data.Typeable (Typeable, typeRep, typeRepFingerprint)
45+import Data.Type.Equality ((:~:) (..), TestEquality (..))
46+import Data.List (intercalate, nubBy)
47+import Data.Maybe (fromMaybe)
48+import Data.Primitive.Array (Array, MutableArray, indexArray, mapArray', sizeofArray,
49+ thawArray, unsafeFreezeArray, writeArray)
50+import Data.Primitive.PrimArray (writePrimArray, newPrimArray, unsafeFreezePrimArray, primArrayToList, MutablePrimArray, primArrayFromListN, PrimArray, indexPrimArray, sizeofPrimArray)
51+import Data.Semigroup (Semigroup (..), All(..))
3652import GHC.Base (Any, Int (..), Int#, (*#), (+#), (<#))
37−import GHC.Exts (inline, sortWith)
53+import GHC.Exts (IsList (..), inline, sortWith)
3854import GHC.Fingerprint (Fingerprint (..))
55+#if WORD_SIZE_IN_BITS >= 64
3956import GHC.Prim (eqWord#, ltWord#)
57+#else
58+import GHC.IntWord64 (eqWord64#, ltWord64#)
59+#define eqWord eqWord64
60+#define ltWord ltWord64
61+#endif
4062import GHC.Word (Word64 (..))
63+import Type.Reflection (SomeTypeRep (..), TypeRep, Typeable, typeRep, withTypeable)
64+import Type.Reflection.Unsafe (typeRepFingerprint)
4165import Unsafe.Coerce (unsafeCoerce)
4266
43−import qualified Data.IntMap.Strict as IM
44−import qualified Data.Map.Strict as Map
4567import qualified GHC.Exts as GHC (fromList, toList)
4668
4769{- |
@@ -69,13 +91,20 @@ data TypeRepMap (f :: k -> Type) =
6991 TypeRepMap
7092 { fingerprintAs :: {-# UNPACK #-} !(PrimArray Word64) -- ^ first components of key fingerprints
7193 , fingerprintBs :: {-# UNPACK #-} !(PrimArray Word64) -- ^ second components of key fingerprints
72− , anys :: {-# UNPACK #-} !(Array Any) -- ^ values stored in the map
94+ , trAnys :: {-# UNPACK #-} !(Array Any) -- ^ values stored in the map
95+ , trKeys :: {-# UNPACK #-} !(Array Any) -- ^ typerep keys
7396 }
7497 -- ^ an unsafe constructor for 'TypeRepMap'
7598
76−-- | Shows only 'Fingerprint's.
99+instance NFData (TypeRepMap f) where
100+ rnf x = rnf (keys x) `seq` ()
101+
102+-- | Shows only keys.
77103instance Show (TypeRepMap f) where
78− show = show . toFingerprints
104+ show TypeRepMap{..} = "TypeRepMap [" ++ showKeys ++ "]"
105+ where
106+ showKeys :: String
107+ showKeys = intercalate ", " $ toList $ mapArray' (show . anyToTypeRep) trKeys
79108
80109-- | Uses 'union' to combine 'TypeRepMap's.
81110instance Semigroup (TypeRepMap f) where
@@ -84,11 +113,36 @@ instance Semigroup (TypeRepMap f) where
84113 {-# INLINE (<>) #-}
85114
86115instance Monoid (TypeRepMap f) where
87− mempty = TypeRepMap mempty mempty mempty
116+ mempty = TypeRepMap mempty mempty mempty mempty
88117 mappend = (<>)
89118 {-# INLINE mempty #-}
90119 {-# INLINE mappend #-}
91120
121+#if __GLASGOW_HASKELL__ >= 806
122+instance (forall a. Typeable a => Eq (f a)) => Eq (TypeRepMap f) where
123+ tm1 == tm2 = size tm1 == size tm2 && go 0
124+ where
125+ go :: Int -> Bool
126+ go i
127+ | i == size tm1 = True
128+ | otherwise = case testEquality tr1i tr2i of
129+ Nothing -> False
130+ Just Refl -> repEq tr1i (fromAny tv1i) (fromAny tv2i) && go (i + 1)
131+ where
132+ tr1i :: TypeRep x
133+ tr1i = anyToTypeRep $ indexArray (trKeys tm1) i
134+
135+ tr2i :: TypeRep y
136+ tr2i = anyToTypeRep $ indexArray (trKeys tm2) i
137+
138+ tv1i, tv2i :: Any
139+ tv1i = indexArray (trAnys tm1) i
140+ tv2i = indexArray (trAnys tm2) i
141+
142+ repEq :: TypeRep x -> f x -> f x -> Bool
143+ repEq tr = withTypeable tr (==)
144+#endif
145+
92146-- | Returns the list of 'Fingerprint's from 'TypeRepMap'.
93147toFingerprints :: TypeRepMap f -> [Fingerprint]
94148toFingerprints TypeRepMap{..} =
@@ -115,33 +169,49 @@ prop> member @a (one (x :: f a)) == True
115169
116170-}
117171one :: forall a f . Typeable a => f a -> TypeRepMap f
118−one x = insert x empty
172+one x = TypeRepMap (primArrayFromListN 1 [fa])
173+ (primArrayFromListN 1 [fb])
174+ (pure @Array v)
175+ (pure @Array k)
176+ where
177+ (Fingerprint fa fb, v, k) = (calcFp @a, toAny x, unsafeCoerce $ typeRep @a)
119178{-# INLINE one #-}
120179
121180{- |
122181
123182Insert a value into a 'TypeRepMap'.
183+TypeRepMap optimizes for fast reads rather than inserts, as a trade-off inserts are @O(n)@.
124184
125185prop> size (insert v tm) >= size tm
126186prop> member @a (insert (x :: f a) tm) == True
127187
128188-}
129189insert :: forall a f . Typeable a => f a -> TypeRepMap f -> TypeRepMap f
130−insert x = fromListPairs . addX . toPairList
190+insert x m
191+ | size m == 0 = one x
192+ | otherwise = case cachedBinarySearch (typeFp @a) (fingerprintAs m) (fingerprintBs m) of
193+ Nothing -> union m $ one x
194+ Just i -> m {trAnys = changeAnyArr i (trAnys m)}
131195 where
132− pairX :: (Fingerprint, Any)
133− pairX@(fpX, _) = (calcFp x, toAny x)
134−
135− addX :: [(Fingerprint, Any)] -> [(Fingerprint, Any)]
136− addX l = pairX : deleteByFst fpX l
196+ changeAnyArr :: Int -> Array Any -> Array Any
197+ changeAnyArr i trAs = runST $ do
198+ let n = sizeofArray trAs
199+ mutArr <- thawArray trAs 0 n
200+ writeArray mutArr i $ toAny x
201+ unsafeFreezeArray mutArr
137202{-# INLINE insert #-}
138203
139204-- Extract the kind of a type. We use it to work around lack of syntax for
140205-- inferred type variables (which are not subject to type applications).
141206type KindOf (a :: k) = k
142207
208+type ArgKindOf (f :: k -> l) = k
209+
143210{- | Delete a value from a 'TypeRepMap'.
144211
212+TypeRepMap optimizes for fast reads rather than modifications, as a trade-off deletes are
213+@O(n)@, with an @O(log(n))@ optimization for when the element is already missing.
214+
145215prop> size (delete @a tm) <= size tm
146216prop> member @a (delete @a tm) == False
147217
@@ -154,9 +224,61 @@ False
154224True
155225-}
156226delete :: forall a (f :: KindOf a -> Type) . Typeable a => TypeRepMap f -> TypeRepMap f
157−delete = fromListPairs . deleteByFst (typeFp @a) . toPairList
227+delete m
228+ -- Lookups are fast, so check if we even have the element first.
229+ | not (member @a m) = m
230+ -- We know we have the element, If the map has exactly one element, we can return the empty map
231+ | size m == 1 = empty
232+ -- Otherwise, filter out the element in linear time.
233+ | otherwise = fromSortedTriples . deleteFirst ((== typeFp @a) . fst3) . toSortedTriples $ m
158234{-# INLINE delete #-}
159235
236+deleteFirst :: (a -> Bool) -> [a] -> [a]
237+deleteFirst _ [] = []
238+deleteFirst p (x : xs) = if p x then xs else x : deleteFirst p xs
239+
240+{- |
241+Update a value at a specific key with the result of the provided function.
242+When the key is not a member of the map, the original map is returned.
243+
244+>>> trmap = fromList @(TypeRepMap Identity) [WrapTypeable $ Identity "a"]
245+>>> lookup @String $ adjust (fmap (++ "ww")) trmap
246+Just (Identity "aww")
247+-}
248+adjust :: forall a f . Typeable a => (f a -> f a) -> TypeRepMap f -> TypeRepMap f
249+adjust fun = alter (fmap fun)
250+{-# INLINE adjust #-}
251+
252+{- |
253+Updates a value at a specific key, whether or not it exists.
254+This can be used to insert, delete, or update a value of a given type in the map.
255+
256+>>> func = (\case Nothing -> Just (Identity "new"); Just (Identity s) -> Just (Identity (reverse s)))
257+>>> lookup @String $ alter func empty
258+Just (Identity "new")
259+>>> trmap = fromList @(TypeRepMap Identity) [WrapTypeable $ Identity "helllo"]
260+>>> lookup @String $ alter func trmap
261+>>> Just (Identity "olleh")
262+-}
263+alter :: forall a f . Typeable a => (Maybe (f a) -> Maybe (f a)) -> TypeRepMap f -> TypeRepMap f
264+alter fun tr = case cachedBinarySearch (typeFp @a) (fingerprintAs tr) (fingerprintBs tr) of
265+ Nothing ->
266+ case (fun Nothing) of
267+ Nothing -> tr
268+ Just v -> insert v tr
269+ Just i ->
270+ case fun (Just . fromAny $ indexArray (trAnys tr) i) of
271+ Nothing -> delete @a tr
272+ Just v -> tr{trAnys = replaceAnyAt i (toAny v) (trAnys tr)}
273+ where
274+ replaceAnyAt :: Int -> Any -> Array Any -> Array Any
275+ replaceAnyAt i v trAs = runST $ do
276+ let n = sizeofArray trAs
277+ mutArr <- thawArray trAs 0 n
278+ writeArray mutArr i v
279+ unsafeFreezeArray mutArr
280+{-# INLINE alter #-}
281+
160282{- | Map over the elements of a 'TypeRepMap'.
161283
162284>>> tm = insert (Identity True) $ one (Identity 'a')
@@ -171,27 +293,107 @@ Just [True]
171293Just "a"
172294-}
173295hoist :: (forall x. f x -> g x) -> TypeRepMap f -> TypeRepMap g
174−hoist f (TypeRepMap as bs ans) = TypeRepMap as bs $ mapArray' (toAny . f . fromAny) ans
296+hoist f (TypeRepMap as bs ans ks) = TypeRepMap as bs (mapArray' (toAny . f . fromAny) ans) ks
175297{-# INLINE hoist #-}
176298
177−-- | The union of two 'TypeRepMap's using a combining function.
178−unionWith :: (forall x. f x -> f x -> f x) -> TypeRepMap f -> TypeRepMap f -> TypeRepMap f
179−unionWith f m1 m2 = fromListPairs
180− $ Map.toList
181− $ Map.unionWith combine
182− (Map.fromList $ toPairList m1)
183− (Map.fromList $ toPairList m2)
299+hoistA :: (Applicative t) => (forall x. f x -> t (g x)) -> TypeRepMap f -> t (TypeRepMap g)
300+hoistA f (TypeRepMap as bs (toList -> ans) ks) = (\l -> TypeRepMap as bs (fromList $ map toAny l) ks)
301+ <$> traverse (f . fromAny) ans
302+{-# INLINE hoistA #-}
303+
304+hoistWithKey :: forall f g. (forall x. Typeable x => f x -> g x) -> TypeRepMap f -> TypeRepMap g
305+hoistWithKey f (TypeRepMap as bs ans ks) = TypeRepMap as bs newAns ks
306+ where
307+ newAns = mapArray' mapAns (mzip ans ks)
308+ mapAns (a, k) = toAny $ withTr (unsafeCoerce k) $ fromAny a
309+
310+ withTr :: forall x. TypeRep x -> f x -> g x
311+ withTr t = withTypeable t f
312+{-# INLINE hoistWithKey #-}
313+
314+-- | The union of two 'TypeRepMap's using a combining function for conflicting entries. @O(n + m)@
315+unionWith :: forall f. (forall x. Typeable x => f x -> f x -> f x) -> TypeRepMap f -> TypeRepMap f -> TypeRepMap f
316+unionWith f ma mb = do
317+ fromSortedTriples $ mergeMaps (toSortedTriples ma) (toSortedTriples mb)
184318 where
185− combine :: Any -> Any -> Any
186− combine a b = toAny $ f (fromAny a) (fromAny b)
319+ f' :: forall x. TypeRep x -> f x -> f x -> f x
320+ f' tr = withTypeable tr f
321+
322+ combine :: (Fingerprint, Any, Any) -> (Fingerprint, Any, Any) -> (Fingerprint, Any, Any)
323+ combine (fp, av, ak) (_, bv, _) = (fp, toAny $ f' (fromAny ak) (fromAny av) (fromAny bv), ak)
324+
325+ -- Merges two typrepmaps into a sorted, dedup'd list of triples.
326+ -- Using 'toSortedTriples' allows us to assume the triples are sorted by fingerprint,
327+ -- Given O(n) performance from 'toSortedTriples', and given that we can merge-sort in
328+ -- O(n + m) time, then can '.fromSortedTriples' back into cachedBinarySearch order in O(n + m)
329+ -- that gives a total of O(n + m).
330+ mergeMaps :: [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
331+ -- We've addressed all elements from both maps
332+ mergeMaps as [] = as
333+ mergeMaps [] bs = bs
334+ -- Merge
335+ mergeMaps (a@(af, _, _) : as) (b@(bf, _, _) : bs) =
336+ case compare af bf of
337+ -- Fingerprints are equal, union the elements using our function
338+ -- If the incoming maps were de-duped, there shouldn't be any other equivalent
339+ -- fingerprints
340+ EQ -> combine a b : mergeMaps as bs
341+ -- First fingerprint must not be in the second map or we would have seen it by now
342+ -- Add it to the result as-is
343+ LT -> a : mergeMaps as (b : bs)
344+ -- Second fingerprint must not be in the first map or we would have seen it by now
345+ -- Add it to the result as-is
346+ GT -> b : mergeMaps (a:as) bs
187347{-# INLINE unionWith #-}
188348
189−-- | The (left-biased) union of two 'TypeRepMap's. It prefers the first map when
349+-- | The (left-biased) union of two 'TypeRepMap's in @O(n + m)@. It prefers the first map when
190350-- duplicate keys are encountered, i.e. @'union' == 'unionWith' const@.
191351union :: TypeRepMap f -> TypeRepMap f -> TypeRepMap f
192352union = unionWith const
193353{-# INLINE union #-}
194354
355+-- | The 'intersection' of two 'TypeRepMap's using a combining function
356+--
357+-- @O(n + m)@
358+intersectionWith :: forall f. (forall x. Typeable x => f x -> f x -> f x) -> TypeRepMap f -> TypeRepMap f -> TypeRepMap f
359+intersectionWith f ma mb =
360+ fromSortedTriples $ mergeMaps (toSortedTriples ma) (toSortedTriples mb)
361+ where
362+ f' :: forall x. TypeRep x -> f x -> f x -> f x
363+ f' tr = withTypeable tr f
364+
365+ combine :: (Fingerprint, Any, Any) -> (Fingerprint, Any, Any) -> (Fingerprint, Any, Any)
366+ combine (fp, av, ak) (_, bv, _) = (fp, toAny $ f' (fromAny ak) (fromAny av) (fromAny bv), ak)
367+
368+ -- Merges two typrepmaps into a sorted, dedup'd list of triples.
369+ mergeMaps :: [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
370+ -- If either list is empty, the intersection must be finished.
371+ mergeMaps _ [] = []
372+ mergeMaps [] _ = []
373+ -- Merge the two maps considering one element at a time.
374+ mergeMaps (a@(af, _, _) : as) (b@(bf, _, _) : bs) =
375+ case compare af bf of
376+ -- Fingerprints are equal, union the elements using our function
377+ -- If the incoming maps were de-duped, there shouldn't be any other equivalent
378+ -- fingerprints
379+ EQ -> combine a b : mergeMaps as bs
380+ -- First fingerprint must not be in the second map or we would have seen it by now
381+ -- Skip it an move on
382+ LT -> mergeMaps as (b : bs)
383+ -- Second fingerprint must not be in the first map or we would have seen it by now
384+ -- Skip it an move on
385+ GT -> mergeMaps (a:as) bs
386+{-# INLINE intersectionWith #-}
387+
388+-- | The intersection of two 'TypeRepMap's.
389+-- It keeps all values from the first map whose keys are present in the second.
390+--
391+-- @O(n + m)@
392+intersection :: TypeRepMap f -> TypeRepMap f -> TypeRepMap f
393+intersection = intersectionWith const
394+{-# INLINE intersection #-}
395+
396+
195397{- | Check if a value of the given type is present in a 'TypeRepMap'.
196398
197399>>> member @Char $ one (Identity 'a')
@@ -214,7 +416,7 @@ Just (Identity 11)
214416Nothing
215417-}
216418lookup :: forall a f . Typeable a => TypeRepMap f -> Maybe (f a)
217−lookup tVect = fromAny . (anys tVect `indexArray`)
419+lookup tVect = fromAny . (trAnys tVect `indexArray`)
218420 <$> cachedBinarySearch (typeFp @a)
219421 (fingerprintAs tVect)
220422 (fingerprintBs tVect)
@@ -225,6 +427,24 @@ size :: TypeRepMap f -> Int
225427size = sizeofPrimArray . fingerprintAs
226428{-# INLINE size #-}
227429
430+-- | Return the list of 'SomeTypeRep' from the keys.
431+keys :: TypeRepMap f -> [SomeTypeRep]
432+keys = keysWith SomeTypeRep
433+{-# INLINE keys #-}
434+
435+-- | Return the list of keys by wrapping them with a user-provided function.
436+keysWith :: (forall (a :: ArgKindOf f). TypeRep a -> r) -> TypeRepMap f -> [r]
437+keysWith f TypeRepMap{..} = f . anyToTypeRep <$> toList trKeys
438+{-# INLINE keysWith #-}
439+
440+-- | Return the list of key-value pairs by wrapping them with a user-provided function.
441+toListWith :: forall f r. (forall (a :: ArgKindOf f). Typeable a => f a -> r) -> TypeRepMap f -> [r]
442+toListWith f = map toF . toTriples
443+ where
444+ withTypeRep :: TypeRep a -> f a -> r
445+ withTypeRep tr an = withTypeable tr $ f an
446+ toF (_, an, k) = withTypeRep (unsafeCoerce k) (fromAny an)
447+
228448-- | Binary searched based on this article
229449-- http://bannalia.blogspot.com/2015/06/cache-friendly-binary-search.html
230450-- with modification for our two-vector search case.
@@ -258,71 +478,193 @@ toAny = unsafeCoerce
258478fromAny :: Any -> f a
259479fromAny = unsafeCoerce
260480
481+anyToTypeRep :: Any -> TypeRep f
482+anyToTypeRep = unsafeCoerce
483+
261484typeFp :: forall a . Typeable a => Fingerprint
262−typeFp = typeRepFingerprint $ typeRep $ Proxy @a
485+typeFp = typeRepFingerprint $ typeRep @a
263486{-# INLINE typeFp #-}
264487
265−toPairList :: TypeRepMap f -> [(Fingerprint, Any)]
266−toPairList tm = zip (toFingerprints tm) (GHC.toList $ anys tm)
488+toTriples :: TypeRepMap f -> [(Fingerprint, Any, Any)]
489+toTriples tm = zip3 (toFingerprints tm) (GHC.toList $ trAnys tm) (GHC.toList $ trKeys tm)
267490
268−deleteByFst :: Eq a => a -> [(a, b)] -> [(a, b)]
269−deleteByFst x = filter ((/= x) . fst)
491+-- | Efficiently get sorted triples from a map in O(n) time
492+--
493+-- We assume the incoming TypeRepMap is already sorted into 'cachedBinarySearch' order using fromSortedList.
494+-- Then we can construct the index mapping from the "cached" ordering into monotonically
495+-- increasing order using 'generateOrderMapping' with the length of the TRM. This takes @O(n).
496+-- We then pull those indexes from the source TRM to get the sorted triples in a total of @O(n).
497+toSortedTriples :: TypeRepMap f -> [(Fingerprint, Any, Any)]
498+toSortedTriples tm = trip <$> ordering
499+ where
500+ trip i = ( Fingerprint (indexPrimArray (fingerprintAs tm) i) (indexPrimArray (fingerprintBs tm) i)
501+ , indexArray (trAnys tm) i
502+ , indexArray (trKeys tm) i)
503+ ordering :: [ Int ]
504+ ordering = generateOrderMapping (size tm)
505+
506+nubByFst :: (Eq a) => [(a, b, c)] -> [(a, b, c)]
507+nubByFst = nubBy ((==) `on` fst3)
270508
271−nubByFst :: (Eq a) => [(a, b)] -> [(a, b)]
272−nubByFst = nubBy ((==) `on` fst)
509+fst3 :: (a, b, c) -> a
510+fst3 (a, _, _) = a
273511
274512----------------------------------------------------------------------------
275513-- Functions for testing and benchmarking
276514----------------------------------------------------------------------------
277515
278516-- | Existential wrapper around 'Typeable' indexed by @f@ type parameter.
279−-- Useful for 'TypeRepMap' structure creation form list of 'TF's.
280−data TF f where
281− TF :: Typeable a => f a -> TF f
517+-- Useful for 'TypeRepMap' structure creation form list of 'WrapTypeable's.
518+data WrapTypeable f where
519+ WrapTypeable :: Typeable a => f a -> WrapTypeable f
282520
283−instance Show (TF f) where
284− show (TF tf) = show $ calcFp tf
521+instance Show (WrapTypeable f) where
522+ show (WrapTypeable (_ :: f a)) = show $ calcFp @a
523+
524+{- |
285525
286−{- | Creates 'TypeRepMap' from a list of 'TF's.
526+prop> fromList . toList == 'id'
527+
528+Creates 'TypeRepMap' from a list of 'WrapTypeable's.
529+
530+>>> show $ fromList [WrapTypeable $ Identity True, WrapTypeable $ Identity 'a']
531+TypeRepMap [Bool, Char]
287532
288−>>> size $ fromList [TF $ Identity True, TF $ Identity 'a']
289−2
290533
291534-}
292−fromList :: forall f . [TF f] -> TypeRepMap f
293−fromList = fromListPairs . map (fp &&& an)
294− where
295− fp :: TF f -> Fingerprint
296− fp (TF x) = calcFp x
535+instance IsList (TypeRepMap f) where
536+ type Item (TypeRepMap f) = WrapTypeable f
297537
298− an :: TF f -> Any
299− an (TF x) = toAny x
538+ fromList :: [WrapTypeable f] -> TypeRepMap f
539+ fromList = fromTriples . map (\x -> (fp x, an x, k x))
540+ where
541+ fp :: WrapTypeable f -> Fingerprint
542+ fp (WrapTypeable (_ :: f a)) = calcFp @a
300543
301−fromF :: Typeable a => f a -> Proxy a
302−fromF _ = Proxy
544+ an :: WrapTypeable f -> Any
545+ an (WrapTypeable x) = toAny x
303546
304−calcFp :: Typeable a => f a -> Fingerprint
305−calcFp = typeRepFingerprint . typeRep . fromF
547+ k :: WrapTypeable f -> Any
548+ k (WrapTypeable (_ :: f a)) = unsafeCoerce $ typeRep @a
306549
307−fromListPairs :: [(Fingerprint, Any)] -> TypeRepMap f
308−fromListPairs kvs = TypeRepMap (GHC.fromList fpAs) (GHC.fromList fpBs) (GHC.fromList ans)
550+ toList :: TypeRepMap f -> [WrapTypeable f]
551+ toList = toListWith WrapTypeable
552+
553+calcFp :: forall a . Typeable a => Fingerprint
554+calcFp = typeRepFingerprint $ typeRep @a
555+
556+fromTriples :: [(Fingerprint, Any, Any)] -> TypeRepMap f
557+fromTriples = fromSortedTriples . sortWith fst3 . nubByFst
558+
559+fromSortedTriples :: [(Fingerprint, Any, Any)] -> TypeRepMap f
560+fromSortedTriples kvs = TypeRepMap (GHC.fromList fpAs) (GHC.fromList fpBs) (GHC.fromList ans) (GHC.fromList ks)
309561 where
310562 (fpAs, fpBs) = unzip $ map (\(Fingerprint a b) -> (a, b)) fps
311− (fps, ans) = unzip $ fromSortedList $ sortWith fst $ nubByFst kvs
563+ (fps, ans, ks) = unzip3 $ fromSortedList kvs
312564
313565----------------------------------------------------------------------------
314566-- Tree-like conversion
315567----------------------------------------------------------------------------
316568
317569fromSortedList :: forall a . [a] -> [a]
318−fromSortedList l = IM.elems $ fst $ go 0 0 mempty (IM.fromList $ zip [0..] l)
570+fromSortedList l = runST $ do
571+ let n = length l
572+ let arrOrigin = fromListN n l
573+ arrResult <- thawArray arrOrigin 0 n
574+ go n arrResult arrOrigin
575+ toList <$> unsafeFreezeArray arrResult
319576 where
320577 -- state monad could be used here, but it's another dependency
321− go :: Int -> Int -> IntMap a -> IntMap a -> (IntMap a, Int)
322− go i first result vector =
323− if i >= IM.size vector
324− then (result, first)
325− else do
326− let (newResult, newFirst) = go (2 * i + 1) first result vector
327− let withCur = IM.insert i (fromJust $ IM.lookup newFirst vector) newResult
328− go (2 * i + 2) (newFirst + 1) withCur vector
578+ go :: forall s . Int -> MutableArray s a -> Array a -> ST s ()
579+ go len result origin = () <$ loop 0 0
580+ where
581+ loop :: Int -> Int -> ST s Int
582+ loop i first =
583+ if i >= len
584+ then pure first
585+ else do
586+ newFirst <- loop (2 * i + 1) first
587+ writeArray result i (indexArray origin newFirst)
588+ loop (2 * i + 2) (newFirst + 1)
589+
590+-- Returns a list of indexes which represents the "sorted" order of an array generated by
591+-- fromSortedList of the provided length.
592+-- I.e. fmap (fromSortedList [1, 2, 3, 4, 5, 6] !!) (generateOrderMapping 6) == [1, 2, 3, 4, 5, 6]
593+--
594+-- >>> generateOrderMapping 6
595+-- [3,1,4,0,5,2]
596+--
597+-- >>> generateOrderMapping 8
598+-- [7,3,1,4,0,5,2,6]
599+generateOrderMapping :: Int -> [Int]
600+generateOrderMapping len = runST $ do
601+ orderMappingArr <- newPrimArray len
602+ _ <- loop orderMappingArr 0 0
603+ primArrayToList <$> unsafeFreezePrimArray orderMappingArr
604+ where
605+ loop :: MutablePrimArray s Int -> Int -> Int -> ST s Int
606+ loop result i first =
607+ if i >= len
608+ then pure first
609+ else do
610+ newFirst <- loop result (2 * i + 1) first
611+ writePrimArray result newFirst i
612+ loop result (2 * i + 2) (newFirst + 1)
613+
614+----------------------------------------------------------------------------
615+-- Helper functions.
616+----------------------------------------------------------------------------
617+
618+-- | Check that invariant of the structure holds.
619+-- The structure maintains the following invariant.
620+-- For each element @A@ at index @i@:
621+--
622+-- 1. if there is an element @B@ at index @2*i+1@,
623+-- then @B < A@.
624+--
625+-- 2. if there is an element @C@ at index @2*i+2@,
626+-- then @A < C@.
627+--
628+invariantCheck :: TypeRepMap f -> Bool
629+invariantCheck TypeRepMap{..} = getAll (check 0)
630+ where
631+ lastMay [] = Nothing
632+ lastMay [x] = Just x
633+ lastMay (_:xs) = lastMay xs
634+ sz = sizeofPrimArray fingerprintAs
635+ check i | i >= sz = All True
636+ | otherwise =
637+ let left = i*2+1
638+ right = i*2+2
639+ -- maximum value in the left branch
640+ leftMax =
641+ fmap (\j -> (indexPrimArray fingerprintAs j, indexPrimArray fingerprintBs j))
642+ $ lastMay
643+ $ takeWhile (<sz)
644+ $ iterate (\j -> j*2+2) left
645+ -- minimum value in the right branch
646+ rightMin =
647+ fmap (\j -> (indexPrimArray fingerprintAs j, indexPrimArray fingerprintBs j))
648+ $ lastMay
649+ $ takeWhile (<sz)
650+ $ iterate (\j -> j*2+1) right
651+ in mconcat
652+ [ All $
653+ if left < sz
654+ then
655+ case indexPrimArray fingerprintAs i `compare` indexPrimArray fingerprintAs left of
656+ LT -> False
657+ EQ -> indexPrimArray fingerprintBs i >= indexPrimArray fingerprintBs left
658+ GT -> True
659+ else True
660+ , All $
661+ if right < sz
662+ then
663+ case indexPrimArray fingerprintAs i `compare` indexPrimArray fingerprintAs right of
664+ LT -> True
665+ EQ -> indexPrimArray fingerprintBs i <= indexPrimArray fingerprintBs right
666+ GT -> False
667+ else True
668+ , All $ fromMaybe True $ (<=) <$> leftMax <*> rightMin
669+ , check (i+1)
670+ ]