Skip to content

Commit 722218d

Browse files
meooow25meooow25
authored
Add an API to merge a Set and Map into a Map (#1227)
This is intended to serve some use cases more efficiently than existing methods, such as partitioning a map according to whether the keys exist in a set. Similar to Map merging, there are two modules, "Data.Map.Merge.Set.Lazy" and "Data.Map.Merge.Set.Strict" exposing the merge functions and tactics. The WhenMissing tactic for missing maps is simply reused from the Map merging interface. A new WhenMatched tactic is added. A new WhenMissingSet tactic is also added, and I chose to added the "Set" suffix to distinguish it from the existing map tactic. Similarly, the functions that create WhenMissingSet tactics are suffixed with "MissingSet".
1 parent dde5527 commit 722218d

8 files changed

Lines changed: 907 additions & 2 deletions

File tree

‎containers-tests/containers-tests.cabal‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,9 @@ library
117117
Data.Map.Lazy
118118
Data.Map.Merge.Lazy
119119
Data.Map.Merge.Strict
120+
Data.Map.Merge.Set.Internal
121+
Data.Map.Merge.Set.Lazy
122+
Data.Map.Merge.Set.Strict
120123
Data.Map.Strict
121124
Data.Map.Strict.Internal
122125
Data.Sequence

‎containers-tests/test-utils/Utils/MergeFunc.hs‎

Lines changed: 43 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
11
module Utils.MergeFunc
22
( WhenMatchedFunc(..)
33
, WhenMissingFunc(..)
4+
, MapSet_WhenMatchedFunc(..)
5+
, MapSet_WhenMissingSetFunc(..)
46
) where
57

6-
import Test.QuickCheck
8+
import Test.QuickCheck (CoArbitrary, Function, Arbitrary(..), oneof)
79
import Utils.Strictness (Func, Func2, Func3)
810

911
-- k: key, x: left map value, y: right map value, z: result map value,
@@ -69,3 +71,43 @@ instance
6971
MapMissingFunc fun -> MapMissingFunc <$> shrink fun
7072
FmapMapMissingFunc fun2 fun1 ->
7173
uncurry FmapMapMissingFunc <$> shrink (fun2, fun1)
74+
75+
-- For Set-Map to Map merge.
76+
-- k: key, a: map value, b: result map value
77+
data MapSet_WhenMatchedFunc k a b
78+
= MapSet_MapMaybeMatchedFunc (Func2 k a (Maybe b))
79+
| MapSet_MapMatchedFunc (Func2 k a b)
80+
deriving Show
81+
82+
instance
83+
( CoArbitrary k, Function k
84+
, CoArbitrary a, Function a
85+
, Arbitrary b
86+
) => Arbitrary (MapSet_WhenMatchedFunc k a b) where
87+
arbitrary = oneof
88+
[ MapSet_MapMaybeMatchedFunc <$> arbitrary
89+
, MapSet_MapMatchedFunc <$> arbitrary
90+
]
91+
shrink wmf = case wmf of
92+
MapSet_MapMaybeMatchedFunc fun -> MapSet_MapMaybeMatchedFunc <$> shrink fun
93+
MapSet_MapMatchedFunc fun -> MapSet_MapMatchedFunc <$> shrink fun
94+
95+
-- For Set-Map to Map merge.
96+
-- k: key, a: result map value
97+
data MapSet_WhenMissingSetFunc k a
98+
= MapSet_GenerateMissingSetFunc (Func k a)
99+
| MapSet_GenerateMaybeMissingSetFunc (Func k (Maybe a))
100+
deriving Show
101+
102+
instance
103+
(CoArbitrary k, Function k, Arbitrary a)
104+
=> Arbitrary (MapSet_WhenMissingSetFunc k a) where
105+
arbitrary = oneof
106+
[ MapSet_GenerateMissingSetFunc <$> arbitrary
107+
, MapSet_GenerateMaybeMissingSetFunc <$> arbitrary
108+
]
109+
shrink wmf = case wmf of
110+
MapSet_GenerateMissingSetFunc fun ->
111+
MapSet_GenerateMissingSetFunc <$> shrink fun
112+
MapSet_GenerateMaybeMissingSetFunc fun ->
113+
MapSet_GenerateMaybeMissingSetFunc <$> shrink fun

‎containers-tests/tests/map-properties.hs‎

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,11 @@
66
#ifdef STRICT
77
import Data.Map.Strict as Data.Map
88
import Data.Map.Merge.Strict
9+
import qualified Data.Map.Merge.Set.Strict as MergeSet
910
#else
1011
import Data.Map.Lazy as Data.Map
1112
import Data.Map.Merge.Lazy
13+
import qualified Data.Map.Merge.Set.Lazy as MergeSet
1214
#endif
1315
import Data.Map.Internal (Map, link2, link)
1416
import Data.Map.Internal.Debug (showTree, showTreeWith, balanced)
@@ -327,6 +329,8 @@ main = defaultMain $ testGroup "map-properties"
327329
, testProperty "mapAccum" prop_mapAccum
328330
, testProperty "mapAccumWithKey" prop_mapAccumWithKey
329331
, testProperty "mapAccumRWithKey" prop_mapAccumRWithKey
332+
, testProperty "Merge.Set.merge" prop_mergeMapSet
333+
, testProperty "Merge.Set.mergeA" prop_mergeAMapSet
330334
, testLaws $ Laws.eqLaws (Proxy :: Proxy (Map Int A))
331335
, testLaws $ Laws.ordLaws (Proxy :: Proxy (Map Int OrdA))
332336
, testLaws $ Laws.showLaws (Proxy :: Proxy (Map Int A))
@@ -1390,6 +1394,129 @@ instance (Arbitrary a, CoArbitrary a, Function a, CoArbitrary k, Function k)
13901394
, ZipWithMaybeMatched <$> arbitrary
13911395
]
13921396

1397+
prop_mergeMapSet
1398+
:: WhenMissingSpec Int A
1399+
-> WhenMissingMapSetSpec Int A
1400+
-> WhenMatchedMapSetSpec Int A
1401+
-> Map Int A
1402+
-> Set Int
1403+
-> Property
1404+
prop_mergeMapSet miss1 miss2 match m1 s2 =
1405+
valid m .&&.
1406+
m ===
1407+
(mapMaybeWithKey (\k x -> runIdentity (runWhenMissing miss1' k x)) m1Only `union`
1408+
mapMaybe id (fromSet (runIdentity . MergeSet.runWhenMissingSet miss2') s2Only) `union`
1409+
mapMaybeWithKey (\k x -> runIdentity (MergeSet.runWhenMatched match' k x)) m12Both)
1410+
where
1411+
miss1' = whenMissing miss1
1412+
miss2' = whenMissingSet miss2
1413+
match' = whenMatched match
1414+
1415+
m = MergeSet.merge miss1' miss2' match' m1 s2
1416+
1417+
m1Only = filterKeys (`Set.notMember` s2) m1
1418+
s2Only = Set.filter (`notMember` m1) s2
1419+
m12Both = filterKeys (`Set.member` s2) m1
1420+
1421+
whenMissing spec = case spec of
1422+
DropMissing -> MergeSet.dropMissing
1423+
PreserveMissing -> MergeSet.preserveMissing
1424+
FilterMissing f -> MergeSet.filterMissing (applyFun2 f)
1425+
MapMissing f -> MergeSet.mapMissing (applyFun2 f)
1426+
MapMaybeMissing f -> MergeSet.mapMaybeMissing (applyFun2 f)
1427+
whenMissingSet spec = case spec of
1428+
DropMissingMapSet -> MergeSet.dropMissingSet
1429+
GenerateMissingMapSet f -> MergeSet.generateMissingSet (applyFun f)
1430+
GenerateMaybeMissingMapSet f -> MergeSet.generateMaybeMissingSet (applyFun f)
1431+
whenMatched spec = case spec of
1432+
FilterMatchedMapSet f -> MergeSet.filterMatched (applyFun2 f)
1433+
MapMatchedMapSet f -> MergeSet.mapMatched (applyFun2 f)
1434+
MapMaybeMatchedMapSet f -> MergeSet.mapMaybeMatched (applyFun2 f)
1435+
1436+
-- This uses the instance
1437+
-- Monoid a => Applicative ((,) a)
1438+
-- to test that effects are sequenced in ascending key order.
1439+
prop_mergeAMapSet
1440+
:: WhenMissingSpec Int A
1441+
-> WhenMissingMapSetSpec Int A
1442+
-> WhenMatchedMapSetSpec Int A
1443+
-> Map Int A
1444+
-> Set Int
1445+
-> Property
1446+
prop_mergeAMapSet miss1 miss2 match m1 s2 =
1447+
valid m .&&.
1448+
m ===
1449+
(mapMaybeWithKey (\k x -> snd (runWhenMissing miss1' k x)) m1Only `union`
1450+
mapMaybe id (fromSet (snd . MergeSet.runWhenMissingSet miss2') s2Only) `union`
1451+
mapMaybeWithKey (\k x -> snd (MergeSet.runWhenMatched match' k x)) m12Both) .&&.
1452+
ks ===
1453+
sort
1454+
(concat
1455+
(fmap (\(k,x) -> fst (runWhenMissing miss1' k x)) (toList m1Only) ++
1456+
fmap (fst . MergeSet.runWhenMissingSet miss2') (Set.toList s2Only) ++
1457+
fmap (\(k,x) -> fst (MergeSet.runWhenMatched match' k x)) (toList m12Both)))
1458+
where
1459+
miss1' = whenMissing miss1
1460+
miss2' = whenMissingSet miss2
1461+
match' = whenMatched match
1462+
1463+
(ks, m) = MergeSet.mergeA miss1' miss2' match' m1 s2
1464+
1465+
m1Only = filterKeys (`Set.notMember` s2) m1
1466+
s2Only = Set.filter (`notMember` m1) s2
1467+
m12Both = filterKeys (`Set.member` s2) m1
1468+
1469+
whenMissing spec = case spec of
1470+
DropMissing -> MergeSet.dropMissing
1471+
PreserveMissing -> MergeSet.preserveMissing
1472+
FilterMissing f -> MergeSet.filterAMissing (\k x -> ([k], applyFun2 f k x))
1473+
MapMissing f -> MergeSet.traverseMissing (\k x -> ([k], applyFun2 f k x))
1474+
MapMaybeMissing f -> MergeSet.traverseMaybeMissing (\k x -> ([k], applyFun2 f k x))
1475+
whenMissingSet spec = case spec of
1476+
DropMissingMapSet -> MergeSet.dropMissingSet
1477+
GenerateMissingMapSet f -> MergeSet.generateAMissingSet (\k -> ([k], applyFun f k))
1478+
GenerateMaybeMissingMapSet f -> MergeSet.generateMaybeAMissingSet (\k -> ([k], applyFun f k))
1479+
whenMatched spec = case spec of
1480+
FilterMatchedMapSet f -> MergeSet.filterAMatched (\k x -> ([k], applyFun2 f k x))
1481+
MapMatchedMapSet f -> MergeSet.traverseMatched (\k x -> ([k], applyFun2 f k x))
1482+
MapMaybeMatchedMapSet f -> MergeSet.traverseMaybeMatched (\k x -> ([k], applyFun2 f k x))
1483+
1484+
data WhenMissingMapSetSpec k a
1485+
= DropMissingMapSet
1486+
| GenerateMissingMapSet (Fun k a)
1487+
| GenerateMaybeMissingMapSet (Fun k (Maybe a))
1488+
deriving Show
1489+
1490+
instance (Arbitrary a, CoArbitrary a, Function a, CoArbitrary k, Function k)
1491+
=> Arbitrary (WhenMissingMapSetSpec k a) where
1492+
arbitrary = oneof
1493+
[ pure DropMissingMapSet
1494+
, GenerateMissingMapSet <$> arbitrary
1495+
, GenerateMaybeMissingMapSet <$> arbitrary
1496+
]
1497+
shrink spec = case spec of
1498+
DropMissingMapSet -> []
1499+
GenerateMissingMapSet f -> GenerateMissingMapSet <$> shrink f
1500+
GenerateMaybeMissingMapSet f -> GenerateMaybeMissingMapSet <$> shrink f
1501+
1502+
data WhenMatchedMapSetSpec k a
1503+
= FilterMatchedMapSet (Fun (k, a) Bool)
1504+
| MapMatchedMapSet (Fun (k, a) a)
1505+
| MapMaybeMatchedMapSet (Fun (k, a) (Maybe a))
1506+
deriving Show
1507+
1508+
instance (Arbitrary a, CoArbitrary a, Function a, CoArbitrary k, Function k)
1509+
=> Arbitrary (WhenMatchedMapSetSpec k a) where
1510+
arbitrary = oneof
1511+
[ FilterMatchedMapSet <$> arbitrary
1512+
, MapMatchedMapSet <$> arbitrary
1513+
, MapMaybeMatchedMapSet <$> arbitrary
1514+
]
1515+
shrink spec = case spec of
1516+
FilterMatchedMapSet f -> FilterMatchedMapSet <$> shrink f
1517+
MapMatchedMapSet f -> MapMatchedMapSet <$> shrink f
1518+
MapMaybeMatchedMapSet f -> MapMaybeMatchedMapSet <$> shrink f
1519+
13931520
----------------------------------------------------------------
13941521

13951522
prop_list :: [Int] -> Bool

0 commit comments

Comments
 (0)