From ba05f939b922d885f5ebda20cf47b0a24d1e93ed Mon Sep 17 00:00:00 2001
From: Harendra Kumar <harendra@composewell.com>
Date: Thu, 31 Mar 2022 16:05:13 +0530
Subject: [PATCH] Use explicit buckets in classifySessions benchmarks

Terminate the session fold after 10 elements
---
 .../Benchmark/Prelude/Serial/NestedFold.hs    | 36 +++++++++++--------
 bin/bench-config.sh                           |  2 +-
 2 files changed, 23 insertions(+), 15 deletions(-)

diff --git a/benchmark/Streamly/Benchmark/Prelude/Serial/NestedFold.hs b/benchmark/Streamly/Benchmark/Prelude/Serial/NestedFold.hs
index d79b54b1c..9685a825d 100644
--- a/benchmark/Streamly/Benchmark/Prelude/Serial/NestedFold.hs
+++ b/benchmark/Streamly/Benchmark/Prelude/Serial/NestedFold.hs
@@ -10,7 +10,6 @@
 {-# LANGUAGE DeriveAnyClass #-}
 {-# LANGUAGE DeriveGeneric #-}
 {-# LANGUAGE RankNTypes #-}
-{-# LANGUAGE TupleSections #-}
 
 module Serial.NestedFold (benchmarks) where
 
@@ -214,35 +213,44 @@ o_n_heap_buffering value =
 -------------------------------------------------------------------------------
 
 {-# INLINE classifySessionsOf #-}
-classifySessionsOf :: (S.MonadAsync m) => SerialT m Int -> m ()
-classifySessionsOf =
+classifySessionsOf :: S.MonadAsync m => (Int -> Int) -> SerialT m Int -> m ()
+classifySessionsOf getKey =
       S.drain
-    . Internal.classifySessionsOf (const (return False)) 3 FL.drain
+    . Internal.classifySessionsOf
+        (const (return False)) 3 (FL.take 10 FL.sum)
     . Internal.timestamped
-    . S.concatMap (\x -> S.map (x,) (S.enumerateFromTo 1 (10 :: Int)))
+    . S.map (\x -> (getKey x, x))
 
 {-# INLINE classifySessionsOfHash #-}
-classifySessionsOfHash :: (S.MonadAsync m) => SerialT m Int -> m ()
-classifySessionsOfHash =
+classifySessionsOfHash :: S.MonadAsync m =>
+    (Int -> Int) -> SerialT m Int -> m ()
+classifySessionsOfHash getKey =
       S.drain
     . Internal.classifySessionsByGeneric
         (Proxy :: Proxy (HashMap k))
-        1 False (const (return False)) 3 FL.drain
+        1 False (const (return False)) 3 (FL.take 10 FL.sum)
     . Internal.timestamped
-    . S.concatMap (\x -> S.map (x,) (S.enumerateFromTo 1 (10 :: Int)))
+    . S.map (\x -> (getKey x, x))
 
 o_n_space_grouping :: Int -> [Benchmark]
 o_n_space_grouping value =
     -- Buffering operations using heap proportional to group/window sizes.
     [ bgroup "grouping"
-        -- We use 10 element stream per input, so div by 10 here
-        [ benchIOSink (value `div` 10) "classifySessionsOf"
-            classifySessionsOf
-        , benchIOSink (value `div` 10) "classifySessionsOfHash"
-            classifySessionsOfHash
+        [ benchIOSink value "classifySessionsOf (10000 buckets)"
+            (classifySessionsOf (getKey 10000))
+        , benchIOSink value "classifySessionsOf (64 buckets)"
+            (classifySessionsOf (getKey 64))
+        , benchIOSink value "classifySessionsOfHash (10000 buckets)"
+            (classifySessionsOfHash (getKey 10000))
+        , benchIOSink value "classifySessionsOfHash (64 buckets)"
+            (classifySessionsOfHash (getKey 64))
         ]
     ]
 
+    where
+
+    getKey n = (`mod` n)
+
 -------------------------------------------------------------------------------
 -- Mixed Transformation
 -------------------------------------------------------------------------------
diff --git a/bin/bench-config.sh b/bin/bench-config.sh
index 1e28d1ab3..70f09a4ad 100644
--- a/bin/bench-config.sh
+++ b/bin/bench-config.sh
@@ -68,7 +68,7 @@ bench_rts_options () {
     Prelude.Serial/o-1-space/mixed/sum-product-fold) echo -n "-K64M" ;;
 
     # XXX These should be moved to o-n-space?
-    Prelude.Serial/o-n-heap/grouping/classifySessionsOf*) echo -n "-K1M -M32M" ;;
+    Prelude.Serial/o-n-heap/grouping/classifySessionsOf*) echo -n "-K4M -M32M" ;;
     Prelude.Serial/o-n-heap/Functor/*) echo -n "-K4M -M32M" ;;
     Prelude.Serial/o-n-heap/transformer/*) echo -n "-K8M -M64M" ;;