diff --git a/src/Data/Map.purs b/src/Data/Map.purs index 945c968..d4dd4e4 100644 --- a/src/Data/Map.purs +++ b/src/Data/Map.purs @@ -5,7 +5,7 @@ module Data.Map import Prelude -import Data.Map.Internal (Map, alter, checkValid, delete, empty, filter, filterKeys, filterWithKey, findMax, findMin, foldSubmap, fromFoldable, fromFoldableWith, fromFoldableWithIndex, insert, insertWith, isEmpty, isSubmap, lookup, lookupGE, lookupGT, lookupLE, lookupLT, member, pop, showTree, singleton, size, submap, toUnfoldable, toUnfoldableUnordered, union, unionWith, unions, difference, update, values, mapMaybeWithKey, mapMaybe) +import Data.Map.Internal (Map, alter, checkValid, delete, empty, filter, filterKeys, filterWithKey, findMax, findMin, foldSubmap, fromFoldable, fromFoldableWith, fromFoldableWithIndex, insert, insertWith, isEmpty, isSubmap, lookup, lookupGE, lookupGT, lookupLE, lookupLT, member, pop, showTree, singleton, size, submap, toUnfoldable, toUnfoldableUnordered, union, unionWith, unions, intersection, intersectionWith, difference, update, values, mapMaybeWithKey, mapMaybe) import Data.Set (Set) import Unsafe.Coerce (unsafeCoerce) diff --git a/src/Data/Map/Internal.purs b/src/Data/Map/Internal.purs index 16346ab..1179165 100644 --- a/src/Data/Map/Internal.purs +++ b/src/Data/Map/Internal.purs @@ -34,6 +34,8 @@ module Data.Map.Internal , union , unionWith , unions + , intersection + , intersectionWith , difference , isSubmap , size @@ -628,6 +630,24 @@ union = unionWith const unions :: forall k v f. Ord k => Foldable f => f (Map k v) -> Map k v unions = foldl union empty +-- | Compute the intersection of two maps, using the specified function +-- | to combine values for duplicate keys. +intersectionWith :: forall k a b c. Ord k => (a -> b -> c) -> Map k a -> Map k b -> Map k c +intersectionWith f m1 m2 = go (toUnfoldable m1 :: List (Tuple k a)) (toUnfoldable m2 :: List (Tuple k b)) empty + where + go Nil _ m = m + go _ Nil m = m + go as@(Cons (Tuple k1 a) ass) bs@(Cons (Tuple k2 b) bss) m = + case compare k1 k2 of + LT -> go ass bs m + EQ -> go ass bss (insert k1 (f a b) m) + GT -> go as bss m + +-- | Compute the intersection of two maps, preferring values from the first map in the case +-- | of duplicate keys. +intersection :: forall k a b. Ord k => Map k a -> Map k b -> Map k a +intersection = intersectionWith const + -- | Difference of two maps. Return elements of the first map where -- | the keys do not exist in the second map. difference :: forall k v w. Ord k => Map k v -> Map k w -> Map k v diff --git a/test/Test/Data/Map.purs b/test/Test/Data/Map.purs index f8a4ae4..37d5661 100644 --- a/test/Test/Data/Map.purs +++ b/test/Test/Data/Map.purs @@ -223,6 +223,23 @@ mapTests = do Just v -> Just v == v2 Nothing -> not (in1 || in2) + log "Lookup from intersection" + quickCheck $ \(TestMap m1) (TestMap m2) k -> + M.lookup (smallKey k) (M.intersection (m1 :: M.Map SmallKey Int) (m2 :: M.Map SmallKey Int)) == (case M.lookup k m2 of + Nothing -> Nothing + Just v -> M.lookup k m1) ("m1: " <> show m1 <> ", m2: " <> show m2 <> ", k: " <> show k <> ", v1: " <> show (M.lookup k m1) <> ", v2: " <> show (M.lookup k m2) <> ", intersection: " <> show (M.intersection m1 m2)) + + log "Intersection is idempotent" + quickCheck $ \(TestMap m1) (TestMap m2) -> ((m1 :: M.Map SmallKey Int) `M.intersection` m2) == ((m1 `M.intersection` m2) `M.intersection` (m2 :: M.Map SmallKey Int)) + + log "intersectionWith" + for_ [(+), (*)] $ \op -> + quickCheck $ \(TestMap m1) (TestMap m2) k -> + let u = M.intersectionWith op m1 m2 :: M.Map SmallKey Int + in case M.lookup k u of + Nothing -> not (M.member k m1 && M.member k m2) + Just v -> Just v == (op <$> M.lookup k m1 <*> M.lookup k m2) + log "difference" quickCheck $ \(TestMap m1) (TestMap m2) -> let d = M.difference (m1 :: M.Map SmallKey Int) (m2 :: M.Map SmallKey String)