Add fromListUpsert for Map and IntMap

containers-tests/tests/intmap-properties.hs

@@ -264,6 +264,7 @@ main = defaultMain $ testGroup "intmap-properties"
              , testProperty "fromDistinctAscList"  prop_fromDistinctAscList
              , testProperty "fromListWith"         prop_fromListWith
              , testProperty "fromListWithKey"      prop_fromListWithKey
+             , testProperty "fromListUpsert"       prop_fromListUpsert
              , testProperty "compareSize"          prop_compareSize
              ]
 
@@ -2174,5 +2175,12 @@ prop_fromListWithKey f kxs =
   where
     m = fromListWithKey (applyFun3 f) kxs
 
+prop_fromListUpsert :: Fun (A, Maybe B) B -> [(Int, A)] -> Property
+prop_fromListUpsert f kxs =
+  valid m' .&&.
+  m' === List.foldl' (\m (kx,x) -> upsert (applyFun2 f x) kx m) empty kxs
+  where
+    m' = fromListUpsert (applyFun2 f) kxs
+
 prop_compareSize :: IntMap A -> Int -> Property
 prop_compareSize t c = compareSize t c === compare (size t) c

containers-tests/tests/intmap-strictness.hs

@@ -141,6 +141,22 @@ prop_lazyFromListWithKey fun kvs = isNotBottomProp (L.fromListWithKey f kvs')
     f = coerce (applyFunc3 fun)
     kvs' = coerce kvs :: [(Key, A)]
 
+prop_strictFromListUpsert
+  :: Func2 A (Maybe B) (Bot B) -> [(Key, Bot A)] -> Property
+prop_strictFromListUpsert fun kvs =
+  isBottom (M.fromListUpsert f kvs') ===
+  isBottom (F.foldl' (\acc (k,x) -> M.upsert (f x) k acc) M.empty kvs')
+  where
+    f = coerce (applyFunc2 fun) :: A -> Maybe B -> B
+    kvs' = coerce kvs :: [(Key, A)]
+
+prop_lazyFromListUpsert
+  :: Func2 A (Maybe B) (Bot B) -> [(Key, Bot A)] -> Property
+prop_lazyFromListUpsert fun kvs = isNotBottomProp (L.fromListUpsert f kvs')
+  where
+    f = coerce (applyFunc2 fun) :: A -> Maybe B -> B
+    kvs' = coerce kvs :: [(Key, A)]
+
 prop_strictFromAscList :: [(Key, Bot A)] -> Property
 prop_strictFromAscList kvs =
   isBottom (M.fromAscList kvs') === isBottom (M.fromList kvs')
@@ -1045,6 +1061,7 @@ tests =
       , testPropStrictLazy "fromList" prop_strictFromList prop_lazyFromList
       , testPropStrictLazy "fromListWith" prop_strictFromListWith prop_lazyFromListWith
       , testPropStrictLazy "fromListWithKey" prop_strictFromListWithKey prop_lazyFromListWithKey
+      , testPropStrictLazy "fromListUpsert" prop_strictFromListUpsert prop_lazyFromListUpsert
       , testPropStrictLazy "fromAscList" prop_strictFromAscList prop_lazyFromAscList
       , testPropStrictLazy "fromAscListWith" prop_strictFromAscListWith prop_lazyFromAscListWith
       , testPropStrictLazy "fromAscListWithKey" prop_strictFromAscListWithKey prop_lazyFromAscListWithKey

containers-tests/tests/map-properties.hs

@@ -208,6 +208,7 @@ main = defaultMain $ testGroup "map-properties"
          , testProperty "fromList"             prop_fromList
          , testProperty "fromListWith"         prop_fromListWith
          , testProperty "fromListWithKey"      prop_fromListWithKey
+         , testProperty "fromListUpsert"       prop_fromListUpsert
          , testProperty "alter"                prop_alter
          , testProperty "alterF/alter"         prop_alterF_alter
          , testProperty "alterF/alter/noRULES" prop_alterF_alter_noRULES
@@ -1457,6 +1458,13 @@ prop_fromListWithKey f kxs =
   fromListWithKey (applyFun3 f) kxs ===
   List.foldl' (\m (kx, x) -> insertWithKey (applyFun3 f) kx x m) empty kxs
 
+prop_fromListUpsert :: Fun (A, Maybe B) B -> [(Int, A)] -> Property
+prop_fromListUpsert f kxs =
+  valid m' .&&.
+  m' === List.foldl' (\m (kx,x) -> upsert (applyFun2 f x) kx m) empty kxs
+  where
+    m' = fromListUpsert (applyFun2 f) kxs
+
 ----------------------------------------------------------------
 
 prop_alter :: UMap -> Int -> Bool

containers-tests/tests/map-strictness.hs

@@ -221,6 +221,22 @@ prop_lazyFromListWithKey fun kvs = isNotBottomProp (L.fromListWithKey f kvs')
     f = coerce (applyFunc3 fun)
     kvs' = coerce kvs :: [(OrdA, A)]
 
+prop_strictFromListUpsert
+  :: Func2 A (Maybe B) (Bot B) -> [(OrdA, Bot A)] -> Property
+prop_strictFromListUpsert fun kvs =
+  isBottom (M.fromListUpsert f kvs') ===
+  isBottom (F.foldl' (\acc (k,x) -> M.upsert (f x) k acc) M.empty kvs')
+  where
+    f = coerce (applyFunc2 fun) :: A -> Maybe B -> B
+    kvs' = coerce kvs :: [(OrdA, A)]
+
+prop_lazyFromListUpsert
+  :: Func2 A (Maybe B) (Bot B) -> [(OrdA, Bot A)] -> Property
+prop_lazyFromListUpsert fun kvs = isNotBottomProp (L.fromListUpsert f kvs')
+  where
+    f = coerce (applyFunc2 fun) :: A -> Maybe B -> B
+    kvs' = coerce kvs :: [(OrdA, A)]
+
 prop_strictFromAscList :: [(OrdA, Bot A)] -> Property
 prop_strictFromAscList kvs =
   isBottom (M.fromAscList kvs') === isBottom (M.fromList kvs')
@@ -1184,6 +1200,7 @@ tests =
       , testPropStrictLazy "fromList" prop_strictFromList prop_lazyFromList
       , testPropStrictLazy "fromListWith" prop_strictFromListWith prop_lazyFromListWith
       , testPropStrictLazy "fromListWithKey" prop_strictFromListWithKey prop_lazyFromListWithKey
+      , testPropStrictLazy "fromListUpsert" prop_strictFromListUpsert prop_lazyFromListUpsert
       , testPropStrictLazy "fromAscList" prop_strictFromAscList prop_lazyFromAscList
       , testPropStrictLazy "fromAscListWith" prop_strictFromAscListWith prop_lazyFromAscListWith
       , testPropStrictLazy "fromAscListWithKey" prop_strictFromAscListWithKey prop_lazyFromAscListWithKey

containers/src/Data/IntMap/Internal.hs

@@ -229,6 +229,7 @@ module Data.IntMap.Internal (
     , fromList
     , fromListWith
     , fromListWithKey
+    , fromListUpsert
 
     -- ** Ordered lists
     , toAscList
@@ -3472,7 +3473,7 @@ fromList :: [(Key,a)] -> IntMap a
 fromList xs = finishB (Foldable.foldl' (\b (kx,x) -> insertB kx x b) emptyB xs)
 {-# INLINE fromList #-} -- Inline for list fusion
 
--- | \(O(n \min(n,W))\). Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
+-- | \(O(n \min(n,W))\). Build a map from a list of key\/value pairs with a combining function.
 --
 -- If the keys are in sorted order, ascending or descending, this function
 -- takes \(O(n)\) time.
@@ -3484,6 +3485,8 @@ fromList xs = finishB (Foldable.foldl' (\b (kx,x) -> insertB kx x b) emptyB xs)
 --
 -- The symmetric combining function @f@ is applied in a left-fold over the list, as @f new old@.
 --
+-- See also: 'fromListUpsert'
+--
 -- === Performance
 --
 -- You should ensure that the given @f@ is fast with this order of arguments.
@@ -3516,7 +3519,7 @@ fromListWith f xs
   = fromListWithKey (\_ x y -> f x y) xs
 {-# INLINE fromListWith #-} -- Inline for list fusion
 
--- | \(O(n \min(n,W))\). Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWithKey'.
+-- | \(O(n \min(n,W))\). Build a map from a list of key\/value pairs with a combining function.
 --
 -- If the keys are in sorted order, ascending or descending, this function
 -- takes \(O(n)\) time.
@@ -3526,12 +3529,36 @@ fromListWith f xs
 -- > fromListWithKey f [] == empty
 --
 -- Also see the performance note on 'fromListWith'.
+--
+-- See also: 'fromListUpsert'
 
 fromListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
 fromListWithKey f xs =
   finishB (Foldable.foldl' (\b (kx,x) -> insertWithB (f kx) kx x b) emptyB xs)
 {-# INLINE fromListWithKey #-} -- Inline for list fusion
 
+-- | \(O(n \min(n,W)\). Build a map from a list of key\/value pairs with a
+-- combining function.
+--
+-- If the keys are in sorted order, ascending or descending, this function
+-- takes \(O(n)\) time.
+--
+-- The result is equivalent to performing an @upsert@ for every key\/value in
+-- the list.
+--
+-- @
+-- fromListUpsert f = foldl' (\\m (k, x) -> 'upsert' (f x) k m) 'empty'
+-- @
+--
+-- > let f x = maybe [x] (x:)
+-- > fromListUpsert f [(5,'a'), (5,'b'), (3,'c'), (3,'d'), (5,'e')] == fromList [(3,"dc"), (5,"eba")]
+--
+-- @since FIXME
+fromListUpsert :: (a -> Maybe b -> b) -> [(Key, a)] -> IntMap b
+fromListUpsert f xs =
+  finishB (Foldable.foldl' (\b (kx, x) -> upsertB (f x) kx b) emptyB xs)
+{-# INLINE fromListUpsert #-}  -- INLINE for fusion
+
 -- | \(O(n)\). Build a map from a list of key\/value pairs where
 -- the keys are in ascending order.
 --
@@ -3716,6 +3743,15 @@ insertWithB f !ky y b = case b of
       Just x' -> BTip ky (f y x') stk'
 {-# INLINE insertWithB #-}
 
+-- Upsert a key-value. The given function is used to generate the value based
+-- on the existing value for the key.
+upsertB :: (Maybe a -> a) -> Key -> IntMapBuilder a -> IntMapBuilder a
+upsertB f !ky b = case b of
+  BNil -> BTip ky (f Nothing) BNada
+  BTip kx x stk -> case moveToB ky kx x stk of
+    MoveResult m stk' -> BTip ky (f m) stk'
+{-# INLINE upsertB #-}
+
 -- GHC >=9.6 supports unpacking sums, so we unpack the Maybe and avoid
 -- allocating Justs. GHC optimizes the workers for moveUpB and moveDownB to
 -- return (# (# (# #) | a #), BStack a #).

containers/src/Data/IntMap/Lazy.hs

@@ -109,6 +109,7 @@ module Data.IntMap.Lazy (
     , fromList
     , fromListWith
     , fromListWithKey
+    , fromListUpsert
 
     -- ** From Ascending Lists
     , fromAscList

containers/src/Data/IntMap/Strict.hs

@@ -127,6 +127,7 @@ module Data.IntMap.Strict (
     , fromList
     , fromListWith
     , fromListWithKey
+    , fromListUpsert
 
     -- ** From Ascending Lists
     , fromAscList

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

@@ -68,6 +68,7 @@ module Data.IntMap.Strict.Internal (
     , fromList
     , fromListWith
     , fromListWithKey
+    , fromListUpsert
 
     -- ** From Ascending Lists
     , fromAscList
@@ -1140,7 +1141,7 @@ fromList :: [(Key,a)] -> IntMap a
 fromList xs = finishB (Foldable.foldl' (\b (kx,!x) -> insertB kx x b) emptyB xs)
 {-# INLINE fromList #-} -- Inline for list fusion
 
--- | \(O(n \min(n,W))\). Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
+-- | \(O(n \min(n,W))\). Build a map from a list of key\/value pairs with a combining function.
 --
 -- If the keys are in sorted order, ascending or descending, this function
 -- takes \(O(n)\) time.
@@ -1152,6 +1153,8 @@ fromList xs = finishB (Foldable.foldl' (\b (kx,!x) -> insertB kx x b) emptyB xs)
 --
 -- The symmetric combining function @f@ is applied in a left-fold over the list, as @f new old@.
 --
+-- See also: 'fromListUpsert'
+--
 -- === Performance
 --
 -- You should ensure that the given @f@ is fast with this order of arguments.
@@ -1184,7 +1187,7 @@ fromListWith f xs
   = fromListWithKey (\_ x y -> f x y) xs
 {-# INLINE fromListWith #-} -- Inline for list fusion
 
--- | \(O(n \min(n,W))\). Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWithKey'.
+-- | \(O(n \min(n,W))\). Build a map from a list of key\/value pairs with a combining function.
 --
 -- If the keys are in sorted order, ascending or descending, this function
 -- takes \(O(n)\) time.
@@ -1194,12 +1197,36 @@ fromListWith f xs
 -- > fromListWithKey f [] == empty
 --
 -- Also see the performance note on 'fromListWith'.
+--
+-- See also: 'fromListUpsert'
 
 fromListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
 fromListWithKey f xs =
   finishB (Foldable.foldl' (\b (kx,x) -> insertWithB (f kx) kx x b) emptyB xs)
 {-# INLINE fromListWithKey #-} -- Inline for list fusion
 
+-- | \(O(n \min(n,W)\). Build a map from a list of key\/value pairs with a
+-- combining function.
+--
+-- If the keys are in sorted order, ascending or descending, this function
+-- takes \(O(n)\) time.
+--
+-- The result is equivalent to performing an @upsert@ for every key\/value in
+-- the list.
+--
+-- @
+-- fromListUpsert f = foldl' (\\m (k, x) -> 'upsert' (f x) k m) 'empty'
+-- @
+--
+-- > let f x = maybe [x] (x:)
+-- > fromListUpsert f [(5,'a'), (5,'b'), (3,'c'), (3,'d'), (5,'e')] == fromList [(3,"dc"), (5,"eba")]
+--
+-- @since FIXME
+fromListUpsert :: (a -> Maybe b -> b) -> [(Key, a)] -> IntMap b
+fromListUpsert f xs =
+  finishB (Foldable.foldl' (\b (kx, x) -> upsertB (f x) kx b) emptyB xs)
+{-# INLINE fromListUpsert #-}  -- INLINE for fusion
+
 -- | \(O(n)\). Build a map from a list of key\/value pairs where
 -- the keys are in ascending order.
 --
@@ -1285,3 +1312,14 @@ insertWithB f !ky y b = case b of
   where
     btip' kx !x = BTip kx x
 {-# INLINE insertWithB #-}
+
+-- Upsert a key-value. The given function is used to generate the value based
+-- on the existing value for the key.
+upsertB :: (Maybe a -> a) -> Key -> IntMapBuilder a -> IntMapBuilder a
+upsertB f !ky b = case b of
+  BNil -> btip' ky (f Nothing) BNada
+  BTip kx x stk -> case moveToB ky kx x stk of
+    MoveResult m stk' -> btip' ky (f m) stk'
+  where
+    btip' kx !x = BTip kx x
+{-# INLINE upsertB #-}

containers/src/Data/Map/Internal.hs

@@ -276,6 +276,7 @@ module Data.Map.Internal (
     , fromList
     , fromListWith
     , fromListWithKey
+    , fromListUpsert
 
     -- ** Ordered lists
     , toAscList
@@ -3457,7 +3458,7 @@ instance (Ord k) => GHCExts.IsList (Map k v) where
   toList   = toList
 #endif
 
--- | \(O(n \log n)\). Build a map from a list of key\/value pairs. See also 'fromAscList'.
+-- | \(O(n \log n)\). Build a map from a list of key\/value pairs.
 -- If the list contains more than one value for the same key, the last value
 -- for the key is retained.
 --
@@ -3471,7 +3472,7 @@ fromList :: Ord k => [(k,a)] -> Map k a
 fromList xs = finishB (Foldable.foldl' (\b (kx, x) -> insertB kx x b) emptyB xs)
 {-# INLINE fromList #-} -- INLINE for fusion
 
--- | \(O(n \log n)\). Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
+-- | \(O(n \log n)\). Build a map from a list of key\/value pairs with a combining function.
 --
 -- If the keys are in non-decreasing order, this function takes \(O(n)\) time.
 --
@@ -3482,6 +3483,8 @@ fromList xs = finishB (Foldable.foldl' (\b (kx, x) -> insertB kx x b) emptyB xs)
 --
 -- The symmetric combining function @f@ is applied in a left-fold over the list, as @f new old@.
 --
+-- See also: 'fromListUpsert'
+--
 -- === Performance
 --
 -- You should ensure that the given @f@ is fast with this order of arguments.
@@ -3514,7 +3517,7 @@ fromListWith f xs =
   finishB (Foldable.foldl' (\b (kx, x) -> insertWithB f kx x b) emptyB xs)
 {-# INLINE fromListWith #-}  -- INLINE for fusion
 
--- | \(O(n \log n)\). Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWithKey'.
+-- | \(O(n \log n)\). Build a map from a list of key\/value pairs with a combining function.
 --
 -- If the keys are in non-decreasing order, this function takes \(O(n)\) time.
 --
@@ -3523,12 +3526,35 @@ fromListWith f xs =
 -- > fromListWithKey f [] == empty
 --
 -- Also see the performance note on 'fromListWith'.
+--
+-- See also: 'fromListUpsert'
 
 fromListWithKey :: Ord k => (k -> a -> a -> a) -> [(k,a)] -> Map k a
 fromListWithKey f xs =
   finishB (Foldable.foldl' (\b (kx, x) -> insertWithB (f kx) kx x b) emptyB xs)
 {-# INLINE fromListWithKey #-}  -- INLINE for fusion
 
+-- | \(O(n \log n)\). Build a map from a list of key\/value pairs with a
+-- combining function.
+--
+-- If the keys are in non-decreasing order, this function takes \(O(n)\) time.
+--
+-- The result is equivalent to performing an @upsert@ for every key\/value in
+-- the list.
+--
+-- @
+-- fromListUpsert f = foldl' (\\m (k, x) -> 'upsert' (f x) k m) 'empty'
+-- @
+--
+-- > let f x = maybe [x] (x:)
+-- > fromListUpsert f [(5,'a'), (5,'b'), (3,'c'), (3,'d'), (5,'e')] == fromList [(3,"dc"), (5,"eba")]
+--
+-- @since FIXME
+fromListUpsert :: Ord k => (a -> Maybe b -> b) -> [(k, a)] -> Map k b
+fromListUpsert f xs =
+  finishB (Foldable.foldl' (\b (kx, x) -> upsertB (f x) kx b) emptyB xs)
+{-# INLINE fromListUpsert #-}  -- INLINE for fusion
+
 -- | \(O(n)\). Convert the map to a list of key\/value pairs. Subject to list fusion.
 --
 -- > toList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")]
@@ -3934,6 +3960,21 @@ insertWithB f !ky y b = case b of
   BMap m -> BMap (insertWith f ky y m)
 {-# INLINE insertWithB #-}
 
+-- Upsert a key-value. The given function is used to generate the value based
+-- on the existing value for the key.
+upsertB :: Ord k => (Maybe a -> a) -> k -> MapBuilder k a -> MapBuilder k a
+upsertB f !ky b = case b of
+  BAsc stk -> case stk of
+    Push kx x l stk' -> case compare ky kx of
+      LT -> BMap (upsert f ky (ascLinkAll stk))
+      EQ -> BAsc (Push ky (f (Just x)) l stk')
+      GT -> case l of
+        Tip -> BAsc (ascLinkTop stk' 1 (singleton kx x) ky (f Nothing))
+        Bin{} -> BAsc (Push ky (f Nothing) Tip stk)
+    Nada -> BAsc (Push ky (f Nothing) Tip Nada)
+  BMap m -> BMap (upsert f ky m)
+{-# INLINE upsertB #-}
+
 -- Finalize the builder into a Map.
 finishB :: MapBuilder k a -> Map k a
 finishB (BAsc stk) = ascLinkAll stk