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.
haskell · Dec 29, 2024 · c1dd12c · c1dd12c
1 parent d4beee8
commit c1dd12c
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 140 deletions.
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./

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