add unionWith

Author: hdgarrood

README.md

@@ -54,6 +54,8 @@
 
     union :: forall k v. (P.Ord k) => Map k v -> Map k v -> Map k v
 
+    unionWith :: forall k v. (P.Ord k) => (v -> v -> v) -> Map k v -> Map k v -> Map k v
+
     unions :: forall k v. (P.Ord k) => [Map k v] -> Map k v
 
     update :: forall k v. (P.Ord k) => (v -> Maybe v) -> k -> Map k v -> Map k v

src/Data/Map.purs

@@ -22,6 +22,7 @@ module Data.Map
     keys,
     values,
     union,
+    unionWith,
     unions,
     map,
     size
@@ -230,8 +231,15 @@ values Leaf = []
 values (Two left _ v right) = values left P.++ [v] P.++ values right
 values (Three left _ v1 mid _ v2 right) = values left P.++ [v1] P.++ values mid P.++ [v2] P.++ values right
 
+-- Computes the union of two maps, except that when a key exists in both maps, its value in the result
+-- is computed by combining them with the supplied function.
+unionWith :: forall k v. (P.Ord k) => (v -> v -> v) -> Map k v -> Map k v -> Map k v
+unionWith f m1 m2 = foldl go m2 (toList m1)
+  where
+  go m (Tuple k v) = alter (Just P.<<< maybe v (f v)) k m
+
 union :: forall k v. (P.Ord k) => Map k v -> Map k v -> Map k v
-union m1 m2 = foldl (\m (Tuple k v) -> insert k v m) m2 (toList m1)
+union = unionWith P.const
 
 unions :: forall k v. (P.Ord k) => [Map k v] -> Map k v
 unions = foldl union empty

tests/Data/Map.purs

@@ -6,7 +6,7 @@ import Data.Maybe
 import Data.Tuple
 import Data.Array (map, length, nubBy)
 import Data.Function (on)
-import Data.Foldable (foldl)
+import Data.Foldable (foldl, for_)
 
 import Test.QuickCheck
 
@@ -166,6 +166,14 @@ mapTests = do
   trace "Union is idempotent"
   quickCheck $ \m1 m2 -> (m1 `M.union` m2) == ((m1 `M.union` m2) `M.union` (m2 :: M.Map SmallKey Number))
 
+  trace "unionWith"
+  for_ [Tuple (+) 0, Tuple (*) 1] $ \(Tuple op ident) ->
+    quickCheck $ \m1 m2 k ->
+      let u = M.unionWith op m1 m2 :: M.Map SmallKey Number
+      in case M.lookup k u of
+           Nothing -> not (M.member k m1 || M.member k m2)
+           Just v -> v == op (fromMaybe ident (M.lookup k m1)) (fromMaybe ident (M.lookup k m2))
+
   trace "size"
   quickCheck $ \xs ->
     let xs' = nubBy ((==) `on` fst) xs