@@ -47,7 +47,7 @@ import Data.List (intercalate, nubBy)
4747import Data.Maybe (fromMaybe)
4848import Data.Primitive.Array (Array, MutableArray, indexArray, mapArray', readArray, sizeofArray,
4949 thawArray, unsafeFreezeArray, writeArray)
50−import Data.Primitive.PrimArray (PrimArray, indexPrimArray, sizeofPrimArray)
50+import Data.Primitive.PrimArray (writePrimArray, newPrimArray, unsafeFreezePrimArray, primArrayToList, MutablePrimArray, primArrayFromListN, PrimArray, indexPrimArray, sizeofPrimArray)
5151import Data.Semigroup (Semigroup (..), All(..))
5252import GHC.Base (Any, Int (..), Int#, (*#), (+#), (<#))
5353import GHC.Exts (IsList (..), inline, sortWith)
@@ -64,7 +64,6 @@ import Type.Reflection (SomeTypeRep (..), TypeRep, Typeable, typeRep, withTypeab
6464import Type.Reflection.Unsafe (typeRepFingerprint)
6565import Unsafe.Coerce (unsafeCoerce)
6666
67−import qualified Data.Map.Strict as Map
6867import qualified GHC.Exts as GHC (fromList, toList)
6968
7069{- |
@@ -170,7 +169,12 @@ prop> member @a (one (x :: f a)) == True
170169
171170-}
172171one :: forall a f . Typeable a => f a -> TypeRepMap f
173−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)
174178{-# INLINE one #-}
175179
176180{- |
@@ -266,28 +270,42 @@ hoistWithKey f (TypeRepMap as bs ans ks) = TypeRepMap as bs newAns ks
266270 withTr t = withTypeable t f
267271{-# INLINE hoistWithKey #-}
268272
269−-- | The union of two 'TypeRepMap's using a combining function.
273+-- | The union of two 'TypeRepMap's using a combining function for conflicting entries. @O(n + m)@
270274unionWith :: forall f. (forall x. Typeable x => f x -> f x -> f x) -> TypeRepMap f -> TypeRepMap f -> TypeRepMap f
271−unionWith f m1 m2 = fromTriples
272− $ toTripleList
273− $ Map.unionWith combine
274− (fromTripleList $ toTriples m1)
275− (fromTripleList $ toTriples m2)
275+unionWith f ma mb = do
276+ fromSortedTriples $ mergeMaps (toSortedTriples ma) (toSortedTriples mb)
276277 where
277278 f' :: forall x. TypeRep x -> f x -> f x -> f x
278279 f' tr = withTypeable tr f
279280
280− combine :: (Any, Any) -> (Any, Any) -> (Any, Any)
281− combine (av, ak) (bv, _) = (toAny $ f' (fromAny ak) (fromAny av) (fromAny bv), ak)
282−
283− fromTripleList :: Ord a => [(a, b, c)] -> Map.Map a (b, c)
284− fromTripleList = Map.fromList . map (\(a, b, c) -> (a, (b, c)))
285−
286− toTripleList :: Map.Map a (b, c) -> [(a, b, c)]
287− toTripleList = map (\(a, (b, c)) -> (a, b, c)) . Map.toList
281+ combine :: (Fingerprint, Any, Any) -> (Fingerprint, Any, Any) -> (Fingerprint, Any, Any)
282+ combine (fp, av, ak) (_, bv, _) = (fp, toAny $ f' (fromAny ak) (fromAny av) (fromAny bv), ak)
283+
284+ -- Merges two typrepmaps into a sorted, dedup'd list of triples.
285+ -- Using 'toSortedTriples' allows us to assume the triples are sorted by fingerprint,
286+ -- Given O(n) performance from 'toSortedTriples', and given that we can merge-sort in
287+ -- O(n + m) time, then can '.fromSortedTriples' back into cachedBinarySearch order in O(n + m)
288+ -- that gives a total of O(n + m).
289+ mergeMaps :: [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
290+ -- We've addressed all elements from both maps
291+ mergeMaps as [] = as
292+ mergeMaps [] bs = bs
293+ -- Merge
294+ mergeMaps (a@(af, _, _) : as) (b@(bf, _, _) : bs) =
295+ case compare af bf of
296+ -- Fingerprints are equal, union the elements using our function
297+ -- If the incoming maps were de-duped, there shouldn't be any other equivalent
298+ -- fingerprints
299+ EQ -> combine a b : mergeMaps as bs
300+ -- First fingerprint must not be in the second map or we would have seen it by now
301+ -- Add it to the result as-is
302+ LT -> a : mergeMaps as (b : bs)
303+ -- Second fingerprint must not be in the first map or we would have seen it by now
304+ -- Add it to the result as-is
305+ GT -> b : mergeMaps (a:as) bs
288306{-# INLINE unionWith #-}
289307
290−-- | The (left-biased) union of two 'TypeRepMap's. It prefers the first map when
308+-- | The (left-biased) union of two 'TypeRepMap's in @O(n + m)@. It prefers the first map when
291309-- duplicate keys are encountered, i.e. @'union' == 'unionWith' const@.
292310union :: TypeRepMap f -> TypeRepMap f -> TypeRepMap f
293311union = unionWith const
@@ -374,6 +392,21 @@ typeFp = typeRepFingerprint $ typeRep @a
374392toTriples :: TypeRepMap f -> [(Fingerprint, Any, Any)]
375393toTriples tm = zip3 (toFingerprints tm) (GHC.toList $ trAnys tm) (GHC.toList $ trKeys tm)
376394
395+-- | Efficiently get sorted triples from a map in O(n) time
396+--
397+-- We assume the incoming TypeRepMap is already sorted into 'cachedBinarySearch' order using fromSortedList.
398+-- Then we can construct the index mapping from the "cached" ordering into monotonically
399+-- increasing order using 'generateOrderMapping' with the length of the TRM. This takes @O(n).
400+-- We then pull those indexes from the source TRM to get the sorted triples in a total of @O(n).
401+toSortedTriples :: TypeRepMap f -> [(Fingerprint, Any, Any)]
402+toSortedTriples tm = trip <$> ordering
403+ where
404+ trip i = ( Fingerprint (indexPrimArray (fingerprintAs tm) i) (indexPrimArray (fingerprintBs tm) i)
405+ , indexArray (trAnys tm) i
406+ , indexArray (trKeys tm) i)
407+ ordering :: [ Int ]
408+ ordering = generateOrderMapping (size tm)
409+
377410deleteByFst :: Eq a => a -> [(a, b, c)] -> [(a, b, c)]
378411deleteByFst x = filter ((/= x) . fst3)
379412
@@ -434,10 +467,13 @@ calcFp :: forall a . Typeable a => Fingerprint
434467calcFp = typeRepFingerprint $ typeRep @a
435468
436469fromTriples :: [(Fingerprint, Any, Any)] -> TypeRepMap f
437−fromTriples kvs = TypeRepMap (GHC.fromList fpAs) (GHC.fromList fpBs) (GHC.fromList ans) (GHC.fromList ks)
470+fromTriples = fromSortedTriples . sortWith fst3 . nubByFst
471+
472+fromSortedTriples :: [(Fingerprint, Any, Any)] -> TypeRepMap f
473+fromSortedTriples kvs = TypeRepMap (GHC.fromList fpAs) (GHC.fromList fpBs) (GHC.fromList ans) (GHC.fromList ks)
438474 where
439475 (fpAs, fpBs) = unzip $ map (\(Fingerprint a b) -> (a, b)) fps
440− (fps, ans, ks) = unzip3 $ fromSortedList $ sortWith fst3 $ nubByFst kvs
476+ (fps, ans, ks) = unzip3 $ fromSortedList kvs
441477
442478----------------------------------------------------------------------------
443479-- Tree-like conversion
@@ -464,11 +500,35 @@ fromSortedList l = runST $ do
464500 writeArray result i (indexArray origin newFirst)
465501 loop (2 * i + 2) (newFirst + 1)
466502
503+-- Returns a list of indexes which represents the "sorted" order of an array generated by
504+-- fromSortedList of the provided length.
505+-- I.e. fmap (fromSortedList [1, 2, 3, 4, 5, 6] !!) (generateOrderMapping 6) == [1, 2, 3, 4, 5, 6]
506+--
507+-- >>> generateOrderMapping 6
508+-- [3,1,4,0,5,2]
509+--
510+-- >>> generateOrderMapping 8
511+-- [7,3,1,4,0,5,2,6]
512+generateOrderMapping :: Int -> [Int]
513+generateOrderMapping len = runST $ do
514+ orderMappingArr <- newPrimArray len
515+ _ <- loop orderMappingArr 0 0
516+ primArrayToList <$> unsafeFreezePrimArray orderMappingArr
517+ where
518+ loop :: MutablePrimArray s Int -> Int -> Int -> ST s Int
519+ loop result i first =
520+ if i >= len
521+ then pure first
522+ else do
523+ newFirst <- loop result (2 * i + 1) first
524+ writePrimArray result newFirst i
525+ loop result (2 * i + 2) (newFirst + 1)
526+
467527----------------------------------------------------------------------------
468528-- Helper functions.
469529----------------------------------------------------------------------------
470530
471−-- | Check that invariant of the structure is hold.
531+-- | Check that invariant of the structure holds.
472532-- The structure maintains the following invariant.
473533-- For each element @A@ at index @i@:
474534--