@@ -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', readArray, sizeofArray,
48+import Data.Primitive.Array (Array, MutableArray, indexArray, mapArray', sizeofArray,
4949 thawArray, unsafeFreezeArray, writeArray)
5050import Data.Primitive.PrimArray (writePrimArray, newPrimArray, unsafeFreezePrimArray, primArrayToList, MutablePrimArray, primArrayFromListN, PrimArray, indexPrimArray, sizeofPrimArray)
5151import Data.Semigroup (Semigroup (..), All(..))
@@ -187,13 +187,26 @@ 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 = union . one
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
191202{-# INLINE insert #-}
192203
193204-- Extract the kind of a type. We use it to work around lack of syntax for
194205-- inferred type variables (which are not subject to type applications).
195206type KindOf (a :: k) = k
196207
208+type ArgKindOf (f :: k -> l) = k
209+
197210{- | Delete a value from a 'TypeRepMap'.
198211
199212TypeRepMap optimizes for fast reads rather than modifications, as a trade-off deletes are
@@ -225,26 +238,46 @@ deleteFirst _ [] = []
225238deleteFirst p (x : xs) = if p x then xs else x : deleteFirst p xs
226239
227240{- |
228−Update a value at a specific key with the result of the provided function. When
229−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.
230243
231244>>> trmap = fromList @(TypeRepMap Identity) [WrapTypeable $ Identity "a"]
232245>>> lookup @String $ adjust (fmap (++ "ww")) trmap
233246Just (Identity "aww")
234247-}
235248adjust :: forall a f . Typeable a => (f a -> f a) -> TypeRepMap f -> TypeRepMap f
236−adjust fun tr = case cachedBinarySearch (typeFp @a) (fingerprintAs tr) (fingerprintBs tr) of
237− Nothing -> tr
238− 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)}
239273 where
240− changeAnyArr :: Int -> Array Any -> Array Any
241− changeAnyArr i trAs = runST $ do
274+ replaceAnyAt :: Int -> Any -> Array Any -> Array Any
275+ replaceAnyAt i v trAs = runST $ do
242276 let n = sizeofArray trAs
243277 mutArr <- thawArray trAs 0 n
244− a <- toAny . fun . fromAny <$> readArray mutArr i
245− writeArray mutArr i a
278+ writeArray mutArr i v
246279 unsafeFreezeArray mutArr
247−{-# INLINE adjust #-}
280+{-# INLINE alter #-}
248281
249282{- | Map over the elements of a 'TypeRepMap'.
250283
@@ -319,6 +352,48 @@ union :: TypeRepMap f -> TypeRepMap f -> TypeRepMap f
319352union = unionWith const
320353{-# INLINE union #-}
321354
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+
322397{- | Check if a value of the given type is present in a 'TypeRepMap'.
323398
324399>>> member @Char $ one (Identity 'a')
@@ -354,9 +429,22 @@ size = sizeofPrimArray . fingerprintAs
354429
355430-- | Return the list of 'SomeTypeRep' from the keys.
356431keys :: TypeRepMap f -> [SomeTypeRep]
357−keys TypeRepMap{..} = SomeTypeRep . anyToTypeRep <$> toList trKeys
432+keys = keysWith SomeTypeRep
358433{-# INLINE keys #-}
359434
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+
360448-- | Binary searched based on this article
361449-- http://bannalia.blogspot.com/2015/06/cache-friendly-binary-search.html
362450-- with modification for our two-vector search case.
@@ -433,9 +521,6 @@ data WrapTypeable f where
433521instance Show (WrapTypeable f) where
434522 show (WrapTypeable (_ :: f a)) = show $ calcFp @a
435523
436−wrapTypeable :: TypeRep a -> f a -> WrapTypeable f
437−wrapTypeable tr = withTypeable tr WrapTypeable
438−
439524{- |
440525
441526prop> fromList . toList == 'id'
@@ -463,10 +548,7 @@ instance IsList (TypeRepMap f) where
463548 k (WrapTypeable (_ :: f a)) = unsafeCoerce $ typeRep @a
464549
465550 toList :: TypeRepMap f -> [WrapTypeable f]
466− toList = map toWrapTypeable . toTriples
467− where
468− toWrapTypeable :: (Fingerprint, Any, Any) -> WrapTypeable f
469− toWrapTypeable (_, an, k) = wrapTypeable (unsafeCoerce k) (fromAny an)
551+ toList = toListWith WrapTypeable
470552
471553calcFp :: forall a . Typeable a => Fingerprint
472554calcFp = typeRepFingerprint $ typeRep @a