Skip to content

Commit c1dd12c

Browse files
meooow25meooow25
committed
Improve fromAscList and friends for Set and Map
* Make fromAscList, fromAscListWith, fromAscListWithKey more efficient by removing the intermediate list and making them good consumers in list fusion. * Update fromDistinct{Asc,Desc}List to take 1 arg for consistent inlining behavior.
1 parent d4beee8 commit c1dd12c

File tree

3 files changed

+65
-140
lines changed

3 files changed

+65
-140
lines changed

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

+24-73
Original file line numberDiff line numberDiff line change
@@ -3702,23 +3702,8 @@ foldlFB = foldlWithKey
37023702
-- > valid (fromAscList [(5,"a"), (3,"b"), (5,"b")]) == False
37033703

37043704
fromAscList :: Eq k => [(k,a)] -> Map k a
3705-
fromAscList xs
3706-
= fromDistinctAscList (combineEq xs)
3707-
where
3708-
-- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs]
3709-
combineEq xs'
3710-
= case xs' of
3711-
[] -> []
3712-
[x] -> [x]
3713-
(x:xx) -> combineEq' x xx
3714-
3715-
combineEq' z [] = [z]
3716-
combineEq' z@(kz,_) (x@(kx,xx):xs')
3717-
| kx==kz = combineEq' (kx,xx) xs'
3718-
| otherwise = z:combineEq' x xs'
3719-
#if __GLASGOW_HASKELL__
3720-
{-# INLINABLE fromAscList #-}
3721-
#endif
3705+
fromAscList xs = fromAscListWithKey (\_ x _ -> x) xs
3706+
{-# INLINE fromAscList #-} -- INLINE for fusion
37223707

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

37333718
fromDescList :: Eq k => [(k,a)] -> Map k a
3734-
fromDescList xs = fromDistinctDescList (combineEq xs)
3735-
where
3736-
-- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs]
3737-
combineEq xs'
3738-
= case xs' of
3739-
[] -> []
3740-
[x] -> [x]
3741-
(x:xx) -> combineEq' x xx
3742-
3743-
combineEq' z [] = [z]
3744-
combineEq' z@(kz,_) (x@(kx,xx):xs')
3745-
| kx==kz = combineEq' (kx,xx) xs'
3746-
| otherwise = z:combineEq' x xs'
3747-
#if __GLASGOW_HASKELL__
3748-
{-# INLINABLE fromDescList #-}
3749-
#endif
3719+
fromDescList xs = fromDescListWithKey (\_ x _ -> x) xs
3720+
{-# INLINE fromDescList #-} -- INLINE for fusion
37503721

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

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

37833750
-- | \(O(n)\). Build a map from an ascending list in linear time with a
37843751
-- combining function for equal keys.
@@ -3792,23 +3759,15 @@ fromDescListWith f xs
37923759
-- Also see the performance note on 'fromListWith'.
37933760

37943761
fromAscListWithKey :: Eq k => (k -> a -> a -> a) -> [(k,a)] -> Map k a
3795-
fromAscListWithKey f xs
3796-
= fromDistinctAscList (combineEq f xs)
3762+
fromAscListWithKey f xs = ascLinkAll (Foldable.foldl' next Nada xs)
37973763
where
3798-
-- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs]
3799-
combineEq _ xs'
3800-
= case xs' of
3801-
[] -> []
3802-
[x] -> [x]
3803-
(x:xx) -> combineEq' x xx
3804-
3805-
combineEq' z [] = [z]
3806-
combineEq' z@(kz,zz) (x@(kx,xx):xs')
3807-
| kx==kz = let yy = f kx xx zz in combineEq' (kx,yy) xs'
3808-
| otherwise = z:combineEq' x xs'
3809-
#if __GLASGOW_HASKELL__
3810-
{-# INLINABLE fromAscListWithKey #-}
3811-
#endif
3764+
next stk (!ky, y) = case stk of
3765+
Push kx x l stk'
3766+
| ky == kx -> Push ky (f ky y x) l stk'
3767+
| Tip <- l -> ascLinkTop stk' 1 (singleton kx x) ky y
3768+
| otherwise -> Push ky y Tip stk
3769+
Nada -> Push ky y Tip stk
3770+
{-# INLINE fromAscListWithKey #-} -- INLINE for fusion
38123771

38133772
-- | \(O(n)\). Build a map from a descending list in linear time with a
38143773
-- combining function for equal keys.
@@ -3822,23 +3781,15 @@ fromAscListWithKey f xs
38223781
-- Also see the performance note on 'fromListWith'.
38233782

38243783
fromDescListWithKey :: Eq k => (k -> a -> a -> a) -> [(k,a)] -> Map k a
3825-
fromDescListWithKey f xs
3826-
= fromDistinctDescList (combineEq f xs)
3784+
fromDescListWithKey f xs = descLinkAll (Foldable.foldl' next Nada xs)
38273785
where
3828-
-- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs]
3829-
combineEq _ xs'
3830-
= case xs' of
3831-
[] -> []
3832-
[x] -> [x]
3833-
(x:xx) -> combineEq' x xx
3834-
3835-
combineEq' z [] = [z]
3836-
combineEq' z@(kz,zz) (x@(kx,xx):xs')
3837-
| kx==kz = let yy = f kx xx zz in combineEq' (kx,yy) xs'
3838-
| otherwise = z:combineEq' x xs'
3839-
#if __GLASGOW_HASKELL__
3840-
{-# INLINABLE fromDescListWithKey #-}
3841-
#endif
3786+
next stk (!ky, y) = case stk of
3787+
Push kx x r stk'
3788+
| ky == kx -> Push ky (f ky y x) r stk'
3789+
| Tip <- r -> descLinkTop ky y 1 (singleton kx x) stk'
3790+
| otherwise -> Push ky y Tip stk
3791+
Nada -> Push ky y Tip stk
3792+
{-# INLINE fromDescListWithKey #-} -- INLINE for fusion
38423793

38433794

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

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

38803831
-- See Note [fromDistinctAscList implementation] in Data.Set.Internal.
38813832
fromDistinctDescList :: [(k,a)] -> Map k a
3882-
fromDistinctDescList = descLinkAll . Foldable.foldl' next Nada
3833+
fromDistinctDescList xs = descLinkAll (Foldable.foldl' next Nada xs)
38833834
where
38843835
next :: Stack k a -> (k, a) -> Stack k a
38853836
next (Push ky y Tip stk) (!kx, x) = descLinkTop kx x 1 (singleton ky y) stk

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

+22-44
Original file line numberDiff line numberDiff line change
@@ -1612,9 +1612,7 @@ fromListWithKey f xs
16121612
fromAscList :: Eq k => [(k,a)] -> Map k a
16131613
fromAscList xs
16141614
= fromAscListWithKey (\_ x _ -> x) xs
1615-
#if __GLASGOW_HASKELL__
1616-
{-# INLINABLE fromAscList #-}
1617-
#endif
1615+
{-# INLINE fromAscList #-} -- INLINE for fusion
16181616

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

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

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

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

16761668
fromAscListWithKey :: Eq k => (k -> a -> a -> a) -> [(k,a)] -> Map k a
1677-
fromAscListWithKey f xs0 = fromDistinctAscList xs1
1669+
fromAscListWithKey f xs = ascLinkAll (Foldable.foldl' next Nada xs)
16781670
where
1679-
xs1 = case xs0 of
1680-
[] -> []
1681-
[x] -> [x]
1682-
x:xs -> combineEq x xs
1683-
1684-
-- We want to have the same strictness as fromListWithKey, which is achieved
1685-
-- with the bang on yy.
1686-
combineEq y@(ky, !yy) xs = case xs of
1687-
[] -> [y]
1688-
x@(kx, xx) : xs'
1689-
| kx == ky -> combineEq (kx, f kx xx yy) xs'
1690-
| otherwise -> y : combineEq x xs'
1691-
#if __GLASGOW_HASKELL__
1692-
{-# INLINABLE fromAscListWithKey #-}
1693-
#endif
1671+
next stk (!ky, y) = case stk of
1672+
Push kx x l stk'
1673+
| ky == kx -> let !y' = f ky y x in Push ky y' l stk'
1674+
| Tip <- l -> y `seq` ascLinkTop stk' 1 (singleton kx x) ky y
1675+
| otherwise -> push ky y Tip stk
1676+
Nada -> push ky y Tip stk
1677+
push kx !x = Push kx x
1678+
{-# INLINE fromAscListWithKey #-} -- INLINE for fusion
16941679

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

17061691
fromDescListWithKey :: Eq k => (k -> a -> a -> a) -> [(k,a)] -> Map k a
1707-
fromDescListWithKey f xs0 = fromDistinctDescList xs1
1692+
fromDescListWithKey f xs = descLinkAll (Foldable.foldl' next Nada xs)
17081693
where
1709-
xs1 = case xs0 of
1710-
[] -> []
1711-
[x] -> [x]
1712-
x:xs -> combineEq x xs
1713-
1714-
-- We want to have the same strictness as fromListWithKey, which is achieved
1715-
-- with the bang on yy.
1716-
combineEq y@(ky, !yy) xs = case xs of
1717-
[] -> [y]
1718-
x@(kx, xx) : xs'
1719-
| kx == ky -> combineEq (kx, f kx xx yy) xs'
1720-
| otherwise -> y : combineEq x xs'
1721-
#if __GLASGOW_HASKELL__
1722-
{-# INLINABLE fromDescListWithKey #-}
1723-
#endif
1694+
next stk (!ky, y) = case stk of
1695+
Push kx x r stk'
1696+
| ky == kx -> let !y' = f ky y x in Push ky y' r stk'
1697+
| Tip <- r -> y `seq` descLinkTop ky y 1 (singleton kx x) stk'
1698+
| otherwise -> push ky y Tip stk
1699+
Nada -> push ky y Tip stk
1700+
push kx !x = Push kx x
1701+
{-# INLINE fromDescListWithKey #-} -- INLINE for fusion
17241702

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

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

+19-23
Original file line numberDiff line numberDiff line change
@@ -1198,41 +1198,37 @@ fromList (x0 : xs0) | not_ordered x0 xs0 = fromList' (Bin 1 x0 Tip Tip) xs0
11981198
-- | \(O(n)\). Build a set from an ascending list in linear time.
11991199
-- /The precondition (input list is ascending) is not checked./
12001200
fromAscList :: Eq a => [a] -> Set a
1201-
fromAscList xs = fromDistinctAscList (combineEq xs)
1202-
#if __GLASGOW_HASKELL__
1203-
{-# INLINABLE fromAscList #-}
1204-
#endif
1201+
fromAscList xs = ascLinkAll (Foldable.foldl' next Nada xs)
1202+
where
1203+
next stk !y = case stk of
1204+
Push x l stk'
1205+
| y == x -> Push y l stk'
1206+
| Tip <- l -> ascLinkTop stk' 1 (singleton x) y
1207+
| otherwise -> Push y Tip stk
1208+
Nada -> Push y Tip stk
1209+
{-# INLINE fromAscList #-} -- INLINE for fusion
12051210

12061211
-- | \(O(n)\). Build a set from a descending list in linear time.
12071212
-- /The precondition (input list is descending) is not checked./
12081213
--
12091214
-- @since 0.5.8
12101215
fromDescList :: Eq a => [a] -> Set a
1211-
fromDescList xs = fromDistinctDescList (combineEq xs)
1212-
#if __GLASGOW_HASKELL__
1213-
{-# INLINABLE fromDescList #-}
1214-
#endif
1215-
1216-
-- [combineEq xs] combines equal elements with [const] in an ordered list [xs]
1217-
--
1218-
-- TODO: combineEq allocates an intermediate list. It *should* be better to
1219-
-- make fromAscListBy and fromDescListBy the fundamental operations, and to
1220-
-- implement the rest using those.
1221-
combineEq :: Eq a => [a] -> [a]
1222-
combineEq [] = []
1223-
combineEq (x : xs) = combineEq' x xs
1216+
fromDescList xs = descLinkAll (Foldable.foldl' next Nada xs)
12241217
where
1225-
combineEq' z [] = [z]
1226-
combineEq' z (y:ys)
1227-
| z == y = combineEq' z ys
1228-
| otherwise = z : combineEq' y ys
1218+
next stk !y = case stk of
1219+
Push x r stk'
1220+
| y == x -> Push y r stk'
1221+
| Tip <- r -> descLinkTop y 1 (singleton x) stk'
1222+
| otherwise -> Push y Tip stk
1223+
Nada -> Push y Tip stk
1224+
{-# INLINE fromDescList #-} -- INLINE for fusion
12291225

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

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

12581254
-- See Note [fromDistinctAscList implementation]
12591255
fromDistinctDescList :: [a] -> Set a
1260-
fromDistinctDescList = descLinkAll . Foldable.foldl' next Nada
1256+
fromDistinctDescList xs = descLinkAll (Foldable.foldl' next Nada xs)
12611257
where
12621258
next :: Stack a -> a -> Stack a
12631259
next (Push y Tip stk) !x = descLinkTop x 1 (singleton y) stk

0 commit comments

Comments
 (0)