Skip to content

Commit

Permalink
Adjust strictness of from{Asc,Desc}List* for maps
Browse files Browse the repository at this point in the history
Make the functions strict in the first value of the run of equal keys.
This makes the strictness match that of the fromList functions.
  • Loading branch information
meooow25 committed Aug 11, 2024
1 parent 96a8dec commit 60e57c2
Show file tree
Hide file tree
Showing 3 changed files with 57 additions and 39 deletions.
12 changes: 12 additions & 0 deletions containers/changelog.md
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,18 @@
`Data.IntMap.Lazy.splitLookup`, `Data.IntMap.Strict.splitLookup` and
`Data.IntSet.splitMember` are now strict in the key. Previously, the key was
ignored for an empty map or set. (Soumik Sarkar)
* The functions below have increased strictness in values:
* `Data.Map.Strict.fromAscList`
* `Data.Map.Strict.fromAscListWith`
* `Data.Map.Strict.fromAscListWithKey`
* `Data.Map.Strict.fromDescList`
* `Data.Map.Strict.fromDescListWith`
* `Data.Map.Strict.fromDescListWithKey`
* `Data.IntMap.Strict.fromAscList`
* `Data.IntMap.Strict.fromAscListWith`
* `Data.IntMap.Strict.fromAscListWithKey`
These functions now match the strictness of the corresponding `fromList`
functions. (Soumik Sarkar)

## Unreleased with `@since` annotation for 0.7.1:

Expand Down
32 changes: 19 additions & 13 deletions containers/src/Data/IntMap/Strict/Internal.hs
Original file line number Diff line number Diff line change
Expand Up @@ -1203,15 +1203,18 @@ fromMonoListWithKey distinct f = go

-- `addAll'` collects all keys equal to `kx` into a single value,
-- and then proceeds with `addAll`.
addAll' !kx vx []
= Tip kx $! vx
addAll' !kx vx ((ky,vy) : zs)
--
-- We want to have the same strictness as fromListWithKey, which is achieved
-- with the bang on vx.
addAll' !kx !vx []
= Tip kx vx
addAll' !kx !vx ((ky,vy) : zs)
| Nondistinct <- distinct, kx == ky
= let !v = f kx vy vx in addAll' ky v zs
-- inlined: | otherwise = addAll kx (Tip kx $! vx) (ky : zs)
= addAll' ky (f kx vy vx) zs
-- inlined: | otherwise = addAll kx (Tip kx vx) (ky : zs)
| m <- branchMask kx ky
, Inserted ty zs' <- addMany' m ky vy zs
= addAll kx (linkWithMask m ky ty kx (Tip kx $! vx)) zs'
= addAll kx (linkWithMask m ky ty kx (Tip kx vx)) zs'

-- for `addAll` and `addMany`, kx is /a/ key inside the tree `tx`
-- `addAll` consumes the rest of the list, adding to the tree `tx`
Expand All @@ -1223,17 +1226,20 @@ fromMonoListWithKey distinct f = go
= addAll kx (linkWithMask m ky ty kx tx) zs'

-- `addMany'` is similar to `addAll'`, but proceeds with `addMany'`.
addMany' !_m !kx vx []
= Inserted (Tip kx $! vx) []
addMany' !m !kx vx zs0@((ky,vy) : zs)
--
-- We want to have the same strictness as fromListWithKey, which is achieved
-- with the bang on vx.
addMany' !_m !kx !vx []
= Inserted (Tip kx vx) []
addMany' !m !kx !vx zs0@((ky,vy) : zs)
| Nondistinct <- distinct, kx == ky
= let !v = f kx vy vx in addMany' m ky v zs
-- inlined: | otherwise = addMany m kx (Tip kx $! vx) (ky : zs)
= addMany' m ky (f kx vy vx) zs
-- inlined: | otherwise = addMany m kx (Tip kx vx) (ky : zs)
| mask kx m /= mask ky m
= Inserted (Tip kx $! vx) zs0
= Inserted (Tip kx vx) zs0
| mxy <- branchMask kx ky
, Inserted ty zs' <- addMany' mxy ky vy zs
= addMany m kx (linkWithMask mxy ky ty kx (Tip kx $! vx)) zs'
= addMany m kx (linkWithMask mxy ky ty kx (Tip kx vx)) zs'

-- `addAll` adds to `tx` all keys whose prefix w.r.t. `m` agrees with `kx`.
addMany !_m !_kx tx []
Expand Down
52 changes: 26 additions & 26 deletions containers/src/Data/Map/Strict/Internal.hs
Original file line number Diff line number Diff line change
Expand Up @@ -1677,20 +1677,20 @@ fromDescListWith f xs
-- Also see the performance note on 'fromListWith'.

fromAscListWithKey :: Eq k => (k -> a -> a -> a) -> [(k,a)] -> Map k a
fromAscListWithKey f xs
= fromDistinctAscList (combineEq f xs)
fromAscListWithKey f xs0 = fromDistinctAscList xs1
where
-- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs]
combineEq _ xs'
= case xs' of
[] -> []
[x] -> [x]
(x:xx) -> combineEq' x xx

combineEq' z [] = [z]
combineEq' z@(kz,zz) (x@(kx,xx):xs')
| kx==kz = let yy = f kx xx zz in yy `seq` combineEq' (kx,yy) xs'
| otherwise = z:combineEq' x xs'
xs1 = case xs0 of
[] -> []
[x] -> [x]
x:xs -> combineEq x xs

-- We want to have the same strictness as fromListWithKey, which is achieved
-- with the bang on yy.
combineEq y@(ky, !yy) xs = case xs of
[] -> [y]
x@(kx, xx) : xs'
| kx == ky -> combineEq (kx, f kx xx yy) xs'
| otherwise -> y : combineEq x xs'
#if __GLASGOW_HASKELL__
{-# INLINABLE fromAscListWithKey #-}
#endif
Expand All @@ -1707,20 +1707,20 @@ fromAscListWithKey f xs
-- Also see the performance note on 'fromListWith'.

fromDescListWithKey :: Eq k => (k -> a -> a -> a) -> [(k,a)] -> Map k a
fromDescListWithKey f xs
= fromDistinctDescList (combineEq f xs)
fromDescListWithKey f xs0 = fromDistinctDescList xs1
where
-- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs]
combineEq _ xs'
= case xs' of
[] -> []
[x] -> [x]
(x:xx) -> combineEq' x xx

combineEq' z [] = [z]
combineEq' z@(kz,zz) (x@(kx,xx):xs')
| kx==kz = let yy = f kx xx zz in yy `seq` combineEq' (kx,yy) xs'
| otherwise = z:combineEq' x xs'
xs1 = case xs0 of
[] -> []
[x] -> [x]
x:xs -> combineEq x xs

-- We want to have the same strictness as fromListWithKey, which is achieved
-- with the bang on yy.
combineEq y@(ky, !yy) xs = case xs of
[] -> [y]
x@(kx, xx) : xs'
| kx == ky -> combineEq (kx, f kx xx yy) xs'
| otherwise -> y : combineEq x xs'
#if __GLASGOW_HASKELL__
{-# INLINABLE fromDescListWithKey #-}
#endif
Expand Down

0 comments on commit 60e57c2

Please sign in to comment.