Add an API to merge a Set and Map into a Map

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

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

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