@@ -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)
@@ -37,20 +45,25 @@ import Data.Kind (Type)
3745import Data.Type.Equality ((:~:) (..), TestEquality (..))
3846import Data.List (intercalate, nubBy)
3947import Data.Maybe (fromMaybe)
40−import Data.Primitive.Array (Array, MutableArray, indexArray, mapArray', readArray, sizeofArray,
48+import Data.Primitive.Array (Array, MutableArray, indexArray, mapArray', sizeofArray,
4149 thawArray, unsafeFreezeArray, writeArray)
42−import Data.Primitive.PrimArray (PrimArray, indexPrimArray, sizeofPrimArray)
50+import Data.Primitive.PrimArray (writePrimArray, newPrimArray, unsafeFreezePrimArray, primArrayToList, MutablePrimArray, primArrayFromListN, PrimArray, indexPrimArray, sizeofPrimArray)
4351import Data.Semigroup (Semigroup (..), All(..))
4452import GHC.Base (Any, Int (..), Int#, (*#), (+#), (<#))
4553import GHC.Exts (IsList (..), inline, sortWith)
4654import GHC.Fingerprint (Fingerprint (..))
55+#if WORD_SIZE_IN_BITS >= 64
4756import GHC.Prim (eqWord#, ltWord#)
57+#else
58+import GHC.IntWord64 (eqWord64#, ltWord64#)
59+#define eqWord eqWord64
60+#define ltWord ltWord64
61+#endif
4862import GHC.Word (Word64 (..))
4963import Type.Reflection (SomeTypeRep (..), TypeRep, Typeable, typeRep, withTypeable)
5064import Type.Reflection.Unsafe (typeRepFingerprint)
5165import Unsafe.Coerce (unsafeCoerce)
5266
53−import qualified Data.Map.Strict as Map
5467import qualified GHC.Exts as GHC (fromList, toList)
5568
5669{- |
@@ -156,33 +169,49 @@ prop> member @a (one (x :: f a)) == True
156169
157170-}
158171one :: forall a f . Typeable a => f a -> TypeRepMap f
159−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)
160178{-# INLINE one #-}
161179
162180{- |
163181
164182Insert a value into a 'TypeRepMap'.
183+TypeRepMap optimizes for fast reads rather than inserts, as a trade-off inserts are @O(n)@.
165184
166185prop> size (insert v tm) >= size tm
167186prop> member @a (insert (x :: f a) tm) == True
168187
169188-}
170189insert :: forall a f . Typeable a => f a -> TypeRepMap f -> TypeRepMap f
171−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)}
172195 where
173− tripleX :: (Fingerprint, Any, Any)
174− tripleX@(fpX, _, _) = (calcFp @a, toAny x, unsafeCoerce $ typeRep @a)
175−
176− addX :: [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
177− 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
178202{-# INLINE insert #-}
179203
180204-- Extract the kind of a type. We use it to work around lack of syntax for
181205-- inferred type variables (which are not subject to type applications).
182206type KindOf (a :: k) = k
183207
208+type ArgKindOf (f :: k -> l) = k
209+
184210{- | Delete a value from a 'TypeRepMap'.
185211
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+
186215prop> size (delete @a tm) <= size tm
187216prop> member @a (delete @a tm) == False
188217
@@ -195,30 +224,60 @@ False
195224True
196225-}
197226delete :: forall a (f :: KindOf a -> Type) . Typeable a => TypeRepMap f -> TypeRepMap f
198−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
199234{-# INLINE delete #-}
200235
236+deleteFirst :: (a -> Bool) -> [a] -> [a]
237+deleteFirst _ [] = []
238+deleteFirst p (x : xs) = if p x then xs else x : deleteFirst p xs
239+
201240{- |
202−Update a value at a specific key with the result of the provided function. When
203−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.
204243
205244>>> trmap = fromList @(TypeRepMap Identity) [WrapTypeable $ Identity "a"]
206245>>> lookup @String $ adjust (fmap (++ "ww")) trmap
207246Just (Identity "aww")
208247-}
209248adjust :: forall a f . Typeable a => (f a -> f a) -> TypeRepMap f -> TypeRepMap f
210−adjust fun tr = case cachedBinarySearch (typeFp @a) (fingerprintAs tr) (fingerprintBs tr) of
211− Nothing -> tr
212− 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)}
213273 where
214− changeAnyArr :: Int -> Array Any -> Array Any
215− changeAnyArr i trAs = runST $ do
274+ replaceAnyAt :: Int -> Any -> Array Any -> Array Any
275+ replaceAnyAt i v trAs = runST $ do
216276 let n = sizeofArray trAs
217277 mutArr <- thawArray trAs 0 n
218− a <- toAny . fun . fromAny <$> readArray mutArr i
219− writeArray mutArr i a
278+ writeArray mutArr i v
220279 unsafeFreezeArray mutArr
221−{-# INLINE adjust #-}
280+{-# INLINE alter #-}
222281
223282{- | Map over the elements of a 'TypeRepMap'.
224283
@@ -252,33 +311,89 @@ hoistWithKey f (TypeRepMap as bs ans ks) = TypeRepMap as bs newAns ks
252311 withTr t = withTypeable t f
253312{-# INLINE hoistWithKey #-}
254313
255−-- | 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)@
256315unionWith :: forall f. (forall x. Typeable x => f x -> f x -> f x) -> TypeRepMap f -> TypeRepMap f -> TypeRepMap f
257−unionWith f m1 m2 = fromTriples
258− $ toTripleList
259− $ Map.unionWith combine
260− (fromTripleList $ toTriples m1)
261− (fromTripleList $ toTriples m2)
316+unionWith f ma mb = do
317+ fromSortedTriples $ mergeMaps (toSortedTriples ma) (toSortedTriples mb)
262318 where
263319 f' :: forall x. TypeRep x -> f x -> f x -> f x
264320 f' tr = withTypeable tr f
265321
266− combine :: (Any, Any) -> (Any, Any) -> (Any, Any)
267− combine (av, ak) (bv, _) = (toAny $ f' (fromAny ak) (fromAny av) (fromAny bv), ak)
268−
269− fromTripleList :: Ord a => [(a, b, c)] -> Map.Map a (b, c)
270− fromTripleList = Map.fromList . map (\(a, b, c) -> (a, (b, c)))
271−
272− toTripleList :: Map.Map a (b, c) -> [(a, b, c)]
273− 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
274347{-# INLINE unionWith #-}
275348
276−-- | 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
277350-- duplicate keys are encountered, i.e. @'union' == 'unionWith' const@.
278351union :: TypeRepMap f -> TypeRepMap f -> TypeRepMap f
279352union = unionWith const
280353{-# INLINE union #-}
281354
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+
282397{- | Check if a value of the given type is present in a 'TypeRepMap'.
283398
284399>>> member @Char $ one (Identity 'a')
@@ -314,9 +429,22 @@ size = sizeofPrimArray . fingerprintAs
314429
315430-- | Return the list of 'SomeTypeRep' from the keys.
316431keys :: TypeRepMap f -> [SomeTypeRep]
317−keys TypeRepMap{..} = SomeTypeRep . anyToTypeRep <$> toList trKeys
432+keys = keysWith SomeTypeRep
318433{-# INLINE keys #-}
319434
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+
320448-- | Binary searched based on this article
321449-- http://bannalia.blogspot.com/2015/06/cache-friendly-binary-search.html
322450-- with modification for our two-vector search case.
@@ -360,8 +488,20 @@ typeFp = typeRepFingerprint $ typeRep @a
360488toTriples :: TypeRepMap f -> [(Fingerprint, Any, Any)]
361489toTriples tm = zip3 (toFingerprints tm) (GHC.toList $ trAnys tm) (GHC.toList $ trKeys tm)
362490
363−deleteByFst :: Eq a => a -> [(a, b, c)] -> [(a, b, c)]
364−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)
365505
366506nubByFst :: (Eq a) => [(a, b, c)] -> [(a, b, c)]
367507nubByFst = nubBy ((==) `on` fst3)
@@ -381,9 +521,6 @@ data WrapTypeable f where
381521instance Show (WrapTypeable f) where
382522 show (WrapTypeable (_ :: f a)) = show $ calcFp @a
383523
384−wrapTypeable :: TypeRep a -> f a -> WrapTypeable f
385−wrapTypeable tr = withTypeable tr WrapTypeable
386−
387524{- |
388525
389526prop> fromList . toList == 'id'
@@ -411,19 +548,19 @@ instance IsList (TypeRepMap f) where
411548 k (WrapTypeable (_ :: f a)) = unsafeCoerce $ typeRep @a
412549
413550 toList :: TypeRepMap f -> [WrapTypeable f]
414− toList = map toWrapTypeable . toTriples
415− where
416− toWrapTypeable :: (Fingerprint, Any, Any) -> WrapTypeable f
417− toWrapTypeable (_, an, k) = wrapTypeable (unsafeCoerce k) (fromAny an)
551+ toList = toListWith WrapTypeable
418552
419553calcFp :: forall a . Typeable a => Fingerprint
420554calcFp = typeRepFingerprint $ typeRep @a
421555
422556fromTriples :: [(Fingerprint, Any, Any)] -> TypeRepMap f
423−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)
424561 where
425562 (fpAs, fpBs) = unzip $ map (\(Fingerprint a b) -> (a, b)) fps
426− (fps, ans, ks) = unzip3 $ fromSortedList $ sortWith fst3 $ nubByFst kvs
563+ (fps, ans, ks) = unzip3 $ fromSortedList kvs
427564
428565----------------------------------------------------------------------------
429566-- Tree-like conversion
@@ -450,18 +587,42 @@ fromSortedList l = runST $ do
450587 writeArray result i (indexArray origin newFirst)
451588 loop (2 * i + 2) (newFirst + 1)
452589
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+
453614----------------------------------------------------------------------------
454615-- Helper functions.
455616----------------------------------------------------------------------------
456617
457−-- | Check that invariant of the structure is hold.
618+-- | Check that invariant of the structure holds.
458619-- The structure maintains the following invariant.
459620-- For each element @A@ at index @i@:
460621--
461−-- 1. if there is an element @B@ at index @2*i+1@,
622+-- 1. if there is an element @B@ at index @2*i+1@,
462623-- then @B < A@.
463624--
464−-- 2. if there is an element @C@ at index @2*i+2@,
625+-- 2. if there is an element @C@ at index @2*i+2@,
465626-- then @A < C@.
466627--
467628invariantCheck :: TypeRepMap f -> Bool
@@ -507,4 +668,3 @@ invariantCheck TypeRepMap{..} = getAll (check 0)
507668 , All $ fromMaybe True $ (<=) <$> leftMax <*> rightMin
508669 , check (i+1)
509670 ]
510−