@@ -1,11 +1,10 @@
1−{-# OPTIONS_GHC -Wno-missing-export-lists #-}
3{-# LANGUAGE AllowAmbiguousTypes #-}
1{-# LANGUAGE AllowAmbiguousTypes #-}
4{-# LANGUAGE BangPatterns #-}
2{-# LANGUAGE BangPatterns #-}
5{-# LANGUAGE CPP #-}
3{-# LANGUAGE CPP #-}
6{-# LANGUAGE DataKinds #-}
4{-# LANGUAGE DataKinds #-}
7{-# LANGUAGE GADTs #-}
5{-# LANGUAGE GADTs #-}
8{-# LANGUAGE InstanceSigs #-}
6{-# LANGUAGE InstanceSigs #-}
7+{-# LANGUAGE KindSignatures #-}
9{-# LANGUAGE MagicHash #-}
8{-# LANGUAGE MagicHash #-}
10{-# LANGUAGE PolyKinds #-}
9{-# LANGUAGE PolyKinds #-}
11{-# LANGUAGE Rank2Types #-}
10{-# LANGUAGE Rank2Types #-}
@@ -45,9 +44,9 @@ import Data.Kind (Type)
45import Data.Type.Equality ((:~:) (..), TestEquality (..))
44import Data.Type.Equality ((:~:) (..), TestEquality (..))
46import Data.List (intercalate, nubBy)
45import Data.List (intercalate, nubBy)
47import Data.Maybe (fromMaybe)
46import Data.Maybe (fromMaybe)
48−import Data.Primitive.Array (Array, MutableArray, indexArray, mapArray', sizeofArray,
47+import Data.Primitive.Array (Array, MutableArray, indexArray, mapArray', readArray, sizeofArray,
49 thawArray, unsafeFreezeArray, writeArray)
48 thawArray, unsafeFreezeArray, writeArray)
50−import Data.Primitive.PrimArray (writePrimArray, newPrimArray, unsafeFreezePrimArray, primArrayToList, MutablePrimArray, primArrayFromListN, PrimArray, indexPrimArray, sizeofPrimArray)
49+import Data.Primitive.PrimArray (PrimArray, indexPrimArray, sizeofPrimArray)
51import Data.Semigroup (Semigroup (..), All(..))
50import Data.Semigroup (Semigroup (..), All(..))
52import GHC.Base (Any, Int (..), Int#, (*#), (+#), (<#))
51import GHC.Base (Any, Int (..), Int#, (*#), (+#), (<#))
53import GHC.Exts (IsList (..), inline, sortWith)
52import GHC.Exts (IsList (..), inline, sortWith)
@@ -64,6 +63,7 @@ import Type.Reflection (SomeTypeRep (..), TypeRep, Typeable, typeRep, withTypeab
64import Type.Reflection.Unsafe (typeRepFingerprint)
63import Type.Reflection.Unsafe (typeRepFingerprint)
65import Unsafe.Coerce (unsafeCoerce)
64import Unsafe.Coerce (unsafeCoerce)
66+import qualified Data.Map.Strict as Map
67import qualified GHC.Exts as GHC (fromList, toList)
67import qualified GHC.Exts as GHC (fromList, toList)
@@ -169,49 +169,33 @@ prop> member @a (one (x :: f a)) == True
171one :: forall a f . Typeable a => f a -> TypeRepMap f
171one :: forall a f . Typeable a => f a -> TypeRepMap f
172−one x = TypeRepMap (primArrayFromListN 1 [fa])
172+one x = insert x empty
173− (primArrayFromListN 1 [fb])
177− (Fingerprint fa fb, v, k) = (calcFp @a, toAny x, unsafeCoerce $ typeRep @a)
178{-# INLINE one #-}
173{-# INLINE one #-}
182Insert a value into a 'TypeRepMap'.
177Insert a value into a 'TypeRepMap'.
183−TypeRepMap optimizes for fast reads rather than inserts, as a trade-off inserts are @O(n)@.
185prop> size (insert v tm) >= size tm
179prop> size (insert v tm) >= size tm
186prop> member @a (insert (x :: f a) tm) == True
180prop> member @a (insert (x :: f a) tm) == True
189insert :: forall a f . Typeable a => f a -> TypeRepMap f -> TypeRepMap f
183insert :: forall a f . Typeable a => f a -> TypeRepMap f -> TypeRepMap f
190−insert x m
184+insert x = fromTriples . addX . toTriples
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)}
196− changeAnyArr :: Int -> Array Any -> Array Any
186+ tripleX :: (Fingerprint, Any, Any)
197− changeAnyArr i trAs = runST $ do
187+ tripleX@(fpX, _, _) = (calcFp @a, toAny x, unsafeCoerce $ typeRep @a)
198− let n = sizeofArray trAs
188+
199− mutArr <- thawArray trAs 0 n
189+ addX :: [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
200− writeArray mutArr i $ toAny x
190+ addX l = tripleX : deleteByFst fpX l
201− unsafeFreezeArray mutArr
202{-# INLINE insert #-}
191{-# INLINE insert #-}
204-- Extract the kind of a type. We use it to work around lack of syntax for
193-- Extract the kind of a type. We use it to work around lack of syntax for
205-- inferred type variables (which are not subject to type applications).
194-- inferred type variables (which are not subject to type applications).
206type KindOf (a :: k) = k
195type KindOf (a :: k) = k
208−type ArgKindOf (f :: k -> l) = k
210{- | Delete a value from a 'TypeRepMap'.
197{- | 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.
215prop> size (delete @a tm) <= size tm
199prop> size (delete @a tm) <= size tm
216prop> member @a (delete @a tm) == False
200prop> member @a (delete @a tm) == False
@@ -224,60 +208,30 @@ False
226delete :: forall a (f :: KindOf a -> Type) . Typeable a => TypeRepMap f -> TypeRepMap f
210delete :: forall a (f :: KindOf a -> Type) . Typeable a => TypeRepMap f -> TypeRepMap f
227−delete m
211+delete = fromTriples . deleteByFst (typeFp @a) . toTriples
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
234{-# INLINE delete #-}
212{-# 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.
215+Update a value at a specific key with the result of the provided function. When
242−When the key is not a member of the map, the original map is returned.
216+the key is not a member of the map, the original map is returned.
244>>> trmap = fromList @(TypeRepMap Identity) [WrapTypeable $ Identity "a"]
218>>> trmap = fromList @(TypeRepMap Identity) [WrapTypeable $ Identity "a"]
245>>> lookup @String $ adjust (fmap (++ "ww")) trmap
219>>> lookup @String $ adjust (fmap (++ "ww")) trmap
246Just (Identity "aww")
220Just (Identity "aww")
248adjust :: forall a f . Typeable a => (f a -> f a) -> TypeRepMap f -> TypeRepMap f
222adjust :: 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
223+adjust fun tr = case cachedBinarySearch (typeFp @a) (fingerprintAs tr) (fingerprintBs tr) of
265− Nothing ->
224+ Nothing -> tr
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)}
225+ Just i -> tr {trAnys = changeAnyArr i (trAnys tr)}
274− replaceAnyAt :: Int -> Any -> Array Any -> Array Any
227+ changeAnyArr :: Int -> Array Any -> Array Any
275− replaceAnyAt i v trAs = runST $ do
228+ changeAnyArr i trAs = runST $ do
276 let n = sizeofArray trAs
229 let n = sizeofArray trAs
277 mutArr <- thawArray trAs 0 n
230 mutArr <- thawArray trAs 0 n
231+ a <- toAny . fun . fromAny <$> readArray mutArr i
278− writeArray mutArr i v
232+ writeArray mutArr i a
279 unsafeFreezeArray mutArr
233 unsafeFreezeArray mutArr
280−{-# INLINE alter #-}
234+{-# INLINE adjust #-}
282{- | Map over the elements of a 'TypeRepMap'.
236{- | Map over the elements of a 'TypeRepMap'.
@@ -311,89 +265,33 @@ hoistWithKey f (TypeRepMap as bs ans ks) = TypeRepMap as bs newAns ks
311 withTr t = withTypeable t f
265 withTr t = withTypeable t f
312{-# INLINE hoistWithKey #-}
266{-# INLINE hoistWithKey #-}
314−-- | The union of two 'TypeRepMap's using a combining function for conflicting entries. @O(n + m)@
268+-- | The union of two 'TypeRepMap's using a combining function.
315unionWith :: forall f. (forall x. Typeable x => f x -> f x -> f x) -> TypeRepMap f -> TypeRepMap f -> TypeRepMap f
269unionWith :: forall f. (forall x. Typeable x => f x -> f x -> f x) -> TypeRepMap f -> TypeRepMap f -> TypeRepMap f
316−unionWith f ma mb = do
270+unionWith f m1 m2 = fromTriples
317− fromSortedTriples $ mergeMaps (toSortedTriples ma) (toSortedTriples mb)
271+ $ toTripleList
272+ $ Map.unionWith combine
273+ (fromTripleList $ toTriples m1)
274+ (fromTripleList $ toTriples m2)
319 f' :: forall x. TypeRep x -> f x -> f x -> f x
276 f' :: forall x. TypeRep x -> f x -> f x -> f x
320 f' tr = withTypeable tr f
277 f' tr = withTypeable tr f
322− combine :: (Fingerprint, Any, Any) -> (Fingerprint, Any, Any) -> (Fingerprint, Any, Any)
279+ combine :: (Any, Any) -> (Any, Any) -> (Any, Any)
323− combine (fp, av, ak) (_, bv, _) = (fp, toAny $ f' (fromAny ak) (fromAny av) (fromAny bv), ak)
280+ combine (av, ak) (bv, _) = (toAny $ f' (fromAny ak) (fromAny av) (fromAny bv), ak)
325− -- Merges two typrepmaps into a sorted, dedup'd list of triples.
282+ fromTripleList :: Ord a => [(a, b, c)] -> Map.Map a (b, c)
326− -- Using 'toSortedTriples' allows us to assume the triples are sorted by fingerprint,
283+ fromTripleList = Map.fromList . map (\(a, b, c) -> (a, (b, c)))
327− -- Given O(n) performance from 'toSortedTriples', and given that we can merge-sort in
284+
328− -- O(n + m) time, then can '.fromSortedTriples' back into cachedBinarySearch order in O(n + m)
285+ toTripleList :: Map.Map a (b, c) -> [(a, b, c)]
329− -- that gives a total of O(n + m).
286+ toTripleList = map (\(a, (b, c)) -> (a, b, c)) . Map.toList
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
347{-# INLINE unionWith #-}
287{-# INLINE unionWith #-}
349−-- | The (left-biased) union of two 'TypeRepMap's in @O(n + m)@. It prefers the first map when
289+-- | The (left-biased) union of two 'TypeRepMap's. It prefers the first map when
350-- duplicate keys are encountered, i.e. @'union' == 'unionWith' const@.
290-- duplicate keys are encountered, i.e. @'union' == 'unionWith' const@.
351union :: TypeRepMap f -> TypeRepMap f -> TypeRepMap f
291union :: TypeRepMap f -> TypeRepMap f -> TypeRepMap f
352union = unionWith const
292union = unionWith const
353{-# INLINE union #-}
293{-# 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 #-}
397{- | Check if a value of the given type is present in a 'TypeRepMap'.
295{- | Check if a value of the given type is present in a 'TypeRepMap'.
399>>> member @Char $ one (Identity 'a')
297>>> member @Char $ one (Identity 'a')
@@ -429,22 +327,9 @@ size = sizeofPrimArray . fingerprintAs
430-- | Return the list of 'SomeTypeRep' from the keys.
328-- | Return the list of 'SomeTypeRep' from the keys.
431keys :: TypeRepMap f -> [SomeTypeRep]
329keys :: TypeRepMap f -> [SomeTypeRep]
432−keys = keysWith SomeTypeRep
330+keys TypeRepMap{..} = SomeTypeRep . anyToTypeRep <$> toList trKeys
433{-# INLINE keys #-}
331{-# 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)
448-- | Binary searched based on this article
333-- | Binary searched based on this article
449-- http://bannalia.blogspot.com/2015/06/cache-friendly-binary-search.html
334-- http://bannalia.blogspot.com/2015/06/cache-friendly-binary-search.html
450-- with modification for our two-vector search case.
335-- with modification for our two-vector search case.
@@ -488,20 +373,8 @@ typeFp = typeRepFingerprint $ typeRep @a
488toTriples :: TypeRepMap f -> [(Fingerprint, Any, Any)]
373toTriples :: TypeRepMap f -> [(Fingerprint, Any, Any)]
489toTriples tm = zip3 (toFingerprints tm) (GHC.toList $ trAnys tm) (GHC.toList $ trKeys tm)
374toTriples tm = zip3 (toFingerprints tm) (GHC.toList $ trAnys tm) (GHC.toList $ trKeys tm)
491−-- | Efficiently get sorted triples from a map in O(n) time
376+deleteByFst :: Eq a => a -> [(a, b, c)] -> [(a, b, c)]
492−--
377+deleteByFst x = filter ((/= x) . fst3)
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)
506nubByFst :: (Eq a) => [(a, b, c)] -> [(a, b, c)]
379nubByFst :: (Eq a) => [(a, b, c)] -> [(a, b, c)]
507nubByFst = nubBy ((==) `on` fst3)
380nubByFst = nubBy ((==) `on` fst3)
@@ -521,6 +394,9 @@ data WrapTypeable f where
521instance Show (WrapTypeable f) where
394instance Show (WrapTypeable f) where
522 show (WrapTypeable (_ :: f a)) = show $ calcFp @a
395 show (WrapTypeable (_ :: f a)) = show $ calcFp @a
397+wrapTypeable :: TypeRep a -> f a -> WrapTypeable f
398+wrapTypeable tr = withTypeable tr WrapTypeable
526prop> fromList . toList == 'id'
402prop> fromList . toList == 'id'
@@ -548,19 +424,19 @@ instance IsList (TypeRepMap f) where
548 k (WrapTypeable (_ :: f a)) = unsafeCoerce $ typeRep @a
424 k (WrapTypeable (_ :: f a)) = unsafeCoerce $ typeRep @a
550 toList :: TypeRepMap f -> [WrapTypeable f]
426 toList :: TypeRepMap f -> [WrapTypeable f]
551− toList = toListWith WrapTypeable
427+ toList = map toWrapTypeable . toTriples
429+ toWrapTypeable :: (Fingerprint, Any, Any) -> WrapTypeable f
430+ toWrapTypeable (_, an, k) = wrapTypeable (unsafeCoerce k) (fromAny an)
553calcFp :: forall a . Typeable a => Fingerprint
432calcFp :: forall a . Typeable a => Fingerprint
554calcFp = typeRepFingerprint $ typeRep @a
433calcFp = typeRepFingerprint $ typeRep @a
556fromTriples :: [(Fingerprint, Any, Any)] -> TypeRepMap f
435fromTriples :: [(Fingerprint, Any, Any)] -> TypeRepMap f
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)
436+fromTriples kvs = TypeRepMap (GHC.fromList fpAs) (GHC.fromList fpBs) (GHC.fromList ans) (GHC.fromList ks)
562 (fpAs, fpBs) = unzip $ map (\(Fingerprint a b) -> (a, b)) fps
438 (fpAs, fpBs) = unzip $ map (\(Fingerprint a b) -> (a, b)) fps
563− (fps, ans, ks) = unzip3 $ fromSortedList kvs
439+ (fps, ans, ks) = unzip3 $ fromSortedList $ sortWith fst3 $ nubByFst kvs
565----------------------------------------------------------------------------
441----------------------------------------------------------------------------
566-- Tree-like conversion
442-- Tree-like conversion
@@ -587,35 +463,11 @@ fromSortedList l = runST $ do
587 writeArray result i (indexArray origin newFirst)
463 writeArray result i (indexArray origin newFirst)
588 loop (2 * i + 2) (newFirst + 1)
464 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----------------------------------------------------------------------------
466----------------------------------------------------------------------------
615-- Helper functions.
467-- Helper functions.
616----------------------------------------------------------------------------
468----------------------------------------------------------------------------
618−-- | Check that invariant of the structure holds.
470+-- | Check that invariant of the structure is hold.
619-- The structure maintains the following invariant.
471-- The structure maintains the following invariant.
620-- For each element @A@ at index @i@:
472-- For each element @A@ at index @i@: