Improve fromAscList and friends for Set and Map (#1083)

* Make fromAscList, fromAscListWith, fromAscListWithKey, fromDescList, fromDescListWith, fromDescListWithKey more efficient by removing the intermediate list and making them good consumers in list fusion. * Update Set's fromAscList and fromDescList to keep the last of duplicates. This makes it consistent with all other fromList functions on Set and Map. * Update fromDistinct{Asc,Desc}List to take 1 arg for consistent inlining behavior.
haskell · Jan 4, 2025 · dfeed8b · dfeed8b
1 parent 2776ace
commit dfeed8b
Show file tree

Hide file tree

Showing 6 changed files with 82 additions and 142 deletions.
diff --git a/containers-tests/benchmarks/Map.hs b/containers-tests/benchmarks/Map.hs
@@ -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

diff --git a/containers-tests/benchmarks/Set.hs b/containers-tests/benchmarks/Set.hs
@@ -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

diff --git a/containers/changelog.md b/containers/changelog.md
@@ -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

diff --git a/containers/src/Data/Map/Internal.hs b/containers/src/Data/Map/Internal.hs
@@ -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./
@@ -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./
@@ -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./
@@ -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.
@@ -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.
@@ -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.
@@ -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
@@ -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

diff --git a/containers/src/Data/Map/Strict/Internal.hs b/containers/src/Data/Map/Strict/Internal.hs
@@ -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./
@@ -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./
@@ -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./
@@ -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.
@@ -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.
@@ -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./
@@ -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
@@ -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

diff --git a/containers/src/Data/Set/Internal.hs b/containers/src/Data/Set/Internal.hs
@@ -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
@@ -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