Merge pull request #970 from haskell/issue-952-keymap-insertWith

Issue #952: KeyMap.insertWith

Author: phadej

changelog.md

@@ -3,6 +3,7 @@ For the latest version of this document, please see [https://github.com/haskell/
 ### 2.1.1.0
 
 - Add `Data.Aeson.KeyMap.!?` (flipped) alias to `Data.Aeson.KeyMap.lookup`.
+- Add `Data.Aeson.KeyMap.insertWith` function.
 - Use `unsafeDupablePerformIO` instead of incorrect `accursedUnutterablePerformIO` in creation of keys in TH serialisation.
   This fixes a bug in TH deriving, e.g. when `Strict` pragma was enabled.
 

src/Data/Aeson/KeyMap.hs

@@ -26,6 +26,7 @@ module Data.Aeson.KeyMap (
 
     -- ** Insertion
     insert,
+    insertWith,
 
     -- * Deletion
     delete,
@@ -191,6 +192,12 @@ lookup t tm = M.lookup t (unKeyMap tm)
 insert :: Key -> v -> KeyMap v -> KeyMap v
 insert k v tm = KeyMap (M.insert k v (unKeyMap tm))
 
+-- | Insert with a function combining new and old values, taken in that order.
+--
+-- @since 2.1.1.0
+insertWith :: (a -> a -> a) -> Key -> a -> KeyMap a -> KeyMap a
+insertWith f k v m = KeyMap (M.insertWith f k v (unKeyMap m))
+
 -- | Map a function over all values in the map.
 map :: (a -> b) -> KeyMap a -> KeyMap b
 map = fmap
@@ -394,6 +401,12 @@ lookup t tm = H.lookup t (unKeyMap tm)
 insert :: Key -> v -> KeyMap v -> KeyMap v
 insert k v tm = KeyMap (H.insert k v (unKeyMap tm))
 
+-- | Insert with a function combining new and old values, taken in that order.
+--
+-- @since 2.1.1.0
+insertWith :: (a -> a -> a) -> Key -> a -> KeyMap a -> KeyMap a
+insertWith f k v m = KeyMap (H.insertWith f k v (unKeyMap m))
+
 -- | Map a function over all values in the map.
 map :: (a -> b) -> KeyMap a -> KeyMap b
 map = fmap

tests/UnitTests.hs

@@ -826,4 +826,9 @@ tests = testGroup "unit" [
     ]
   , monadFixTests
   , issue967
+  , testCase "KeyMap.insertWith" $ do
+      KM.insertWith (-)        "a" 2 (KM.fromList [("a", 1)]) @?= KM.fromList [("a",1 :: Int)]
+      KM.insertWith (flip (-)) "a" 2 (KM.fromList [("a", 1)]) @?= KM.fromList [("a",-1 :: Int)]
+      KM.insertWith (-)        "b" 2 (KM.fromList [("a", 1)]) @?= KM.fromList [("a",1),("b",2 :: Int)]
+      KM.insertWith (-)        "b" 2 KM.empty                 @?= KM.fromList [("b",2 :: Int)]
   ]