Skip to content

Commit d82a404

Browse files
meooow25meooow25
authored
Add from{Asc,Desc}ListUpsert for Map and IntMap (#1199)
As with fromList, this is a less error-prone alternative to from{Asc,Desc}ListWith functions.
1 parent 4a0ae42 commit d82a404

12 files changed

Lines changed: 330 additions & 0 deletions

File tree

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

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -265,8 +265,10 @@ main = defaultMain $ testGroup "intmap-properties"
265265
, testProperty "fromAscList" prop_fromAscList
266266
, testProperty "fromAscListWith" prop_fromAscListWith
267267
, testProperty "fromAscListWithKey" prop_fromAscListWithKey
268+
, testProperty "fromAscListUpsert" prop_fromAscListUpsert
268269
, testProperty "fromDistinctAscList" prop_fromDistinctAscList
269270
, testProperty "fromDescList" prop_fromDescList
271+
, testProperty "fromDescListUpsert" prop_fromDescListUpsert
270272
, testProperty "fromListWith" prop_fromListWith
271273
, testProperty "fromListWithKey" prop_fromListWithKey
272274
, testProperty "fromListUpsert" prop_fromListUpsert
@@ -2173,6 +2175,14 @@ prop_fromAscListWithKey f kxs =
21732175
sortedKxs = List.sortBy (comparing fst) kxs
21742176
t = fromAscListWithKey (applyFun3 f) sortedKxs
21752177

2178+
prop_fromAscListUpsert :: Fun (A, Maybe B) B -> [(Int, A)] -> Property
2179+
prop_fromAscListUpsert f kxs =
2180+
valid t .&&.
2181+
t === fromListUpsert (applyFun2 f) sortedKxs
2182+
where
2183+
sortedKxs = List.sortBy (comparing fst) kxs
2184+
t = fromAscListUpsert (applyFun2 f) sortedKxs
2185+
21762186
prop_fromDistinctAscList :: [(Int, A)] -> Property
21772187
prop_fromDistinctAscList kxs =
21782188
valid t .&&.
@@ -2192,6 +2202,14 @@ prop_fromDescList kxs =
21922202
sortedKxs = List.sortBy (comparing (Down . fst)) kxs
21932203
t = fromDescList sortedKxs
21942204

2205+
prop_fromDescListUpsert :: Fun (A, Maybe B) B -> [(Int, A)] -> Property
2206+
prop_fromDescListUpsert f kxs =
2207+
valid t .&&.
2208+
t === fromListUpsert (applyFun2 f) sortedKxs
2209+
where
2210+
sortedKxs = List.sortBy (comparing (Down . fst)) kxs
2211+
t = fromDescListUpsert (applyFun2 f) sortedKxs
2212+
21952213
prop_fromListWith :: Fun (A, A) A -> [(Int, A)] -> Property
21962214
prop_fromListWith f kxs =
21972215
valid m .&&.

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

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -203,6 +203,22 @@ prop_lazyFromAscListWithKey fun kvs =
203203
f = coerce (applyFunc3 fun)
204204
kvs' = List.sortBy (comparing fst) (coerce kvs) :: [(Key, A)]
205205

206+
prop_strictFromAscListUpsert
207+
:: Func2 A (Maybe B) (Bot B) -> [(Key, Bot A)] -> Property
208+
prop_strictFromAscListUpsert fun kvs =
209+
isBottom (M.fromAscListUpsert f kvs') === isBottom (M.fromListUpsert f kvs')
210+
where
211+
f = coerce (applyFunc2 fun) :: A -> Maybe B -> B
212+
kvs' = List.sortBy (comparing fst) (coerce kvs) :: [(Key, A)]
213+
214+
prop_lazyFromAscListUpsert
215+
:: Func2 A (Maybe B) (Bot B) -> [(Key, Bot A)] -> Property
216+
prop_lazyFromAscListUpsert fun kvs =
217+
isNotBottomProp (L.fromAscListUpsert f kvs')
218+
where
219+
f = coerce (applyFunc2 fun) :: A -> Maybe B -> B
220+
kvs' = List.sortBy (comparing fst) (coerce kvs) :: [(Key, A)]
221+
206222
prop_strictFromDistinctAscList :: [(Key, Bot A)] -> Property
207223
prop_strictFromDistinctAscList kvs =
208224
isBottom (M.fromDistinctAscList kvs') === isBottom (M.fromList kvs')
@@ -225,6 +241,22 @@ prop_lazyFromDescList kvs = isNotBottomProp (L.fromDescList kvs')
225241
where
226242
kvs' = List.sortBy (comparing (Down . fst)) (coerce kvs) :: [(Key, A)]
227243

244+
prop_strictFromDescListUpsert
245+
:: Func2 A (Maybe B) (Bot B) -> [(Key, Bot A)] -> Property
246+
prop_strictFromDescListUpsert fun kvs =
247+
isBottom (M.fromDescListUpsert f kvs') === isBottom (M.fromListUpsert f kvs')
248+
where
249+
f = coerce (applyFunc2 fun) :: A -> Maybe B -> B
250+
kvs' = List.sortBy (comparing (Down . fst)) (coerce kvs) :: [(Key, A)]
251+
252+
prop_lazyFromDescListUpsert
253+
:: Func2 A (Maybe B) (Bot B) -> [(Key, Bot A)] -> Property
254+
prop_lazyFromDescListUpsert fun kvs =
255+
isNotBottomProp (L.fromDescListUpsert f kvs')
256+
where
257+
f = coerce (applyFunc2 fun) :: A -> Maybe B -> B
258+
kvs' = List.sortBy (comparing (Down . fst)) (coerce kvs) :: [(Key, A)]
259+
228260
prop_strictInsert :: Key -> Bot A -> IntMap A -> Property
229261
prop_strictInsert k (Bot x) m = isBottom (M.insert k x m) === isBottom x
230262

@@ -1080,9 +1112,11 @@ tests =
10801112
, testPropStrictLazy "fromListUpsert" prop_strictFromListUpsert prop_lazyFromListUpsert
10811113
, testPropStrictLazy "fromAscList" prop_strictFromAscList prop_lazyFromAscList
10821114
, testPropStrictLazy "fromAscListWith" prop_strictFromAscListWith prop_lazyFromAscListWith
1115+
, testPropStrictLazy "fromAscListUpsert" prop_strictFromAscListUpsert prop_lazyFromAscListUpsert
10831116
, testPropStrictLazy "fromAscListWithKey" prop_strictFromAscListWithKey prop_lazyFromAscListWithKey
10841117
, testPropStrictLazy "fromDistinctAscList" prop_strictFromDistinctAscList prop_lazyFromDistinctAscList
10851118
, testPropStrictLazy "fromDescList" prop_strictFromDescList prop_lazyFromDescList
1119+
, testPropStrictLazy "fromDescListUpsert" prop_strictFromDescListUpsert prop_lazyFromDescListUpsert
10861120
, testPropStrictLazy "insert" prop_strictInsert prop_lazyInsert
10871121
, testPropStrictLazy "insertWith" prop_strictInsertWith prop_lazyInsertWith
10881122
, testPropStrictLazy "insertWithKey" prop_strictInsertWithKey prop_lazyInsertWithKey

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

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -201,10 +201,12 @@ main = defaultMain $ testGroup "map-properties"
201201
, testProperty "fromAscList" prop_fromAscList
202202
, testProperty "fromAscListWith" prop_fromAscListWith
203203
, testProperty "fromAscListWithKey" prop_fromAscListWithKey
204+
, testProperty "fromAscListUpsert" prop_fromAscListUpsert
204205
, testProperty "fromDistinctAscList" prop_fromDistinctAscList
205206
, testProperty "fromDescList" prop_fromDescList
206207
, testProperty "fromDescListWith" prop_fromDescListWith
207208
, testProperty "fromDescListWithKey" prop_fromDescListWithKey
209+
, testProperty "fromDescListUpsert" prop_fromDescListUpsert
208210
, testProperty "fromDistinctDescList" prop_fromDistinctDescList
209211
, testProperty "fromList then toList" prop_list
210212
, testProperty "toDescList" prop_descList
@@ -1419,6 +1421,14 @@ prop_fromDistinctDescList kxs =
14191421
List.sortBy (comparing (Down . fst)) kxs
14201422
t = fromDistinctDescList nubDownSortedKxs
14211423

1424+
prop_fromDescListUpsert :: Fun (A, Maybe B) B -> [(Int, A)] -> Property
1425+
prop_fromDescListUpsert f kxs =
1426+
valid t .&&.
1427+
t === fromListUpsert (applyFun2 f) downSortedKxs
1428+
where
1429+
downSortedKxs = List.sortBy (comparing (Down . fst)) kxs
1430+
t = fromDescListUpsert (applyFun2 f) downSortedKxs
1431+
14221432
prop_ascDescList :: [Int] -> Bool
14231433
prop_ascDescList xs = toAscList m == reverse (toDescList m)
14241434
where m = fromList $ zip xs $ repeat ()
@@ -1454,6 +1464,14 @@ prop_fromAscListWithKey f kxs =
14541464
sortedKxs = List.sortBy (comparing fst) kxs
14551465
t = fromAscListWithKey (apply3 f) sortedKxs
14561466

1467+
prop_fromAscListUpsert :: Fun (A, Maybe B) B -> [(Int, A)] -> Property
1468+
prop_fromAscListUpsert f kxs =
1469+
valid t .&&.
1470+
t === fromListUpsert (applyFun2 f) sortedKxs
1471+
where
1472+
sortedKxs = List.sortBy (comparing fst) kxs
1473+
t = fromAscListUpsert (applyFun2 f) sortedKxs
1474+
14571475
prop_fromDistinctAscList :: [(Int, A)] -> Property
14581476
prop_fromDistinctAscList kxs =
14591477
valid t .&&.

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

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -283,6 +283,22 @@ prop_lazyFromAscListWithKey fun kvs =
283283
f = coerce (applyFunc3 fun)
284284
kvs' = List.sortBy (comparing fst) (coerce kvs) :: [(OrdA, A)]
285285

286+
prop_strictFromAscListUpsert
287+
:: Func2 A (Maybe B) (Bot B) -> [(OrdA, Bot A)] -> Property
288+
prop_strictFromAscListUpsert fun kvs =
289+
isBottom (M.fromAscListUpsert f kvs') === isBottom (M.fromListUpsert f kvs')
290+
where
291+
f = coerce (applyFunc2 fun) :: A -> Maybe B -> B
292+
kvs' = List.sortBy (comparing fst) (coerce kvs) :: [(OrdA, A)]
293+
294+
prop_lazyFromAscListUpsert
295+
:: Func2 A (Maybe B) (Bot B) -> [(OrdA, Bot A)] -> Property
296+
prop_lazyFromAscListUpsert fun kvs =
297+
isNotBottomProp (L.fromAscListUpsert f kvs')
298+
where
299+
f = coerce (applyFunc2 fun) :: A -> Maybe B -> B
300+
kvs' = List.sortBy (comparing fst) (coerce kvs) :: [(OrdA, A)]
301+
286302
prop_strictFromDistinctAscList :: [(OrdA, Bot A)] -> Property
287303
prop_strictFromDistinctAscList kvs =
288304
isBottom (M.fromDistinctAscList kvs') === isBottom (M.fromList kvs')
@@ -335,6 +351,22 @@ prop_lazyFromDescListWithKey fun kvs =
335351
f = coerce (applyFunc3 fun)
336352
kvs' = List.sortBy (comparing (Down . fst)) (coerce kvs) :: [(OrdA, A)]
337353

354+
prop_strictFromDescListUpsert
355+
:: Func2 A (Maybe B) (Bot B) -> [(OrdA, Bot A)] -> Property
356+
prop_strictFromDescListUpsert fun kvs =
357+
isBottom (M.fromDescListUpsert f kvs') === isBottom (M.fromListUpsert f kvs')
358+
where
359+
f = coerce (applyFunc2 fun) :: A -> Maybe B -> B
360+
kvs' = List.sortBy (comparing (Down . fst)) (coerce kvs) :: [(OrdA, A)]
361+
362+
prop_lazyFromDescListUpsert
363+
:: Func2 A (Maybe B) (Bot B) -> [(OrdA, Bot A)] -> Property
364+
prop_lazyFromDescListUpsert fun kvs =
365+
isNotBottomProp (L.fromDescListUpsert f kvs')
366+
where
367+
f = coerce (applyFunc2 fun) :: A -> Maybe B -> B
368+
kvs' = List.sortBy (comparing (Down . fst)) (coerce kvs) :: [(OrdA, A)]
369+
338370
prop_strictFromDistinctDescList :: [(OrdA, Bot A)] -> Property
339371
prop_strictFromDistinctDescList kvs =
340372
isBottom (M.fromDistinctDescList kvs') === isBottom (M.fromList kvs')
@@ -1209,10 +1241,12 @@ tests =
12091241
, testPropStrictLazy "fromAscList" prop_strictFromAscList prop_lazyFromAscList
12101242
, testPropStrictLazy "fromAscListWith" prop_strictFromAscListWith prop_lazyFromAscListWith
12111243
, testPropStrictLazy "fromAscListWithKey" prop_strictFromAscListWithKey prop_lazyFromAscListWithKey
1244+
, testPropStrictLazy "fromAscListUpsert" prop_strictFromAscListUpsert prop_lazyFromAscListUpsert
12121245
, testPropStrictLazy "fromDistinctAscList" prop_strictFromDistinctAscList prop_lazyFromDistinctAscList
12131246
, testPropStrictLazy "fromDescList" prop_strictFromDescList prop_lazyFromDescList
12141247
, testPropStrictLazy "fromDescListWith" prop_strictFromDescListWith prop_lazyFromDescListWith
12151248
, testPropStrictLazy "fromDescListWithKey" prop_strictFromDescListWithKey prop_lazyFromDescListWithKey
1249+
, testPropStrictLazy "fromDescListUpsert" prop_strictFromDescListUpsert prop_lazyFromDescListUpsert
12161250
, testPropStrictLazy "fromDistinctDescList" prop_strictFromDistinctDescList prop_lazyFromDistinctDescList
12171251
, testPropStrictLazy "insert" prop_strictInsert prop_lazyInsert
12181252
, testPropStrictLazy "insertWith" prop_strictInsertWith prop_lazyInsertWith

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

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -237,8 +237,10 @@ module Data.IntMap.Internal (
237237
, fromAscList
238238
, fromAscListWith
239239
, fromAscListWithKey
240+
, fromAscListUpsert
240241
, fromDistinctAscList
241242
, fromDescList
243+
, fromDescListUpsert
242244

243245
-- * Filter
244246
, filter
@@ -3588,6 +3590,8 @@ fromAscList xs =
35883590
-- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")]
35893591
--
35903592
-- Also see the performance note on 'fromListWith'.
3593+
--
3594+
-- See also: 'fromAscListUpsert'
35913595

35923596
fromAscListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a
35933597
fromAscListWith f xs = fromAscListWithKey (\_ x y -> f x y) xs
@@ -3605,6 +3609,8 @@ fromAscListWith f xs = fromAscListWithKey (\_ x y -> f x y) xs
36053609
-- > fromAscListWithKey f [] == empty
36063610
--
36073611
-- Also see the performance note on 'fromListWith'.
3612+
--
3613+
-- See also: 'fromAscListUpsert'
36083614

36093615
-- See Note [fromAscList implementation]
36103616
fromAscListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a
@@ -3618,6 +3624,29 @@ fromAscListWithKey f xs = ascLinkAll (Foldable.foldl' next MSNada xs)
36183624
in MSPush ky y (ascLinkTop stk kx (Tip kx x) m)
36193625
{-# INLINE fromAscListWithKey #-} -- Inline for list fusion
36203626

3627+
-- | \(O(n)\). Build a map from an ascending list in linear time with a
3628+
-- combining function for equal keys.
3629+
--
3630+
-- __Warning__: This function should be used only if the keys are in
3631+
-- non-decreasing order. This precondition is not checked. Use 'fromListUpsert'
3632+
-- if the precondition may not hold.
3633+
--
3634+
-- > let f x = maybe [x] (x:)
3635+
-- > fromAscListUpsert f [(3,'a'), (3,'b'), (5,'c'), (5,'d'), (5,'e')] == fromList [(3,"ba"), (5,"edc")]
3636+
--
3637+
-- @since FIXME
3638+
fromAscListUpsert :: (a -> Maybe b -> b) -> [(Key, a)] -> IntMap b
3639+
fromAscListUpsert f xs = ascLinkAll (Foldable.foldl' next MSNada xs)
3640+
where
3641+
next s (!ky, y) = case s of
3642+
MSNada -> MSPush ky (f y Nothing) Nada
3643+
MSPush kx x stk
3644+
| kx == ky -> MSPush ky (f y (Just x)) stk
3645+
| otherwise ->
3646+
let m = branchMask kx ky
3647+
in MSPush ky (f y Nothing) (ascLinkTop stk kx (Tip kx x) m)
3648+
{-# INLINE fromAscListUpsert #-} -- Inline for list fusion
3649+
36213650
-- | \(O(n)\). Build a map from a list of key\/value pairs where
36223651
-- the keys are in ascending order and all distinct.
36233652
--
@@ -3652,6 +3681,29 @@ fromDescList xs =
36523681
descLinkAll (Foldable.foldl' (\s (ky, y) -> descInsert ky y s) MSNada xs)
36533682
{-# INLINE fromDescList #-} -- Inline for list fusion
36543683

3684+
-- | \(O(n)\). Build a map from a descending list in linear time with a
3685+
-- combining function for equal keys.
3686+
--
3687+
-- __Warning__: This function should be used only if the keys are in
3688+
-- non-increasing order. This precondition is not checked. Use 'fromListUpsert'
3689+
-- if the precondition may not hold.
3690+
--
3691+
-- > let f x = maybe [x] (x:)
3692+
-- > fromDescListUpsert f [(5,'a'), (5,'b'), (5,'c'), (3,'d'), (3,'e')] == fromList [(3,"ed"), (5,"cba")]
3693+
--
3694+
-- @since FIXME
3695+
fromDescListUpsert :: (a -> Maybe b -> b) -> [(Key, a)] -> IntMap b
3696+
fromDescListUpsert f xs = descLinkAll (Foldable.foldl' next MSNada xs)
3697+
where
3698+
next s (!ky, y) = case s of
3699+
MSNada -> MSPush ky (f y Nothing) Nada
3700+
MSPush kx x stk
3701+
| kx == ky -> MSPush ky (f y (Just x)) stk
3702+
| otherwise ->
3703+
let m = branchMask kx ky
3704+
in MSPush ky (f y Nothing) (descLinkTop kx (Tip kx x) m stk)
3705+
{-# INLINE fromDescListUpsert #-} -- Inline for list fusion
3706+
36553707
data Stack a
36563708
= Nada
36573709
| Push {-# UNPACK #-} !Int !(IntMap a) !(Stack a)

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -115,8 +115,10 @@ module Data.IntMap.Lazy (
115115
, fromAscList
116116
, fromAscListWith
117117
, fromAscListWithKey
118+
, fromAscListUpsert
118119
, fromDistinctAscList
119120
, fromDescList
121+
, fromDescListUpsert
120122

121123
-- * Insertion
122124
, insert

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -133,8 +133,10 @@ module Data.IntMap.Strict (
133133
, fromAscList
134134
, fromAscListWith
135135
, fromAscListWithKey
136+
, fromAscListUpsert
136137
, fromDistinctAscList
137138
, fromDescList
139+
, fromDescListUpsert
138140

139141
-- * Insertion
140142
, insert

0 commit comments

Comments
 (0)