Implementation for function insertWithFun for strict and lazy maps.

containers-tests/tests/map-properties.hs

@@ -66,6 +66,7 @@ main = defaultMain
          , testCase "singleton" test_singleton
          , testCase "insert" test_insert
          , testCase "insertWith" test_insertWith
+         , testCase "insertWithFun" test_insertWithFun
          , testCase "insertWithKey" test_insertWithKey
          , testCase "insertLookupWithKey" test_insertLookupWithKey
          , testCase "delete" test_delete
@@ -465,6 +466,12 @@ test_insertWith = do
     insertWith (++) 7 "xxx" (fromList [(5,"a"), (3,"b")]) @?= fromList [(3, "b"), (5, "a"), (7, "xxx")]
     insertWith (++) 5 "xxx" empty                         @?= singleton 5 "xxx"
 
+test_insertWithFun :: Assertion
+test_insertWithFun = do
+    insertWithFun (:) (: []) 5 "x" (fromList [(5,["a"]), (3,["b"])]) @?= fromList [(3, ["b"]), (5, ["x", "a"])]
+    insertWithFun (:) (: []) 7 "x" (fromList [(5,["a"]), (3,["b"])]) @?= fromList [(3, ["b"]), (5, ["a"]), (7, ["x"])]
+    insertWithFun (:) (: []) 5 "x" empty                             @?= singleton 5 ["x"]
+
 test_insertWithKey :: Assertion
 test_insertWithKey = do
     insertWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) @?= fromList [(3, "b"), (5, "5:xxx|a")]

containers-tests/tests/map-strictness.hs

@@ -71,6 +71,17 @@ pInsertWithValueStrict f k v m
                      not (isBottom $ M.insertWith (const2 1) k bottom m)
     | otherwise    = isBottom $ M.insertWith (apply2 f) k bottom m
 
+pInsertWithFunKeyStrict :: Fun (Int, Int) Int -> Fun Int Int -> Int -> Map Int Int -> Bool
+pInsertWithFunKeyStrict f g v m = isBottom $ M.insertWithFun (apply2 f) (apply g) bottom v m
+
+pInsertWithFunValueStrict :: Fun (Int, Int) Int -> Fun Int Int -> Int -> Int -> Map Int Int
+                             -> Bool
+pInsertWithFunValueStrict f g k v m
+    | M.member k m = (isBottom $ M.insertWithFun (const2 bottom) (const bottom) k v m) &&
+                     not (isBottom $ M.insertWithFun (const2 1) (const bottom) k bottom m)
+    | otherwise    = (isBottom $ M.insertWithFun (apply2 f) (apply g) k bottom m) &&
+                     not (isBottom $ M.insertWithFun (const2 bottom) (const 1) k bottom m)
+
 pInsertLookupWithKeyKeyStrict :: Fun (Int, Int, Int) Int -> Int
                               -> Map Int Int -> Bool
 pInsertLookupWithKeyKeyStrict f v m = isBottom $ M.insertLookupWithKey (apply3 f) bottom v m
@@ -94,19 +105,25 @@ tExtraThunksM :: Test
 tExtraThunksM = testGroup "Map.Strict - extra thunks" $
     if not isUnitSupported then [] else
     -- for strict maps, all the values should be evaluated to ()
-    [ check "singleton"           $ m0
-    , check "insert"              $ M.insert 42 () m0
-    , check "insertWith"          $ M.insertWith const 42 () m0
-    , check "fromList"            $ M.fromList [(42,()),(42,())]
-    , check "fromListWith"        $ M.fromListWith const [(42,()),(42,())]
-    , check "fromAscList"         $ M.fromAscList [(42,()),(42,())]
-    , check "fromAscListWith"     $ M.fromAscListWith const [(42,()),(42,())]
-    , check "fromDistinctAscList" $ M.fromAscList [(42,())]
+    [ check  "singleton"           $ m0
+    , check  "insert"              $ M.insert 42 () m0
+    , check  "insertWith"          $ M.insertWith const 42 () m0
+    , check  "insertWithFun_f"     $ M.insertWithFun const (const ()) 42 () m0
+    , check' "insertWithFun_g" 21  $ M.insertWithFun const (const ()) 21 () m0
+    , check  "fromList"            $ M.fromList [(42,()),(42,())]
+    , check  "fromListWith"        $ M.fromListWith const [(42,()),(42,())]
+    , check  "fromAscList"         $ M.fromAscList [(42,()),(42,())]
+    , check  "fromAscListWith"     $ M.fromAscListWith const [(42,()),(42,())]
+    , check  "fromDistinctAscList" $ M.fromAscList [(42,())]
     ]
   where
     m0 = M.singleton 42 ()
+
     check :: TestName -> M.Map Int () -> Test
-    check n m = testCase n $ case M.lookup 42 m of
+    check n m = check' n 42 m
+
+    check' :: TestName -> Int -> M.Map Int () -> Test
+    check' n k m = testCase n $ case M.lookup k m of
         Just v -> assertBool msg (isUnit v)
         _      -> assertString "key not found"
       where
@@ -117,19 +134,25 @@ tExtraThunksL = testGroup "Map.Lazy - extra thunks" $
     if not isUnitSupported then [] else
     -- for lazy maps, the *With functions should leave `const () ()` thunks,
     -- but the other functions should produce fully evaluated ().
-    [ check "singleton"       True  $ m0
-    , check "insert"          True  $ L.insert 42 () m0
-    , check "insertWith"      False $ L.insertWith const 42 () m0
-    , check "fromList"        True  $ L.fromList [(42,()),(42,())]
-    , check "fromListWith"    False $ L.fromListWith const [(42,()),(42,())]
-    , check "fromAscList"     True  $ L.fromAscList [(42,()),(42,())]
-    , check "fromAscListWith" False $ L.fromAscListWith const [(42,()),(42,())]
-    , check "fromDistinctAscList" True $ L.fromAscList [(42,())]
+    [ check  "singleton"           True     $ m0
+    , check  "insert"              True     $ L.insert 42 () m0
+    , check  "insertWith"          False    $ L.insertWith const 42 () m0
+    , check  "insertWithFun_f"     False    $ L.insertWithFun const (const ()) 42 () m0
+    , check' "insertWithFun_g"     False 21 $ L.insertWithFun const (const ()) 21 () m0
+    , check  "fromList"            True     $ L.fromList [(42,()),(42,())]
+    , check  "fromListWith"        False    $ L.fromListWith const [(42,()),(42,())]
+    , check  "fromAscList"         True     $ L.fromAscList [(42,()),(42,())]
+    , check  "fromAscListWith"     False    $ L.fromAscListWith const [(42,()),(42,())]
+    , check  "fromDistinctAscList" True     $ L.fromAscList [(42,())]
     ]
   where
     m0 = L.singleton 42 ()
+
     check :: TestName -> Bool -> L.Map Int () -> Test
-    check n e m = testCase n $ case L.lookup 42 m of
+    check n e m = check' n e 42 m
+
+    check' :: TestName -> Bool -> Int -> L.Map Int () -> Test
+    check' n e k m = testCase n $ case L.lookup k m of
         Just v -> assertBool msg (e == isUnit v)
         _      -> assertString "key not found"
       where
@@ -157,7 +180,9 @@ tests =
       , testProperty "insert is key-strict" pInsertKeyStrict
       , testProperty "insert is value-strict" pInsertValueStrict
       , testProperty "insertWith is key-strict" pInsertWithKeyStrict
+      , testProperty "insertWithFun is key-strict" pInsertWithFunKeyStrict
       , testProperty "insertWith is value-strict" pInsertWithValueStrict
+      , testProperty "insertWithFun is value-strict" pInsertWithFunValueStrict
       , testProperty "insertLookupWithKey is key-strict"
         pInsertLookupWithKeyKeyStrict
       , testProperty "insertLookupWithKey is value-strict"

containers/changelog.md

@@ -9,6 +9,8 @@
 
 ### Additions
 
+* [Add `insertKeyWithFun` to `Data.Map.Strict` and `Data.Map.Lazy`]() (Thanks, Neophytos Michael)
+
 * [Add `reverseTopSort` to `Data.Graph`](https://github.com/haskell/containers/pull/638) (Thanks, James Parker)
 
 * [Expose `traverseMaybeWithKey` from `Data.IntMap.{Lazy,Strict}`](https://github.com/haskell/containers/pull/743) (Thanks, Simon

containers/src/Data/Map/Internal.hs

@@ -157,6 +157,7 @@ module Data.Map.Internal (
     -- ** Insertion
     , insert
     , insertWith
+    , insertWithFun
     , insertWithKey
     , insertLookupWithKey
 
@@ -886,6 +887,37 @@ insertWith = go
 {-# INLINE insertWith #-}
 #endif
 
+-- | /O(log n)/. Insert with two functions: 'f' to combine a value (of type 'a')
+-- with the Map value (of type 'b') to produce a new Map value, and 'g' to inject the
+-- input value of type 'a' into type 'b' if key is not present.
+-- This version can be used to avoid unnecessary boxing when 'b' is for example
+-- some sort of container containing elements of type 'a'.
+-- @'insertWithFun' f g key value mp@
+-- will insert the pair (key, g value) into @mp@ if key does
+-- not exist in the map. If the key does exist, the function will
+-- insert the pair @(key, f new_value old_value)@.
+--
+-- > insertWithFun (:) (: []) 5 "x" (fromList [(5,["a"]), (3,["b"])]) == fromList [(3, ["b"]), (5, ["x", "a"])]
+-- > insertWithFun (:) (: []) 7 "x" (fromList [(5,["a"]), (3,["b"])]) == fromList [(3, ["b"]), (5, ["a"]), (7, ["x"])]
+-- > insertWithFun (:) (: []) 5 "x" empty                             == singleton 5 ["x"]
+
+insertWithFun :: Ord k => (a -> b -> b) -> (a -> b) -> k -> a -> Map k b -> Map k b
+insertWithFun = go
+  where
+    go :: Ord k => (a -> b -> b) -> (a -> b) -> k -> a -> Map k b -> Map k b
+    go _ g !kx x Tip = singleton kx (g x)
+    go f g !kx x (Bin sy ky y l r) =
+        case compare kx ky of
+            LT -> balanceL ky y (go f g kx x l) r
+            GT -> balanceR ky y l (go f g kx x r)
+            EQ -> Bin sy kx (f x y) l r
+
+#if __GLASGOW_HASKELL__
+{-# INLINABLE insertWithFun #-}
+#else
+{-# INLINE insertWithFun #-}
+#endif
+
 -- | A helper function for 'unionWith'. When the key is already in
 -- the map, the key is left alone, not replaced. The combining
 -- function is flipped--it is applied to the old value and then the

containers/src/Data/Map/Lazy.hs

@@ -113,6 +113,7 @@ module Data.Map.Lazy (
     -- * Insertion
     , insert
     , insertWith
+    , insertWithFun
     , insertWithKey
     , insertLookupWithKey
 

containers/src/Data/Map/Strict.hs

@@ -129,6 +129,7 @@ module Data.Map.Strict
     -- * Insertion
     , insert
     , insertWith
+    , insertWithFun
     , insertWithKey
     , insertLookupWithKey
 

containers/src/Data/Map/Strict/Internal.hs

@@ -110,6 +110,7 @@ module Data.Map.Strict.Internal
     -- ** Insertion
     , insert
     , insertWith
+    , insertWithFun
     , insertWithKey
     , insertLookupWithKey
 
@@ -559,6 +560,37 @@ insertWith = go
 {-# INLINE insertWith #-}
 #endif
 
+-- | /O(log n)/. Insert with two functions: 'f' to combine a value (of type 'a')
+-- with the Map value (of type 'b') to produce a new Map value, and 'g' to inject the
+-- input value of type 'a' into type 'b' if key is not present.
+-- This version can be used to avoid unnecessary boxing when 'b' is for example
+-- some sort of container containing elements of type 'a'.
+-- @'insertWithFun' f g key value mp@
+-- will insert the pair (key, g value) into @mp@ if key does
+-- not exist in the map. If the key does exist, the function will
+-- insert the pair @(key, f new_value old_value)@.
+--
+-- > insertWithFun (:) (: []) 5 "x" (fromList [(5, ["a"]), (3, ["b"])]) == fromList [(3, ["b"]), (5, ["x", "a"])]
+-- > insertWithFun (:) (: []) 7 "x" (fromList [(5, ["a"]), (3, ["b"])]) == fromList [(3, ["b"]), (5, ["a"]), (7, ["x"])]
+-- > insertWithFun (:) (: []) 5 "x" empty                               == singleton 5 ["x"]
+
+insertWithFun :: Ord k => (a -> b -> b) -> (a -> b) -> k -> a -> Map k b -> Map k b
+insertWithFun = go
+  where
+    go :: Ord k => (a -> b -> b) -> (a -> b) -> k -> a -> Map k b -> Map k b
+    go _ g !kx x Tip = let !y' = g x in singleton kx y'
+    go f g !kx x (Bin sy ky y l r) =
+        case compare kx ky of
+            LT -> balanceL ky y (go f g kx x l) r
+            GT -> balanceR ky y l (go f g kx x r)
+            EQ -> let !y' = f x y in Bin sy kx y' l r
+
+#if __GLASGOW_HASKELL__
+{-# INLINABLE insertWithFun #-}
+#else
+{-# INLINE insertWithFun #-}
+#endif
+
 insertWithR :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
 insertWithR = go
   where