Add differenceWith function.

Author: mschristiansen

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

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