Make Int{Map,Set} folds friendlier to optimizations (#1149)

Move the Nil branch to the top-level and error on Nil in go.
The Note [IntMap folds] added in Data.IntMap.Internal explains the
details.

For the "maximum" benchmarks added here, the time improves by ~40% for
IntMap and ~30% for sparse IntSets. "traverseSum" and "cpsSum" benchmarks
also improve by 30-60% for IntMaps and sparse IntSets. Dense IntSets are
barely affected.

Author: meooow25

containers-tests/benchmarks/Utils/Fold.hs

@@ -44,9 +44,11 @@ foldBenchmarks foldr foldl foldr' foldl' foldMap xs =
 
     -- foldr'
   , bench "foldr'_sum" $ whnf (foldr' (+) 0) xs
+  , bench "foldr'_maximum" $ whnf foldr'_maximum xs
 
     -- foldl'
   , bench "foldl'_sum" $ whnf (foldl' (+) 0) xs
+  , bench "foldl'_maximum" $ whnf foldl'_maximum xs
 
     -- foldMap
   , bench "foldMap_elem" $ whnf foldMap_elem xs
@@ -81,6 +83,12 @@ foldBenchmarks foldr foldl foldr' foldl' foldMap xs =
     foldl_traverseSum xs =
       execState (foldl (\z x -> modify' (+x) *> z) (pure ()) xs) 0
 
+    foldr'_maximum :: f -> Maybe Int
+    foldr'_maximum = foldr' (\x z -> Just $! maybe x (max x) z) Nothing
+
+    foldl'_maximum :: f -> Maybe Int
+    foldl'_maximum = foldl' (\z x -> Just $! maybe x (max x) z) Nothing
+
     foldMap_elem :: f -> Any
     foldMap_elem = foldMap (\x -> Any (x == minBound))
 
@@ -138,9 +146,12 @@ instance Applicative f => Monoid (Effect f) where
 --   Folding with an effect. In practice:
 --   * Folds defined using foldr, such as Data.Foldable.traverse_ and friends
 --
--- foldl', foldr'
+-- foldl'_sum, foldr'_sum
 --   Strict folds.
 --
+-- foldl'_maximum, foldr'_maximum
+--  Strict folds with a `Maybe` as accumulator which could be optimized away.
+--
 -- foldMap_elem
 --   Simple lazy fold that visits every element. In practice:
 --   * Worst case for lazy folds defined using foldMap, such as

containers/src/Data/IntMap/Internal.hs

@@ -326,7 +326,7 @@ import Data.Bits
 import qualified Data.Foldable as Foldable
 import Data.Maybe (fromMaybe)
 import Utils.Containers.Internal.Prelude hiding
-  (lookup, map, filter, foldr, foldl, foldl', null)
+  (lookup, map, filter, foldr, foldl, foldl', foldMap, null)
 import Prelude ()
 
 import Data.IntSet.Internal (IntSet)
@@ -470,23 +470,13 @@ instance Semigroup (IntMap a) where
 
 -- | Folds in order of increasing key.
 instance Foldable.Foldable IntMap where
-  fold = go
-    where go Nil = mempty
-          go (Tip _ v) = v
-          go (Bin p l r)
-            | signBranch p = go r `mappend` go l
-            | otherwise = go l `mappend` go r
+  fold = foldMap id
   {-# INLINABLE fold #-}
   foldr = foldr
   {-# INLINE foldr #-}
   foldl = foldl
   {-# INLINE foldl #-}
-  foldMap f t = go t
-    where go Nil = mempty
-          go (Tip _ v) = f v
-          go (Bin p l r)
-            | signBranch p = go r `mappend` go l
-            | otherwise = go l `mappend` go r
+  foldMap = foldMap
   {-# INLINE foldMap #-}
   foldl' = foldl'
   {-# INLINE foldl' #-}
@@ -3033,31 +3023,37 @@ splitLookup k t =
 --
 -- > let f a len = len + (length a)
 -- > foldr f 0 (fromList [(5,"a"), (3,"bbb")]) == 4
+
+-- See Note [IntMap folds]
 foldr :: (a -> b -> b) -> b -> IntMap a -> b
 foldr f z = \t ->      -- Use lambda t to be inlinable with two arguments only.
   case t of
+    Nil -> z
     Bin p l r
       | signBranch p -> go (go z l) r -- put negative numbers before
       | otherwise -> go (go z r) l
     _ -> go z t
   where
-    go z' Nil         = z'
+    go _ Nil          = error "foldr.go: Nil"
     go z' (Tip _ x)   = f x z'
     go z' (Bin _ l r) = go (go z' r) l
 {-# INLINE foldr #-}
 
 -- | \(O(n)\). A strict version of 'foldr'. Each application of the operator is
 -- evaluated before using the result in the next application. This
 -- function is strict in the starting value.
+
+-- See Note [IntMap folds]
 foldr' :: (a -> b -> b) -> b -> IntMap a -> b
 foldr' f z = \t ->      -- Use lambda t to be inlinable with two arguments only.
   case t of
+    Nil -> z
     Bin p l r
       | signBranch p -> go (go z l) r -- put negative numbers before
       | otherwise -> go (go z r) l
     _ -> go z t
   where
-    go !z' Nil        = z'
+    go !_ Nil         = error "foldr'.go: Nil"
     go z' (Tip _ x)   = f x z'
     go z' (Bin _ l r) = go (go z' r) l
 {-# INLINE foldr' #-}
@@ -3071,35 +3067,65 @@ foldr' f z = \t ->      -- Use lambda t to be inlinable with two arguments only.
 --
 -- > let f len a = len + (length a)
 -- > foldl f 0 (fromList [(5,"a"), (3,"bbb")]) == 4
+
+-- See Note [IntMap folds]
 foldl :: (a -> b -> a) -> a -> IntMap b -> a
 foldl f z = \t ->      -- Use lambda t to be inlinable with two arguments only.
   case t of
+    Nil -> z
     Bin p l r
       | signBranch p -> go (go z r) l -- put negative numbers before
       | otherwise -> go (go z l) r
     _ -> go z t
   where
-    go z' Nil         = z'
+    go _ Nil          = error "foldl.go: Nil"
     go z' (Tip _ x)   = f z' x
     go z' (Bin _ l r) = go (go z' l) r
 {-# INLINE foldl #-}
 
 -- | \(O(n)\). A strict version of 'foldl'. Each application of the operator is
 -- evaluated before using the result in the next application. This
 -- function is strict in the starting value.
+
+-- See Note [IntMap folds]
 foldl' :: (a -> b -> a) -> a -> IntMap b -> a
 foldl' f z = \t ->      -- Use lambda t to be inlinable with two arguments only.
   case t of
+    Nil -> z
     Bin p l r
       | signBranch p -> go (go z r) l -- put negative numbers before
       | otherwise -> go (go z l) r
     _ -> go z t
   where
-    go !z' Nil        = z'
+    go !_ Nil         = error "foldl'.go: Nil"
     go z' (Tip _ x)   = f z' x
     go z' (Bin _ l r) = go (go z' l) r
 {-# INLINE foldl' #-}
 
+-- See Note [IntMap folds]
+foldMap :: Monoid m => (a -> m) -> IntMap a -> m
+foldMap f = \t -> -- Use lambda to be inlinable with two arguments.
+  case t of
+    Nil -> mempty
+    Bin p l r
+#if MIN_VERSION_base(4,11,0)
+      | signBranch p -> go r <> go l
+      | otherwise -> go l <> go r
+#else
+      | signBranch p -> go r `mappend` go l
+      | otherwise -> go l `mappend` go r
+#endif
+    _ -> go t
+  where
+    go Nil = error "foldMap.go: Nil"
+    go (Tip _ x) = f x
+#if MIN_VERSION_base(4,11,0)
+    go (Bin _ l r) = go l <> go r
+#else
+    go (Bin _ l r) = go l `mappend` go r
+#endif
+{-# INLINE foldMap #-}
+
 -- | \(O(n)\). Fold the keys and values in the map using the given right-associative
 -- binary operator, such that
 -- @'foldrWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@.
@@ -3110,31 +3136,37 @@ foldl' f z = \t ->      -- Use lambda t to be inlinable with two arguments only.
 --
 -- > let f k a result = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")"
 -- > foldrWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (5:a)(3:b)"
+
+-- See Note [IntMap folds]
 foldrWithKey :: (Key -> a -> b -> b) -> b -> IntMap a -> b
 foldrWithKey f z = \t ->      -- Use lambda t to be inlinable with two arguments only.
   case t of
+    Nil -> z
     Bin p l r
       | signBranch p -> go (go z l) r -- put negative numbers before
       | otherwise -> go (go z r) l
     _ -> go z t
   where
-    go z' Nil         = z'
+    go _ Nil          = error "foldrWithKey.go: Nil"
     go z' (Tip kx x)  = f kx x z'
     go z' (Bin _ l r) = go (go z' r) l
 {-# INLINE foldrWithKey #-}
 
 -- | \(O(n)\). A strict version of 'foldrWithKey'. Each application of the operator is
 -- evaluated before using the result in the next application. This
 -- function is strict in the starting value.
+
+-- See Note [IntMap folds]
 foldrWithKey' :: (Key -> a -> b -> b) -> b -> IntMap a -> b
 foldrWithKey' f z = \t ->      -- Use lambda t to be inlinable with two arguments only.
   case t of
+    Nil -> z
     Bin p l r
       | signBranch p -> go (go z l) r -- put negative numbers before
       | otherwise -> go (go z r) l
     _ -> go z t
   where
-    go !z' Nil        = z'
+    go !_ Nil         = error "foldrWithKey'.go: Nil"
     go z' (Tip kx x)  = f kx x z'
     go z' (Bin _ l r) = go (go z' r) l
 {-# INLINE foldrWithKey' #-}
@@ -3149,31 +3181,37 @@ foldrWithKey' f z = \t ->      -- Use lambda t to be inlinable with two argument
 --
 -- > let f result k a = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")"
 -- > foldlWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (3:b)(5:a)"
+
+-- See Note [IntMap folds]
 foldlWithKey :: (a -> Key -> b -> a) -> a -> IntMap b -> a
 foldlWithKey f z = \t ->      -- Use lambda t to be inlinable with two arguments only.
   case t of
+    Nil -> z
     Bin p l r
       | signBranch p -> go (go z r) l -- put negative numbers before
       | otherwise -> go (go z l) r
     _ -> go z t
   where
-    go z' Nil         = z'
+    go _ Nil          = error "foldlWithKey.go: Nil"
     go z' (Tip kx x)  = f z' kx x
     go z' (Bin _ l r) = go (go z' l) r
 {-# INLINE foldlWithKey #-}
 
 -- | \(O(n)\). A strict version of 'foldlWithKey'. Each application of the operator is
 -- evaluated before using the result in the next application. This
 -- function is strict in the starting value.
+
+-- See Note [IntMap folds]
 foldlWithKey' :: (a -> Key -> b -> a) -> a -> IntMap b -> a
 foldlWithKey' f z = \t ->      -- Use lambda t to be inlinable with two arguments only.
   case t of
+    Nil -> z
     Bin p l r
       | signBranch p -> go (go z r) l -- put negative numbers before
       | otherwise -> go (go z l) r
     _ -> go z t
   where
-    go !z' Nil        = z'
+    go !_ Nil         = error "foldlWithKey'.go: Nil"
     go z' (Tip kx x)  = f z' kx x
     go z' (Bin _ l r) = go (go z' l) r
 {-# INLINE foldlWithKey' #-}
@@ -3185,14 +3223,29 @@ foldlWithKey' f z = \t ->      -- Use lambda t to be inlinable with two argument
 -- This can be an asymptotically faster than 'foldrWithKey' or 'foldlWithKey' for some monoids.
 --
 -- @since 0.5.4
+
+-- See Note [IntMap folds]
 foldMapWithKey :: Monoid m => (Key -> a -> m) -> IntMap a -> m
-foldMapWithKey f = go
+foldMapWithKey f = \t -> -- Use lambda to be inlinable with two arguments.
+  case t of
+    Nil -> mempty
+    Bin p l r
+#if MIN_VERSION_base(4,11,0)
+      | signBranch p -> go r <> go l
+      | otherwise -> go l <> go r
+#else
+      | signBranch p -> go r `mappend` go l
+      | otherwise -> go l `mappend` go r
+#endif
+    _ -> go t
   where
-    go Nil           = mempty
-    go (Tip kx x)    = f kx x
-    go (Bin p l r)
-      | signBranch p = go r `mappend` go l
-      | otherwise = go l `mappend` go r
+    go Nil = error "foldMap.go: Nil"
+    go (Tip kx x) = f kx x
+#if MIN_VERSION_base(4,11,0)
+    go (Bin _ l r) = go l <> go r
+#else
+    go (Bin _ l r) = go l `mappend` go r
+#endif
 {-# INLINE foldMapWithKey #-}
 
 {--------------------------------------------------------------------
@@ -4069,3 +4122,40 @@ withEmpty bars = "   ":bars
 --
 -- The implementation is defined as a foldl' over the input list, which makes
 -- it a good consumer in list fusion.
+
+-- Note [IntMap folds]
+-- ~~~~~~~~~~~~~~~~~~~
+-- Folds on IntMap are defined in a particular way for a few reasons.
+--
+-- foldl' :: (a -> b -> a) -> a -> IntMap b -> a
+-- foldl' f z = \t ->
+--   case t of
+--     Nil -> z
+--     Bin p l r
+--       | signBranch p -> go (go z r) l
+--       | otherwise -> go (go z l) r
+--     _ -> go z t
+--   where
+--     go !_ Nil         = error "foldl'.go: Nil"
+--     go z' (Tip _ x)   = f z' x
+--     go z' (Bin _ l r) = go (go z' l) r
+-- {-# INLINE foldl' #-}
+--
+-- 1. We first check if the Bin separates negative and positive keys, and fold
+--    over the children accordingly. This check is not inside `go` because it
+--    can only happen at the top level and we don't need to check every Bin.
+-- 2. We also check for Nil at the top level instead of, say, `go z Nil = z`.
+--    That's because `Nil` is also allowed only at the top-level, but more
+--    importantly it allows for better optimizations if the `Nil` branch errors
+--    in `go`. For example, if we have
+--      maximum :: Ord a => IntMap a -> Maybe a
+--      maximum = foldl' (\m x -> Just $! maybe x (max x) m) Nothing
+--    because `go` certainly returns a `Just` (or errors), CPR analysis will
+--    optimize it to return `(# a #)` instead of `Maybe a`. This makes it
+--    satisfy the conditions for SpecConstr, which generates two specializations
+--    of `go` for `Nothing` and `Just` inputs. Now both `Maybe`s have been
+--    optimized out of `go`.
+-- 3. The `Tip` is not matched on at the top-level to avoid using `f` more than
+--    once. This allows `f` to be inlined into `go` even if `f` is big, since
+--    it's likely to be the only place `f` is used, and not inlining `f` means
+--    missing out on optimizations. See GHC #25259 for more on this.