@@ -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 #-}
@@ -18,13 +19,20 @@
1819
1920-- {-# OPTIONS_GHC -ddump-simpl -dsuppress-idinfo -dsuppress-coercions -dsuppress-type-applications -dsuppress-uniques -dsuppress-module-prefixes #-}
2021
21−{- | Internal API for 'TypeRepMap' and operations on it. The functions here do
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
2228not have any stability guarantees and can change between minor versions.
2329
2430If you need to use this module for purposes other than tests,
2531create an issue.
2632-}
2733
34+#include "MachDeps.h"
35+
2836module Data.TypeRepMap.Internal where
2937
3038import Prelude hiding (lookup)
@@ -36,20 +44,26 @@ import Data.Function (on)
3644import Data.Kind (Type)
3745import Data.Type.Equality ((:~:) (..), TestEquality (..))
3846import Data.List (intercalate, nubBy)
39−import Data.Primitive.Array (Array, MutableArray, indexArray, mapArray', readArray, sizeofArray,
47+import Data.Maybe (fromMaybe)
48+import Data.Primitive.Array (Array, MutableArray, indexArray, mapArray', sizeofArray,
4049 thawArray, unsafeFreezeArray, writeArray)
41−import Data.Primitive.PrimArray (PrimArray, indexPrimArray, sizeofPrimArray)
42−import Data.Semigroup (Semigroup (..))
50+import Data.Primitive.PrimArray (writePrimArray, newPrimArray, unsafeFreezePrimArray, primArrayToList, MutablePrimArray, primArrayFromListN, PrimArray, indexPrimArray, sizeofPrimArray)
51+import Data.Semigroup (Semigroup (..), All(..))
4352import GHC.Base (Any, Int (..), Int#, (*#), (+#), (<#))
4453import GHC.Exts (IsList (..), inline, sortWith)
4554import GHC.Fingerprint (Fingerprint (..))
55+#if WORD_SIZE_IN_BITS >= 64
4656import GHC.Prim (eqWord#, ltWord#)
57+#else
58+import GHC.IntWord64 (eqWord64#, ltWord64#)
59+#define eqWord eqWord64
60+#define ltWord ltWord64
61+#endif
4762import GHC.Word (Word64 (..))
4863import Type.Reflection (SomeTypeRep (..), TypeRep, Typeable, typeRep, withTypeable)
4964import Type.Reflection.Unsafe (typeRepFingerprint)
5065import Unsafe.Coerce (unsafeCoerce)
5166
52−import qualified Data.Map.Strict as Map
5367import qualified GHC.Exts as GHC (fromList, toList)
5468
5569{- |
@@ -155,33 +169,49 @@ prop> member @a (one (x :: f a)) == True
155169
156170-}
157171one :: forall a f . Typeable a => f a -> TypeRepMap f
158−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)
159178{-# INLINE one #-}
160179
161180{- |
162181
163182Insert a value into a 'TypeRepMap'.
183+TypeRepMap optimizes for fast reads rather than inserts, as a trade-off inserts are @O(n)@.
164184
165185prop> size (insert v tm) >= size tm
166186prop> member @a (insert (x :: f a) tm) == True
167187
168188-}
169189insert :: forall a f . Typeable a => f a -> TypeRepMap f -> TypeRepMap f
170−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)}
171195 where
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
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
177202{-# INLINE insert #-}
178203
179204-- Extract the kind of a type. We use it to work around lack of syntax for
180205-- inferred type variables (which are not subject to type applications).
181206type KindOf (a :: k) = k
182207
208+type ArgKindOf (f :: k -> l) = k
209+
183210{- | Delete a value from a 'TypeRepMap'.
184211
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+
185215prop> size (delete @a tm) <= size tm
186216prop> member @a (delete @a tm) == False
187217
@@ -194,30 +224,60 @@ False
194224True
195225-}
196226delete :: forall a (f :: KindOf a -> Type) . Typeable a => TypeRepMap f -> TypeRepMap f
197−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
198234{-# INLINE delete #-}
199235
236+deleteFirst :: (a -> Bool) -> [a] -> [a]
237+deleteFirst _ [] = []
238+deleteFirst p (x : xs) = if p x then xs else x : deleteFirst p xs
239+
200240{- |
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.
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.
203243
204244>>> trmap = fromList @(TypeRepMap Identity) [WrapTypeable $ Identity "a"]
205245>>> lookup @String $ adjust (fmap (++ "ww")) trmap
206246Just (Identity "aww")
207247-}
208248adjust :: forall a f . Typeable a => (f a -> f a) -> TypeRepMap f -> TypeRepMap f
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)}
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)}
212273 where
213− changeAnyArr :: Int -> Array Any -> Array Any
214− changeAnyArr i trAs = runST $ do
274+ replaceAnyAt :: Int -> Any -> Array Any -> Array Any
275+ replaceAnyAt i v trAs = runST $ do
215276 let n = sizeofArray trAs
216277 mutArr <- thawArray trAs 0 n
217− a <- toAny . fun . fromAny <$> readArray mutArr i
218− writeArray mutArr i a
278+ writeArray mutArr i v
219279 unsafeFreezeArray mutArr
220−{-# INLINE adjust #-}
280+{-# INLINE alter #-}
221281
222282{- | Map over the elements of a 'TypeRepMap'.
223283
@@ -251,33 +311,89 @@ hoistWithKey f (TypeRepMap as bs ans ks) = TypeRepMap as bs newAns ks
251311 withTr t = withTypeable t f
252312{-# INLINE hoistWithKey #-}
253313
254−-- | 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)@
255315unionWith :: forall f. (forall x. Typeable x => f x -> f x -> f x) -> TypeRepMap f -> TypeRepMap f -> TypeRepMap f
256−unionWith f m1 m2 = fromTriples
257− $ toTripleList
258− $ Map.unionWith combine
259− (fromTripleList $ toTriples m1)
260− (fromTripleList $ toTriples m2)
316+unionWith f ma mb = do
317+ fromSortedTriples $ mergeMaps (toSortedTriples ma) (toSortedTriples mb)
261318 where
262319 f' :: forall x. TypeRep x -> f x -> f x -> f x
263320 f' tr = withTypeable tr f
264321
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
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
273347{-# INLINE unionWith #-}
274348
275−-- | 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
276350-- duplicate keys are encountered, i.e. @'union' == 'unionWith' const@.
277351union :: TypeRepMap f -> TypeRepMap f -> TypeRepMap f
278352union = unionWith const
279353{-# INLINE union #-}
280354
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+
281397{- | Check if a value of the given type is present in a 'TypeRepMap'.
282398
283399>>> member @Char $ one (Identity 'a')
@@ -313,9 +429,22 @@ size = sizeofPrimArray . fingerprintAs
313429
314430-- | Return the list of 'SomeTypeRep' from the keys.
315431keys :: TypeRepMap f -> [SomeTypeRep]
316−keys TypeRepMap{..} = SomeTypeRep . anyToTypeRep <$> toList trKeys
432+keys = keysWith SomeTypeRep
317433{-# INLINE keys #-}
318434
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+
319448-- | Binary searched based on this article
320449-- http://bannalia.blogspot.com/2015/06/cache-friendly-binary-search.html
321450-- with modification for our two-vector search case.
@@ -359,8 +488,20 @@ typeFp = typeRepFingerprint $ typeRep @a
359488toTriples :: TypeRepMap f -> [(Fingerprint, Any, Any)]
360489toTriples tm = zip3 (toFingerprints tm) (GHC.toList $ trAnys tm) (GHC.toList $ trKeys tm)
361490
362−deleteByFst :: Eq a => a -> [(a, b, c)] -> [(a, b, c)]
363−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)
364505
365506nubByFst :: (Eq a) => [(a, b, c)] -> [(a, b, c)]
366507nubByFst = nubBy ((==) `on` fst3)
@@ -380,9 +521,6 @@ data WrapTypeable f where
380521instance Show (WrapTypeable f) where
381522 show (WrapTypeable (_ :: f a)) = show $ calcFp @a
382523
383−wrapTypeable :: TypeRep a -> f a -> WrapTypeable f
384−wrapTypeable tr = withTypeable tr WrapTypeable
385−
386524{- |
387525
388526prop> fromList . toList == 'id'
@@ -410,19 +548,19 @@ instance IsList (TypeRepMap f) where
410548 k (WrapTypeable (_ :: f a)) = unsafeCoerce $ typeRep @a
411549
412550 toList :: TypeRepMap f -> [WrapTypeable f]
413− toList = map toWrapTypeable . toTriples
414− where
415− toWrapTypeable :: (Fingerprint, Any, Any) -> WrapTypeable f
416− toWrapTypeable (_, an, k) = wrapTypeable (unsafeCoerce k) (fromAny an)
551+ toList = toListWith WrapTypeable
417552
418553calcFp :: forall a . Typeable a => Fingerprint
419554calcFp = typeRepFingerprint $ typeRep @a
420555
421556fromTriples :: [(Fingerprint, Any, Any)] -> TypeRepMap f
422−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)
423561 where
424562 (fpAs, fpBs) = unzip $ map (\(Fingerprint a b) -> (a, b)) fps
425− (fps, ans, ks) = unzip3 $ fromSortedList $ sortWith fst3 $ nubByFst kvs
563+ (fps, ans, ks) = unzip3 $ fromSortedList kvs
426564
427565----------------------------------------------------------------------------
428566-- Tree-like conversion
@@ -448,3 +586,85 @@ fromSortedList l = runST $ do
448586 newFirst <- loop (2 * i + 1) first
449587 writeArray result i (indexArray origin newFirst)
450588 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+ ]