Implement intersectionWithKey for HashMap

Author: RyanGlScott

Data/HashMap/Base.hs

@@ -45,6 +45,7 @@ module Data.HashMap.Base
     , difference
     , intersection
     , intersectionWith
+    , intersectionWithKey
 
       -- * Folds
     , foldl'
@@ -801,6 +802,18 @@ intersectionWith f a b = foldlWithKey' go empty a
                  _      -> m
 {-# INLINABLE intersectionWith #-}
 
+-- | /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.
+intersectionWithKey :: (Eq k, Hashable k) => (k -> v1 -> v2 -> v3)
+                    -> HashMap k v1 -> HashMap k v2 -> HashMap k v3
+intersectionWithKey f a b = foldlWithKey' go empty a
+  where
+    go m k v = case lookup k b of
+                 Just w -> insert k (f k v w) m
+                 _      -> m
+{-# INLINABLE intersectionWithKey #-}
+
 ------------------------------------------------------------------------
 -- * Folds
 

Data/HashMap/Lazy.hs

@@ -65,6 +65,7 @@ module Data.HashMap.Lazy
     , difference
     , intersection
     , intersectionWith
+    , intersectionWithKey
 
       -- * Folds
     , foldl'

Data/HashMap/Strict.hs

@@ -65,6 +65,7 @@ module Data.HashMap.Strict
     , difference
     , intersection
     , intersectionWith
+    , intersectionWithKey
 
       -- * Folds
     , foldl'
@@ -97,7 +98,8 @@ 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,
-    map, mapWithKey, mapMaybe, mapMaybeWithKey, singleton, update, unionWith)
+    intersectionWithKey, map, mapWithKey, mapMaybe, mapMaybeWithKey, singleton,
+    update, unionWith)
 import Data.HashMap.Unsafe (runST)
 
 -- $strictness
@@ -396,6 +398,18 @@ intersectionWith f a b = foldlWithKey' go empty a
                  _      -> m
 {-# INLINABLE intersectionWith #-}
 
+-- | /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.
+intersectionWithKey :: (Eq k, Hashable k) => (k -> v1 -> v2 -> v3)
+                    -> HashMap k v1 -> HashMap k v2 -> HashMap k v3
+intersectionWithKey f a b = foldlWithKey' go empty a
+  where
+    go m k v = case HM.lookup k b of
+                 Just w -> insert k (f k v w) m
+                 _      -> m
+{-# INLINABLE intersectionWithKey #-}
+
 ------------------------------------------------------------------------
 -- ** Lists
 

tests/HashMapProperties.hs

@@ -149,6 +149,13 @@ pIntersectionWith :: [(Key, Int)] -> [(Key, Int)] -> Bool
 pIntersectionWith xs ys = M.intersectionWith (-) (M.fromList xs) `eq_`
                           HM.intersectionWith (-) (HM.fromList xs) $ ys
 
+pIntersectionWithKey :: [(Key, Int)] -> [(Key, Int)] -> Bool
+pIntersectionWithKey xs ys = M.intersectionWithKey go (M.fromList xs) `eq_`
+                             HM.intersectionWithKey go (HM.fromList xs) $ ys
+  where
+    go :: Key -> Int -> Int -> Int
+    go (K k) i1 i2 = k - i1 - i2
+
 ------------------------------------------------------------------------
 -- ** Folds
 
@@ -255,6 +262,7 @@ tests =
       [ testProperty "difference" pDifference
       , testProperty "intersection" pIntersection
       , testProperty "intersectionWith" pIntersectionWith
+      , testProperty "intersectionWithKey" pIntersectionWithKey
       ]
     -- Filter
     , testGroup "filter"