Add insertWith for Map

Author: garyb

src/Data/Map.purs

@@ -5,7 +5,7 @@ module Data.Map
 
 import Prelude
 
-import Data.Map.Internal (Map, alter, checkValid, delete, empty, filter, filterKeys, filterWithKey, findMax, findMin, foldSubmap, fromFoldable, fromFoldableWith, fromFoldableWithIndex, insert, isEmpty, isSubmap, lookup, lookupGE, lookupGT, lookupLE, lookupLT, member, pop, showTree, singleton, size, submap, toUnfoldable, toUnfoldableUnordered, union, unionWith, unions, difference, update, values, mapMaybeWithKey, mapMaybe)
+import Data.Map.Internal (Map, alter, checkValid, delete, empty, filter, filterKeys, filterWithKey, findMax, findMin, foldSubmap, fromFoldable, fromFoldableWith, fromFoldableWithIndex, insert, insertWith, isEmpty, isSubmap, lookup, lookupGE, lookupGT, lookupLE, lookupLT, member, pop, showTree, singleton, size, submap, toUnfoldable, toUnfoldableUnordered, union, unionWith, unions, difference, update, values, mapMaybeWithKey, mapMaybe)
 import Data.Set (Set)
 import Unsafe.Coerce (unsafeCoerce)
 

src/Data/Map/Internal.purs

@@ -9,6 +9,7 @@ module Data.Map.Internal
   , singleton
   , checkValid
   , insert
+  , insertWith
   , lookup
   , lookupLE
   , lookupLT
@@ -455,6 +456,13 @@ insert k v = down Nil
       ThreeMiddle a k1 v1 k2 v2 d, KickUp b k' v' c -> up ctx (KickUp (Two a k1 v1 b) k' v' (Two c k2 v2 d))
       ThreeRight a k1 v1 b k2 v2, KickUp c k' v' d -> up ctx (KickUp (Two a k1 v1 b) k2 v2 (Two c k' v' d))
 
+-- | Inserts or updates a value with the given function.
+-- |
+-- | The combining function is called with the existing value as the first
+-- | argument and the new value as the second argument.
+insertWith :: forall k v. Ord k => (v -> v -> v) -> k -> v -> Map k v -> Map k v
+insertWith f k v = alter (Just <<< maybe v (flip f v)) k
+
 -- | Delete a key and its corresponding value from a map.
 delete :: forall k v. Ord k => k -> Map k v -> Map k v
 delete k m = maybe m snd (pop k m)

test/Test/Data/Map.purs

@@ -83,6 +83,14 @@ mapTests = do
   quickCheck $ \k v1 v2 ->
     M.lookup (smallKey k) (M.insert k v2 (M.insert k v1 M.empty)) == Just (number v2)
 
+  log "Test insertWith combining values"
+  quickCheck $ \k v1 v2 ->
+    M.lookup (smallKey k) (M.insertWith (+) k v2 (M.insert k v1 M.empty)) == Just (number (v1 + v2))
+
+  log "Test insertWith passes the first value as the first argument to the combining function"
+  quickCheck $ \k v1 v2 ->
+    M.lookup (smallKey k) (M.insertWith const k v2 (M.insert k v1 M.empty)) == Just (number v1)
+
   log "Test delete after inserting"
   quickCheck $ \k v -> M.isEmpty (M.delete (smallKey k) (M.insert k (number v) M.empty))
     <?> ("k: " <> show k <> ", v: " <> show v)