Add benchmarks for fold functions

containers-tests/benchmarks/IntMap.hs

@@ -3,13 +3,15 @@ module Main where
 
 import Control.DeepSeq (rnf)
 import Control.Exception (evaluate)
-import Test.Tasty.Bench (bench, defaultMain, whnf)
+import Test.Tasty.Bench (bench, bgroup, defaultMain, whnf)
 import Data.List (foldl')
 import qualified Data.IntMap as M
 import qualified Data.IntMap.Strict as MS
 import Data.Maybe (fromMaybe)
 import Prelude hiding (lookup)
 
+import Utils.Fold (foldBenchmarks, foldWithKeyBenchmarks)
+
 main = do
     let m     = M.fromAscList elems_hits   :: M.IntMap Int
     let m'    = M.fromAscList elems_mid    :: M.IntMap Int
@@ -36,9 +38,6 @@ main = do
         , bench "insertLookupWithKey update" $ whnf (insLookupWithKey elems) m
         , bench "map" $ whnf (M.map (+ 1)) m
         , bench "mapWithKey" $ whnf (M.mapWithKey (+)) m
-        , bench "foldlWithKey" $ whnf (ins elems) m
-        , bench "foldlWithKey'" $ whnf (M.foldlWithKey' sum 0) m
-        , bench "foldrWithKey" $ whnf (M.foldrWithKey consPair []) m
         , bench "delete" $ whnf (del keys) m
         , bench "update" $ whnf (upd keys) m
         , bench "updateLookupWithKey" $ whnf (upd' keys) m
@@ -54,6 +53,9 @@ main = do
         , bench "split" $ whnf (M.split key_mid) m
         , bench "splitLookup" $ whnf (M.splitLookup key_mid) m
         , bench "eq" $ whnf (\m' -> m' == m') m -- worst case, compares everything
+        , bgroup "folds" $ foldBenchmarks M.foldr M.foldl M.foldr' M.foldl' foldMap m
+        , bgroup "folds with key" $
+            foldWithKeyBenchmarks M.foldrWithKey M.foldlWithKey M.foldrWithKey' M.foldlWithKey' M.foldMapWithKey m
         ]
   where
     elems = elems_hits

containers-tests/benchmarks/IntSet.hs

@@ -5,7 +5,7 @@ module Main where
 
 import Control.DeepSeq (rnf)
 import Control.Exception (evaluate)
-import Test.Tasty.Bench (bench, defaultMain, whnf)
+import Test.Tasty.Bench (bench, bgroup, defaultMain, whnf)
 import Data.List (foldl')
 import Data.Monoid (Sum(..), All(..))
 import qualified Data.IntSet as IS
@@ -15,6 +15,8 @@ import qualified Data.Set as S
 import qualified Data.IntMap as IM
 import qualified Data.Map.Strict as M
 
+import Utils.Fold (foldBenchmarks)
+
 main = do
     let s = IS.fromAscList elems :: IS.IntSet
         s_even = IS.fromAscList elems_even :: IS.IntSet
@@ -56,8 +58,8 @@ main = do
         , bench "splitMember:dense" $ whnf (IS.splitMember elem_mid) s
         , bench "splitMember:sparse" $ whnf (IS.splitMember elem_sparse_mid) s_sparse
         , bench "eq" $ whnf (\s' -> s' == s') s -- worst case, compares everything
-        , bench "foldMap:dense" $ whnf (IS.foldMap (All . (>0))) s
-        , bench "foldMap:sparse" $ whnf (IS.foldMap (All . (>0))) s_sparse
+        , bgroup "folds:dense" $ foldBenchmarks IS.foldr IS.foldl IS.foldr' IS.foldl' IS.foldMap s
+        , bgroup "folds:sparse" $ foldBenchmarks IS.foldr IS.foldl IS.foldr' IS.foldl' IS.foldMap s_sparse
         ]
   where
     bound = 2^12

containers-tests/benchmarks/Map.hs

@@ -5,7 +5,7 @@ module Main where
 import Control.Applicative (Const(Const, getConst), pure)
 import Control.DeepSeq (rnf)
 import Control.Exception (evaluate)
-import Test.Tasty.Bench (bench, defaultMain, whnf, nf)
+import Test.Tasty.Bench (bench, bgroup, defaultMain, whnf, nf)
 import Data.Functor.Identity (Identity(..))
 import Data.List (foldl')
 import qualified Data.Map as M
@@ -16,6 +16,8 @@ import Data.Functor ((<$))
 import Data.Coerce
 import Prelude hiding (lookup)
 
+import Utils.Fold (foldBenchmarks, foldWithKeyBenchmarks)
+
 main = do
     let m = M.fromAscList elems :: M.Map Int Int
         m_even = M.fromAscList elems_even :: M.Map Int Int
@@ -70,10 +72,6 @@ main = do
         , bench "insertLookupWithKey' absent" $ whnf (insLookupWithKey' elems_even) m_odd
         , bench "insertLookupWithKey' present" $ whnf (insLookupWithKey' elems_even) m_even
         , bench "mapWithKey" $ whnf (M.mapWithKey (+)) m
-        , bench "foldlWithKey" $ whnf (ins elems) m
-        , bench "foldlWithKey'" $ whnf (M.foldlWithKey' sumkv 0) m
-        , bench "foldrWithKey" $ whnf (M.foldrWithKey consPair []) m
-        , bench "foldrWithKey'" $ whnf (M.foldrWithKey' consPair []) m
         , bench "update absent" $ whnf (upd Just evens) m_odd
         , bench "update present" $ whnf (upd Just evens) m_even
         , bench "update delete" $ whnf (upd (const Nothing) evens) m
@@ -102,6 +100,9 @@ main = do
         , bench "minView" $ whnf (\m' -> case M.minViewWithKey m' of {Nothing -> 0; Just ((k,v),m'') -> k+v+M.size m''}) (M.fromAscList $ zip [1..10::Int] [100..110::Int])
         , bench "eq" $ whnf (\m' -> m' == m') m -- worst case, compares everything
         , bench "compare" $ whnf (\m' -> compare m' m') m -- worst case, compares everything
+        , bgroup "folds" $ foldBenchmarks M.foldr M.foldl M.foldr' M.foldl' foldMap m
+        , bgroup "folds with key" $
+            foldWithKeyBenchmarks M.foldrWithKey M.foldlWithKey M.foldrWithKey' M.foldlWithKey' M.foldMapWithKey m
         ]
   where
     bound = 2^12

containers-tests/benchmarks/Sequence.hs

@@ -11,6 +11,8 @@ import qualified Data.Foldable
 import Data.Traversable (traverse, sequenceA)
 import System.Random (mkStdGen, randoms)
 
+import Utils.Fold (foldBenchmarks)
+
 main = do
     let s10 = S.fromList [1..10] :: S.Seq Int
         s100 = S.fromList [1..100] :: S.Seq Int
@@ -53,18 +55,6 @@ main = do
          , bench "1000" $ nf (S.partition even) s1000
          , bench "10000" $ nf (S.partition even) s10000
          ]
-      , bgroup "foldl'"
-         [ bench "10" $ nf (foldl' (+) 0) s10
-         , bench "100" $ nf (foldl' (+) 0) s100
-         , bench "1000" $ nf (foldl' (+) 0) s1000
-         , bench "10000" $ nf (foldl' (+) 0) s10000
-         ]
-      , bgroup "foldr'"
-         [ bench "10" $ nf (foldr' (+) 0) s10
-         , bench "100" $ nf (foldr' (+) 0) s100
-         , bench "1000" $ nf (foldr' (+) 0) s1000
-         , bench "10000" $ nf (foldr' (+) 0) s10000
-         ]
       , bgroup "update"
          [ bench "10" $ nf (updatePoints r10 10) s10
          , bench "100" $ nf (updatePoints r100 10) s100
@@ -184,6 +174,8 @@ main = do
         , bench "100/10000" $  whnf (uncurry compare) (s100, s10000)
         , bench "10000/100" $  whnf (uncurry compare) (s10000, s100)
         ]
+      , bgroup "folds 10" $ foldBenchmarks foldr foldl foldr' foldl' foldMap s10
+      , bgroup "folds 10000" $ foldBenchmarks foldr foldl foldr' foldl' foldMap s10000
       ]
 
 {-

containers-tests/benchmarks/Set.hs

@@ -4,10 +4,12 @@ module Main where
 
 import Control.DeepSeq (rnf)
 import Control.Exception (evaluate)
-import Test.Tasty.Bench (bench, defaultMain, whnf)
+import Test.Tasty.Bench (bench, bgroup, defaultMain, whnf)
 import Data.List (foldl')
 import qualified Data.Set as S
 
+import Utils.Fold (foldBenchmarks)
+
 main = do
     let s = S.fromAscList elems :: S.Set Int
         s_even = S.fromAscList elems_even :: S.Set Int
@@ -55,6 +57,7 @@ main = do
         , bench "member.powerSet (15)" $ whnf (\ s -> all (flip S.member s) s) (S.powerSet (S.fromList [1..15]))
         , bench "eq" $ whnf (\s' -> s' == s') s -- worst case, compares everything
         , bench "compare" $ whnf (\s' -> compare s' s') s -- worst case, compares everything
+        , bgroup "folds" $ foldBenchmarks S.foldr S.foldl S.foldr' S.foldl' foldMap s
         ]
   where
     bound = 2^12

containers-tests/benchmarks/Tree.hs

@@ -4,7 +4,7 @@ module Main where
 import Control.DeepSeq (NFData, rnf)
 import Control.Exception (evaluate)
 import Data.Coerce (coerce)
-import Data.Foldable (fold, foldl', toList)
+import qualified Data.Foldable as F
 import Data.Monoid (All(..))
 #if MIN_VERSION_base(4,18,0)
 import Data.Monoid (Sum(..))
@@ -13,20 +13,20 @@ import qualified Data.Foldable1 as Foldable1
 import Test.Tasty.Bench (Benchmark, Benchmarkable, bench, bgroup, defaultMain, whnf, nf)
 import qualified Data.Tree as T
 
+import Utils.Fold (foldBenchmarks)
+
 main :: IO ()
 main = do
   evaluate $ rnf ts `seq` rnf tsBool
   defaultMain
     [ bgroup "Foldable"
-      [ bgroup "fold" $ forTs tsBool $ whnf fold . (coerce :: T.Tree Bool -> T.Tree All)
-      , bgroup "foldMap" $ forTs tsBool $ whnf (foldMap All)
-      , bgroup "foldr_1" $ forTs tsBool $ whnf (foldr (&&) True)
-      , bgroup "foldr_2" $ forTs ts $ whnf (length . foldr (:) [])
-      , bgroup "foldr_3" $ forTs ts $ whnf (\t -> foldr (\x k acc -> if acc < 0 then acc else k $! acc + x) id t 0)
-      , bgroup "foldl'" $ forTs ts $ whnf (foldl' (+) 0)
+      [ bgroup "folds"
+        [ bgroup label $ foldBenchmarks foldr foldl F.foldr' F.foldl' foldMap t
+        | Tree label t <- ts
+        ]
       , bgroup "foldr1" $ forTs tsBool $ whnf (foldr1 (&&))
       , bgroup "foldl1" $ forTs ts $ whnf (foldl1 (+))
-      , bgroup "toList" $ forTs ts $ nf toList
+      , bgroup "toList" $ forTs ts $ nf F.toList
       , bgroup "elem" $ forTs ts $ whnf (elem 0)
       , bgroup "maximum" $ forTs ts $ whnf maximum
       , bgroup "sum" $ forTs ts $ whnf sum

containers-tests/benchmarks/Utils/Fold.hs

@@ -0,0 +1,142 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+
+module Utils.Fold
+  ( foldBenchmarks
+  , foldWithKeyBenchmarks
+  ) where
+
+import Control.Monad.Trans.State.Strict
+import Prelude hiding (Foldable(..))
+import Test.Tasty.Bench (Benchmark, bench, defaultMain, whnf, nf)
+import qualified GHC.Exts
+
+-- | Benchmarks for folds on a structure of @Int@s.
+
+-- See Note [Choice of benchmarks]
+foldBenchmarks
+  :: forall f.
+     (forall b. (Int -> b -> b) -> b -> f -> b)
+  -> (forall b. (b -> Int -> b) -> b -> f -> b)
+  -> (forall b. (Int -> b -> b) -> b -> f -> b)
+  -> (forall b. (b -> Int -> b) -> b -> f -> b)
+  -> (forall m. Monoid m => (Int -> m) -> f -> m)
+  -> f
+  -> [Benchmark]
+foldBenchmarks foldr foldl foldr' foldl' foldMap xs =
+  [-- foldr
+    bench "foldr_skip" $ whnf (foldr (\_ z -> z) ()) xs
+  , bench "foldr_cpsSum" $ whnf foldr_cpsSum xs
+  , bench "foldr_cpsOneShotSum" $ whnf foldr_cpsOneShotSum xs
+  , bench "foldr_traverseSum" $ whnf foldr_traverseSum xs
+
+    -- foldl
+  , bench "foldl_skip" $ whnf (foldl (\z _ -> z) ()) xs
+  , bench "foldl_cpsSum" $ whnf foldl_cpsSum xs
+  , bench "foldl_cpsOneShotSum" $ whnf foldl_cpsOneShotSum xs
+  , bench "foldl_traverseSum" $ whnf foldl_traverseSum xs
+
+    -- foldr'
+  , bench "foldr'_sum" $ whnf (foldr' (+) 0) xs
+
+    -- foldl'
+  , bench "foldl'_sum" $ whnf (foldl' (+) 0) xs
+
+    -- foldMap
+  , bench "foldMap_seq" $ whnf (foldMap (\_ -> Unit ())) xs
+  , bench "foldMap_traverseSum" $ whnf foldMap_traverseSum xs
+  ]
+  where
+    foldr_cpsSum :: f -> Int
+    foldr_cpsSum xs = foldr (\x k !acc -> k (x + acc)) id xs 0
+
+    foldr_cpsOneShotSum :: f -> Int
+    foldr_cpsOneShotSum xs =
+      foldr (\x k -> GHC.Exts.oneShot (\ !acc -> k (x + acc))) id xs 0
+
+    foldr_traverseSum :: f -> Int
+    foldr_traverseSum xs =
+      execState (foldr (\x z -> modify' (+x) *> z) (pure ()) xs) 0
+
+    foldl_cpsSum :: f -> Int
+    foldl_cpsSum xs = foldl (\k x !acc -> k (x + acc)) id xs 0
+
+    foldl_cpsOneShotSum :: f -> Int
+    foldl_cpsOneShotSum xs =
+      foldl (\k x -> GHC.Exts.oneShot (\ !acc -> k (x + acc))) id xs 0
+
+    foldl_traverseSum :: f -> Int
+    foldl_traverseSum xs =
+      execState (foldl (\z x -> modify' (+x) *> z) (pure ()) xs) 0
+
+    foldMap_traverseSum :: f -> Int
+    foldMap_traverseSum xs =
+      execState (runEffect (foldMap (\x -> Effect (modify' (+x))) xs)) 0
+{-# INLINE foldBenchmarks #-}
+
+-- | Benchmarks for folds on a structure of @Int@ keys and @Int@ values.
+foldWithKeyBenchmarks
+  :: (forall b. (Int -> Int -> b -> b) -> b -> f -> b)
+  -> (forall b. (b -> Int -> Int -> b) -> b -> f -> b)
+  -> (forall b. (Int -> Int -> b -> b) -> b -> f -> b)
+  -> (forall b. (b -> Int -> Int -> b) -> b -> f -> b)
+  -> (forall m. Monoid m => (Int -> Int -> m) -> f -> m)
+  -> f
+  -> [Benchmark]
+foldWithKeyBenchmarks
+  foldrWithKey foldlWithKey foldrWithKey' foldlWithKey' foldMapWithKey =
+  foldBenchmarks
+    (\f -> foldrWithKey (\k x z -> f (k + x) z))
+    (\f -> foldlWithKey (\z k x -> f z (k + x)))
+    (\f -> foldrWithKey' (\k x z -> f (k + x) z))
+    (\f -> foldlWithKey' (\z k x -> f z (k + x)))
+    (\f -> foldMapWithKey (\k x -> f (k + x)))
+{-# INLINE foldWithKeyBenchmarks #-}
+
+newtype Effect f = Effect { runEffect :: f () }
+
+instance Applicative f => Semigroup (Effect f) where
+  Effect f1 <> Effect f2 = Effect (f1 *> f2)
+
+instance Applicative f => Monoid (Effect f) where
+  mempty = Effect (pure ())
+
+newtype Unit = Unit ()
+
+instance Semigroup Unit where
+  (<>) = seq
+
+instance Monoid Unit where
+  mempty = Unit ()
+
+
+-- Note [Choice of benchmarks]
+-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~
+--
+-- foldr_skip, foldl_skip
+--   Simplest folds that visit every element. In practice:
+--   * Worst case for lazy folds like `foldr (&&) True`
+--   * Data.Foldable.toList
+--
+-- foldr_cpsSum, foldr_cpsOneShotSum, foldl_cpsSum, foldl_cpsOneShotSum
+--   The well-known foldl'-via-foldr pattern. GHC.Exts.oneShot is used to help
+--   GHC with optimizations. In practice:
+--   * Used for early-return with an accumulator
+--   * Used by the foldl library
+--
+-- foldr_traverseSum, foldr_traverseSum
+--   Folding with an effect. In practice:
+--   * Folds defined using foldr, such as Data.Foldable.traverse_ and friends
+--
+-- foldl', foldr'
+--   Strict folds.
+--
+-- foldMap_seq
+--   Simplest fold that visits every element. In practice:
+--   * Worst case for lazy folds defined using foldMap, such as
+--     Data.Foldable.any, Data.Foldable.find, etc.
+--
+-- foldMap_traverseSum
+--   Folding with an effect. In practice:
+--   * With the lens library, using traverseOf_ on a foldMap based fold.