Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Improve fromAscList and friends for Set and Map #1083

Merged
merged 3 commits into from
Jan 4, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions containers-tests/benchmarks/Map.hs
Original file line number Diff line number Diff line change
Expand Up @@ -90,11 +90,19 @@ main = do
, bench "fromList" $ whnf M.fromList elems
, bench "fromList-desc" $ whnf M.fromList elems_desc
, bench "fromAscList" $ whnf M.fromAscList elems_asc
, bench "fromAscList:fusion" $
whnf (\n -> M.fromAscList [(i `div` 2, i) | i <- [1..n]]) bound
, bench "fromAscListWithKey" $
whnf (M.fromAscListWithKey sumkv) elems_asc
, bench "fromAscListWithKey:fusion" $
whnf (\n -> M.fromAscListWithKey sumkv [(i `div` 2, i) | i <- [1..n]]) bound
, bench "fromDescList" $ whnf M.fromDescList elems_desc
, bench "fromDescList:fusion" $
whnf (\n -> M.fromDescList [(i `div` 2, i) | i <- [n,n-1..1]]) bound
, bench "fromDescListWithKey" $
whnf (M.fromDescListWithKey sumkv) elems_desc
, bench "fromDescListWithKey:fusion" $
whnf (\n -> M.fromDescListWithKey sumkv [(i `div` 2, i) | i <- [n,n-1..1]]) bound
, bench "fromDistinctAscList" $ whnf M.fromDistinctAscList elems_distinct_asc
, bench "fromDistinctAscList:fusion" $ whnf (\n -> M.fromDistinctAscList [(i,i) | i <- [1..n]]) bound
, bench "fromDistinctDescList" $ whnf M.fromDistinctDescList elems_distinct_desc
Expand Down
4 changes: 4 additions & 0 deletions containers-tests/benchmarks/Set.hs
Original file line number Diff line number Diff line change
Expand Up @@ -37,9 +37,13 @@ main = do
, bench "fromList" $ whnf S.fromList elems
, bench "fromList-desc" $ whnf S.fromList elems_desc
, bench "fromAscList" $ whnf S.fromAscList elems_asc
, bench "fromAscList:fusion" $
whnf (\n -> S.fromAscList [i `div` 2 | i <- [1..n]]) bound
, bench "fromDistinctAscList" $ whnf S.fromDistinctAscList elems_distinct_asc
, bench "fromDistinctAscList:fusion" $ whnf (\n -> S.fromDistinctAscList [1..n]) bound
, bench "fromDescList" $ whnf S.fromDescList elems_desc
, bench "fromDescList:fusion" $
whnf (\n -> S.fromDescList [i `div` 2 | i <- [n,n-1..1]]) bound
, bench "fromDistinctDescList" $ whnf S.fromDistinctDescList elems_distinct_desc
, bench "fromDistinctDescList:fusion" $ whnf (\n -> S.fromDistinctDescList [n,n-1..1]) bound
, bench "disjoint:false" $ whnf (S.disjoint s) s_even
Expand Down
3 changes: 3 additions & 0 deletions containers/changelog.md
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,9 @@

* Improved performance for `Data.Intset`'s `foldr`, `foldl'`, `foldl`, `foldr'`.

* Improved performance for `Data.Set` and `Data.Map`'s `fromAscList*` and
`fromDescList*` functions.

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

### Additions
Expand Down
97 changes: 24 additions & 73 deletions containers/src/Data/Map/Internal.hs
Original file line number Diff line number Diff line change
Expand Up @@ -3702,23 +3702,8 @@ foldlFB = foldlWithKey
-- > valid (fromAscList [(5,"a"), (3,"b"), (5,"b")]) == False

fromAscList :: Eq k => [(k,a)] -> Map k a
fromAscList xs
= fromDistinctAscList (combineEq xs)
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,_) (x@(kx,xx):xs')
| kx==kz = combineEq' (kx,xx) xs'
| otherwise = z:combineEq' x xs'
#if __GLASGOW_HASKELL__
{-# INLINABLE fromAscList #-}
#endif
fromAscList xs = fromAscListWithKey (\_ x _ -> x) xs
{-# INLINE fromAscList #-} -- INLINE for fusion

-- | \(O(n)\). Build a map from a descending list in linear time.
-- /The precondition (input list is descending) is not checked./
Expand All @@ -3731,22 +3716,8 @@ fromAscList xs
-- @since 0.5.8

fromDescList :: Eq k => [(k,a)] -> Map k a
fromDescList xs = fromDistinctDescList (combineEq xs)
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,_) (x@(kx,xx):xs')
| kx==kz = combineEq' (kx,xx) xs'
| otherwise = z:combineEq' x xs'
#if __GLASGOW_HASKELL__
{-# INLINABLE fromDescList #-}
#endif
fromDescList xs = fromDescListWithKey (\_ x _ -> x) xs
{-# INLINE fromDescList #-} -- INLINE for fusion

-- | \(O(n)\). Build a map from an ascending list in linear time with a combining function for equal keys.
-- /The precondition (input list is ascending) is not checked./
Expand All @@ -3758,9 +3729,7 @@ fromDescList xs = fromDistinctDescList (combineEq xs)
fromAscListWith :: Eq k => (a -> a -> a) -> [(k,a)] -> Map k a
fromAscListWith f xs
= fromAscListWithKey (\_ x y -> f x y) xs
#if __GLASGOW_HASKELL__
{-# INLINABLE fromAscListWith #-}
#endif
{-# INLINE fromAscListWith #-} -- INLINE for fusion

-- | \(O(n)\). Build a map from a descending list in linear time with a combining function for equal keys.
-- /The precondition (input list is descending) is not checked./
Expand All @@ -3776,9 +3745,7 @@ fromAscListWith f xs
fromDescListWith :: Eq k => (a -> a -> a) -> [(k,a)] -> Map k a
fromDescListWith f xs
= fromDescListWithKey (\_ x y -> f x y) xs
#if __GLASGOW_HASKELL__
{-# INLINABLE fromDescListWith #-}
#endif
{-# INLINE fromDescListWith #-} -- INLINE for fusion

-- | \(O(n)\). Build a map from an ascending list in linear time with a
-- combining function for equal keys.
Expand All @@ -3792,23 +3759,15 @@ 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 xs = ascLinkAll (Foldable.foldl' next Nada xs)
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 combineEq' (kx,yy) xs'
| otherwise = z:combineEq' x xs'
#if __GLASGOW_HASKELL__
{-# INLINABLE fromAscListWithKey #-}
#endif
next stk (!ky, y) = case stk of
Push kx x l stk'
| ky == kx -> Push ky (f ky y x) l stk'
| Tip <- l -> ascLinkTop stk' 1 (singleton kx x) ky y
| otherwise -> Push ky y Tip stk
Nada -> Push ky y Tip stk
{-# INLINE fromAscListWithKey #-} -- INLINE for fusion

-- | \(O(n)\). Build a map from a descending list in linear time with a
-- combining function for equal keys.
Expand All @@ -3822,23 +3781,15 @@ 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 xs = descLinkAll (Foldable.foldl' next Nada xs)
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 combineEq' (kx,yy) xs'
| otherwise = z:combineEq' x xs'
#if __GLASGOW_HASKELL__
{-# INLINABLE fromDescListWithKey #-}
#endif
next stk (!ky, y) = case stk of
Push kx x r stk'
| ky == kx -> Push ky (f ky y x) r stk'
| Tip <- r -> descLinkTop ky y 1 (singleton kx x) stk'
| otherwise -> Push ky y Tip stk
Nada -> Push ky y Tip stk
{-# INLINE fromDescListWithKey #-} -- INLINE for fusion


-- | \(O(n)\). Build a map from an ascending list of distinct elements in linear time.
Expand All @@ -3850,7 +3801,7 @@ fromDescListWithKey f xs

-- See Note [fromDistinctAscList implementation] in Data.Set.Internal.
fromDistinctAscList :: [(k,a)] -> Map k a
fromDistinctAscList = ascLinkAll . Foldable.foldl' next Nada
fromDistinctAscList xs = ascLinkAll (Foldable.foldl' next Nada xs)
where
next :: Stack k a -> (k, a) -> Stack k a
next (Push kx x Tip stk) (!ky, y) = ascLinkTop stk 1 (singleton kx x) ky y
Expand Down Expand Up @@ -3879,7 +3830,7 @@ ascLinkAll stk = foldl'Stack (\r kx x l -> link kx x l r) Tip stk

-- See Note [fromDistinctAscList implementation] in Data.Set.Internal.
fromDistinctDescList :: [(k,a)] -> Map k a
fromDistinctDescList = descLinkAll . Foldable.foldl' next Nada
fromDistinctDescList xs = descLinkAll (Foldable.foldl' next Nada xs)
where
next :: Stack k a -> (k, a) -> Stack k a
next (Push ky y Tip stk) (!kx, x) = descLinkTop kx x 1 (singleton ky y) stk
Expand Down
70 changes: 24 additions & 46 deletions containers/src/Data/Map/Strict/Internal.hs
Original file line number Diff line number Diff line change
Expand Up @@ -1612,9 +1612,7 @@ fromListWithKey f xs
fromAscList :: Eq k => [(k,a)] -> Map k a
fromAscList xs
= fromAscListWithKey (\_ x _ -> x) xs
#if __GLASGOW_HASKELL__
{-# INLINABLE fromAscList #-}
#endif
{-# INLINE fromAscList #-} -- INLINE for fusion

-- | \(O(n)\). Build a map from a descending list in linear time.
-- /The precondition (input list is descending) is not checked./
Expand All @@ -1626,9 +1624,7 @@ fromAscList xs
fromDescList :: Eq k => [(k,a)] -> Map k a
fromDescList xs
= fromDescListWithKey (\_ x _ -> x) xs
#if __GLASGOW_HASKELL__
{-# INLINABLE fromDescList #-}
#endif
{-# INLINE fromDescList #-} -- INLINE for fusion

-- | \(O(n)\). Build a map from an ascending list in linear time with a combining function for equal keys.
-- /The precondition (input list is ascending) is not checked./
Expand All @@ -1642,9 +1638,7 @@ fromDescList xs
fromAscListWith :: Eq k => (a -> a -> a) -> [(k,a)] -> Map k a
fromAscListWith f xs
= fromAscListWithKey (\_ x y -> f x y) xs
#if __GLASGOW_HASKELL__
{-# INLINABLE fromAscListWith #-}
#endif
{-# INLINE fromAscListWith #-} -- INLINE for fusion

-- | \(O(n)\). Build a map from a descending list in linear time with a combining function for equal keys.
-- /The precondition (input list is descending) is not checked./
Expand All @@ -1658,9 +1652,7 @@ fromAscListWith f xs
fromDescListWith :: Eq k => (a -> a -> a) -> [(k,a)] -> Map k a
fromDescListWith f xs
= fromDescListWithKey (\_ x y -> f x y) xs
#if __GLASGOW_HASKELL__
{-# INLINABLE fromDescListWith #-}
#endif
{-# INLINE fromDescListWith #-} -- INLINE for fusion

-- | \(O(n)\). Build a map from an ascending list in linear time with a
-- combining function for equal keys.
Expand All @@ -1674,23 +1666,16 @@ fromDescListWith f xs
-- Also see the performance note on 'fromListWith'.

fromAscListWithKey :: Eq k => (k -> a -> a -> a) -> [(k,a)] -> Map k a
fromAscListWithKey f xs0 = fromDistinctAscList xs1
fromAscListWithKey f xs = ascLinkAll (Foldable.foldl' next Nada xs)
where
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
next stk (!ky, y) = case stk of
Push kx x l stk'
| ky == kx -> let !y' = f ky y x in Push ky y' l stk'
| Tip <- l -> y `seq` ascLinkTop stk' 1 (singleton kx x) ky y
| otherwise -> push ky y Tip stk
Nada -> push ky y Tip stk
push kx !x = Push kx x
{-# INLINE fromAscListWithKey #-} -- INLINE for fusion

-- | \(O(n)\). Build a map from a descending list in linear time with a
-- combining function for equal keys.
Expand All @@ -1704,23 +1689,16 @@ fromAscListWithKey f xs0 = fromDistinctAscList xs1
-- Also see the performance note on 'fromListWith'.

fromDescListWithKey :: Eq k => (k -> a -> a -> a) -> [(k,a)] -> Map k a
fromDescListWithKey f xs0 = fromDistinctDescList xs1
fromDescListWithKey f xs = descLinkAll (Foldable.foldl' next Nada xs)
where
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
next stk (!ky, y) = case stk of
Push kx x r stk'
| ky == kx -> let !y' = f ky y x in Push ky y' r stk'
| Tip <- r -> y `seq` descLinkTop ky y 1 (singleton kx x) stk'
| otherwise -> push ky y Tip stk
Nada -> push ky y Tip stk
push kx !x = Push kx x
{-# INLINE fromDescListWithKey #-} -- INLINE for fusion

-- | \(O(n)\). Build a map from an ascending list of distinct elements in linear time.
-- /The precondition is not checked./
Expand All @@ -1731,7 +1709,7 @@ fromDescListWithKey f xs0 = fromDistinctDescList xs1

-- See Note [fromDistinctAscList implementation] in Data.Set.Internal.
fromDistinctAscList :: [(k,a)] -> Map k a
fromDistinctAscList = ascLinkAll . Foldable.foldl' next Nada
fromDistinctAscList xs = ascLinkAll (Foldable.foldl' next Nada xs)
where
next :: Stack k a -> (k, a) -> Stack k a
next (Push kx x Tip stk) (!ky, !y) = ascLinkTop stk 1 (singleton kx x) ky y
Expand All @@ -1747,7 +1725,7 @@ fromDistinctAscList = ascLinkAll . Foldable.foldl' next Nada

-- See Note [fromDistinctAscList implementation] in Data.Set.Internal.
fromDistinctDescList :: [(k,a)] -> Map k a
fromDistinctDescList = descLinkAll . Foldable.foldl' next Nada
fromDistinctDescList xs = descLinkAll (Foldable.foldl' next Nada xs)
where
next :: Stack k a -> (k, a) -> Stack k a
next (Push ky y Tip stk) (!kx, !x) = descLinkTop kx x 1 (singleton ky y) stk
Expand Down
42 changes: 19 additions & 23 deletions containers/src/Data/Set/Internal.hs
Original file line number Diff line number Diff line change
Expand Up @@ -1198,41 +1198,37 @@ fromList (x0 : xs0) | not_ordered x0 xs0 = fromList' (Bin 1 x0 Tip Tip) xs0
-- | \(O(n)\). Build a set from an ascending list in linear time.
-- /The precondition (input list is ascending) is not checked./
fromAscList :: Eq a => [a] -> Set a
fromAscList xs = fromDistinctAscList (combineEq xs)
#if __GLASGOW_HASKELL__
{-# INLINABLE fromAscList #-}
#endif
fromAscList xs = ascLinkAll (Foldable.foldl' next Nada xs)
where
next stk !y = case stk of
Push x l stk'
| y == x -> Push y l stk'
| Tip <- l -> ascLinkTop stk' 1 (singleton x) y
| otherwise -> Push y Tip stk
Nada -> Push y Tip stk
{-# INLINE fromAscList #-} -- INLINE for fusion

-- | \(O(n)\). Build a set from a descending list in linear time.
-- /The precondition (input list is descending) is not checked./
--
-- @since 0.5.8
fromDescList :: Eq a => [a] -> Set a
fromDescList xs = fromDistinctDescList (combineEq xs)
#if __GLASGOW_HASKELL__
{-# INLINABLE fromDescList #-}
#endif

-- [combineEq xs] combines equal elements with [const] in an ordered list [xs]
--
-- TODO: combineEq allocates an intermediate list. It *should* be better to
-- make fromAscListBy and fromDescListBy the fundamental operations, and to
-- implement the rest using those.
combineEq :: Eq a => [a] -> [a]
combineEq [] = []
combineEq (x : xs) = combineEq' x xs
fromDescList xs = descLinkAll (Foldable.foldl' next Nada xs)
where
combineEq' z [] = [z]
combineEq' z (y:ys)
| z == y = combineEq' z ys
| otherwise = z : combineEq' y ys
next stk !y = case stk of
Push x r stk'
| y == x -> Push y r stk'
| Tip <- r -> descLinkTop y 1 (singleton x) stk'
| otherwise -> Push y Tip stk
Nada -> Push y Tip stk
{-# INLINE fromDescList #-} -- INLINE for fusion

-- | \(O(n)\). Build a set from an ascending list of distinct elements in linear time.
-- /The precondition (input list is strictly ascending) is not checked./

-- See Note [fromDistinctAscList implementation]
fromDistinctAscList :: [a] -> Set a
fromDistinctAscList = ascLinkAll . Foldable.foldl' next Nada
fromDistinctAscList xs = ascLinkAll (Foldable.foldl' next Nada xs)
where
next :: Stack a -> a -> Stack a
next (Push x Tip stk) !y = ascLinkTop stk 1 (singleton x) y
Expand All @@ -1257,7 +1253,7 @@ ascLinkAll stk = foldl'Stack (\r x l -> link x l r) Tip stk

-- See Note [fromDistinctAscList implementation]
fromDistinctDescList :: [a] -> Set a
fromDistinctDescList = descLinkAll . Foldable.foldl' next Nada
fromDistinctDescList xs = descLinkAll (Foldable.foldl' next Nada xs)
where
next :: Stack a -> a -> Stack a
next (Push y Tip stk) !x = descLinkTop x 1 (singleton y) stk
Expand Down
Loading