@@ -17,23 +17,102 @@
1717{-# OPTIONS_GHC -fplugin Test.Inspection.Plugin #-}
1818#endif
1919
20−module Main (main) where
21−
2220import Control.DeepSeq (NFData(..))
21+import Control.Monad.IO.Class (MonadIO)
22+import Data.Functor ((<&>))
2323import Streamly.Prelude (MonadAsync, SerialT, IsStream)
24+import Streamly.Internal.Data.Stream.Serial (SerialT(..))
2425import System.Random (randomRIO)
26+import Prelude
27+ ( IO
28+ , Int
29+ , Integral(..)
30+ , Maybe(..)
31+ , Monad(..)
32+ , Num(..)
33+ , Ord(..)
34+ , String
35+ , ($)
36+ , (.)
37+ , (||)
38+ , concat
39+ , const
40+ , id
41+ , undefined
42+ )
2543
26−import qualified Streamly.Prelude as Stream
2744import qualified Streamly.Internal.Data.Array.Foreign as Array
2845import qualified Streamly.Internal.Data.Array.Foreign.Mut as MArray
46+import qualified Streamly.Internal.Data.Stream.StreamD as StreamD
47+import qualified Streamly.Prelude as Stream
48+
49+import Gauge
50+import Streamly.Benchmark.Common hiding (benchPureSrc)
51+
52+type Stream = MArray.Array
53+
54+-------------------------------------------------------------------------------
55+-- Benchmark helpers
56+-------------------------------------------------------------------------------
57+
58+{-# INLINE benchIO #-}
59+benchIO :: NFData b => String -> (Int -> IO a) -> (a -> b) -> Benchmark
60+benchIO name src f = bench name $ nfIO $
61+ (randomRIO (1,1) >>= src) <&> f
62+
63+{-# INLINE benchPureSink #-}
64+benchPureSink :: NFData b => Int -> String -> (Stream Int -> b) -> Benchmark
65+benchPureSink value name = benchIO name (sourceIntFromTo value)
2966
30−import Gauge hiding (env)
31−import Streamly.Benchmark.Common
67+{-# INLINE benchIO' #-}
68+benchIO' :: NFData b => String -> (Int -> IO a) -> (a -> IO b) -> Benchmark
69+benchIO' name src f = bench name $ nfIO $
70+ randomRIO (1,1) >>= src >>= f
71+
72+{-# INLINE benchIOSink #-}
73+benchIOSink :: NFData b => Int -> String -> (Stream Int -> IO b) -> Benchmark
74+benchIOSink value name = benchIO' name (sourceIntFromTo value)
75+
76+-- Drain a source that generates an array in the IO monad
77+{-# INLINE benchIOSrc #-}
78+benchIOSrc :: String -> (Int -> IO (Stream a)) -> Benchmark
79+benchIOSrc name src = benchIO name src id
3280
3381-------------------------------------------------------------------------------
34−-- Utilities
82+-- Bench Ops
3583-------------------------------------------------------------------------------
3684
85+{-# INLINE sourceUnfoldr #-}
86+sourceUnfoldr :: MonadIO m => Int -> Int -> m (Stream Int)
87+sourceUnfoldr value n =
88+ Stream.fold (MArray.writeN value) $ Stream.unfoldr step n
89+
90+ where
91+
92+ step cnt =
93+ if cnt > n + value
94+ then Nothing
95+ else Just (cnt, cnt + 1)
96+
97+{-# INLINE sourceIntFromTo #-}
98+sourceIntFromTo :: MonadIO m => Int -> Int -> m (Stream Int)
99+sourceIntFromTo value n =
100+ Stream.fold (MArray.writeN value) $ Stream.enumerateFromTo n (n + value)
101+
102+{-# INLINE sourceFromList #-}
103+sourceFromList :: MonadIO m => Int -> Int -> m (Stream Int)
104+sourceFromList value n =
105+ Stream.fold (MArray.writeN value) $ Stream.fromList [n .. n + value]
106+
107+{-# INLINE sourceIntFromToFromList #-}
108+sourceIntFromToFromList :: MonadIO m => Int -> Int -> m (Stream Int)
109+sourceIntFromToFromList value n = MArray.fromListN value [n..n + value]
110+
111+{-# INLINE sourceIntFromToFromStream #-}
112+sourceIntFromToFromStream :: MonadIO m => Int -> Int -> m (Stream Int)
113+sourceIntFromToFromStream value n =
114+ Stream.fold MArray.write $ Stream.enumerateFromTo n (n + value)
115+
37116{-# INLINE sourceUnfoldrM #-}
38117sourceUnfoldrM :: (IsStream t, MonadAsync m) => Int -> Int -> t m Int
39118sourceUnfoldrM value n = Stream.unfoldrM step n
@@ -43,29 +122,140 @@ sourceUnfoldrM value n = Stream.unfoldrM step n
43122 then return Nothing
44123 else return (Just (cnt, cnt + 1))
45124
46−{-# INLINE benchIO #-}
47−benchIO
48− :: NFData b
49− => String -> (Int -> a) -> (a -> IO b) -> Benchmark
50−benchIO name src sink =
51− bench name $ nfIO $ randomRIO (1,1) >>= sink . src
125+-------------------------------------------------------------------------------
126+-- Transformation
127+-------------------------------------------------------------------------------
128+
129+{-# INLINE composeN #-}
130+composeN :: Monad m
131+ => Int -> (Stream Int -> m (Stream Int)) -> Stream Int -> m (Stream Int)
132+composeN n f x =
133+ case n of
134+ 1 -> f x
135+ 2 -> f x >>= f
136+ 3 -> f x >>= f >>= f
137+ 4 -> f x >>= f >>= f >>= f
138+ _ -> undefined
139+
140+{-# INLINE scanl' #-}
141+scanl' :: MonadIO m => Int -> Int -> Stream Int -> m (Stream Int)
142+scanl' value n = composeN n $ onArray value $ Stream.scanl' (+) 0
143+
144+{-# INLINE scanl1' #-}
145+scanl1' :: MonadIO m => Int -> Int -> Stream Int -> m (Stream Int)
146+scanl1' value n = composeN n $ onArray value $ Stream.scanl1' (+)
147+
148+{-# INLINE map #-}
149+map :: MonadIO m => Int -> Int -> Stream Int -> m (Stream Int)
150+map value n = composeN n $ onArray value $ Stream.map (+ 1)
151+
152+{-# INLINE onArray #-}
153+onArray
154+ :: MonadIO m => Int -> (Stream.SerialT m Int -> Stream.SerialT m Int)
155+ -> Stream Int
156+ -> m (Stream Int)
157+onArray value f arr =
158+ Stream.fold (MArray.writeN value) $ f $ Stream.unfold MArray.read arr
159+
160+-------------------------------------------------------------------------------
161+-- Elimination
162+-------------------------------------------------------------------------------
163+
164+{-# INLINE unfoldReadDrain #-}
165+unfoldReadDrain :: MonadIO m => Stream Int -> m ()
166+unfoldReadDrain = Stream.drain . Stream.unfold MArray.read
167+
168+{-# INLINE unfoldReadRevDrain #-}
169+unfoldReadRevDrain :: MonadIO m => Stream Int -> m ()
170+unfoldReadRevDrain = Stream.drain . Stream.unfold MArray.readRev
171+
172+{-# INLINE toStreamDRevDrain #-}
173+toStreamDRevDrain :: MonadIO m => Stream Int -> m ()
174+toStreamDRevDrain =
175+ Stream.drain . SerialT . StreamD.toStreamK . MArray.toStreamDRev
176+
177+{-# INLINE toStreamDDrain #-}
178+toStreamDDrain :: MonadIO m => Stream Int -> m ()
179+toStreamDDrain = Stream.drain . SerialT . StreamD.toStreamK . MArray.toStreamD
180+
181+{-# INLINE unfoldFold #-}
182+unfoldFold :: MonadIO m => Stream Int -> m Int
183+unfoldFold = Stream.foldl' (+) 0 . Stream.unfold MArray.read
184+
185+-------------------------------------------------------------------------------
186+-- Bench groups
187+-------------------------------------------------------------------------------
188+
189+o_1_space_generation :: Int -> [Benchmark]
190+o_1_space_generation value =
191+ [ bgroup
192+ "generation"
193+ [ benchIOSrc "writeN . intFromTo" (sourceIntFromTo value)
194+ , benchIOSrc
195+ "fromList . intFromTo"
196+ (sourceIntFromToFromList value)
197+ , benchIOSrc "writeN . unfoldr" (sourceUnfoldr value)
198+ , benchIOSrc "writeN . fromList" (sourceFromList value)
199+ , benchIOSrc "write . intFromTo" (sourceIntFromToFromStream value)
200+ ]
201+ ]
202+
203+o_1_space_elimination :: Int -> [Benchmark]
204+o_1_space_elimination value =
205+ [ bgroup "elimination"
206+ [ benchPureSink value "id" id
207+ , benchIOSink value "foldl'" unfoldFold
208+ , benchIOSink value "read" unfoldReadDrain
209+ , benchIOSink value "readRev" unfoldReadRevDrain
210+ , benchIOSink value "toStream" toStreamDDrain
211+ , benchIOSink value "toStreamRev" toStreamDRevDrain
212+ ]
213+ ]
214+
215+o_n_heap_serial :: Int -> [Benchmark]
216+o_n_heap_serial value =
217+ [ bgroup "elimination"
218+ [
219+ -- Converting the stream to an array
220+ benchFold "writeN" (Stream.fold (MArray.writeN value))
221+ (sourceUnfoldrM value)
222+ ]
223+ ]
224+
225+o_1_space_transformation :: Int -> [Benchmark]
226+o_1_space_transformation value =
227+ [ bgroup "transformation"
228+ [ benchIOSink value "scanl'" (scanl' value 1)
229+ , benchIOSink value "scanl1'" (scanl1' value 1)
230+ , benchIOSink value "map" (map value 1)
231+ ]
232+ ]
233+
234+o_1_space_transformationX4 :: Int -> [Benchmark]
235+o_1_space_transformationX4 value =
236+ [ bgroup "transformationX4"
237+ [ benchIOSink value "scanl'" (scanl' value 4)
238+ , benchIOSink value "scanl1'" (scanl1' value 4)
239+ , benchIOSink value "map" (map value 4)
240+ ]
241+ ]
52242
53243o_1_space_serial_marray ::
54244 Int -> (MArray.Array Int, Array.Array Int) -> [Benchmark]
55245o_1_space_serial_marray value ~(array, indices) =
56− [ benchIO "partitionBy (< 0)" (const array)
246+ [ benchIO' "partitionBy (< 0)" (const (return array))
57247 $ MArray.partitionBy (< 0)
58− , benchIO "partitionBy (> 0)" (const array)
248+ , benchIO' "partitionBy (> 0)" (const (return array))
59249 $ MArray.partitionBy (> 0)
60− , benchIO "partitionBy (< value/2)" (const array)
250+ , benchIO' "partitionBy (< value/2)" (const (return array))
61251 $ MArray.partitionBy (< (value `div` 2))
62− , benchIO "partitionBy (> value/2)" (const array)
252+ , benchIO' "partitionBy (> value/2)" (const (return array))
63253 $ MArray.partitionBy (> (value `div` 2))
64− , benchIO "strip (< value/2 || > value/2)" (const array)
254+ , benchIO' "strip (< value/2 || > value/2)" (const (return array))
65255 $ MArray.strip (\x -> x < value `div` 2 || x > value `div` 2)
66− , benchIO "strip (> 0)" (const array)
256+ , benchIO' "strip (> 0)" (const (return array))
67257 $ MArray.strip (> 0)
68− , benchIO "modifyIndices (+ 1)" (const indices)
258+ , benchIO' "modifyIndices (+ 1)" (const (return indices))
69259 $ Stream.fold (MArray.modifyIndices (\_idx val -> val + 1) array)
70260 . Stream.unfold Array.read
71261 ]
@@ -89,6 +279,13 @@ main = do
89279 return (marr, indices)
90280
91281 allBenchmarks array value =
92− [ bgroup (o_1_space_prefix moduleName) $
93− o_1_space_serial_marray value array
282+ [ bgroup (o_1_space_prefix moduleName)
283+ $ concat
284+ [ o_1_space_serial_marray value array
285+ , o_1_space_generation value
286+ , o_1_space_elimination value
287+ , o_1_space_transformation value
288+ , o_1_space_transformationX4 value
289+ ]
290+ , bgroup (o_n_heap_prefix moduleName) (o_n_heap_serial value)
94291 ]