diff --git a/containers-tests/containers-tests.cabal b/containers-tests/containers-tests.cabal index 0cf2f5eac..fbcf186ce 100644 --- a/containers-tests/containers-tests.cabal +++ b/containers-tests/containers-tests.cabal @@ -117,6 +117,9 @@ library Data.Map.Lazy Data.Map.Merge.Lazy Data.Map.Merge.Strict + Data.Map.Merge.Set.Internal + Data.Map.Merge.Set.Lazy + Data.Map.Merge.Set.Strict Data.Map.Strict Data.Map.Strict.Internal Data.Sequence diff --git a/containers-tests/test-utils/Utils/MergeFunc.hs b/containers-tests/test-utils/Utils/MergeFunc.hs index 5e6ca9811..8d9f95f94 100644 --- a/containers-tests/test-utils/Utils/MergeFunc.hs +++ b/containers-tests/test-utils/Utils/MergeFunc.hs @@ -1,9 +1,11 @@ module Utils.MergeFunc ( WhenMatchedFunc(..) , WhenMissingFunc(..) + , MapSet_WhenMatchedFunc(..) + , MapSet_WhenMissingSetFunc(..) ) where -import Test.QuickCheck +import Test.QuickCheck (CoArbitrary, Function, Arbitrary(..), oneof) import Utils.Strictness (Func, Func2, Func3) -- k: key, x: left map value, y: right map value, z: result map value, @@ -69,3 +71,43 @@ instance MapMissingFunc fun -> MapMissingFunc <$> shrink fun FmapMapMissingFunc fun2 fun1 -> uncurry FmapMapMissingFunc <$> shrink (fun2, fun1) + +-- For Set-Map to Map merge. +-- k: key, a: map value, b: result map value +data MapSet_WhenMatchedFunc k a b + = MapSet_MapMaybeMatchedFunc (Func2 k a (Maybe b)) + | MapSet_MapMatchedFunc (Func2 k a b) + deriving Show + +instance + ( CoArbitrary k, Function k + , CoArbitrary a, Function a + , Arbitrary b + ) => Arbitrary (MapSet_WhenMatchedFunc k a b) where + arbitrary = oneof + [ MapSet_MapMaybeMatchedFunc <$> arbitrary + , MapSet_MapMatchedFunc <$> arbitrary + ] + shrink wmf = case wmf of + MapSet_MapMaybeMatchedFunc fun -> MapSet_MapMaybeMatchedFunc <$> shrink fun + MapSet_MapMatchedFunc fun -> MapSet_MapMatchedFunc <$> shrink fun + +-- For Set-Map to Map merge. +-- k: key, a: result map value +data MapSet_WhenMissingSetFunc k a + = MapSet_GenerateMissingSetFunc (Func k a) + | MapSet_GenerateMaybeMissingSetFunc (Func k (Maybe a)) + deriving Show + +instance + (CoArbitrary k, Function k, Arbitrary a) + => Arbitrary (MapSet_WhenMissingSetFunc k a) where + arbitrary = oneof + [ MapSet_GenerateMissingSetFunc <$> arbitrary + , MapSet_GenerateMaybeMissingSetFunc <$> arbitrary + ] + shrink wmf = case wmf of + MapSet_GenerateMissingSetFunc fun -> + MapSet_GenerateMissingSetFunc <$> shrink fun + MapSet_GenerateMaybeMissingSetFunc fun -> + MapSet_GenerateMaybeMissingSetFunc <$> shrink fun diff --git a/containers-tests/tests/map-properties.hs b/containers-tests/tests/map-properties.hs index 8bd02ca81..62b2a4787 100644 --- a/containers-tests/tests/map-properties.hs +++ b/containers-tests/tests/map-properties.hs @@ -6,9 +6,11 @@ #ifdef STRICT import Data.Map.Strict as Data.Map import Data.Map.Merge.Strict +import qualified Data.Map.Merge.Set.Strict as MergeSet #else import Data.Map.Lazy as Data.Map import Data.Map.Merge.Lazy +import qualified Data.Map.Merge.Set.Lazy as MergeSet #endif import Data.Map.Internal (Map, link2, link) import Data.Map.Internal.Debug (showTree, showTreeWith, balanced) @@ -327,6 +329,8 @@ main = defaultMain $ testGroup "map-properties" , testProperty "mapAccum" prop_mapAccum , testProperty "mapAccumWithKey" prop_mapAccumWithKey , testProperty "mapAccumRWithKey" prop_mapAccumRWithKey + , testProperty "Merge.Set.merge" prop_mergeMapSet + , testProperty "Merge.Set.mergeA" prop_mergeAMapSet , testLaws $ Laws.eqLaws (Proxy :: Proxy (Map Int A)) , testLaws $ Laws.ordLaws (Proxy :: Proxy (Map Int OrdA)) , testLaws $ Laws.showLaws (Proxy :: Proxy (Map Int A)) @@ -1390,6 +1394,129 @@ instance (Arbitrary a, CoArbitrary a, Function a, CoArbitrary k, Function k) , ZipWithMaybeMatched <$> arbitrary ] +prop_mergeMapSet + :: WhenMissingSpec Int A + -> WhenMissingMapSetSpec Int A + -> WhenMatchedMapSetSpec Int A + -> Map Int A + -> Set Int + -> Property +prop_mergeMapSet miss1 miss2 match m1 s2 = + valid m .&&. + m === + (mapMaybeWithKey (\k x -> runIdentity (runWhenMissing miss1' k x)) m1Only `union` + mapMaybe id (fromSet (runIdentity . MergeSet.runWhenMissingSet miss2') s2Only) `union` + mapMaybeWithKey (\k x -> runIdentity (MergeSet.runWhenMatched match' k x)) m12Both) + where + miss1' = whenMissing miss1 + miss2' = whenMissingSet miss2 + match' = whenMatched match + + m = MergeSet.merge miss1' miss2' match' m1 s2 + + m1Only = filterKeys (`Set.notMember` s2) m1 + s2Only = Set.filter (`notMember` m1) s2 + m12Both = filterKeys (`Set.member` s2) m1 + + whenMissing spec = case spec of + DropMissing -> MergeSet.dropMissing + PreserveMissing -> MergeSet.preserveMissing + FilterMissing f -> MergeSet.filterMissing (applyFun2 f) + MapMissing f -> MergeSet.mapMissing (applyFun2 f) + MapMaybeMissing f -> MergeSet.mapMaybeMissing (applyFun2 f) + whenMissingSet spec = case spec of + DropMissingMapSet -> MergeSet.dropMissingSet + GenerateMissingMapSet f -> MergeSet.generateMissingSet (applyFun f) + GenerateMaybeMissingMapSet f -> MergeSet.generateMaybeMissingSet (applyFun f) + whenMatched spec = case spec of + FilterMatchedMapSet f -> MergeSet.filterMatched (applyFun2 f) + MapMatchedMapSet f -> MergeSet.mapMatched (applyFun2 f) + MapMaybeMatchedMapSet f -> MergeSet.mapMaybeMatched (applyFun2 f) + +-- This uses the instance +-- Monoid a => Applicative ((,) a) +-- to test that effects are sequenced in ascending key order. +prop_mergeAMapSet + :: WhenMissingSpec Int A + -> WhenMissingMapSetSpec Int A + -> WhenMatchedMapSetSpec Int A + -> Map Int A + -> Set Int + -> Property +prop_mergeAMapSet miss1 miss2 match m1 s2 = + valid m .&&. + m === + (mapMaybeWithKey (\k x -> snd (runWhenMissing miss1' k x)) m1Only `union` + mapMaybe id (fromSet (snd . MergeSet.runWhenMissingSet miss2') s2Only) `union` + mapMaybeWithKey (\k x -> snd (MergeSet.runWhenMatched match' k x)) m12Both) .&&. + ks === + sort + (concat + (fmap (\(k,x) -> fst (runWhenMissing miss1' k x)) (toList m1Only) ++ + fmap (fst . MergeSet.runWhenMissingSet miss2') (Set.toList s2Only) ++ + fmap (\(k,x) -> fst (MergeSet.runWhenMatched match' k x)) (toList m12Both))) + where + miss1' = whenMissing miss1 + miss2' = whenMissingSet miss2 + match' = whenMatched match + + (ks, m) = MergeSet.mergeA miss1' miss2' match' m1 s2 + + m1Only = filterKeys (`Set.notMember` s2) m1 + s2Only = Set.filter (`notMember` m1) s2 + m12Both = filterKeys (`Set.member` s2) m1 + + whenMissing spec = case spec of + DropMissing -> MergeSet.dropMissing + PreserveMissing -> MergeSet.preserveMissing + FilterMissing f -> MergeSet.filterAMissing (\k x -> ([k], applyFun2 f k x)) + MapMissing f -> MergeSet.traverseMissing (\k x -> ([k], applyFun2 f k x)) + MapMaybeMissing f -> MergeSet.traverseMaybeMissing (\k x -> ([k], applyFun2 f k x)) + whenMissingSet spec = case spec of + DropMissingMapSet -> MergeSet.dropMissingSet + GenerateMissingMapSet f -> MergeSet.generateAMissingSet (\k -> ([k], applyFun f k)) + GenerateMaybeMissingMapSet f -> MergeSet.generateMaybeAMissingSet (\k -> ([k], applyFun f k)) + whenMatched spec = case spec of + FilterMatchedMapSet f -> MergeSet.filterAMatched (\k x -> ([k], applyFun2 f k x)) + MapMatchedMapSet f -> MergeSet.traverseMatched (\k x -> ([k], applyFun2 f k x)) + MapMaybeMatchedMapSet f -> MergeSet.traverseMaybeMatched (\k x -> ([k], applyFun2 f k x)) + +data WhenMissingMapSetSpec k a + = DropMissingMapSet + | GenerateMissingMapSet (Fun k a) + | GenerateMaybeMissingMapSet (Fun k (Maybe a)) + deriving Show + +instance (Arbitrary a, CoArbitrary a, Function a, CoArbitrary k, Function k) + => Arbitrary (WhenMissingMapSetSpec k a) where + arbitrary = oneof + [ pure DropMissingMapSet + , GenerateMissingMapSet <$> arbitrary + , GenerateMaybeMissingMapSet <$> arbitrary + ] + shrink spec = case spec of + DropMissingMapSet -> [] + GenerateMissingMapSet f -> GenerateMissingMapSet <$> shrink f + GenerateMaybeMissingMapSet f -> GenerateMaybeMissingMapSet <$> shrink f + +data WhenMatchedMapSetSpec k a + = FilterMatchedMapSet (Fun (k, a) Bool) + | MapMatchedMapSet (Fun (k, a) a) + | MapMaybeMatchedMapSet (Fun (k, a) (Maybe a)) + deriving Show + +instance (Arbitrary a, CoArbitrary a, Function a, CoArbitrary k, Function k) + => Arbitrary (WhenMatchedMapSetSpec k a) where + arbitrary = oneof + [ FilterMatchedMapSet <$> arbitrary + , MapMatchedMapSet <$> arbitrary + , MapMaybeMatchedMapSet <$> arbitrary + ] + shrink spec = case spec of + FilterMatchedMapSet f -> FilterMatchedMapSet <$> shrink f + MapMatchedMapSet f -> MapMatchedMapSet <$> shrink f + MapMaybeMatchedMapSet f -> MapMaybeMatchedMapSet <$> shrink f + ---------------------------------------------------------------- prop_list :: [Int] -> Bool diff --git a/containers-tests/tests/map-strictness.hs b/containers-tests/tests/map-strictness.hs index 2ecd7ab82..2daaafc48 100644 --- a/containers-tests/tests/map-strictness.hs +++ b/containers-tests/tests/map-strictness.hs @@ -36,10 +36,17 @@ import Data.Map.Merge.Lazy (WhenMatched, WhenMissing) import qualified Data.Map.Merge.Lazy as LMerge import Data.Set (Set) import qualified Data.Set as Set +import qualified Data.Map.Merge.Set.Lazy as MergeSetLazy +import qualified Data.Map.Merge.Set.Strict as MergeSetStrict import Utils.ArbitrarySetMap (setFromList, mapFromKeysList) import Utils.QuickCheck (NubSorted(..), NubSortedOnFst(..), SortedOnFst(..)) -import Utils.MergeFunc (WhenMatchedFunc(..), WhenMissingFunc(..)) +import Utils.MergeFunc + ( WhenMatchedFunc(..) + , WhenMissingFunc(..) + , MapSet_WhenMatchedFunc(..) + , MapSet_WhenMissingSetFunc(..) + ) import Utils.Strictness (Bot(..), Func, Func2, Func3, applyFunc, applyFunc2, applyFunc3) @@ -1069,6 +1076,86 @@ prop_lazyMergeA misfun1 misfun2 matfun m1 m2 = mis2 = toLazyWhenMissingA (coerce misfun2 :: WhenMissingFunc OrdA B C C C) mat = toLazyWhenMatchedA (coerce matfun :: WhenMatchedFunc OrdA A B C C C) +prop_strictMapSetMerge + :: WhenMissingFunc OrdA A (Bot B) B (Bot B) + -> MapSet_WhenMissingSetFunc OrdA (Bot B) + -> MapSet_WhenMatchedFunc OrdA A (Bot B) + -> Map OrdA A + -> Set OrdA + -> Property +prop_strictMapSetMerge misfun1 misfun2 matfun m1 s2 = + isBottom (MergeSetStrict.merge mis1 mis2 mat m1 s2) === + any isBottom + (catMaybes $ concat + [ [matf k x | (k,x) <- M.toList m1, k `Set.member` s2] + , [misf1 k x | (k,x) <- M.toList m1, k `Set.notMember` s2] + , [misf2 k | k <- Set.toList s2, k `M.notMember` m1] + ]) + where + misfun1' = coerce misfun1 :: WhenMissingFunc OrdA A B B B + misfun2' = coerce misfun2 :: MapSet_WhenMissingSetFunc OrdA B + matfun' = coerce matfun :: MapSet_WhenMatchedFunc OrdA A B + mis1 = toStrictWhenMissing misfun1' + mis2 = toStrictMapSetWhenMissingSet misfun2' + mat = toStrictMapSetWhenMatched matfun' + misf1 = whenMissingApplyStrict misfun1' + misf2 = mapSetWhenMissingSetApply misfun2' + matf = mapSetWhenMatchedApply matfun' + +prop_lazyMapSetMerge + :: WhenMissingFunc OrdA A (Bot B) B (Bot B) + -> MapSet_WhenMissingSetFunc OrdA (Bot B) + -> MapSet_WhenMatchedFunc OrdA A (Bot B) + -> Map OrdA A + -> Set OrdA + -> Property +prop_lazyMapSetMerge misfun1 misfun2 matfun m1 m2 = + isNotBottomProp (MergeSetLazy.merge mis1 mis2 mat m1 m2) + where + mis1 = toLazyWhenMissing (coerce misfun1 :: WhenMissingFunc OrdA A B B B) + mis2 = toLazyMapSetWhenMissingSet (coerce misfun2 :: MapSet_WhenMissingSetFunc OrdA B) + mat = toLazyMapSetWhenMatched (coerce matfun :: MapSet_WhenMatchedFunc OrdA A B) + +prop_strictMapSetMergeA + :: WhenMissingFunc OrdA A (Bot B) B (Bot B) + -> MapSet_WhenMissingSetFunc OrdA (Bot B) + -> MapSet_WhenMatchedFunc OrdA A (Bot B) + -> Map OrdA A + -> Set OrdA + -> Property +prop_strictMapSetMergeA misfun1 misfun2 matfun m1 s2 = + isBottom (runIdentity (MergeSetStrict.mergeA mis1 mis2 mat m1 s2)) === + any isBottom + (catMaybes $ concat + [ [matf k x | (k,x) <- M.toList m1, k `Set.member` s2] + , [misf1 k x | (k,x) <- M.toList m1, k `Set.notMember` s2] + , [misf2 k | k <- Set.toList s2, k `M.notMember` m1] + ]) + where + misfun1' = coerce misfun1 :: WhenMissingFunc OrdA A B B B + misfun2' = coerce misfun2 :: MapSet_WhenMissingSetFunc OrdA B + matfun' = coerce matfun :: MapSet_WhenMatchedFunc OrdA A B + mis1 = toStrictWhenMissingA misfun1' + mis2 = toStrictMapSetWhenMissingSetA misfun2' + mat = toStrictMapSetWhenMatchedA matfun' + misf1 = whenMissingApplyStrict misfun1' + misf2 = mapSetWhenMissingSetApply misfun2' + matf = mapSetWhenMatchedApply matfun' + +prop_lazyMapSetMergeA + :: WhenMissingFunc OrdA A (Bot B) B (Bot B) + -> MapSet_WhenMissingSetFunc OrdA (Bot B) + -> MapSet_WhenMatchedFunc OrdA A (Bot B) + -> Map OrdA A + -> Set OrdA + -> Property +prop_lazyMapSetMergeA misfun1 misfun2 matfun m1 m2 = + isNotBottomProp (runIdentity (MergeSetLazy.mergeA mis1 mis2 mat m1 m2)) + where + mis1 = toLazyWhenMissingA (coerce misfun1 :: WhenMissingFunc OrdA A B B B) + mis2 = toLazyMapSetWhenMissingSetA (coerce misfun2 :: MapSet_WhenMissingSetFunc OrdA B) + mat = toLazyMapSetWhenMatchedA (coerce matfun :: MapSet_WhenMatchedFunc OrdA A B) + ------------------------------------------------------------------------ -- ** Strict module @@ -1315,6 +1402,8 @@ tests = , testPropStrictLazy "updateMaxWithKey" prop_strictUpdateMaxWithKey prop_lazyUpdateMaxWithKey , testPropStrictLazy "merge" prop_strictMerge prop_lazyMerge , testPropStrictLazy "mergeA" prop_strictMergeA prop_lazyMergeA + , testPropStrictLazy "Merge.Set.merge" prop_strictMapSetMerge prop_lazyMapSetMerge + , testPropStrictLazy "Merge.Set.mergeA" prop_strictMapSetMergeA prop_lazyMapSetMergeA ] ] @@ -1471,3 +1560,75 @@ whenMissingApplyStrict wmf = case wmf of \k x -> Just $ applyFunc fun2 $! -- Strict in the intermediate result applyFunc2 fun1 k x + +toStrictMapSetWhenMatched + :: MapSet_WhenMatchedFunc k a b -> MergeSetStrict.WhenMatched Identity k a b +toStrictMapSetWhenMatched wmf = case wmf of + MapSet_MapMaybeMatchedFunc fun -> + MergeSetStrict.mapMaybeMatched (applyFunc2 fun) + MapSet_MapMatchedFunc fun -> MergeSetStrict.mapMatched (applyFunc2 fun) + +toStrictMapSetWhenMatchedA + :: MapSet_WhenMatchedFunc k a b -> MergeSetStrict.WhenMatched Identity k a b +toStrictMapSetWhenMatchedA wmf = case wmf of + MapSet_MapMaybeMatchedFunc fun -> + MergeSetStrict.traverseMaybeMatched (coerce (applyFunc2 fun)) + MapSet_MapMatchedFunc fun -> + MergeSetStrict.traverseMatched (coerce (applyFunc2 fun)) + +toLazyMapSetWhenMatched + :: MapSet_WhenMatchedFunc k a b -> MergeSetLazy.WhenMatched Identity k a b +toLazyMapSetWhenMatched wmf = case wmf of + MapSet_MapMaybeMatchedFunc fun -> + MergeSetLazy.mapMaybeMatched (applyFunc2 fun) + MapSet_MapMatchedFunc fun -> MergeSetLazy.mapMatched (applyFunc2 fun) + +toLazyMapSetWhenMatchedA + :: MapSet_WhenMatchedFunc k a b -> MergeSetLazy.WhenMatched Identity k a b +toLazyMapSetWhenMatchedA wmf = case wmf of + MapSet_MapMaybeMatchedFunc fun -> + MergeSetLazy.traverseMaybeMatched (coerce (applyFunc2 fun)) + MapSet_MapMatchedFunc fun -> + MergeSetLazy.traverseMatched (coerce (applyFunc2 fun)) + +mapSetWhenMatchedApply :: MapSet_WhenMatchedFunc k a b -> k -> a -> Maybe b +mapSetWhenMatchedApply wmf = case wmf of + MapSet_MapMaybeMatchedFunc fun -> applyFunc2 fun + MapSet_MapMatchedFunc fun -> \k x -> Just (applyFunc2 fun k x) + +toStrictMapSetWhenMissingSet + :: MapSet_WhenMissingSetFunc k a -> MergeSetStrict.WhenMissingSet Identity k a +toStrictMapSetWhenMissingSet wmf = case wmf of + MapSet_GenerateMissingSetFunc fun -> + MergeSetStrict.generateMissingSet (applyFunc fun) + MapSet_GenerateMaybeMissingSetFunc fun -> + MergeSetStrict.generateMaybeMissingSet (applyFunc fun) + +toStrictMapSetWhenMissingSetA + :: MapSet_WhenMissingSetFunc k a -> MergeSetStrict.WhenMissingSet Identity k a +toStrictMapSetWhenMissingSetA wmf = case wmf of + MapSet_GenerateMissingSetFunc fun -> + MergeSetStrict.generateAMissingSet (coerce (applyFunc fun)) + MapSet_GenerateMaybeMissingSetFunc fun -> + MergeSetStrict.generateMaybeAMissingSet (coerce (applyFunc fun)) + +toLazyMapSetWhenMissingSet + :: MapSet_WhenMissingSetFunc k a -> MergeSetLazy.WhenMissingSet Identity k a +toLazyMapSetWhenMissingSet wmf = case wmf of + MapSet_GenerateMissingSetFunc fun -> + MergeSetLazy.generateMissingSet (applyFunc fun) + MapSet_GenerateMaybeMissingSetFunc fun -> + MergeSetLazy.generateMaybeMissingSet (applyFunc fun) + +toLazyMapSetWhenMissingSetA + :: MapSet_WhenMissingSetFunc k a -> MergeSetLazy.WhenMissingSet Identity k a +toLazyMapSetWhenMissingSetA wmf = case wmf of + MapSet_GenerateMissingSetFunc fun -> + MergeSetLazy.generateAMissingSet (coerce (applyFunc fun)) + MapSet_GenerateMaybeMissingSetFunc fun -> + MergeSetLazy.generateMaybeAMissingSet (coerce (applyFunc fun)) + +mapSetWhenMissingSetApply :: MapSet_WhenMissingSetFunc k a -> k -> Maybe a +mapSetWhenMissingSetApply wmf = case wmf of + MapSet_GenerateMissingSetFunc fun -> Just . applyFunc fun + MapSet_GenerateMaybeMissingSetFunc fun -> applyFunc fun diff --git a/containers/containers.cabal b/containers/containers.cabal index ec494b983..2e01e6138 100644 --- a/containers/containers.cabal +++ b/containers/containers.cabal @@ -73,6 +73,9 @@ Library Data.Map.Strict.Internal Data.Map.Strict Data.Map.Merge.Strict + Data.Map.Merge.Set.Internal + Data.Map.Merge.Set.Lazy + Data.Map.Merge.Set.Strict Data.Map.Internal Data.Map.Internal.Debug Data.Set.Internal diff --git a/containers/src/Data/Map/Merge/Set/Internal.hs b/containers/src/Data/Map/Merge/Set/Internal.hs new file mode 100644 index 000000000..fe7b18014 --- /dev/null +++ b/containers/src/Data/Map/Merge/Set/Internal.hs @@ -0,0 +1,254 @@ +{-# OPTIONS_HADDOCK not-home #-} + +-- | +-- = WARNING +-- +-- This module is considered __internal__. +-- +-- The Package Versioning Policy __does not apply__. +-- +-- The contents of this module may change __in any way whatsoever__ +-- and __without any warning__ between minor versions of this package. +-- +-- Authors importing this module are expected to track development +-- closely. +-- +-- = Description +-- +-- This module defines common constructs used by both "Data.Map.Merge.Set.Lazy" +-- and "Data.Map.Merge.Set.Strict". +-- +-- @since FIXME +-- +module Data.Map.Merge.Set.Internal + ( WhenMatched(..) + , SimpleWhenMatched + , filterMatched + , filterAMatched + + , WhenMissingSet(..) + , SimpleWhenMissingSet + , dropMissingSet + + , merge + , mergeA + + , runWhenMatched + , runWhenMissingSet + ) where + +import Control.Applicative (liftA3) +import Data.Functor.Identity (Identity(..)) +import Data.Set (Set) +import qualified Data.Set.Internal as S +import Data.Map (Map) +import qualified Data.Map.Internal as M + +-- | A tactic for dealing with keys present in both the set and the map in +-- 'merge' or 'mergeA'. +-- +-- A tactic of type @WhenMatched f k a b@ is an abstract representation of +-- a function of type @k -> a -> f (Maybe b)@. +-- +-- @since FIXME +newtype WhenMatched f k a b = WhenMatched + { matchedKey :: k -> a -> f (Maybe b) + } + +-- | Run @WhenMatched@. +-- +-- @since FIXME +runWhenMatched :: WhenMatched f k a b -> k -> a -> f (Maybe b) +runWhenMatched = matchedKey + +-- | A tactic for dealing with keys present in both the set and the map in +-- 'merge'. +-- +-- A tactic of type @SimpleWhenMatched k a b@ is an abstract representation of +-- a function of type @k -> a -> Maybe b@. +-- +-- @since FIXME +type SimpleWhenMatched = WhenMatched Identity + +-- | When a key is found in both the map and the set, apply a function to the +-- key and the value in the map and keep the value in the merged map if the +-- result is @True@. +-- +-- @since FIXME +filterMatched :: Applicative f => (k -> a -> Bool) -> WhenMatched f k a a +filterMatched f = + WhenMatched (\k x -> if f k x then pure (Just x) else pure Nothing) +{-# INLINE filterMatched #-} + +-- | When a key is found in both the map and the set, apply a function to the +-- key and the value in the map and keep the value in the merged map if the +-- result of the action is @True@. +-- +-- @since FIXME +filterAMatched :: Functor f => (k -> a -> f Bool) -> WhenMatched f k a a +filterAMatched f = + WhenMatched (\k x -> (\b -> if b then Just x else Nothing) <$> f k x) +{-# INLINE filterAMatched #-} + +-- | A tactic for dealing with keys present in the set but not in the map in +-- 'merge' or 'mergeA'. +-- +-- A tactic of type @WhenMissingSet f k a@ is an abstract representation of +-- a function of type @k -> f (Maybe a)@. +-- +-- @since FIXME +data WhenMissingSet f k a = WhenMissingSet + { missingSubtree :: Set k -> f (Map k a) + , missingKey :: k -> f (Maybe a) + } + +-- | Run @WhenMissingSet@. +-- +-- @since FIXME +runWhenMissingSet :: WhenMissingSet f k a -> k -> f (Maybe a) +runWhenMissingSet = missingKey + +-- | A tactic for dealing with keys present in the set but not in the map in +-- 'merge'. +-- +-- A tactic of type @SimpleWhenMissingSet k a@ is an abstract representation of +-- a function of type @k -> Maybe a@. +-- +-- @since FIXME +type SimpleWhenMissingSet = WhenMissingSet Identity + +-- | Drop keys that are present in the set but missing from the map. +-- +-- @since FIXME +dropMissingSet :: Applicative f => WhenMissingSet f k a +dropMissingSet = WhenMissingSet + { missingSubtree = \_ -> pure M.empty + , missingKey = \_ -> pure Nothing + } +{-# INLINE dropMissingSet #-} + +-- | Merge a map and a set into a map. +-- +-- 'merge' takes a 'M.SimpleWhenMissing' tactic, a 'SimpleWhenMissingSet' +-- tactic, a 'SimpleWhenMatched' tactic, a map and a set. It uses the tactics to +-- merge the map and the set into a map. +-- +-- Its behavior is best understood via the tactics @mapMaybeMissing@, +-- @generateMaybeMissingSet@, and @mapMaybeMatched@. Consider +-- +-- @ +-- merge (mapMaybeMissing g1) (generateMaybeMissingSet g2) (mapMaybeMatched f) m1 s2 +-- @ +-- +-- @ +-- g1 k x = if k == 2 then Just ("1" ++ x) else Nothing +-- g2 k = if k == 3 then Just "2" else Nothing +-- f k x = if k == 6 then Just ("3" ++ x) else Nothing +-- m1 = fromList [(2,"a"), (4,"b"), (6,"c"), (8,"d"), (10,"e"), (12,"f")] +-- s2 = fromList [3, 6, 9, 12] +-- @ +-- +-- 'merge' will pass the keys and values to @g1@, @g2@, or @f@ as appropriate, +-- producing a @Maybe@ for each element. +-- +-- @ +-- m1: [ (2, "a"), (4, "b"), (6, "c"), (8, "d"), (10, "e"), (12, "f")] +-- s2: [ 3, 6, 9, 12] +-- result: [ g1 2 "a", g2 3, g1 4 "b", f 6 "c", g1 8 "d", g2 9, g1 10 "e", f 12 "f"] +-- = [Just "1a", Just "2", Nothing, Just "3c", Nothing, Nothing, Nothing, Nothing] +-- @ +-- +-- The result map contains the @Just@ values. +-- +-- >>> merge (mapMaybeMissing g1) (generateMaybeMissingSet g2) (mapMaybeMatched f) m1 s2 +-- fromList [(2,"1a"), (3,"2g"), (6,"3c")] +-- +-- When 'merge' is given three arguments, it is inlined at the call +-- site. To prevent excessive inlining, you should typically use 'merge' +-- to define your custom combining functions. +-- +-- @since FIXME +merge + :: Ord k + => M.SimpleWhenMissing k a b -- ^ What to do with keys in @m1@ but not @s2@ + -> SimpleWhenMissingSet k b -- ^ What to do with keys in @s2@ but not @m1@ + -> SimpleWhenMatched k a b -- ^ What to do with keys in both @m1@ and @s2@ + -> Map k a -- ^ Map @m1@ + -> Set k -- ^ Set @s2@ + -> Map k b +merge miss1 miss2 match = \t1 t2 -> runIdentity (mergeA miss1 miss2 match t1 t2) +{-# INLINE merge #-} + +-- | Merge a map and a set into a map. Applicative version of 'merge'. +-- +-- 'mergeA' takes a 'M.WhenMissing' tactic, a 'WhenMissingSet' tactic, a +-- 'WhenMatched' tactic, a map and a set. It uses the tactics to merge the map +-- and the set into a map. +-- +-- Behaves just like 'merge' while allowing @Applicative@ effects. Effects are +-- performed in increasing order of keys. +-- +-- Consider +-- +-- @ +-- mergeA (traverseMaybeMissing g1) +-- (generateMaybeAMissingSet g2) +-- (traverseMaybeMatched f) +-- m1 +-- s2 +-- @ +-- +-- @ +-- g1 k x = let z = if k == 2 then Just ("1" ++ x) else Nothing +-- in z <$ putStrLn ("g1 " ++ show (k, x)) +-- g2 k = let z = if k == 3 then Just "2" else Nothing +-- in z <$ putStrLn ("g2 " ++ show k) +-- f k x = let z = if k == 6 then Just ("3" ++ x) else Nothing +-- in z <$ putStrLn ("f " ++ show (k, x)) +-- m1 = fromList [(2,"a"), (4,"b"), (6,"c"), (8,"d"), (10,"e"), (12,"f")] +-- m2 = fromList [3, 6, 9, 12] +-- @ +-- +-- As with 'merge', the result map is @[(2,"1a"), (3,"2"), (6,"3c")]@. +-- Additionally, @g1@, @g2@, and @f@ perform @IO@ effects, printing in +-- increasing order of key. +-- +-- >>> mergeA (traverseMaybeMissing g1) (generateMaybeAMissingSet g2) (traverseMaybeMatched f) m1 s2 +-- g1 (2,"a") +-- g2 3 +-- g1 (4,"b") +-- f (6,"c") +-- g1 (8,"d") +-- g2 9 +-- g1 (10,"e") +-- f (12,"f") +-- fromList [(2,"1a"),(3,"2"),(6,"3c")] +-- +-- When 'mergeA' is given three arguments, it is inlined at the call +-- site. To prevent excessive inlining, you should generally only use +-- 'mergeA' to define custom combining functions. +-- +-- @since FIXME +mergeA + :: (Applicative f, Ord k) + => M.WhenMissing f k a b -- ^ What to do with keys in @m1@ but not @s2@ + -> WhenMissingSet f k b -- ^ What to do with keys in @s2@ but not @m1@ + -> WhenMatched f k a b -- ^ What to do with keys in both @m1@ and @s2@ + -> Map k a -- ^ Map @m1@ + -> Set k -- ^ Set @s2@ + -> f (Map k b) +mergeA + M.WhenMissing{M.missingSubtree = g1t, M.missingKey = g1k} + WhenMissingSet{missingSubtree = g2t} + WhenMatched{matchedKey = f} = go + where + go t1 S.Tip = g1t t1 + go M.Tip t2 = g2t t2 + go (M.Bin _ k1 x1 l1 r1) t2 = case S.splitMember k1 t2 of + (l2, found, r2) -> + liftA3 + (\l' mx' r' -> maybe M.link2 (M.link k1) mx' l' r') + (go l1 l2) + (if found then f k1 x1 else g1k k1 x1) + (go r1 r2) +{-# INLINE mergeA #-} diff --git a/containers/src/Data/Map/Merge/Set/Lazy.hs b/containers/src/Data/Map/Merge/Set/Lazy.hs new file mode 100644 index 000000000..53efec88d --- /dev/null +++ b/containers/src/Data/Map/Merge/Set/Lazy.hs @@ -0,0 +1,153 @@ +-- | +-- This module defines an API for writing functions that merge a map and a +-- set into a map. The key functions are 'Data.Map.Merge.Set.Lazy.merge' and +-- 'Data.Map.Merge.Set.Lazy.mergeA'. Each of these can be used with several +-- different \"merge tactics\". +-- +-- The @merge@ and @mergeA@ functions are shared by the lazy and strict +-- modules. Only the choice of merge tactics determines strictness. If you +-- use 'Data.Map.Merge.Set.Strict.mapMissing' from "Data.Map.Merge.Set.Strict" +-- then the results will be forced before they are inserted. If you use +-- 'Data.Map.Merge.Set.Lazy.mapMissing' from this module then they will not. +-- +-- @since FIXME +-- +module Data.Map.Merge.Set.Lazy + ( + -- ** Simple merge tactic types + M.SimpleWhenMissing + , Internal.SimpleWhenMissingSet + , Internal.SimpleWhenMatched + + -- ** General combining function + , Internal.merge + + -- *** @WhenMatched@ tactics + , Internal.filterMatched + , mapMatched + , mapMaybeMatched + + -- *** @WhenMissing@ tactics + , M.dropMissing + , M.preserveMissing + , M.mapMissing + , M.filterMissing + , M.mapMaybeMissing + + -- *** @WhenMissingSet@ tactics + , Internal.dropMissingSet + , generateMissingSet + , generateMaybeMissingSet + + -- ** Applicative merge tactic types + , M.WhenMissing + , Internal.WhenMissingSet + , Internal.WhenMatched + + -- ** General combining function + , Internal.mergeA + + -- *** @WhenMatched@ tactics + , Internal.filterAMatched + , traverseMatched + , traverseMaybeMatched + + -- *** @WhenMissing@ tactics + , M.filterAMissing + , M.traverseMissing + , M.traverseMaybeMissing + + -- *** @WhenMissingSet@ tactics + , generateAMissingSet + , generateMaybeAMissingSet + + -- ** Miscellaneous + , Internal.runWhenMatched + , Internal.runWhenMissingSet + ) where + +import qualified Data.Map.Internal as M +import qualified Data.Map.Merge.Set.Internal as Internal +import Data.Map.Merge.Set.Internal (WhenMatched(..), WhenMissingSet(..)) + +-- | When a key is found in both the map and the set, apply a function to the +-- key and the value in the map and use the result as the value for the merged +-- map. +-- +-- @since FIXME +mapMatched :: Applicative f => (k -> a -> b) -> WhenMatched f k a b +mapMatched f = WhenMatched (\k x -> pure (Just (f k x))) +{-# INLINE mapMatched #-} + +-- | When a key is found in both the map and the set, apply a function to the +-- key and the value in the map and maybe use the result as the value for the +-- merged map. +-- +-- @since FIXME +mapMaybeMatched :: Applicative f => (k -> a -> Maybe b) -> WhenMatched f k a b +mapMaybeMatched f = WhenMatched (\k x -> pure (f k x)) +{-# INLINE mapMaybeMatched #-} + +-- | When a key is found in both the map and the set, apply a function to the +-- key and the value in the map, and use the result of the action as the value +-- for the merged map. +-- +-- @since FIXME +traverseMatched :: Functor f => (k -> a -> f b) -> WhenMatched f k a b +traverseMatched f = WhenMatched (\k x -> Just <$> f k x) +{-# INLINE traverseMatched #-} + +-- | When a key is found in both the map and the set, apply a function to the +-- key and the value in the map, and maybe use the result of the action as the +-- value for the merged map. +-- +-- @since FIXME +traverseMaybeMatched :: (k -> a -> f (Maybe b)) -> WhenMatched f k a b +traverseMaybeMatched = WhenMatched + +-- | For keys that are present in the set but missing from the map, apply a +-- function and use the result as the value for the merged map. +-- +-- @since FIXME +generateMissingSet :: Applicative f => (k -> a) -> WhenMissingSet f k a +generateMissingSet f = WhenMissingSet + { missingSubtree = \s -> pure (M.fromSet f s) + , missingKey = \k -> pure (Just (f k)) + } +{-# INLINE generateMissingSet #-} + +-- | For keys that are present in the set but missing from the map, apply a +-- function, and use the result of the action as the value for the merged map. +-- +-- @since FIXME +generateAMissingSet :: Applicative f => (k -> f a) -> WhenMissingSet f k a +generateAMissingSet f = WhenMissingSet + { missingSubtree = M.fromSetA f + , missingKey = \k -> Just <$> f k + } +{-# INLINE generateAMissingSet #-} + +-- | For keys that are present in the set but missing from the map, apply a +-- function and maybe use the result as the value for the merged map. +-- +-- @since FIXME +generateMaybeMissingSet + :: Applicative f => (k -> Maybe a) -> WhenMissingSet f k a +generateMaybeMissingSet f = WhenMissingSet + { missingSubtree = \s -> pure (M.fromSetMaybe f s) + , missingKey = \k -> pure (f k) + } +{-# INLINE generateMaybeMissingSet #-} + +-- | For keys that are present in the set but missing from the map, apply a +-- function, and maybe use the result of the action as the value for the merged +-- map. +-- +-- @since FIXME +generateMaybeAMissingSet + :: Applicative f => (k -> f (Maybe a)) -> WhenMissingSet f k a +generateMaybeAMissingSet f = WhenMissingSet + { missingSubtree = M.fromSetMaybeA f + , missingKey = f + } +{-# INLINE generateMaybeAMissingSet #-} diff --git a/containers/src/Data/Map/Merge/Set/Strict.hs b/containers/src/Data/Map/Merge/Set/Strict.hs new file mode 100644 index 000000000..eaedd6ef9 --- /dev/null +++ b/containers/src/Data/Map/Merge/Set/Strict.hs @@ -0,0 +1,162 @@ +{-# LANGUAGE BangPatterns #-} + +-- | +-- This module defines an API for writing functions that merge a map and a +-- set into a map. The key functions are 'Data.Map.Merge.Set.Strict.merge' and +-- 'Data.Map.Merge.Set.Strict.mergeA'. Each of these can be used with several +-- different \"merge tactics\". +-- +-- The @merge@ and @mergeA@ functions are shared by the lazy and strict +-- modules. Only the choice of merge tactics determines strictness. +-- If you use 'Data.Map.Merge.Set.Strict.mapMissing' from this module +-- then the results will be forced before they are inserted. If you use +-- 'Data.Map.Merge.Set.Lazy.mapMissing' from "Data.Map.Merge.Set.Lazy" then they +-- will not. +-- +-- @since FIXME +-- +module Data.Map.Merge.Set.Strict + ( + -- ** Simple merge tactic types + MS.SimpleWhenMissing + , Internal.SimpleWhenMissingSet + , Internal.SimpleWhenMatched + + -- ** General combining function + , Internal.merge + + -- *** @WhenMatched@ tactics + , Internal.filterMatched + , mapMatched + , mapMaybeMatched + + -- *** @WhenMissing@ tactics + , MS.dropMissing + , MS.preserveMissing + , MS.mapMissing + , MS.filterMissing + , MS.mapMaybeMissing + + -- *** @WhenMissingSet@ tactics + , Internal.dropMissingSet + , generateMissingSet + , generateMaybeMissingSet + + -- ** Applicative merge tactic types + , MS.WhenMissing + , Internal.WhenMissingSet + , Internal.WhenMatched + + -- ** General combining function + , Internal.mergeA + + -- *** @WhenMatched@ tactics + , Internal.filterAMatched + , traverseMatched + , traverseMaybeMatched + + -- *** @WhenMissing@ tactics + , MS.filterAMissing + , MS.traverseMissing + , MS.traverseMaybeMissing + + -- *** @WhenMissingSet@ tactics + , generateAMissingSet + , generateMaybeAMissingSet + + -- ** Miscellaneous + , Internal.runWhenMatched + , Internal.runWhenMissingSet + ) where + +import qualified Data.Map.Strict.Internal as MS +import qualified Data.Map.Merge.Set.Internal as Internal +import Data.Map.Merge.Set.Internal (WhenMatched(..), WhenMissingSet(..)) + +-- | When the key is found in both the map and the set, apply a function to the +-- key and the value in the map and use the result as the value for the result +-- map. +-- +-- @since FIXME +mapMatched :: Applicative f => (k -> a -> b) -> WhenMatched f k a b +mapMatched f = WhenMatched (\k x -> pure (Just $! f k x)) +{-# INLINE mapMatched #-} + +-- | When a key is found in both the map and the set, apply a function to the +-- key and the value in the map and maybe use the result as the value for the +-- merged map. +-- +-- @since FIXME +mapMaybeMatched :: Applicative f => (k -> a -> Maybe b) -> WhenMatched f k a b +mapMaybeMatched f = WhenMatched (\k x -> pure (forceMaybe (f k x))) +{-# INLINE mapMaybeMatched #-} + +-- | When a key is found in both the map and the set, apply a function to the +-- key and the value in the map, and use the result of the action as the value +-- for the merged map. +-- +-- @since FIXME +traverseMatched :: Functor f => (k -> a -> f b) -> WhenMatched f k a b +traverseMatched f = WhenMatched (\k x -> (Just $!) <$> f k x) +{-# INLINE traverseMatched #-} + +-- | When a key is found in both the map and the set, apply a function to the +-- key and the value in the map, and maybe use the result of the action as the +-- value for the merged map. +-- +-- @since FIXME +traverseMaybeMatched + :: Functor f => (k -> a -> f (Maybe b)) -> WhenMatched f k a b +traverseMaybeMatched f = WhenMatched (\k x -> forceMaybe <$> f k x) +{-# INLINE traverseMaybeMatched #-} + +-- | For keys that are present in the set but missing from the map, apply a +-- function and use the result as the value for the merge map. +-- +-- @since FIXME +generateMissingSet :: Applicative f => (k -> a) -> WhenMissingSet f k a +generateMissingSet f = WhenMissingSet + { missingSubtree = \s -> pure (MS.fromSet f s) + , missingKey = \k -> pure (Just $! f k) + } +{-# INLINE generateMissingSet #-} + +-- | For keys that are present in the set but missing from the map, apply a +-- function, and use the result of the action as the value for the merged map. +-- +-- @since FIXME +generateAMissingSet :: Applicative f => (k -> f a) -> WhenMissingSet f k a +generateAMissingSet f = WhenMissingSet + { missingSubtree = MS.fromSetA f + , missingKey = \k -> (Just $!) <$> f k + } +{-# INLINE generateAMissingSet #-} + +-- | For keys that are present in the set but missing from the map, apply a +-- function and maybe use the result as the value for the merged map. +-- +-- @since FIXME +generateMaybeMissingSet + :: Applicative f => (k -> Maybe a) -> WhenMissingSet f k a +generateMaybeMissingSet f = WhenMissingSet + { missingSubtree = \s -> pure (MS.fromSetMaybe f s) + , missingKey = \k -> pure (f k) + } +{-# INLINE generateMaybeMissingSet #-} + +-- | For keys that are present in the set but missing from the map, apply a +-- function, and maybe use the result of the action as the value for the merged +-- map. +-- +-- @since FIXME +generateMaybeAMissingSet + :: Applicative f => (k -> f (Maybe a)) -> WhenMissingSet f k a +generateMaybeAMissingSet f = WhenMissingSet + { missingSubtree = MS.fromSetMaybeA f + , missingKey = \k -> forceMaybe <$> f k + } +{-# INLINE generateMaybeAMissingSet #-} + +forceMaybe :: Maybe a -> Maybe a +forceMaybe Nothing = Nothing +forceMaybe m@(Just !_) = m