@@ -45,7 +45,7 @@ import Data.Kind (Type)
45import Data.Type.Equality ((:~:) (..), TestEquality (..))
45import Data.Type.Equality ((:~:) (..), TestEquality (..))
46import Data.List (intercalate, nubBy)
46import Data.List (intercalate, nubBy)
47import Data.Maybe (fromMaybe)
47import Data.Maybe (fromMaybe)
48−import Data.Primitive.Array (Array, MutableArray, indexArray, mapArray', readArray, sizeofArray,
48+import Data.Primitive.Array (Array, MutableArray, indexArray, mapArray', sizeofArray,
49 thawArray, unsafeFreezeArray, writeArray)
49 thawArray, unsafeFreezeArray, writeArray)
50import Data.Primitive.PrimArray (writePrimArray, newPrimArray, unsafeFreezePrimArray, primArrayToList, MutablePrimArray, primArrayFromListN, PrimArray, indexPrimArray, sizeofPrimArray)
50import Data.Primitive.PrimArray (writePrimArray, newPrimArray, unsafeFreezePrimArray, primArrayToList, MutablePrimArray, primArrayFromListN, PrimArray, indexPrimArray, sizeofPrimArray)
51import Data.Semigroup (Semigroup (..), All(..))
51import Data.Semigroup (Semigroup (..), All(..))
@@ -187,13 +187,26 @@ prop> member @a (insert (x :: f a) tm) == True
189insert :: forall a f . Typeable a => f a -> TypeRepMap f -> TypeRepMap f
189insert :: 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)}
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
191{-# INLINE insert #-}
202{-# INLINE insert #-}
193-- Extract the kind of a type. We use it to work around lack of syntax for
204-- Extract the kind of a type. We use it to work around lack of syntax for
194-- inferred type variables (which are not subject to type applications).
205-- inferred type variables (which are not subject to type applications).
195type KindOf (a :: k) = k
206type KindOf (a :: k) = k
208+type ArgKindOf (f :: k -> l) = k
197{- | Delete a value from a 'TypeRepMap'.
210{- | Delete a value from a 'TypeRepMap'.
199TypeRepMap optimizes for fast reads rather than modifications, as a trade-off deletes are
212TypeRepMap optimizes for fast reads rather than modifications, as a trade-off deletes are
@@ -225,26 +238,46 @@ deleteFirst _ [] = []
225deleteFirst p (x : xs) = if p x then xs else x : deleteFirst p xs
238deleteFirst p (x : xs) = if p x then xs else x : deleteFirst p xs
228−Update a value at a specific key with the result of the provided function. When
241+Update a value at a specific key with the result of the provided function.
229−the key is not a member of the map, the original map is returned.
242+When the key is not a member of the map, the original map is returned.
231>>> trmap = fromList @(TypeRepMap Identity) [WrapTypeable $ Identity "a"]
244>>> trmap = fromList @(TypeRepMap Identity) [WrapTypeable $ Identity "a"]
232>>> lookup @String $ adjust (fmap (++ "ww")) trmap
245>>> lookup @String $ adjust (fmap (++ "ww")) trmap
233Just (Identity "aww")
246Just (Identity "aww")
235adjust :: forall a f . Typeable a => (f a -> f a) -> TypeRepMap f -> TypeRepMap f
248adjust :: 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
236−adjust fun tr = case cachedBinarySearch (typeFp @a) (fingerprintAs tr) (fingerprintBs tr) of
264+alter fun tr = case cachedBinarySearch (typeFp @a) (fingerprintAs tr) (fingerprintBs tr) of
237− Nothing -> tr
265+ Nothing ->
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
238− Just i -> tr {trAnys = changeAnyArr i (trAnys tr)}
272+ Just v -> tr{trAnys = replaceAnyAt i (toAny v) (trAnys tr)}
240− changeAnyArr :: Int -> Array Any -> Array Any
274+ replaceAnyAt :: Int -> Any -> Array Any -> Array Any
241− changeAnyArr i trAs = runST $ do
275+ replaceAnyAt i v trAs = runST $ do
242 let n = sizeofArray trAs
276 let n = sizeofArray trAs
243 mutArr <- thawArray trAs 0 n
277 mutArr <- thawArray trAs 0 n
244− a <- toAny . fun . fromAny <$> readArray mutArr i
245− writeArray mutArr i a
278+ writeArray mutArr i v
246 unsafeFreezeArray mutArr
279 unsafeFreezeArray mutArr
247−{-# INLINE adjust #-}
280+{-# INLINE alter #-}
249{- | Map over the elements of a 'TypeRepMap'.
282{- | Map over the elements of a 'TypeRepMap'.
@@ -319,6 +352,48 @@ union :: TypeRepMap f -> TypeRepMap f -> TypeRepMap f
319union = unionWith const
352union = unionWith const
320{-# INLINE union #-}
353{-# 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 #-}
322{- | Check if a value of the given type is present in a 'TypeRepMap'.
397{- | Check if a value of the given type is present in a 'TypeRepMap'.
324>>> member @Char $ one (Identity 'a')
399>>> member @Char $ one (Identity 'a')
@@ -354,9 +429,22 @@ size = sizeofPrimArray . fingerprintAs
355-- | Return the list of 'SomeTypeRep' from the keys.
430-- | Return the list of 'SomeTypeRep' from the keys.
356keys :: TypeRepMap f -> [SomeTypeRep]
431keys :: TypeRepMap f -> [SomeTypeRep]
357−keys TypeRepMap{..} = SomeTypeRep . anyToTypeRep <$> toList trKeys
432+keys = keysWith SomeTypeRep
358{-# INLINE keys #-}
433{-# 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)
360-- | Binary searched based on this article
448-- | Binary searched based on this article
361-- http://bannalia.blogspot.com/2015/06/cache-friendly-binary-search.html
449-- http://bannalia.blogspot.com/2015/06/cache-friendly-binary-search.html
362-- with modification for our two-vector search case.
450-- with modification for our two-vector search case.
@@ -433,9 +521,6 @@ data WrapTypeable f where
433instance Show (WrapTypeable f) where
521instance Show (WrapTypeable f) where
434 show (WrapTypeable (_ :: f a)) = show $ calcFp @a
522 show (WrapTypeable (_ :: f a)) = show $ calcFp @a
436−wrapTypeable :: TypeRep a -> f a -> WrapTypeable f
437−wrapTypeable tr = withTypeable tr WrapTypeable
441prop> fromList . toList == 'id'
526prop> fromList . toList == 'id'
@@ -463,10 +548,7 @@ instance IsList (TypeRepMap f) where
463 k (WrapTypeable (_ :: f a)) = unsafeCoerce $ typeRep @a
548 k (WrapTypeable (_ :: f a)) = unsafeCoerce $ typeRep @a
465 toList :: TypeRepMap f -> [WrapTypeable f]
550 toList :: TypeRepMap f -> [WrapTypeable f]
466− toList = map toWrapTypeable . toTriples
551+ toList = toListWith WrapTypeable
468− toWrapTypeable :: (Fingerprint, Any, Any) -> WrapTypeable f
469− toWrapTypeable (_, an, k) = wrapTypeable (unsafeCoerce k) (fromAny an)
471calcFp :: forall a . Typeable a => Fingerprint
553calcFp :: forall a . Typeable a => Fingerprint
472calcFp = typeRepFingerprint $ typeRep @a
554calcFp = typeRepFingerprint $ typeRep @a