Skip to content

Commit 316cde6

Browse files
meooow25meooow25
authored
Add upsert for Map and IntMap (#1145)
This is the missing member of the group of functions for modifying a single entry. * adjust takes (a -> a) * update takes (a -> Maybe a) * upsert takes (Maybe a -> a) * alter takes (Maybe a -> Maybe a)
1 parent 180c65e commit 316cde6

File tree

12 files changed

+124
-0
lines changed

12 files changed

+124
-0
lines changed

‎containers-tests/tests/intmap-properties.hs‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -229,6 +229,7 @@ main = defaultMain $ testGroup "intmap-properties"
229229
, testProperty "adjustWithKey" prop_adjustWithKey
230230
, testProperty "update" prop_update
231231
, testProperty "updateWithKey" prop_updateWithKey
232+
, testProperty "upsert" prop_upsert
232233
, testProperty "updateLookupWithKey" prop_updateLookupWithKey
233234
, testProperty "differenceWith" prop_differenceWith
234235
, testProperty "differenceWithKey" prop_differenceWithKey
@@ -1824,6 +1825,12 @@ prop_updateWithKey f k m = valid m' .&&. m' === m''
18241825
Nothing -> delete k m
18251826
Just x' -> insert k x' m
18261827

1828+
prop_upsert :: Fun (Maybe A) A -> Int -> IntMap A -> Property
1829+
prop_upsert f k m = valid m' .&&. m' == m''
1830+
where
1831+
m' = upsert (applyFun f) k m
1832+
m'' = insert k (applyFun f (lookup k m)) m
1833+
18271834
prop_updateLookupWithKey
18281835
:: Fun (Int, A) (Maybe A) -> Int -> IntMap A -> Property
18291836
prop_updateLookupWithKey f k m = valid m' .&&. r === (lookup k m, m'')

‎containers-tests/tests/intmap-strictness.hs‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -288,6 +288,17 @@ prop_lazyUpdateWithKey fun k m = isNotBottomProp (L.updateWithKey f k m)
288288
where
289289
f = coerce (applyFunc2 fun)
290290

291+
prop_strictUpsert :: Func (Maybe A) (Bot A) -> Int -> IntMap A -> Property
292+
prop_strictUpsert fun k m =
293+
isBottom (M.upsert f k m) === isBottom (M.alter (Just . f) k m)
294+
where
295+
f = coerce (applyFunc fun)
296+
297+
prop_lazyUpsert :: Func (Maybe A) (Bot A) -> Int -> IntMap A -> Property
298+
prop_lazyUpsert fun k m = isNotBottomProp (L.upsert f k m)
299+
where
300+
f = coerce (applyFunc fun)
301+
291302
prop_strictUpdateLookupWithKey
292303
:: Func2 Key A (Maybe (Bot A)) -> Key -> IntMap A -> Property
293304
prop_strictUpdateLookupWithKey fun k m =
@@ -1018,6 +1029,7 @@ tests =
10181029
, testPropStrictLazy "adjust" prop_strictAdjust prop_lazyAdjust
10191030
, testPropStrictLazy "adjustWithKey" prop_strictAdjustWithKey prop_lazyAdjustWithKey
10201031
, testPropStrictLazy "update" prop_strictUpdate prop_lazyUpdate
1032+
, testPropStrictLazy "upsert" prop_strictUpsert prop_lazyUpsert
10211033
, testPropStrictLazy "updateWithKey" prop_strictUpdateWithKey prop_lazyUpdateWithKey
10221034
, testPropStrictLazy "updateLookupWithKey" prop_strictUpdateLookupWithKey prop_lazyUpdateLookupWithKey
10231035
, testPropStrictLazy "alter" prop_strictAlter prop_lazyAlter

‎containers-tests/tests/map-properties.hs‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -276,6 +276,7 @@ main = defaultMain $ testGroup "map-properties"
276276
, testProperty "adjustWithKey" prop_adjustWithKey
277277
, testProperty "update" prop_update
278278
, testProperty "updateWithKey" prop_updateWithKey
279+
, testProperty "upsert" prop_upsert
279280
, testProperty "updateLookupWithKey" prop_updateLookupWithKey
280281
, testProperty "elemAt" prop_elemAt
281282
, testProperty "updateAt" prop_updateAt
@@ -1808,6 +1809,12 @@ prop_updateWithKey f k m = valid m' .&&. m' === m''
18081809
Nothing -> delete k m
18091810
Just x' -> insert k x' m
18101811

1812+
prop_upsert :: Fun (Maybe A) A -> Int -> Map Int A -> Property
1813+
prop_upsert f k m = valid m' .&&. m' == m''
1814+
where
1815+
m' = upsert (applyFun f) k m
1816+
m'' = insert k (applyFun f (lookup k m)) m
1817+
18111818
prop_updateLookupWithKey
18121819
:: Fun (Int, A) (Maybe A) -> Int -> Map Int A -> Property
18131820
prop_updateLookupWithKey f k m =

‎containers-tests/tests/map-strictness.hs‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -425,6 +425,17 @@ prop_lazyUpdateWithKey fun k m = isNotBottomProp (L.updateWithKey f k m)
425425
where
426426
f = coerce (applyFunc2 fun)
427427

428+
prop_strictUpsert :: Func (Maybe A) (Bot A) -> OrdA -> Map OrdA A -> Property
429+
prop_strictUpsert fun k m =
430+
isBottom (M.upsert f k m) === isBottom (M.alter (Just . f) k m)
431+
where
432+
f = coerce (applyFunc fun)
433+
434+
prop_lazyUpsert :: Func (Maybe A) (Bot A) -> OrdA -> Map OrdA A -> Property
435+
prop_lazyUpsert fun k m = isNotBottomProp (L.upsert f k m)
436+
where
437+
f = coerce (applyFunc fun)
438+
428439
prop_strictUpdateLookupWithKey
429440
:: Func2 OrdA A (Maybe (Bot A)) -> OrdA -> Map OrdA A -> Property
430441
prop_strictUpdateLookupWithKey fun k m =
@@ -1162,6 +1173,7 @@ tests =
11621173
, testPropStrictLazy "adjustWithKey" prop_strictAdjustWithKey prop_lazyAdjustWithKey
11631174
, testPropStrictLazy "update" prop_strictUpdate prop_lazyUpdate
11641175
, testPropStrictLazy "updateWithKey" prop_strictUpdateWithKey prop_lazyUpdateWithKey
1176+
, testPropStrictLazy "upsert" prop_strictUpsert prop_lazyUpsert
11651177
, testPropStrictLazy "updateLookupWithKey" prop_strictUpdateLookupWithKey prop_lazyUpdateLookupWithKey
11661178
, testPropStrictLazy "alter" prop_strictAlter prop_lazyAlter
11671179
, testPropStrictLazy "alterF" prop_strictAlterF prop_lazyAlterF

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

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,7 @@ module Data.IntMap.Internal (
122122
, adjustWithKey
123123
, update
124124
, updateWithKey
125+
, upsert
125126
, updateLookupWithKey
126127
, alter
127128
, alterF
@@ -1036,6 +1037,26 @@ updateWithKey f k t@(Tip ky y)
10361037
| otherwise = t
10371038
updateWithKey _ _ Nil = Nil
10381039

1040+
-- | \(O(\min(n,W))\). Update the value at a key or insert a value if the key is
1041+
-- not in the map.
1042+
--
1043+
-- @
1044+
-- let inc = maybe 1 (+1)
1045+
-- upsert inc 100 (fromList [(100,1),(300,2)]) == fromList [(100,2),(300,2)]
1046+
-- upsert inc 200 (fromList [(100,1),(300,2)]) == fromList [(100,1),(200,1),(300,2)]
1047+
-- @
1048+
--
1049+
-- @since FIXME
1050+
upsert :: (Maybe a -> a) -> Key -> IntMap a -> IntMap a
1051+
upsert f !k t@(Bin p l r)
1052+
| nomatch k p = linkKey k (Tip k (f Nothing)) p t
1053+
| left k p = Bin p (upsert f k l) r
1054+
| otherwise = Bin p l (upsert f k r)
1055+
upsert f !k t@(Tip ky y)
1056+
| k == ky = Tip ky (f (Just y))
1057+
| otherwise = link k (Tip k (f Nothing)) ky t
1058+
upsert f !k Nil = Tip k (f Nothing)
1059+
10391060
-- | \(O(\min(n,W))\). Look up and update.
10401061
-- This function returns the original value, if it is updated.
10411062
-- This is different behavior than 'Data.Map.updateLookupWithKey'.

‎containers/src/Data/IntMap/Lazy.hs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -127,6 +127,7 @@ module Data.IntMap.Lazy (
127127
, adjustWithKey
128128
, update
129129
, updateWithKey
130+
, upsert
130131
, updateLookupWithKey
131132
, alter
132133
, alterF

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -145,6 +145,7 @@ module Data.IntMap.Strict (
145145
, adjustWithKey
146146
, update
147147
, updateWithKey
148+
, upsert
148149
, updateLookupWithKey
149150
, alter
150151
, alterF

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

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,7 @@ module Data.IntMap.Strict.Internal (
8888
, adjustWithKey
8989
, update
9090
, updateWithKey
91+
, upsert
9192
, updateLookupWithKey
9293
, alter
9394
, alterF
@@ -514,6 +515,26 @@ updateWithKey f !k t =
514515
| otherwise -> t
515516
Nil -> Nil
516517

518+
-- | \(O(\min(n,W))\). Update the value at a key or insert a value if the key is
519+
-- not in the map.
520+
--
521+
-- @
522+
-- let inc = maybe 1 (+1)
523+
-- upsert inc 100 (fromList [(100,1),(300,2)]) == fromList [(100,2),(300,2)]
524+
-- upsert inc 200 (fromList [(100,1),(300,2)]) == fromList [(100,1),(200,1),(300,2)]
525+
-- @
526+
--
527+
-- @since FIXME
528+
upsert :: (Maybe a -> a) -> Key -> IntMap a -> IntMap a
529+
upsert f !k t@(Bin p l r)
530+
| nomatch k p = linkKey k (Tip k $! f Nothing) p t
531+
| left k p = Bin p (upsert f k l) r
532+
| otherwise = Bin p l (upsert f k r)
533+
upsert f !k t@(Tip ky y)
534+
| k == ky = Tip ky $! f (Just y)
535+
| otherwise = link k (Tip k $! f Nothing) ky t
536+
upsert f !k Nil = Tip k $! f Nothing
537+
517538
-- | \(O(\min(n,W))\). Look up and update.
518539
-- The function returns original value, if it is updated.
519540
-- This is different behavior than 'Data.Map.updateLookupWithKey'.

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

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -163,6 +163,7 @@ module Data.Map.Internal (
163163
, adjustWithKey
164164
, update
165165
, updateWithKey
166+
, upsert
166167
, updateLookupWithKey
167168
, alter
168169
, alterF
@@ -1120,6 +1121,25 @@ updateWithKey = go
11201121
{-# INLINE updateWithKey #-}
11211122
#endif
11221123

1124+
-- | \(O(\log n)\). Update the value at a key or insert a value if the key is
1125+
-- not in the map.
1126+
--
1127+
-- @
1128+
-- let inc = maybe 1 (+1)
1129+
-- upsert inc \'a\' (fromList [(\'a\',1),(\'c\',2)]) == fromList [(\'a\',2),(\'c\',2)]
1130+
-- upsert inc \'b\' (fromList [(\'a\',1),(\'c\',2)]) == fromList [(\'a\',1),(\'b\',1),(\'c\',2)]
1131+
-- @
1132+
--
1133+
-- @since FIXME
1134+
upsert :: Ord k => (Maybe a -> a) -> k -> Map k a -> Map k a
1135+
upsert f !k (Bin sz kx x l r) =
1136+
case compare k kx of
1137+
LT -> balanceL kx x (upsert f k l) r
1138+
EQ -> Bin sz kx (f (Just x)) l r
1139+
GT -> balanceR kx x l (upsert f k r)
1140+
upsert f !k Tip = singleton k (f Nothing)
1141+
{-# INLINABLE upsert #-}
1142+
11231143
-- | \(O(\log n)\). Look up and update. See also 'updateWithKey'.
11241144
-- This function returns the changed value, if it is updated.
11251145
-- Returns the original key value if the map entry is deleted.

‎containers/src/Data/Map/Lazy.hs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,7 @@ module Data.Map.Lazy (
137137
, adjustWithKey
138138
, update
139139
, updateWithKey
140+
, upsert
140141
, updateLookupWithKey
141142
, alter
142143
, alterF

0 commit comments

Comments
 (0)