git.iko.soy/streamly
on efecb26 help sign in
streamly/benchmark/Streamly/Benchmark/Data/Stream/Common.hs
516 lines 14 kB history raw
1{-# LANGUAGE ConstraintKinds #-}
2{-# LANGUAGE CPP #-}
3{-# LANGUAGE FlexibleContexts #-}
4{-# LANGUAGE ScopedTypeVariables #-}
6-- |
7-- Module : Stream.Common
8-- Copyright : (c) 2018 Composewell Technologies
9-- License : BSD-3-Clause
10-- Maintainer : streamly@composewell.com
12#ifdef USE_PRELUDE
13{-# OPTIONS_GHC -Wno-deprecations #-}
14#endif
16module Stream.Common
17 ( MonadAsync
19 , fromStream
20 , toStream
22 -- Generation
23 , fromListM
24 , fromFoldableM
26 , append
27 , append2
29 -- Elimination
30 , drain
31 , foldl'
32 , scanl'
34 -- Benchmark stream generation
35 , sourceUnfoldr
36 , sourceUnfoldrM
37 , sourceUnfoldrAction
38 , sourceConcatMapId
39 , sourceFromFoldable
41 -- Benchmark stream elimination
42 , benchIOSink
43 , benchIOSrc
44 , benchIO
46 -- Benchmarking functions
47#ifdef USE_STREAMK
48 , concatStreamsWith
49 , mergeMapWith
50#endif
51 , apDiscardFst
52 , apDiscardSnd
53 , apLiftA2
54 , toNullAp
55 , monadThen
56 , toNullM
57 , toNullM3
58 , filterAllOutM
59 , filterAllInM
60 , filterSome
61 , breakAfterSome
62 , toListM
63 , toListSome
64 , composeN
65 , mapN
66 , mapM
67 , transformMapM
68 , transformComposeMapM
69 , transformTeeMapM
70 , transformZipMapM
71 )
72where
74import Control.Applicative (liftA2)
75import Control.DeepSeq (NFData)
76import Control.Exception (try)
77import GHC.Exception (ErrorCall)
78import System.Random (randomRIO)
80import qualified Streamly.Internal.Data.Fold as Fold
81import qualified Streamly.Internal.Data.Pipe as Pipe
83#ifdef USE_PRELUDE
84import Streamly.Prelude (foldl', scanl')
85import qualified Streamly.Internal.Data.Stream.IsStream as Stream
86import qualified Streamly.Prelude as Stream
87import qualified Streamly.Prelude as StreamK
88import Streamly.Benchmark.Prelude
89 ( composeN, sourceConcatMapId, benchIOSink
90 , concatStreamsWith
91 )
92#else
93import Streamly.Internal.Data.Stream (Stream)
94import qualified Streamly.Internal.Data.Stream as D
95#ifdef USE_STREAMK
96import Streamly.Internal.Data.StreamK (StreamK)
97import qualified Streamly.Internal.Data.Stream as Stream
98import qualified Streamly.Internal.Data.StreamK as StreamK
99#else
100import qualified Streamly.Internal.Data.Stream as Stream
101import qualified Streamly.Internal.Data.Stream as StreamK
102#endif
103#endif
105import Gauge
106import Prelude hiding (Foldable(..), mapM, replicate)
108#ifdef USE_STREAMK
109toStream :: Applicative m => StreamK m a -> Stream m a
110toStream = StreamK.toStream
112fromStream :: Monad m => Stream m a -> StreamK m a
113fromStream = StreamK.fromStream
114#else
115fromStream :: a -> a
116fromStream = id
118toStream :: a -> a
119toStream = id
120#endif
122#ifdef USE_PRELUDE
123type Stream = Stream.SerialT
124type MonadAsync m = Stream.MonadAsync m
125mkCross = id
126unCross = id
127#else
128type MonadAsync = Monad
130#ifdef USE_STREAMK
131mkCross :: StreamK m a -> StreamK.CrossStreamK m a
132mkCross = StreamK.mkCross
134unCross :: StreamK.CrossStreamK m a -> StreamK m a
135unCross = StreamK.unCross
136#else
137mkCross :: Stream m a -> Stream.CrossStream m a
138mkCross = Stream.mkCross
140unCross :: Stream.CrossStream m a -> Stream m a
141unCross = Stream.unCross
142#endif
143#endif
145#ifdef USE_PRELUDE
146{-# INLINE append #-}
147append :: Monad m => Stream m a -> Stream m a -> Stream m a
148append = Stream.serial
149#else
150#ifdef USE_STREAMK
151append :: StreamK m a -> StreamK m a -> StreamK m a
152append = StreamK.append
153#else
154append :: Monad m => Stream m a -> Stream m a -> Stream m a
155append = Stream.append
156#endif
157#endif
159{-# INLINE append2 #-}
160append2 :: Monad m => Stream m a -> Stream m a -> Stream m a
161#ifdef USE_PRELUDE
162append2 = Stream.append
163#else
164append2 = D.append
165#endif
167{-# INLINE drain #-}
168drain :: Monad m => Stream m a -> m ()
169{-# INLINE toList #-}
170toList :: Monad m => Stream m a -> m [a]
171#ifdef USE_PRELUDE
172drain = Stream.drain
173toList = Stream.toList
174#else
175drain = Stream.fold Fold.drain
176toList = Stream.fold Fold.toList
177#endif
179{-# INLINE fromListM #-}
180fromListM :: MonadAsync m => [m a] -> Stream m a
181#ifdef USE_PRELUDE
182fromListM = Stream.fromListM
183#else
184fromListM = Stream.sequence . Stream.fromList
185#endif
187{-# INLINE fromFoldableM #-}
188fromFoldableM :: MonadAsync m => [m a] -> Stream m a
189#ifdef USE_PRELUDE
190fromFoldableM = Stream.fromFoldableM
191#else
192fromFoldableM = Stream.sequence . toStream . StreamK.fromFoldable
193#endif
195{-# INLINE sourceUnfoldrM #-}
196sourceUnfoldrM :: MonadAsync m => Int -> Int -> Stream m Int
197sourceUnfoldrM count start = Stream.unfoldrM step start
199 where
201 step cnt =
202 if cnt > start + count
203 then return Nothing
204 else return (Just (cnt, cnt + 1))
206{-# INLINE sourceUnfoldr #-}
207sourceUnfoldr :: Monad m => Int -> Int -> Stream m Int
208sourceUnfoldr count start = Stream.unfoldr step start
210 where
212 step cnt =
213 if cnt > start + count
214 then Nothing
215 else Just (cnt, cnt + 1)
217{-# INLINE sourceUnfoldrAction #-}
218sourceUnfoldrAction :: (Monad m1, Monad m) => Int -> Int -> Stream m (m1 Int)
219sourceUnfoldrAction value n = Stream.unfoldr step n
221 where
223 step cnt =
224 if cnt > n + value
225 then Nothing
226 else Just (return cnt, cnt + 1)
228{-# INLINE sourceFromFoldable #-}
229sourceFromFoldable :: Monad m => Int -> Int -> Stream m Int
230sourceFromFoldable value n = toStream $ StreamK.fromFoldable [n..n+value]
232#ifndef USE_PRELUDE
233{-# INLINE benchIOSink #-}
234benchIOSink
235 :: (NFData b)
236 => Int -> String -> (Stream IO Int -> IO b) -> Benchmark
237benchIOSink value name f =
238 bench name $ nfIO $ randomRIO (1,1) >>= f . sourceUnfoldrM value
239#endif
241-- | Takes a source, and uses it with a default drain/fold method.
242{-# INLINE benchIOSrc #-}
243benchIOSrc
244 :: String
245 -> (Int -> Stream IO a)
246 -> Benchmark
247benchIOSrc name f =
248 bench name $ nfIO $ randomRIO (1,1) >>= drain . f
250{-# NOINLINE benchIO #-}
251benchIO :: (NFData b) => String -> (Int -> IO b) -> Benchmark
252benchIO name f = bench name $ nfIO $ randomRIO (1,1) >>= f
254#ifndef USE_PRELUDE
255#ifdef USE_STREAMK
256{-# INLINE concatStreamsWith #-}
257concatStreamsWith
258 :: (StreamK IO Int -> StreamK IO Int -> StreamK IO Int)
259 -> Int
260 -> Int
261 -> Int
262 -> IO ()
263concatStreamsWith op outer inner n =
264 drain $ toStream $ StreamK.concatMapWith op
265 (fromStream . sourceUnfoldrM inner)
266 (fromStream $ sourceUnfoldrM outer n)
268{-# INLINE mergeMapWith #-}
269mergeMapWith
270 :: (StreamK IO Int -> StreamK IO Int -> StreamK IO Int)
271 -> Int
272 -> Int
273 -> Int
274 -> IO ()
275mergeMapWith op outer inner n =
276 drain $ toStream $ StreamK.mergeMapWith op
277 (fromStream . sourceUnfoldrM inner)
278 (fromStream $ sourceUnfoldrM outer n)
279#endif
281{-# INLINE sourceConcatMapId #-}
282sourceConcatMapId :: (Monad m)
283 => Int -> Int -> Stream m (Stream m Int)
284sourceConcatMapId value n =
285 Stream.fromList $ fmap (D.fromEffect . return) [n..n+value]
286#endif
288{-# INLINE apDiscardFst #-}
289apDiscardFst :: MonadAsync m =>
290 Int -> Int -> m ()
291apDiscardFst linearCount start = drain $ toStream $ unCross $
292 mkCross (fromStream $ sourceUnfoldrM nestedCount2 start)
293 *> mkCross (fromStream $ sourceUnfoldrM nestedCount2 start)
295 where
297 nestedCount2 = round (fromIntegral linearCount**(1/2::Double))
299{-# INLINE apDiscardSnd #-}
300apDiscardSnd :: MonadAsync m => Int -> Int -> m ()
301apDiscardSnd linearCount start = drain $ toStream $ unCross $
302 mkCross (fromStream $ sourceUnfoldrM nestedCount2 start)
303 <* mkCross (fromStream $ sourceUnfoldrM nestedCount2 start)
305 where
307 nestedCount2 = round (fromIntegral linearCount**(1/2::Double))
309{-# INLINE apLiftA2 #-}
310apLiftA2 :: MonadAsync m => Int -> Int -> m ()
311apLiftA2 linearCount start = drain $ toStream $ unCross $
312 liftA2 (+) (mkCross (fromStream $ sourceUnfoldrM nestedCount2 start))
313 (mkCross (fromStream $ sourceUnfoldrM nestedCount2 start))
315 where
317 nestedCount2 = round (fromIntegral linearCount**(1/2::Double))
319{-# INLINE toNullAp #-}
320toNullAp :: MonadAsync m => Int -> Int -> m ()
321toNullAp linearCount start = drain $ toStream $ unCross $
322 (+) <$> mkCross (fromStream $ sourceUnfoldrM nestedCount2 start)
323 <*> mkCross (fromStream $ sourceUnfoldrM nestedCount2 start)
325 where
327 nestedCount2 = round (fromIntegral linearCount**(1/2::Double))
329{-# INLINE monadThen #-}
330monadThen :: MonadAsync m => Int -> Int -> m ()
331monadThen linearCount start = drain $ toStream $ unCross $ do
332 mkCross (fromStream $ sourceUnfoldrM nestedCount2 start) >>
333 mkCross (fromStream $ sourceUnfoldrM nestedCount2 start)
335 where
337 nestedCount2 = round (fromIntegral linearCount**(1/2::Double))
339{-# INLINE toNullM #-}
340toNullM :: MonadAsync m => Int -> Int -> m ()
341toNullM linearCount start = drain $ toStream $ unCross $ do
342 x <- mkCross (fromStream $ sourceUnfoldrM nestedCount2 start)
343 y <- mkCross (fromStream $ sourceUnfoldrM nestedCount2 start)
344 return $ x + y
346 where
348 nestedCount2 = round (fromIntegral linearCount**(1/2::Double))
350{-# INLINE toNullM3 #-}
351toNullM3 :: MonadAsync m => Int -> Int -> m ()
352toNullM3 linearCount start = drain $ toStream $ unCross $ do
353 x <- mkCross (fromStream $ sourceUnfoldrM nestedCount3 start)
354 y <- mkCross (fromStream $ sourceUnfoldrM nestedCount3 start)
355 z <- mkCross (fromStream $ sourceUnfoldrM nestedCount3 start)
356 return $ x + y + z
357 where
358 nestedCount3 = round (fromIntegral linearCount**(1/3::Double))
360{-# INLINE filterAllOutM #-}
361filterAllOutM :: MonadAsync m => Int -> Int -> m ()
362filterAllOutM linearCount start = drain $ toStream $ unCross $ do
363 x <- mkCross (fromStream $ sourceUnfoldrM nestedCount2 start)
364 y <- mkCross (fromStream $ sourceUnfoldrM nestedCount2 start)
365 let s = x + y
366 if s < 0
367 then return s
368 else mkCross StreamK.nil
369 where
370 nestedCount2 = round (fromIntegral linearCount**(1/2::Double))
372{-# INLINE filterAllInM #-}
373filterAllInM :: MonadAsync m => Int -> Int -> m ()
374filterAllInM linearCount start = drain $ toStream $ unCross $ do
375 x <- mkCross (fromStream $ sourceUnfoldrM nestedCount2 start)
376 y <- mkCross (fromStream $ sourceUnfoldrM nestedCount2 start)
377 let s = x + y
378 if s > 0
379 then return s
380 else mkCross StreamK.nil
381 where
382 nestedCount2 = round (fromIntegral linearCount**(1/2::Double))
384{-# INLINE filterSome #-}
385filterSome :: MonadAsync m => Int -> Int -> m ()
386filterSome linearCount start = drain $ toStream $ unCross $ do
387 x <- mkCross (fromStream $ sourceUnfoldrM nestedCount2 start)
388 y <- mkCross (fromStream $ sourceUnfoldrM nestedCount2 start)
389 let s = x + y
390 if s > 1100000
391 then return s
392 else mkCross StreamK.nil
393 where
394 nestedCount2 = round (fromIntegral linearCount**(1/2::Double))
396{-# INLINE breakAfterSome #-}
397breakAfterSome :: Int -> Int -> IO ()
398breakAfterSome linearCount start = do
399 (_ :: Either ErrorCall ()) <- try $ drain $ toStream $ unCross $ do
400 x <- mkCross (fromStream $ sourceUnfoldrM nestedCount2 start)
401 y <- mkCross (fromStream $ sourceUnfoldrM nestedCount2 start)
402 let s = x + y
403 if s > 1100000
404 then error "break"
405 else return s
406 return ()
407 where
408 nestedCount2 = round (fromIntegral linearCount**(1/2::Double))
410{-# INLINE toListM #-}
411toListM :: MonadAsync m => Int -> Int -> m [Int]
412toListM linearCount start = toList $ toStream $ unCross $ do
413 x <- mkCross (fromStream $ sourceUnfoldrM nestedCount2 start)
414 y <- mkCross (fromStream $ sourceUnfoldrM nestedCount2 start)
415 return $ x + y
416 where
417 nestedCount2 = round (fromIntegral linearCount**(1/2::Double))
419-- Taking a specified number of elements is very expensive in logict so we have
420-- a test to measure the same.
421{-# INLINE toListSome #-}
422toListSome :: MonadAsync m => Int -> Int -> m [Int]
423toListSome linearCount start =
424 toList $ Stream.take 10000 $ toStream $ unCross $ do
425 x <- mkCross (fromStream $ sourceUnfoldrM nestedCount2 start)
426 y <- mkCross (fromStream $ sourceUnfoldrM nestedCount2 start)
427 return $ x + y
428 where
429 nestedCount2 = round (fromIntegral linearCount**(1/2::Double))
431#ifndef USE_PRELUDE
432{-# INLINE composeN #-}
433composeN ::
434 (Monad m)
435 => Int
436 -> (Stream m Int -> Stream m Int)
437 -> Stream m Int
438 -> m ()
439composeN n f =
440 case n of
441 1 -> drain . f
442 2 -> drain . f . f
443 3 -> drain . f . f . f
444 4 -> drain . f . f . f . f
445 _ -> undefined
446#endif
448{-# INLINE mapN #-}
449mapN ::
450 Monad m
451 => Int
452 -> Stream m Int
453 -> m ()
454mapN n = composeN n $ fmap (+ 1)
456{-# INLINE mapM #-}
457mapM ::
458 MonadAsync m
459 => Int
460 -> Stream m Int
461 -> m ()
462mapM n = composeN n $ Stream.mapM return
464#ifndef USE_PRELUDE
465foldl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> m b
466foldl' f z = Stream.fold (Fold.foldl' f z)
468scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
469scanl' f z = Stream.scan (Fold.foldl' f z)
470#endif
472{-# INLINE transformMapM #-}
473transformMapM ::
474 (Monad m)
475 => Int
476 -> Stream m Int
477 -> m ()
478transformMapM n = composeN n $ Stream.transform (Pipe.mapM return)
480{-# INLINE transformComposeMapM #-}
481transformComposeMapM ::
482 (Monad m)
483 => Int
484 -> Stream m Int
485 -> m ()
486transformComposeMapM n =
487 composeN n $
488 Stream.transform
489 (Pipe.mapM (\x -> return (x + 1)) `Pipe.compose`
490 Pipe.mapM (\x -> return (x + 2)))
492{-# INLINE transformTeeMapM #-}
493transformTeeMapM ::
494 (Monad m)
495 => Int
496 -> Stream m Int
497 -> m ()
498transformTeeMapM n =
499 composeN n $
500 Stream.transform
501 (Pipe.mapM (\x -> return (x + 1)) `Pipe.tee`
502 Pipe.mapM (\x -> return (x + 2)))
504{-# INLINE transformZipMapM #-}
505transformZipMapM ::
506 (Monad m)
507 => Int
508 -> Stream m Int
509 -> m ()
510transformZipMapM n =
511 composeN n $
512 Stream.transform
513 (Pipe.zipWith
514 (+)
515 (Pipe.mapM (\x -> return (x + 1)))
516 (Pipe.mapM (\x -> return (x + 2))))