Merge pull request #146 from mschristiansen/diffwith

Add `differenceWith` function.

Author: tibbe

Data/HashMap/Base.hs

@@ -45,6 +45,7 @@ module Data.HashMap.Base
 
       -- * Difference and intersection
     , difference
+    , differenceWith
     , intersection
     , intersectionWith
     , intersectionWithKey
@@ -934,6 +935,18 @@ difference a b = foldlWithKey' go empty a
                  _       -> m
 {-# INLINABLE difference #-}
 
+-- | /O(n*log m)/ Difference with a combining function. When two equal keys are
+-- encountered, the combining function is applied to the values of these keys.
+-- If it returns 'Nothing', the element is discarded (proper set difference). If
+-- it returns (@'Just' y@), the element is updated with a new value @y@.
+differenceWith :: (Eq k, Hashable k) => (v -> w -> Maybe v) -> HashMap k v -> HashMap k w -> HashMap k v
+differenceWith f a b = foldlWithKey' go empty a
+  where
+    go m k v = case lookup k b of
+                 Nothing -> insert k v m
+                 Just w  -> maybe m (\y -> insert k y m) (f v w)
+{-# INLINABLE differenceWith #-}
+
 -- | /O(n*log m)/ Intersection of two maps. Return elements of the first
 -- map for keys existing in the second.
 intersection :: (Eq k, Hashable k) => HashMap k v -> HashMap k w -> HashMap k v

Data/HashMap/Lazy.hs

@@ -64,6 +64,7 @@ module Data.HashMap.Lazy
 
       -- * Difference and intersection
     , difference
+    , differenceWith
     , intersection
     , intersectionWith
     , intersectionWithKey

Data/HashMap/Strict.hs

@@ -64,6 +64,7 @@ module Data.HashMap.Strict
 
       -- * Difference and intersection
     , difference
+    , differenceWith
     , intersection
     , intersectionWith
     , intersectionWithKey
@@ -98,9 +99,9 @@ import Prelude hiding (map)
 import qualified Data.HashMap.Array as A
 import qualified Data.HashMap.Base as HM
 import Data.HashMap.Base hiding (
-    alter, adjust, fromList, fromListWith, insert, insertWith, intersectionWith,
-    intersectionWithKey, map, mapWithKey, mapMaybe, mapMaybeWithKey, singleton,
-    update, unionWith, unionWithKey)
+    alter, adjust, fromList, fromListWith, insert, insertWith, differenceWith,
+    intersectionWith, intersectionWithKey, map, mapWithKey, mapMaybe,
+    mapMaybeWithKey, singleton, update, unionWith, unionWithKey)
 import Data.HashMap.Unsafe (runST)
 
 -- $strictness
@@ -394,6 +395,18 @@ mapMaybe f = mapMaybeWithKey (const f)
 ------------------------------------------------------------------------
 -- * Difference and intersection
 
+-- | /O(n*log m)/ Difference with a combining function. When two equal keys are
+-- encountered, the combining function is applied to the values of these keys.
+-- If it returns 'Nothing', the element is discarded (proper set difference). If
+-- it returns (@'Just' y@), the element is updated with a new value @y@.
+differenceWith :: (Eq k, Hashable k) => (v -> w -> Maybe v) -> HashMap k v -> HashMap k w -> HashMap k v
+differenceWith f a b = foldlWithKey' go empty a
+  where
+    go m k v = case HM.lookup k b of
+                 Nothing -> insert k v m
+                 Just w  -> maybe m (\y -> insert k y m) (f v w)
+{-# INLINABLE differenceWith #-}
+
 -- | /O(n+m)/ Intersection of two maps. If a key occurs in both maps
 -- the provided function is used to combine the values from the two
 -- maps.

tests/HashMapProperties.hs

@@ -163,6 +163,12 @@ pDifference :: [(Key, Int)] -> [(Key, Int)] -> Bool
 pDifference xs ys = M.difference (M.fromList xs) `eq_`
                     HM.difference (HM.fromList xs) $ ys
 
+pDifferenceWith :: [(Key, Int)] -> [(Key, Int)] -> Bool
+pDifferenceWith xs ys = M.differenceWith f (M.fromList xs) `eq_`
+                        HM.differenceWith f (HM.fromList xs) $ ys
+  where
+    f x y = if x == 0 then Nothing else Just (x - y)
+
 pIntersection :: [(Key, Int)] -> [(Key, Int)] -> Bool
 pIntersection xs ys = M.intersection (M.fromList xs) `eq_`
                       HM.intersection (HM.fromList xs) $ ys
@@ -284,6 +290,7 @@ tests =
       ]
     , testGroup "difference and intersection"
       [ testProperty "difference" pDifference
+      , testProperty "differenceWith" pDifferenceWith
       , testProperty "intersection" pIntersection
       , testProperty "intersectionWith" pIntersectionWith
       , testProperty "intersectionWithKey" pIntersectionWithKey