@@ -45,7 +45,7 @@ import Data.Kind (Type)
4545import Data.Type.Equality ((:~:) (..), TestEquality (..))
4646import Data.List (intercalate, nubBy)
4747import Data.Maybe (fromMaybe)
48−import Data.Primitive.Array (Array, MutableArray, indexArray, mapArray', sizeofArray,
48+import Data.Primitive.Array (Array, MutableArray, indexArray, mapArray', readArray, sizeofArray,
4949 thawArray, unsafeFreezeArray, writeArray)
5050import Data.Primitive.PrimArray (writePrimArray, newPrimArray, unsafeFreezePrimArray, primArrayToList, MutablePrimArray, primArrayFromListN, PrimArray, indexPrimArray, sizeofPrimArray)
5151import Data.Semigroup (Semigroup (..), All(..))
@@ -187,26 +187,13 @@ prop> member @a (insert (x :: f a) tm) == True
187187
188188-}
189189insert :: 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)}
195− 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
190+insert = union . one
202191{-# INLINE insert #-}
203192
204193-- Extract the kind of a type. We use it to work around lack of syntax for
205194-- inferred type variables (which are not subject to type applications).
206195type KindOf (a :: k) = k
207196
208−type ArgKindOf (f :: k -> l) = k
209−
210197{- | Delete a value from a 'TypeRepMap'.
211198
212199TypeRepMap optimizes for fast reads rather than modifications, as a trade-off deletes are
@@ -230,54 +217,30 @@ delete m
230217 -- We know we have the element, If the map has exactly one element, we can return the empty map
231218 | size m == 1 = empty
232219 -- Otherwise, filter out the element in linear time.
233− | otherwise = fromSortedTriples . deleteFirst ((== typeFp @a) . fst3) . toSortedTriples $ m
220+ | otherwise = fromSortedTriples . filter ((/= typeFp @a) . fst3) . toSortedTriples $ m
234221{-# INLINE delete #-}
235222
236−deleteFirst :: (a -> Bool) -> [a] -> [a]
237−deleteFirst _ [] = []
238−deleteFirst p (x : xs) = if p x then xs else x : deleteFirst p xs
239−
240223{- |
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.
224+Update a value at a specific key with the result of the provided function. When
225+the key is not a member of the map, the original map is returned.
243226
244227>>> trmap = fromList @(TypeRepMap Identity) [WrapTypeable $ Identity "a"]
245228>>> lookup @String $ adjust (fmap (++ "ww")) trmap
246229Just (Identity "aww")
247230-}
248231adjust :: 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)}
232+adjust fun tr = case cachedBinarySearch (typeFp @a) (fingerprintAs tr) (fingerprintBs tr) of
233+ Nothing -> tr
234+ Just i -> tr {trAnys = changeAnyArr i (trAnys tr)}
273235 where
274− replaceAnyAt :: Int -> Any -> Array Any -> Array Any
275− replaceAnyAt i v trAs = runST $ do
236+ changeAnyArr :: Int -> Array Any -> Array Any
237+ changeAnyArr i trAs = runST $ do
276238 let n = sizeofArray trAs
277239 mutArr <- thawArray trAs 0 n
278− writeArray mutArr i v
240+ a <- toAny . fun . fromAny <$> readArray mutArr i
241+ writeArray mutArr i a
279242 unsafeFreezeArray mutArr
280−{-# INLINE alter #-}
243+{-# INLINE adjust #-}
281244
282245{- | Map over the elements of a 'TypeRepMap'.
283246
@@ -352,48 +315,6 @@ union :: TypeRepMap f -> TypeRepMap f -> TypeRepMap f
352315union = unionWith const
353316{-# INLINE union #-}
354317
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−
397318{- | Check if a value of the given type is present in a 'TypeRepMap'.
398319
399320>>> member @Char $ one (Identity 'a')
@@ -429,22 +350,9 @@ size = sizeofPrimArray . fingerprintAs
429350
430351-- | Return the list of 'SomeTypeRep' from the keys.
431352keys :: TypeRepMap f -> [SomeTypeRep]
432−keys = keysWith SomeTypeRep
353+keys TypeRepMap{..} = SomeTypeRep . anyToTypeRep <$> toList trKeys
433354{-# INLINE keys #-}
434355
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−
448356-- | Binary searched based on this article
449357-- http://bannalia.blogspot.com/2015/06/cache-friendly-binary-search.html
450358-- with modification for our two-vector search case.
@@ -521,6 +429,9 @@ data WrapTypeable f where
521429instance Show (WrapTypeable f) where
522430 show (WrapTypeable (_ :: f a)) = show $ calcFp @a
523431
432+wrapTypeable :: TypeRep a -> f a -> WrapTypeable f
433+wrapTypeable tr = withTypeable tr WrapTypeable
434+
524435{- |
525436
526437prop> fromList . toList == 'id'
@@ -548,7 +459,10 @@ instance IsList (TypeRepMap f) where
548459 k (WrapTypeable (_ :: f a)) = unsafeCoerce $ typeRep @a
549460
550461 toList :: TypeRepMap f -> [WrapTypeable f]
551− toList = toListWith WrapTypeable
462+ toList = map toWrapTypeable . toTriples
463+ where
464+ toWrapTypeable :: (Fingerprint, Any, Any) -> WrapTypeable f
465+ toWrapTypeable (_, an, k) = wrapTypeable (unsafeCoerce k) (fromAny an)
552466
553467calcFp :: forall a . Typeable a => Fingerprint
554468calcFp = typeRepFingerprint $ typeRep @a