Merge pull request #24 from hdgarrood/unionWith

add `unionWith`

Author: paf31

README.md

@@ -58,6 +58,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
@@ -237,8 +238,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

@@ -2,11 +2,12 @@ module Tests.Data.Map where
 
 import Debug.Trace
 
+import Control.Alt ((<|>))
 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 +167,30 @@ 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 "Union prefers left"
+  quickCheck $ \m1 m2 k -> M.lookup k (M.union m1 (m2 :: M.Map SmallKey Number)) == (M.lookup k m1 <|> M.lookup k m2)
+
+  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 "unionWith argument order"
+  quickCheck $ \m1 m2 k ->
+    let u   = M.unionWith (-) m1 m2 :: M.Map SmallKey Number
+        in1 = M.member k m1
+        v1  = M.lookup k m1
+        in2 = M.member k m2
+        v2  = M.lookup k m2
+    in case M.lookup k u of
+          Just v | in1 && in2 -> Just v == ((-) <$> v1 <*> v2)
+          Just v | in1        -> Just v == v1
+          Just v              -> Just v == v2
+          Nothing             -> not (in1 || in2)
+
   trace "size"
   quickCheck $ \xs ->
     let xs' = nubBy ((==) `on` fst) xs