@@ -6,11 +6,15 @@
66-- Maintainer : streamly@composewell.com
77
88{-# LANGUAGE CPP #-}
9+{-# LANGUAGE FlexibleContexts #-}
910
1011import Control.DeepSeq (NFData(..))
1112import System.Random (randomRIO)
1213
1314import qualified Streamly.Benchmark.Data.ArrayOps as Ops
15+import qualified Streamly.Data.Array.Storable.Foreign as A
16+import qualified Streamly.Internal.Data.Array.Storable.Foreign as IA
17+import qualified Streamly.Prelude as S
1418
1519import Gauge
1620import Streamly.Benchmark.Common hiding (benchPureSrc)
@@ -66,6 +70,26 @@ benchIO' name src f = bench name $ nfIO $
6670benchIOSink :: NFData b => Int -> String -> (Ops.Stream Int -> IO b) -> Benchmark
6771benchIOSink value name f = benchIO' name (Ops.sourceIntFromTo value) f
6872
73+{-# INLINE sourceUnfoldrM #-}
74+sourceUnfoldrM :: (S.IsStream t, S.MonadAsync m) => Int -> Int -> t m Int
75+sourceUnfoldrM value n = S.unfoldrM step n
76+ where
77+ step cnt =
78+ if cnt > n + value
79+ then return Nothing
80+ else return (Just (cnt, cnt + 1))
81+
82+{-# INLINE source #-}
83+source :: (S.MonadAsync m, S.IsStream t) => Int -> Int -> t m Int
84+source = sourceUnfoldrM
85+
86+ -- | Takes a fold method, and uses it with a default source.
87+{-# INLINE benchIOSinkF #-}
88+benchIOSinkF
89+ :: (S.IsStream t, NFData b)
90+ => Int -> String -> (t IO Int -> IO b) -> Benchmark
91+benchIOSinkF value name f = bench name $ nfIO $ randomRIO (1,1) >>= f . source value
92+
6993o_1_space_generation :: Int -> [Benchmark]
7094o_1_space_generation value =
7195 [ bgroup
@@ -115,6 +139,8 @@ o_1_space_elimination value =
115139 , benchIOSink value "foldl'" Ops.pureFoldl'
116140 , benchIOSink value "read" Ops.unfoldReadDrain
117141 , benchIOSink value "toStreamRev" Ops.toStreamRevDrain
142+ , benchIOSinkF value "lastN.1" (S.fold (IA.lastN 1))
143+ , benchIOSinkF value "lastN.10" (S.fold (IA.lastN 10))
118144#if !defined(DATA_ARRAY_PRIM) && !defined(DATA_ARRAY_PRIM_PINNED)
119145#ifdef DEVBUILD
120146 , benchPureSink value "foldable/foldl'" Ops.foldableFoldl'
@@ -124,6 +150,16 @@ o_1_space_elimination value =
124150 ]
125151 ]
126152
153+o_n_heap_serial :: Int -> [Benchmark]
154+o_n_heap_serial value =
155+ [ bgroup "elimination"
156+ [
157+ -- Converting the stream to an array
158+ benchIOSinkF value "lastN.Max" (S.fold (IA.lastN (value + 1)))
159+ , benchIOSinkF value "writeN" (S.fold (A.writeN value))
160+ ]
161+ ]
162+
127163o_1_space_transformation :: Int -> [Benchmark]
128164o_1_space_transformation value =
129165 [ bgroup "transformation"
@@ -176,5 +212,6 @@ main = do
176212 , o_1_space_elimination size
177213 , o_1_space_transformation size
178214 , o_1_space_transformationX4 size
215+ , o_n_heap_serial size
179216 ]
180217 ]