@@ -1,10 +1,11 @@
1+{-# OPTIONS_GHC -Wno-missing-export-lists #-}
2+
13{-# LANGUAGE AllowAmbiguousTypes #-}
24{-# LANGUAGE BangPatterns #-}
35{-# LANGUAGE CPP #-}
46{-# LANGUAGE DataKinds #-}
57{-# LANGUAGE GADTs #-}
68{-# LANGUAGE InstanceSigs #-}
7−{-# LANGUAGE KindSignatures #-}
89{-# LANGUAGE MagicHash #-}
910{-# LANGUAGE PolyKinds #-}
1011{-# LANGUAGE Rank2Types #-}
@@ -44,9 +45,9 @@ import Data.Kind (Type)
4445import Data.Type.Equality ((:~:) (..), TestEquality (..))
4546import Data.List (intercalate, nubBy)
4647import Data.Maybe (fromMaybe)
47−import Data.Primitive.Array (Array, MutableArray, indexArray, mapArray', readArray, sizeofArray,
48+import Data.Primitive.Array (Array, MutableArray, indexArray, mapArray', sizeofArray,
4849 thawArray, unsafeFreezeArray, writeArray)
49−import Data.Primitive.PrimArray (PrimArray, indexPrimArray, sizeofPrimArray)
50+import Data.Primitive.PrimArray (writePrimArray, newPrimArray, unsafeFreezePrimArray, primArrayToList, MutablePrimArray, primArrayFromListN, PrimArray, indexPrimArray, sizeofPrimArray)
5051import Data.Semigroup (Semigroup (..), All(..))
5152import GHC.Base (Any, Int (..), Int#, (*#), (+#), (<#))
5253import GHC.Exts (IsList (..), inline, sortWith)
@@ -63,7 +64,6 @@ import Type.Reflection (SomeTypeRep (..), TypeRep, Typeable, typeRep, withTypeab
6364import Type.Reflection.Unsafe (typeRepFingerprint)
6465import Unsafe.Coerce (unsafeCoerce)
6566
66−import qualified Data.Map.Strict as Map
6767import qualified GHC.Exts as GHC (fromList, toList)
6868
6969{- |
@@ -169,33 +169,49 @@ prop> member @a (one (x :: f a)) == True
169169
170170-}
171171one :: forall a f . Typeable a => f a -> TypeRepMap f
172−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)
173178{-# INLINE one #-}
174179
175180{- |
176181
177182Insert a value into a 'TypeRepMap'.
183+TypeRepMap optimizes for fast reads rather than inserts, as a trade-off inserts are @O(n)@.
178184
179185prop> size (insert v tm) >= size tm
180186prop> member @a (insert (x :: f a) tm) == True
181187
182188-}
183189insert :: forall a f . Typeable a => f a -> TypeRepMap f -> TypeRepMap f
184−insert x = fromTriples . addX . toTriples
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)}
185195 where
186− tripleX :: (Fingerprint, Any, Any)
187− tripleX@(fpX, _, _) = (calcFp @a, toAny x, unsafeCoerce $ typeRep @a)
188−
189− addX :: [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
190− addX l = tripleX : 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
191202{-# INLINE insert #-}
192203
193204-- Extract the kind of a type. We use it to work around lack of syntax for
194205-- inferred type variables (which are not subject to type applications).
195206type KindOf (a :: k) = k
196207
208+type ArgKindOf (f :: k -> l) = k
209+
197210{- | Delete a value from a 'TypeRepMap'.
198211
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+
199215prop> size (delete @a tm) <= size tm
200216prop> member @a (delete @a tm) == False
201217
@@ -208,30 +224,60 @@ False
208224True
209225-}
210226delete :: forall a (f :: KindOf a -> Type) . Typeable a => TypeRepMap f -> TypeRepMap f
211−delete = fromTriples . deleteByFst (typeFp @a) . toTriples
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
212234{-# INLINE delete #-}
213235
236+deleteFirst :: (a -> Bool) -> [a] -> [a]
237+deleteFirst _ [] = []
238+deleteFirst p (x : xs) = if p x then xs else x : deleteFirst p xs
239+
214240{- |
215−Update a value at a specific key with the result of the provided function. When
216−the key is not a member of the map, the original map is returned.
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.
217243
218244>>> trmap = fromList @(TypeRepMap Identity) [WrapTypeable $ Identity "a"]
219245>>> lookup @String $ adjust (fmap (++ "ww")) trmap
220246Just (Identity "aww")
221247-}
222248adjust :: forall a f . Typeable a => (f a -> f a) -> TypeRepMap f -> TypeRepMap f
223−adjust fun tr = case cachedBinarySearch (typeFp @a) (fingerprintAs tr) (fingerprintBs tr) of
224− Nothing -> tr
225− Just i -> tr {trAnys = changeAnyArr i (trAnys tr)}
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)}
226273 where
227− changeAnyArr :: Int -> Array Any -> Array Any
228− changeAnyArr i trAs = runST $ do
274+ replaceAnyAt :: Int -> Any -> Array Any -> Array Any
275+ replaceAnyAt i v trAs = runST $ do
229276 let n = sizeofArray trAs
230277 mutArr <- thawArray trAs 0 n
231− a <- toAny . fun . fromAny <$> readArray mutArr i
232− writeArray mutArr i a
278+ writeArray mutArr i v
233279 unsafeFreezeArray mutArr
234−{-# INLINE adjust #-}
280+{-# INLINE alter #-}
235281
236282{- | Map over the elements of a 'TypeRepMap'.
237283
@@ -265,33 +311,89 @@ hoistWithKey f (TypeRepMap as bs ans ks) = TypeRepMap as bs newAns ks
265311 withTr t = withTypeable t f
266312{-# INLINE hoistWithKey #-}
267313
268−-- | The union of two 'TypeRepMap's using a combining function.
314+-- | The union of two 'TypeRepMap's using a combining function for conflicting entries. @O(n + m)@
269315unionWith :: forall f. (forall x. Typeable x => f x -> f x -> f x) -> TypeRepMap f -> TypeRepMap f -> TypeRepMap f
270−unionWith f m1 m2 = fromTriples
271− $ toTripleList
272− $ Map.unionWith combine
273− (fromTripleList $ toTriples m1)
274− (fromTripleList $ toTriples m2)
316+unionWith f ma mb = do
317+ fromSortedTriples $ mergeMaps (toSortedTriples ma) (toSortedTriples mb)
275318 where
276319 f' :: forall x. TypeRep x -> f x -> f x -> f x
277320 f' tr = withTypeable tr f
278321
279− combine :: (Any, Any) -> (Any, Any) -> (Any, Any)
280− combine (av, ak) (bv, _) = (toAny $ f' (fromAny ak) (fromAny av) (fromAny bv), ak)
281−
282− fromTripleList :: Ord a => [(a, b, c)] -> Map.Map a (b, c)
283− fromTripleList = Map.fromList . map (\(a, b, c) -> (a, (b, c)))
284−
285− toTripleList :: Map.Map a (b, c) -> [(a, b, c)]
286− toTripleList = map (\(a, (b, c)) -> (a, b, c)) . Map.toList
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
287347{-# INLINE unionWith #-}
288348
289−-- | 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
290350-- duplicate keys are encountered, i.e. @'union' == 'unionWith' const@.
291351union :: TypeRepMap f -> TypeRepMap f -> TypeRepMap f
292352union = unionWith const
293353{-# INLINE union #-}
294354
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+
295397{- | Check if a value of the given type is present in a 'TypeRepMap'.
296398
297399>>> member @Char $ one (Identity 'a')
@@ -327,9 +429,22 @@ size = sizeofPrimArray . fingerprintAs
327429
328430-- | Return the list of 'SomeTypeRep' from the keys.
329431keys :: TypeRepMap f -> [SomeTypeRep]
330−keys TypeRepMap{..} = SomeTypeRep . anyToTypeRep <$> toList trKeys
432+keys = keysWith SomeTypeRep
331433{-# INLINE keys #-}
332434
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+
333448-- | Binary searched based on this article
334449-- http://bannalia.blogspot.com/2015/06/cache-friendly-binary-search.html
335450-- with modification for our two-vector search case.
@@ -373,8 +488,20 @@ typeFp = typeRepFingerprint $ typeRep @a
373488toTriples :: TypeRepMap f -> [(Fingerprint, Any, Any)]
374489toTriples tm = zip3 (toFingerprints tm) (GHC.toList $ trAnys tm) (GHC.toList $ trKeys tm)
375490
376−deleteByFst :: Eq a => a -> [(a, b, c)] -> [(a, b, c)]
377−deleteByFst x = filter ((/= x) . fst3)
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)
378505
379506nubByFst :: (Eq a) => [(a, b, c)] -> [(a, b, c)]
380507nubByFst = nubBy ((==) `on` fst3)
@@ -394,9 +521,6 @@ data WrapTypeable f where
394521instance Show (WrapTypeable f) where
395522 show (WrapTypeable (_ :: f a)) = show $ calcFp @a
396523
397−wrapTypeable :: TypeRep a -> f a -> WrapTypeable f
398−wrapTypeable tr = withTypeable tr WrapTypeable
399−
400524{- |
401525
402526prop> fromList . toList == 'id'
@@ -424,19 +548,19 @@ instance IsList (TypeRepMap f) where
424548 k (WrapTypeable (_ :: f a)) = unsafeCoerce $ typeRep @a
425549
426550 toList :: TypeRepMap f -> [WrapTypeable f]
427− toList = map toWrapTypeable . toTriples
428− where
429− toWrapTypeable :: (Fingerprint, Any, Any) -> WrapTypeable f
430− toWrapTypeable (_, an, k) = wrapTypeable (unsafeCoerce k) (fromAny an)
551+ toList = toListWith WrapTypeable
431552
432553calcFp :: forall a . Typeable a => Fingerprint
433554calcFp = typeRepFingerprint $ typeRep @a
434555
435556fromTriples :: [(Fingerprint, Any, Any)] -> TypeRepMap f
436−fromTriples kvs = TypeRepMap (GHC.fromList fpAs) (GHC.fromList fpBs) (GHC.fromList ans) (GHC.fromList ks)
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)
437561 where
438562 (fpAs, fpBs) = unzip $ map (\(Fingerprint a b) -> (a, b)) fps
439− (fps, ans, ks) = unzip3 $ fromSortedList $ sortWith fst3 $ nubByFst kvs
563+ (fps, ans, ks) = unzip3 $ fromSortedList kvs
440564
441565----------------------------------------------------------------------------
442566-- Tree-like conversion
@@ -463,11 +587,35 @@ fromSortedList l = runST $ do
463587 writeArray result i (indexArray origin newFirst)
464588 loop (2 * i + 2) (newFirst + 1)
465589
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+
466614----------------------------------------------------------------------------
467615-- Helper functions.
468616----------------------------------------------------------------------------
469617
470−-- | Check that invariant of the structure is hold.
618+-- | Check that invariant of the structure holds.
471619-- The structure maintains the following invariant.
472620-- For each element @A@ at index @i@:
473621--