Add upsert for Map and IntMap

containers-tests/tests/intmap-properties.hs

@@ -229,6 +229,7 @@ main = defaultMain $ testGroup "intmap-properties"
              , testProperty "adjustWithKey"        prop_adjustWithKey
              , testProperty "update"               prop_update
              , testProperty "updateWithKey"        prop_updateWithKey
+             , testProperty "upsert"               prop_upsert
              , testProperty "updateLookupWithKey"  prop_updateLookupWithKey
              , testProperty "differenceWith"       prop_differenceWith
              , testProperty "differenceWithKey"    prop_differenceWithKey
@@ -1824,6 +1825,12 @@ prop_updateWithKey f k m = valid m' .&&. m' === m''
         Nothing -> delete k m
         Just x' -> insert k x' m
 
+prop_upsert :: Fun (Maybe A) A -> Int -> IntMap A -> Property
+prop_upsert f k m = valid m' .&&. m' == m''
+  where
+    m' = upsert (applyFun f) k m
+    m'' = insert k (applyFun f (lookup k m)) m
+
 prop_updateLookupWithKey
   :: Fun (Int, A) (Maybe A) -> Int -> IntMap A -> Property
 prop_updateLookupWithKey f k m = valid m' .&&. r === (lookup k m, m'')

containers-tests/tests/intmap-strictness.hs

@@ -288,6 +288,17 @@ prop_lazyUpdateWithKey fun k m = isNotBottomProp (L.updateWithKey f k m)
   where
     f = coerce (applyFunc2 fun)
 
+prop_strictUpsert :: Func (Maybe A) (Bot A) -> Int -> IntMap A -> Property
+prop_strictUpsert fun k m =
+  isBottom (M.upsert f k m) === isBottom (M.alter (Just . f) k m)
+  where
+    f = coerce (applyFunc fun)
+
+prop_lazyUpsert :: Func (Maybe A) (Bot A) -> Int -> IntMap A -> Property
+prop_lazyUpsert fun k m = isNotBottomProp (L.upsert f k m)
+  where
+    f = coerce (applyFunc fun)
+
 prop_strictUpdateLookupWithKey
   :: Func2 Key A (Maybe (Bot A)) -> Key -> IntMap A -> Property
 prop_strictUpdateLookupWithKey fun k m =
@@ -1018,6 +1029,7 @@ tests =
       , testPropStrictLazy "adjust" prop_strictAdjust prop_lazyAdjust
       , testPropStrictLazy "adjustWithKey" prop_strictAdjustWithKey prop_lazyAdjustWithKey
       , testPropStrictLazy "update" prop_strictUpdate prop_lazyUpdate
+      , testPropStrictLazy "upsert" prop_strictUpsert prop_lazyUpsert
       , testPropStrictLazy "updateWithKey" prop_strictUpdateWithKey prop_lazyUpdateWithKey
       , testPropStrictLazy "updateLookupWithKey" prop_strictUpdateLookupWithKey prop_lazyUpdateLookupWithKey
       , testPropStrictLazy "alter" prop_strictAlter prop_lazyAlter

containers-tests/tests/map-properties.hs

@@ -276,6 +276,7 @@ main = defaultMain $ testGroup "map-properties"
          , testProperty "adjustWithKey"        prop_adjustWithKey
          , testProperty "update"               prop_update
          , testProperty "updateWithKey"        prop_updateWithKey
+         , testProperty "upsert"               prop_upsert
          , testProperty "updateLookupWithKey"  prop_updateLookupWithKey
          , testProperty "elemAt"               prop_elemAt
          , testProperty "updateAt"             prop_updateAt
@@ -1808,6 +1809,12 @@ prop_updateWithKey f k m = valid m' .&&. m' === m''
         Nothing -> delete k m
         Just x' -> insert k x' m
 
+prop_upsert :: Fun (Maybe A) A -> Int -> Map Int A -> Property
+prop_upsert f k m = valid m' .&&. m' == m''
+  where
+    m' = upsert (applyFun f) k m
+    m'' = insert k (applyFun f (lookup k m)) m
+
 prop_updateLookupWithKey
   :: Fun (Int, A) (Maybe A) -> Int -> Map Int A -> Property
 prop_updateLookupWithKey f k m =

containers-tests/tests/map-strictness.hs

@@ -425,6 +425,17 @@ prop_lazyUpdateWithKey fun k m = isNotBottomProp (L.updateWithKey f k m)
   where
     f = coerce (applyFunc2 fun)
 
+prop_strictUpsert :: Func (Maybe A) (Bot A) -> OrdA -> Map OrdA A -> Property
+prop_strictUpsert fun k m =
+  isBottom (M.upsert f k m) === isBottom (M.alter (Just . f) k m)
+  where
+    f = coerce (applyFunc fun)
+
+prop_lazyUpsert :: Func (Maybe A) (Bot A) -> OrdA -> Map OrdA A -> Property
+prop_lazyUpsert fun k m = isNotBottomProp (L.upsert f k m)
+  where
+    f = coerce (applyFunc fun)
+
 prop_strictUpdateLookupWithKey
   :: Func2 OrdA A (Maybe (Bot A)) -> OrdA -> Map OrdA A -> Property
 prop_strictUpdateLookupWithKey fun k m =
@@ -1162,6 +1173,7 @@ tests =
       , testPropStrictLazy "adjustWithKey" prop_strictAdjustWithKey prop_lazyAdjustWithKey
       , testPropStrictLazy "update" prop_strictUpdate prop_lazyUpdate
       , testPropStrictLazy "updateWithKey" prop_strictUpdateWithKey prop_lazyUpdateWithKey
+      , testPropStrictLazy "upsert" prop_strictUpsert prop_lazyUpsert
       , testPropStrictLazy "updateLookupWithKey" prop_strictUpdateLookupWithKey prop_lazyUpdateLookupWithKey
       , testPropStrictLazy "alter" prop_strictAlter prop_lazyAlter
       , testPropStrictLazy "alterF" prop_strictAlterF prop_lazyAlterF

containers/src/Data/IntMap/Internal.hs

@@ -122,6 +122,7 @@ module Data.IntMap.Internal (
     , adjustWithKey
     , update
     , updateWithKey
+    , upsert
     , updateLookupWithKey
     , alter
     , alterF
@@ -1036,6 +1037,26 @@ updateWithKey f k t@(Tip ky y)
   | otherwise     = t
 updateWithKey _ _ Nil = Nil
 
+-- | \(O(\min(n,W))\). Update the value at a key or insert a value if the key is
+-- not in the map.
+--
+-- @
+-- let inc = maybe 1 (+1)
+-- upsert inc 100 (fromList [(100,1),(300,2)]) == fromList [(100,2),(300,2)]
+-- upsert inc 200 (fromList [(100,1),(300,2)]) == fromList [(100,1),(200,1),(300,2)]
+-- @
+--
+-- @since FIXME
+upsert :: (Maybe a -> a) -> Key -> IntMap a -> IntMap a
+upsert f !k t@(Bin p l r)
+  | nomatch k p = linkKey k (Tip k (f Nothing)) p t
+  | left k p = Bin p (upsert f k l) r
+  | otherwise = Bin p l (upsert f k r)
+upsert f !k t@(Tip ky y)
+  | k == ky = Tip ky (f (Just y))
+  | otherwise = link k (Tip k (f Nothing)) ky t
+upsert f !k Nil = Tip k (f Nothing)
+
 -- | \(O(\min(n,W))\). Look up and update.
 -- This function returns the original value, if it is updated.
 -- This is different behavior than 'Data.Map.updateLookupWithKey'.

containers/src/Data/IntMap/Lazy.hs

@@ -127,6 +127,7 @@ module Data.IntMap.Lazy (
     , adjustWithKey
     , update
     , updateWithKey
+    , upsert
     , updateLookupWithKey
     , alter
     , alterF

containers/src/Data/IntMap/Strict.hs

@@ -145,6 +145,7 @@ module Data.IntMap.Strict (
     , adjustWithKey
     , update
     , updateWithKey
+    , upsert
     , updateLookupWithKey
     , alter
     , alterF

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

@@ -88,6 +88,7 @@ module Data.IntMap.Strict.Internal (
     , adjustWithKey
     , update
     , updateWithKey
+    , upsert
     , updateLookupWithKey
     , alter
     , alterF
@@ -514,6 +515,26 @@ updateWithKey f !k t =
       | otherwise     -> t
     Nil -> Nil
 
+-- | \(O(\min(n,W))\). Update the value at a key or insert a value if the key is
+-- not in the map.
+--
+-- @
+-- let inc = maybe 1 (+1)
+-- upsert inc 100 (fromList [(100,1),(300,2)]) == fromList [(100,2),(300,2)]
+-- upsert inc 200 (fromList [(100,1),(300,2)]) == fromList [(100,1),(200,1),(300,2)]
+-- @
+--
+-- @since FIXME
+upsert :: (Maybe a -> a) -> Key -> IntMap a -> IntMap a
+upsert f !k t@(Bin p l r)
+  | nomatch k p = linkKey k (Tip k $! f Nothing) p t
+  | left k p = Bin p (upsert f k l) r
+  | otherwise = Bin p l (upsert f k r)
+upsert f !k t@(Tip ky y)
+  | k == ky = Tip ky $! f (Just y)
+  | otherwise = link k (Tip k $! f Nothing) ky t
+upsert f !k Nil = Tip k $! f Nothing
+
 -- | \(O(\min(n,W))\). Look up and update.
 -- The function returns original value, if it is updated.
 -- This is different behavior than 'Data.Map.updateLookupWithKey'.

containers/src/Data/Map/Internal.hs

@@ -163,6 +163,7 @@ module Data.Map.Internal (
     , adjustWithKey
     , update
     , updateWithKey
+    , upsert
     , updateLookupWithKey
     , alter
     , alterF
@@ -1120,6 +1121,25 @@ updateWithKey = go
 {-# INLINE updateWithKey #-}
 #endif
 
+-- | \(O(\log n)\). Update the value at a key or insert a value if the key is
+-- not in the map.
+--
+-- @
+-- let inc = maybe 1 (+1)
+-- upsert inc \'a\' (fromList [(\'a\',1),(\'c\',2)]) == fromList [(\'a\',2),(\'c\',2)]
+-- upsert inc \'b\' (fromList [(\'a\',1),(\'c\',2)]) == fromList [(\'a\',1),(\'b\',1),(\'c\',2)]
+-- @
+--
+-- @since FIXME
+upsert :: Ord k => (Maybe a -> a) -> k -> Map k a -> Map k a
+upsert f !k (Bin sz kx x l r) =
+  case compare k kx of
+    LT -> balanceL kx x (upsert f k l) r
+    EQ -> Bin sz kx (f (Just x)) l r
+    GT -> balanceR kx x l (upsert f k r)
+upsert f !k Tip = singleton k (f Nothing)
+{-# INLINABLE upsert #-}
+
 -- | \(O(\log n)\). Look up and update. See also 'updateWithKey'.
 -- This function returns the changed value, if it is updated.
 -- Returns the original key value if the map entry is deleted.

containers/src/Data/Map/Lazy.hs

@@ -137,6 +137,7 @@ module Data.Map.Lazy (
     , adjustWithKey
     , update
     , updateWithKey
+    , upsert
     , updateLookupWithKey
     , alter
     , alterF

containers/src/Data/Map/Strict.hs

@@ -151,6 +151,7 @@ module Data.Map.Strict
     , adjustWithKey
     , update
     , updateWithKey
+    , upsert
     , updateLookupWithKey
     , alter
     , alterF

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

@@ -101,6 +101,7 @@ module Data.Map.Strict.Internal
     , adjustWithKey
     , update
     , updateWithKey
+    , upsert
     , updateLookupWithKey
     , alter
     , alterF
@@ -705,6 +706,25 @@ updateWithKey = go
 {-# INLINE updateWithKey #-}
 #endif
 
+-- | \(O(\log n)\). Update the value at a key or insert a value if the key is
+-- not in the map.
+--
+-- @
+-- let inc = maybe 1 (+1)
+-- upsert inc \'a\' (fromList [(\'a\',1),(\'c\',2)]) == fromList [(\'a\',2),(\'c\',2)]
+-- upsert inc \'b\' (fromList [(\'a\',1),(\'c\',2)]) == fromList [(\'a\',1),(\'b\',1),(\'c\',2)]
+-- @
+--
+-- @since FIXME
+upsert :: Ord k => (Maybe a -> a) -> k -> Map k a -> Map k a
+upsert f !k (Bin sz kx x l r) =
+  case compare k kx of
+    LT -> balanceL kx x (upsert f k l) r
+    EQ -> let !x' = f (Just x) in Bin sz kx x' l r
+    GT -> balanceR kx x l (upsert f k r)
+upsert f !k Tip = singleton k (f Nothing)
+{-# INLINABLE upsert #-}
+
 -- | \(O(\log n)\). Look up and update. See also 'updateWithKey'.
 -- This function returns the changed value, if it is updated.
 -- Returns the original key value if the map entry is deleted.