@@ -1,48 +1,69 @@
1+{-# OPTIONS_GHC -Wno-missing-export-lists #-}
1−{-# LANGUAGE AllowAmbiguousTypes #-}
3+{-# LANGUAGE AllowAmbiguousTypes #-}
2−{-# LANGUAGE BangPatterns #-}
4+{-# LANGUAGE BangPatterns #-}
3−{-# LANGUAGE DataKinds #-}
6+{-# LANGUAGE DataKinds #-}
4−{-# LANGUAGE GADTs #-}
7+{-# LANGUAGE GADTs #-}
5−{-# LANGUAGE InstanceSigs #-}
8+{-# LANGUAGE InstanceSigs #-}
6−{-# LANGUAGE KindSignatures #-}
7−{-# LANGUAGE MagicHash #-}
9+{-# LANGUAGE MagicHash #-}
8−{-# LANGUAGE PolyKinds #-}
10+{-# LANGUAGE PolyKinds #-}
9−{-# LANGUAGE Rank2Types #-}
11+{-# LANGUAGE Rank2Types #-}
10−{-# LANGUAGE TypeFamilies #-}
12+{-# LANGUAGE TypeFamilies #-}
11−{-# LANGUAGE TypeInType #-}
13+{-# LANGUAGE TypeInType #-}
12−{-# LANGUAGE ViewPatterns #-}
14+{-# LANGUAGE ViewPatterns #-}
16+#if __GLASGOW_HASKELL__ >= 806
17+{-# LANGUAGE QuantifiedConstraints #-}
14-- {-# OPTIONS_GHC -ddump-simpl -dsuppress-idinfo -dsuppress-coercions -dsuppress-type-applications -dsuppress-uniques -dsuppress-module-prefixes #-}
20-- {-# OPTIONS_GHC -ddump-simpl -dsuppress-idinfo -dsuppress-coercions -dsuppress-type-applications -dsuppress-uniques -dsuppress-module-prefixes #-}
23+Copyright: (c) 2017-2020 Kowainik
24+SPDX-License-Identifier: MPL-2.0
25+Maintainer: Kowainik <xrom.xkov@gmail.com>
16−-- | Internal API for 'TypeRepMap' and operations on it. The functions here do
27+Internal API for 'TypeRepMap' and operations on it. The functions here do
17−-- not have any stability guarantees and can change between minor versions.
28+not have any stability guarantees and can change between minor versions.
19−-- If you need to use this module for purposes other than tests,
30+If you need to use this module for purposes other than tests,
20−-- create an issue.
31+create an issue.
22module Data.TypeRepMap.Internal where
36module Data.TypeRepMap.Internal where
24import Prelude hiding (lookup)
38import Prelude hiding (lookup)
40+import Control.DeepSeq
41+import Control.Monad.ST (ST, runST)
26import Control.Monad.Zip (mzip)
42import Control.Monad.Zip (mzip)
27import Data.Function (on)
43import Data.Function (on)
28−import Data.IntMap.Strict (IntMap)
29import Data.Kind (Type)
44import Data.Kind (Type)
45+import Data.Type.Equality ((:~:) (..), TestEquality (..))
30import Data.List (intercalate, nubBy)
46import Data.List (intercalate, nubBy)
31−import Data.Maybe (fromJust)
47+import Data.Maybe (fromMaybe)
32−import Data.Primitive.Array (Array, indexArray, mapArray')
48+import Data.Primitive.Array (Array, MutableArray, indexArray, mapArray', sizeofArray,
49+ thawArray, unsafeFreezeArray, writeArray)
33−import Data.Primitive.PrimArray (PrimArray, indexPrimArray, sizeofPrimArray)
50+import Data.Primitive.PrimArray (writePrimArray, newPrimArray, unsafeFreezePrimArray, primArrayToList, MutablePrimArray, primArrayFromListN, PrimArray, indexPrimArray, sizeofPrimArray)
34−import Data.Semigroup (Semigroup (..))
51+import Data.Semigroup (Semigroup (..), All(..))
35import GHC.Base (Any, Int (..), Int#, (*#), (+#), (<#))
52import GHC.Base (Any, Int (..), Int#, (*#), (+#), (<#))
36import GHC.Exts (IsList (..), inline, sortWith)
53import GHC.Exts (IsList (..), inline, sortWith)
37import GHC.Fingerprint (Fingerprint (..))
54import GHC.Fingerprint (Fingerprint (..))
55+#if WORD_SIZE_IN_BITS >= 64
38import GHC.Prim (eqWord#, ltWord#)
56import GHC.Prim (eqWord#, ltWord#)
58+import GHC.IntWord64 (eqWord64#, ltWord64#)
59+#define eqWord eqWord64
60+#define ltWord ltWord64
39import GHC.Word (Word64 (..))
62import GHC.Word (Word64 (..))
40import Type.Reflection (SomeTypeRep (..), TypeRep, Typeable, typeRep, withTypeable)
63import Type.Reflection (SomeTypeRep (..), TypeRep, Typeable, typeRep, withTypeable)
41import Type.Reflection.Unsafe (typeRepFingerprint)
64import Type.Reflection.Unsafe (typeRepFingerprint)
42import Unsafe.Coerce (unsafeCoerce)
65import Unsafe.Coerce (unsafeCoerce)
44−import qualified Data.IntMap.Strict as IM
45−import qualified Data.Map.Strict as Map
46import qualified GHC.Exts as GHC (fromList, toList)
67import qualified GHC.Exts as GHC (fromList, toList)
@@ -75,6 +96,9 @@ data TypeRepMap (f :: k -> Type) =
76 -- ^ an unsafe constructor for 'TypeRepMap'
97 -- ^ an unsafe constructor for 'TypeRepMap'
99+instance NFData (TypeRepMap f) where
100+ rnf x = rnf (keys x) `seq` ()
78-- | Shows only keys.
102-- | Shows only keys.
79instance Show (TypeRepMap f) where
103instance Show (TypeRepMap f) where
80 show TypeRepMap{..} = "TypeRepMap [" ++ showKeys ++ "]"
104 show TypeRepMap{..} = "TypeRepMap [" ++ showKeys ++ "]"
@@ -94,6 +118,31 @@ instance Monoid (TypeRepMap f) where
94 {-# INLINE mempty #-}
118 {-# INLINE mempty #-}
95 {-# INLINE mappend #-}
119 {-# INLINE mappend #-}
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
127+ | i == size tm1 = True
128+ | otherwise = case testEquality tr1i tr2i of
130+ Just Refl -> repEq tr1i (fromAny tv1i) (fromAny tv2i) && go (i + 1)
133+ tr1i = anyToTypeRep $ indexArray (trKeys tm1) i
136+ tr2i = anyToTypeRep $ indexArray (trKeys tm2) i
139+ tv1i = indexArray (trAnys tm1) i
140+ tv2i = indexArray (trAnys tm2) i
142+ repEq :: TypeRep x -> f x -> f x -> Bool
143+ repEq tr = withTypeable tr (==)
97-- | Returns the list of 'Fingerprint's from 'TypeRepMap'.
146-- | Returns the list of 'Fingerprint's from 'TypeRepMap'.
98toFingerprints :: TypeRepMap f -> [Fingerprint]
147toFingerprints :: TypeRepMap f -> [Fingerprint]
99toFingerprints TypeRepMap{..} =
148toFingerprints TypeRepMap{..} =
@@ -120,33 +169,49 @@ prop> member @a (one (x :: f a)) == True
122one :: forall a f . Typeable a => f a -> TypeRepMap f
171one :: forall a f . Typeable a => f a -> TypeRepMap f
123−one x = insert x empty
172+one x = TypeRepMap (primArrayFromListN 1 [fa])
173+ (primArrayFromListN 1 [fb])
177+ (Fingerprint fa fb, v, k) = (calcFp @a, toAny x, unsafeCoerce $ typeRep @a)
124{-# INLINE one #-}
178{-# INLINE one #-}
128Insert a value into a 'TypeRepMap'.
182Insert a value into a 'TypeRepMap'.
183+TypeRepMap optimizes for fast reads rather than inserts, as a trade-off inserts are @O(n)@.
130prop> size (insert v tm) >= size tm
185prop> size (insert v tm) >= size tm
131prop> member @a (insert (x :: f a) tm) == True
186prop> member @a (insert (x :: f a) tm) == True
134insert :: forall a f . Typeable a => f a -> TypeRepMap f -> TypeRepMap f
189insert :: forall a f . Typeable a => f a -> TypeRepMap f -> TypeRepMap f
135−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)}
137− tripleX :: (Fingerprint, Any, Any)
196+ changeAnyArr :: Int -> Array Any -> Array Any
138− tripleX@(fpX, _, _) = (calcFp @a, toAny x, unsafeCoerce $ typeRep @a)
197+ changeAnyArr i trAs = runST $ do
139−
198+ let n = sizeofArray trAs
140− addX :: [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
199+ mutArr <- thawArray trAs 0 n
141− addX l = tripleX : deleteByFst fpX l
200+ writeArray mutArr i $ toAny x
201+ unsafeFreezeArray mutArr
142{-# INLINE insert #-}
202{-# INLINE insert #-}
144-- Extract the kind of a type. We use it to work around lack of syntax for
204-- Extract the kind of a type. We use it to work around lack of syntax for
145-- inferred type variables (which are not subject to type applications).
205-- inferred type variables (which are not subject to type applications).
146type KindOf (a :: k) = k
206type KindOf (a :: k) = k
208+type ArgKindOf (f :: k -> l) = k
148{- | Delete a value from a 'TypeRepMap'.
210{- | Delete a value from a 'TypeRepMap'.
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.
150prop> size (delete @a tm) <= size tm
215prop> size (delete @a tm) <= size tm
151prop> member @a (delete @a tm) == False
216prop> member @a (delete @a tm) == False
@@ -159,9 +224,61 @@ False
161delete :: forall a (f :: KindOf a -> Type) . Typeable a => TypeRepMap f -> TypeRepMap f
226delete :: forall a (f :: KindOf a -> Type) . Typeable a => TypeRepMap f -> TypeRepMap f
162−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
163{-# INLINE delete #-}
234{-# INLINE delete #-}
236+deleteFirst :: (a -> Bool) -> [a] -> [a]
237+deleteFirst _ [] = []
238+deleteFirst p (x : xs) = if p x then xs else x : deleteFirst p xs
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.
244+>>> trmap = fromList @(TypeRepMap Identity) [WrapTypeable $ Identity "a"]
245+>>> lookup @String $ adjust (fmap (++ "ww")) trmap
246+Just (Identity "aww")
248+adjust :: forall a f . Typeable a => (f a -> f a) -> TypeRepMap f -> TypeRepMap f
249+adjust fun = alter (fmap fun)
250+{-# INLINE adjust #-}
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.
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")
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
266+ case (fun Nothing) of
268+ Just v -> insert v tr
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)}
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
165{- | Map over the elements of a 'TypeRepMap'.
282{- | Map over the elements of a 'TypeRepMap'.
167>>> tm = insert (Identity True) $ one (Identity 'a')
284>>> tm = insert (Identity True) $ one (Identity 'a')
@@ -194,30 +311,89 @@ hoistWithKey f (TypeRepMap as bs ans ks) = TypeRepMap as bs newAns ks
194 withTr t = withTypeable t f
311 withTr t = withTypeable t f
195{-# INLINE hoistWithKey #-}
312{-# INLINE hoistWithKey #-}
197−-- | 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)@
198−unionWith :: (forall x. f x -> f x -> f x) -> TypeRepMap f -> TypeRepMap f -> TypeRepMap f
315+unionWith :: forall f. (forall x. Typeable x => f x -> f x -> f x) -> TypeRepMap f -> TypeRepMap f -> TypeRepMap f
199−unionWith f m1 m2 = fromTriples
316+unionWith f ma mb = do
200− $ toTripleList
317+ fromSortedTriples $ mergeMaps (toSortedTriples ma) (toSortedTriples mb)
201− $ Map.unionWith combine
202− (fromTripleList $ toTriples m1)
203− (fromTripleList $ toTriples m2)
205− combine :: (Any, Any) -> (Any, Any) -> (Any, Any)
319+ f' :: forall x. TypeRep x -> f x -> f x -> f x
206− combine (av, ak) (bv, _) = (toAny $ f (fromAny av) (fromAny bv), ak)
320+ f' tr = withTypeable tr f
208− fromTripleList :: Ord a => [(a, b, c)] -> Map.Map a (b, c)
322+ combine :: (Fingerprint, Any, Any) -> (Fingerprint, Any, Any) -> (Fingerprint, Any, Any)
209− fromTripleList = Map.fromList . map (\(a, b, c) -> (a, (b, c)))
323+ combine (fp, av, ak) (_, bv, _) = (fp, toAny $ f' (fromAny ak) (fromAny av) (fromAny bv), ak)
211− toTripleList :: Map.Map a (b, c) -> [(a, b, c)]
325+ -- Merges two typrepmaps into a sorted, dedup'd list of triples.
212− toTripleList = map (\(a, (b, c)) -> (a, b, c)) . Map.toList
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
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
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
213{-# INLINE unionWith #-}
347{-# INLINE unionWith #-}
215−-- | 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
216-- duplicate keys are encountered, i.e. @'union' == 'unionWith' const@.
350-- duplicate keys are encountered, i.e. @'union' == 'unionWith' const@.
217union :: TypeRepMap f -> TypeRepMap f -> TypeRepMap f
351union :: TypeRepMap f -> TypeRepMap f -> TypeRepMap f
218union = unionWith const
352union = unionWith const
219{-# INLINE union #-}
353{-# INLINE union #-}
355+-- | The 'intersection' of two 'TypeRepMap's using a combining function
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)
362+ f' :: forall x. TypeRep x -> f x -> f x -> f x
363+ f' tr = withTypeable tr f
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)
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.
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
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 #-}
388+-- | The intersection of two 'TypeRepMap's.
389+-- It keeps all values from the first map whose keys are present in the second.
392+intersection :: TypeRepMap f -> TypeRepMap f -> TypeRepMap f
393+intersection = intersectionWith const
394+{-# INLINE intersection #-}
221{- | Check if a value of the given type is present in a 'TypeRepMap'.
397{- | Check if a value of the given type is present in a 'TypeRepMap'.
223>>> member @Char $ one (Identity 'a')
399>>> member @Char $ one (Identity 'a')
@@ -253,9 +429,22 @@ size = sizeofPrimArray . fingerprintAs
254-- | Return the list of 'SomeTypeRep' from the keys.
430-- | Return the list of 'SomeTypeRep' from the keys.
255keys :: TypeRepMap f -> [SomeTypeRep]
431keys :: TypeRepMap f -> [SomeTypeRep]
256−keys TypeRepMap{..} = SomeTypeRep . anyToTypeRep <$> toList trKeys
432+keys = keysWith SomeTypeRep
257{-# INLINE keys #-}
433{-# INLINE keys #-}
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 #-}
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
444+ withTypeRep :: TypeRep a -> f a -> r
445+ withTypeRep tr an = withTypeable tr $ f an
446+ toF (_, an, k) = withTypeRep (unsafeCoerce k) (fromAny an)
259-- | Binary searched based on this article
448-- | Binary searched based on this article
260-- http://bannalia.blogspot.com/2015/06/cache-friendly-binary-search.html
449-- http://bannalia.blogspot.com/2015/06/cache-friendly-binary-search.html
261-- with modification for our two-vector search case.
450-- with modification for our two-vector search case.
@@ -299,8 +488,20 @@ typeFp = typeRepFingerprint $ typeRep @a
299toTriples :: TypeRepMap f -> [(Fingerprint, Any, Any)]
488toTriples :: TypeRepMap f -> [(Fingerprint, Any, Any)]
300toTriples tm = zip3 (toFingerprints tm) (GHC.toList $ trAnys tm) (GHC.toList $ trKeys tm)
489toTriples tm = zip3 (toFingerprints tm) (GHC.toList $ trAnys tm) (GHC.toList $ trKeys tm)
302−deleteByFst :: Eq a => a -> [(a, b, c)] -> [(a, b, c)]
491+-- | Efficiently get sorted triples from a map in O(n) time
303−deleteByFst x = filter ((/= x) . fst3)
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
500+ trip i = ( Fingerprint (indexPrimArray (fingerprintAs tm) i) (indexPrimArray (fingerprintBs tm) i)
501+ , indexArray (trAnys tm) i
502+ , indexArray (trKeys tm) i)
504+ ordering = generateOrderMapping (size tm)
305nubByFst :: (Eq a) => [(a, b, c)] -> [(a, b, c)]
506nubByFst :: (Eq a) => [(a, b, c)] -> [(a, b, c)]
306nubByFst = nubBy ((==) `on` fst3)
507nubByFst = nubBy ((==) `on` fst3)
@@ -347,33 +548,123 @@ instance IsList (TypeRepMap f) where
347 k (WrapTypeable (_ :: f a)) = unsafeCoerce $ typeRep @a
548 k (WrapTypeable (_ :: f a)) = unsafeCoerce $ typeRep @a
349 toList :: TypeRepMap f -> [WrapTypeable f]
550 toList :: TypeRepMap f -> [WrapTypeable f]
350− toList = map toWrapTypeable . toTriples
551+ toList = toListWith WrapTypeable
352− toWrapTypeable :: (Fingerprint, Any, Any) -> WrapTypeable f
353− toWrapTypeable (_, an, k) = withTypeable (unsafeCoerce k) $ fromAny an
355calcFp :: forall a . Typeable a => Fingerprint
553calcFp :: forall a . Typeable a => Fingerprint
356calcFp = typeRepFingerprint $ typeRep @a
554calcFp = typeRepFingerprint $ typeRep @a
358fromTriples :: [(Fingerprint, Any, Any)] -> TypeRepMap f
556fromTriples :: [(Fingerprint, Any, Any)] -> TypeRepMap f
359−fromTriples kvs = TypeRepMap (GHC.fromList fpAs) (GHC.fromList fpBs) (GHC.fromList ans) (GHC.fromList ks)
557+fromTriples = fromSortedTriples . sortWith fst3 . nubByFst
559+fromSortedTriples :: [(Fingerprint, Any, Any)] -> TypeRepMap f
560+fromSortedTriples kvs = TypeRepMap (GHC.fromList fpAs) (GHC.fromList fpBs) (GHC.fromList ans) (GHC.fromList ks)
361 (fpAs, fpBs) = unzip $ map (\(Fingerprint a b) -> (a, b)) fps
562 (fpAs, fpBs) = unzip $ map (\(Fingerprint a b) -> (a, b)) fps
362− (fps, ans, ks) = unzip3 $ fromSortedList $ sortWith fst3 $ nubByFst kvs
563+ (fps, ans, ks) = unzip3 $ fromSortedList kvs
364----------------------------------------------------------------------------
565----------------------------------------------------------------------------
365-- Tree-like conversion
566-- Tree-like conversion
366----------------------------------------------------------------------------
567----------------------------------------------------------------------------
368fromSortedList :: forall a . [a] -> [a]
569fromSortedList :: forall a . [a] -> [a]
369−fromSortedList l = IM.elems $ fst $ go 0 0 mempty (IM.fromList $ zip [0..] l)
570+fromSortedList l = runST $ do
572+ let arrOrigin = fromListN n l
573+ arrResult <- thawArray arrOrigin 0 n
574+ go n arrResult arrOrigin
575+ toList <$> unsafeFreezeArray arrResult
371 -- state monad could be used here, but it's another dependency
577 -- state monad could be used here, but it's another dependency
372− go :: Int -> Int -> IntMap a -> IntMap a -> (IntMap a, Int)
578+ go :: forall s . Int -> MutableArray s a -> Array a -> ST s ()
373− go i first result vector =
579+ go len result origin = () <$ loop 0 0
374− if i >= IM.size vector
580+ where
375− then (result, first)
581+ loop :: Int -> Int -> ST s Int
376− else do
582+ loop i first =
377− let (newResult, newFirst) = go (2 * i + 1) first result vector
583+ if i >= len
378− let withCur = IM.insert i (fromJust $ IM.lookup newFirst vector) newResult
584+ then pure first
379− go (2 * i + 2) (newFirst + 1) withCur vector
585+ else do
586+ newFirst <- loop (2 * i + 1) first
587+ writeArray result i (indexArray origin newFirst)
588+ loop (2 * i + 2) (newFirst + 1)
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]
594+-- >>> generateOrderMapping 6
597+-- >>> generateOrderMapping 8
599+generateOrderMapping :: Int -> [Int]
600+generateOrderMapping len = runST $ do
601+ orderMappingArr <- newPrimArray len
602+ _ <- loop orderMappingArr 0 0
603+ primArrayToList <$> unsafeFreezePrimArray orderMappingArr
605+ loop :: MutablePrimArray s Int -> Int -> Int -> ST s Int
606+ loop result i first =
610+ newFirst <- loop result (2 * i + 1) first
611+ writePrimArray result newFirst i
612+ loop result (2 * i + 2) (newFirst + 1)
614+----------------------------------------------------------------------------
616+----------------------------------------------------------------------------
618+-- | Check that invariant of the structure holds.
619+-- The structure maintains the following invariant.
620+-- For each element @A@ at index @i@:
622+-- 1. if there is an element @B@ at index @2*i+1@,
625+-- 2. if there is an element @C@ at index @2*i+2@,
628+invariantCheck :: TypeRepMap f -> Bool
629+invariantCheck TypeRepMap{..} = getAll (check 0)
631+ lastMay [] = Nothing
632+ lastMay [x] = Just x
633+ lastMay (_:xs) = lastMay xs
634+ sz = sizeofPrimArray fingerprintAs
635+ check i | i >= sz = All True
639+ -- maximum value in the left branch
641+ fmap (\j -> (indexPrimArray fingerprintAs j, indexPrimArray fingerprintBs j))
644+ $ iterate (\j -> j*2+2) left
645+ -- minimum value in the right branch
647+ fmap (\j -> (indexPrimArray fingerprintAs j, indexPrimArray fingerprintBs j))
650+ $ iterate (\j -> j*2+1) right
655+ case indexPrimArray fingerprintAs i `compare` indexPrimArray fingerprintAs left of
657+ EQ -> indexPrimArray fingerprintBs i >= indexPrimArray fingerprintBs left
663+ case indexPrimArray fingerprintAs i `compare` indexPrimArray fingerprintAs right of
665+ EQ -> indexPrimArray fingerprintBs i <= indexPrimArray fingerprintBs right
668+ , All $ fromMaybe True $ (<=) <$> leftMax <*> rightMin