Merge pull request #15 from karshan/master

Add intersection,intersectionWith for Data.Map

Author: garyb

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)
 

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

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)