@@ -1,11 +1,10 @@
1−{-# OPTIONS_GHC -Wno-missing-export-lists #-}
2−
31{-# LANGUAGE AllowAmbiguousTypes #-}
42{-# LANGUAGE BangPatterns #-}
53{-# LANGUAGE CPP #-}
64{-# LANGUAGE DataKinds #-}
75{-# LANGUAGE GADTs #-}
86{-# LANGUAGE InstanceSigs #-}
7+{-# LANGUAGE KindSignatures #-}
98{-# LANGUAGE MagicHash #-}
109{-# LANGUAGE PolyKinds #-}
1110{-# LANGUAGE Rank2Types #-}
@@ -19,20 +18,13 @@
1918
2019-- {-# OPTIONS_GHC -ddump-simpl -dsuppress-idinfo -dsuppress-coercions -dsuppress-type-applications -dsuppress-uniques -dsuppress-module-prefixes #-}
2120
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
21+{- | Internal API for 'TypeRepMap' and operations on it. The functions here do
2822not have any stability guarantees and can change between minor versions.
2923
3024If you need to use this module for purposes other than tests,
3125create an issue.
3226-}
3327
34−#include "MachDeps.h"
35−
3628module Data.TypeRepMap.Internal where
3729
3830import Prelude hiding (lookup)
@@ -44,26 +36,20 @@ import Data.Function (on)
4436import Data.Kind (Type)
4537import Data.Type.Equality ((:~:) (..), TestEquality (..))
4638import Data.List (intercalate, nubBy)
47−import Data.Maybe (fromMaybe)
48−import Data.Primitive.Array (Array, MutableArray, indexArray, mapArray', sizeofArray,
39+import Data.Primitive.Array (Array, MutableArray, indexArray, mapArray', readArray, sizeofArray,
4940 thawArray, unsafeFreezeArray, writeArray)
50−import Data.Primitive.PrimArray (writePrimArray, newPrimArray, unsafeFreezePrimArray, primArrayToList, MutablePrimArray, primArrayFromListN, PrimArray, indexPrimArray, sizeofPrimArray)
51−import Data.Semigroup (Semigroup (..), All(..))
41+import Data.Primitive.PrimArray (PrimArray, indexPrimArray, sizeofPrimArray)
42+import Data.Semigroup (Semigroup (..))
5243import GHC.Base (Any, Int (..), Int#, (*#), (+#), (<#))
5344import GHC.Exts (IsList (..), inline, sortWith)
5445import GHC.Fingerprint (Fingerprint (..))
55−#if WORD_SIZE_IN_BITS >= 64
5646import GHC.Prim (eqWord#, ltWord#)
57−#else
58−import GHC.IntWord64 (eqWord64#, ltWord64#)
59−#define eqWord eqWord64
60−#define ltWord ltWord64
61−#endif
6247import GHC.Word (Word64 (..))
6348import Type.Reflection (SomeTypeRep (..), TypeRep, Typeable, typeRep, withTypeable)
6449import Type.Reflection.Unsafe (typeRepFingerprint)
6550import Unsafe.Coerce (unsafeCoerce)
6651
52+import qualified Data.Map.Strict as Map
6753import qualified GHC.Exts as GHC (fromList, toList)
6854
6955{- |
@@ -169,49 +155,33 @@ prop> member @a (one (x :: f a)) == True
169155
170156-}
171157one :: forall a f . Typeable a => f a -> TypeRepMap f
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)
158+one x = insert x empty
178159{-# INLINE one #-}
179160
180161{- |
181162
182163Insert a value into a 'TypeRepMap'.
183−TypeRepMap optimizes for fast reads rather than inserts, as a trade-off inserts are @O(n)@.
184164
185165prop> size (insert v tm) >= size tm
186166prop> member @a (insert (x :: f a) tm) == True
187167
188168-}
189169insert :: forall a f . Typeable a => f a -> TypeRepMap f -> TypeRepMap f
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)}
170+insert x = fromTriples . addX . toTriples
195171 where
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
172+ tripleX :: (Fingerprint, Any, Any)
173+ tripleX@(fpX, _, _) = (calcFp @a, toAny x, unsafeCoerce $ typeRep @a)
174+
175+ addX :: [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
176+ addX l = tripleX : deleteByFst fpX l
202177{-# INLINE insert #-}
203178
204179-- Extract the kind of a type. We use it to work around lack of syntax for
205180-- inferred type variables (which are not subject to type applications).
206181type KindOf (a :: k) = k
207182
208−type ArgKindOf (f :: k -> l) = k
209−
210183{- | Delete a value from a 'TypeRepMap'.
211184
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−
215185prop> size (delete @a tm) <= size tm
216186prop> member @a (delete @a tm) == False
217187
@@ -224,60 +194,30 @@ False
224194True
225195-}
226196delete :: forall a (f :: KindOf a -> Type) . Typeable a => TypeRepMap f -> TypeRepMap f
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
197+delete = fromTriples . deleteByFst (typeFp @a) . toTriples
234198{-# INLINE delete #-}
235199
236−deleteFirst :: (a -> Bool) -> [a] -> [a]
237−deleteFirst _ [] = []
238−deleteFirst p (x : xs) = if p x then xs else x : deleteFirst p xs
239−
240200{- |
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.
201+Update a value at a specific key with the result of the provided function. When
202+the key is not a member of the map, the original map is returned.
243203
244204>>> trmap = fromList @(TypeRepMap Identity) [WrapTypeable $ Identity "a"]
245205>>> lookup @String $ adjust (fmap (++ "ww")) trmap
246206Just (Identity "aww")
247207-}
248208adjust :: 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)}
209+adjust fun tr = case cachedBinarySearch (typeFp @a) (fingerprintAs tr) (fingerprintBs tr) of
210+ Nothing -> tr
211+ Just i -> tr {trAnys = changeAnyArr i (trAnys tr)}
273212 where
274− replaceAnyAt :: Int -> Any -> Array Any -> Array Any
275− replaceAnyAt i v trAs = runST $ do
213+ changeAnyArr :: Int -> Array Any -> Array Any
214+ changeAnyArr i trAs = runST $ do
276215 let n = sizeofArray trAs
277216 mutArr <- thawArray trAs 0 n
278− writeArray mutArr i v
217+ a <- toAny . fun . fromAny <$> readArray mutArr i
218+ writeArray mutArr i a
279219 unsafeFreezeArray mutArr
280−{-# INLINE alter #-}
220+{-# INLINE adjust #-}
281221
282222{- | Map over the elements of a 'TypeRepMap'.
283223
@@ -311,89 +251,33 @@ hoistWithKey f (TypeRepMap as bs ans ks) = TypeRepMap as bs newAns ks
311251 withTr t = withTypeable t f
312252{-# INLINE hoistWithKey #-}
313253
314−-- | The union of two 'TypeRepMap's using a combining function for conflicting entries. @O(n + m)@
254+-- | The union of two 'TypeRepMap's using a combining function.
315255unionWith :: 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)
256+unionWith f m1 m2 = fromTriples
257+ $ toTripleList
258+ $ Map.unionWith combine
259+ (fromTripleList $ toTriples m1)
260+ (fromTripleList $ toTriples m2)
318261 where
319262 f' :: forall x. TypeRep x -> f x -> f x -> f x
320263 f' tr = withTypeable tr f
321264
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
265+ combine :: (Any, Any) -> (Any, Any) -> (Any, Any)
266+ combine (av, ak) (bv, _) = (toAny $ f' (fromAny ak) (fromAny av) (fromAny bv), ak)
267+
268+ fromTripleList :: Ord a => [(a, b, c)] -> Map.Map a (b, c)
269+ fromTripleList = Map.fromList . map (\(a, b, c) -> (a, (b, c)))
270+
271+ toTripleList :: Map.Map a (b, c) -> [(a, b, c)]
272+ toTripleList = map (\(a, (b, c)) -> (a, b, c)) . Map.toList
347273{-# INLINE unionWith #-}
348274
349−-- | The (left-biased) union of two 'TypeRepMap's in @O(n + m)@. It prefers the first map when
275+-- | The (left-biased) union of two 'TypeRepMap's. It prefers the first map when
350276-- duplicate keys are encountered, i.e. @'union' == 'unionWith' const@.
351277union :: TypeRepMap f -> TypeRepMap f -> TypeRepMap f
352278union = unionWith const
353279{-# INLINE union #-}
354280
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−
397281{- | Check if a value of the given type is present in a 'TypeRepMap'.
398282
399283>>> member @Char $ one (Identity 'a')
@@ -429,22 +313,9 @@ size = sizeofPrimArray . fingerprintAs
429313
430314-- | Return the list of 'SomeTypeRep' from the keys.
431315keys :: TypeRepMap f -> [SomeTypeRep]
432−keys = keysWith SomeTypeRep
316+keys TypeRepMap{..} = SomeTypeRep . anyToTypeRep <$> toList trKeys
433317{-# INLINE keys #-}
434318
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−
448319-- | Binary searched based on this article
449320-- http://bannalia.blogspot.com/2015/06/cache-friendly-binary-search.html
450321-- with modification for our two-vector search case.
@@ -488,20 +359,8 @@ typeFp = typeRepFingerprint $ typeRep @a
488359toTriples :: TypeRepMap f -> [(Fingerprint, Any, Any)]
489360toTriples tm = zip3 (toFingerprints tm) (GHC.toList $ trAnys tm) (GHC.toList $ trKeys tm)
490361
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)
362+deleteByFst :: Eq a => a -> [(a, b, c)] -> [(a, b, c)]
363+deleteByFst x = filter ((/= x) . fst3)
505364
506365nubByFst :: (Eq a) => [(a, b, c)] -> [(a, b, c)]
507366nubByFst = nubBy ((==) `on` fst3)
@@ -521,6 +380,9 @@ data WrapTypeable f where
521380instance Show (WrapTypeable f) where
522381 show (WrapTypeable (_ :: f a)) = show $ calcFp @a
523382
383+wrapTypeable :: TypeRep a -> f a -> WrapTypeable f
384+wrapTypeable tr = withTypeable tr WrapTypeable
385+
524386{- |
525387
526388prop> fromList . toList == 'id'
@@ -548,19 +410,19 @@ instance IsList (TypeRepMap f) where
548410 k (WrapTypeable (_ :: f a)) = unsafeCoerce $ typeRep @a
549411
550412 toList :: TypeRepMap f -> [WrapTypeable f]
551− toList = toListWith WrapTypeable
413+ toList = map toWrapTypeable . toTriples
414+ where
415+ toWrapTypeable :: (Fingerprint, Any, Any) -> WrapTypeable f
416+ toWrapTypeable (_, an, k) = wrapTypeable (unsafeCoerce k) (fromAny an)
552417
553418calcFp :: forall a . Typeable a => Fingerprint
554419calcFp = typeRepFingerprint $ typeRep @a
555420
556421fromTriples :: [(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)
422+fromTriples kvs = TypeRepMap (GHC.fromList fpAs) (GHC.fromList fpBs) (GHC.fromList ans) (GHC.fromList ks)
561423 where
562424 (fpAs, fpBs) = unzip $ map (\(Fingerprint a b) -> (a, b)) fps
563− (fps, ans, ks) = unzip3 $ fromSortedList kvs
425+ (fps, ans, ks) = unzip3 $ fromSortedList $ sortWith fst3 $ nubByFst kvs
564426
565427----------------------------------------------------------------------------
566428-- Tree-like conversion
@@ -586,85 +448,3 @@ fromSortedList l = runST $ do
586448 newFirst <- loop (2 * i + 1) first
587449 writeArray result i (indexArray origin newFirst)
588450 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− ]