on 4eaeb10 help sign in
451 lines 11 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
12module Stream.Common
13 ( MonadAsync
15 -- Generation
16 , fromListM
17 , fromFoldableM
19 , append
20 , append2
22 -- Elimination
23 , drain
24 , foldl'
25 , scanl'
27 -- Benchmark stream generation
28 , sourceUnfoldr
29 , sourceUnfoldrM
30 , sourceUnfoldrAction
31 , sourceConcatMapId
32 , sourceFromFoldable
34 -- Benchmark stream elimination
35 , benchIOSink
36 , benchIOSrc
37 , benchIO
39 -- Benchmarking functions
40 , concatStreamsWith
41 , concatPairsWith
42 , apDiscardFst
43 , apDiscardSnd
44 , apLiftA2
45 , toNullAp
46 , monadThen
47 , toNullM
48 , toNullM3
49 , filterAllOutM
50 , filterAllInM
51 , filterSome
52 , breakAfterSome
53 , toListM
54 , toListSome
55 , composeN
56 , mapN
57 , mapM
58 , transformMapM
59 , transformComposeMapM
60 , transformTeeMapM
61 , transformZipMapM
62 )
63where
65import Control.Applicative (liftA2)
66import Control.DeepSeq (NFData)
67import Control.Exception (try)
68import GHC.Exception (ErrorCall)
69import Streamly.Internal.Data.Stream (Stream)
70import System.Random (randomRIO)
72import qualified Streamly.Internal.Data.Fold as Fold
73import qualified Streamly.Internal.Data.Pipe as Pipe
75#ifdef USE_PRELUDE
76import Streamly.Prelude (foldl', scanl')
77import qualified Streamly.Internal.Data.Stream.IsStream as Stream
78import qualified Streamly.Prelude as Stream
79import Streamly.Benchmark.Prelude
80 ( composeN, sourceConcatMapId, benchIOSink
81 , concatStreamsWith, concatPairsWith
82 )
83#else
84import qualified Streamly.Internal.Data.Stream as Stream
85#endif
87import Gauge
88import Prelude hiding (mapM, replicate)
90#ifdef USE_PRELUDE
91type MonadAsync m = Stream.MonadAsync m
92#else
93type MonadAsync = Monad
94#endif
96{-# INLINE append #-}
97append :: Stream m a -> Stream m a -> Stream m a
98#ifdef USE_PRELUDE
99append = Stream.serial
100#else
101append = Stream.append
102#endif
104{-# INLINE append2 #-}
105append2 :: Monad m => Stream m a -> Stream m a -> Stream m a
106#ifdef USE_PRELUDE
107append2 = Stream.append
108#else
109append2 = Stream.append2
110#endif
112{-# INLINE drain #-}
113drain :: Monad m => Stream m a -> m ()
114drain = Stream.fold Fold.drain
116{-# INLINE fromListM #-}
117fromListM :: MonadAsync m => [m a] -> Stream m a
118#ifdef USE_PRELUDE
119fromListM = Stream.fromListM
120#else
121fromListM = Stream.sequence . Stream.fromList
122#endif
124{-# INLINE fromFoldableM #-}
125fromFoldableM :: MonadAsync m => [m a] -> Stream m a
126#ifdef USE_PRELUDE
127fromFoldableM = Stream.fromFoldableM
128#else
129fromFoldableM = Stream.sequence . Stream.fromFoldable
130#endif
132{-# INLINE sourceUnfoldrM #-}
133sourceUnfoldrM :: MonadAsync m => Int -> Int -> Stream m Int
134sourceUnfoldrM count start = Stream.unfoldrM step start
136 where
138 step cnt =
139 if cnt > start + count
140 then return Nothing
141 else return (Just (cnt, cnt + 1))
143{-# INLINE sourceUnfoldr #-}
144sourceUnfoldr :: Monad m => Int -> Int -> Stream m Int
145sourceUnfoldr count start = Stream.unfoldr step start
147 where
149 step cnt =
150 if cnt > start + count
151 then Nothing
152 else Just (cnt, cnt + 1)
154{-# INLINE sourceUnfoldrAction #-}
155sourceUnfoldrAction :: (Monad m1, Monad m) => Int -> Int -> Stream m (m1 Int)
156sourceUnfoldrAction value n = Stream.unfoldr step n
158 where
160 step cnt =
161 if cnt > n + value
162 then Nothing
163 else Just (return cnt, cnt + 1)
165{-# INLINE sourceFromFoldable #-}
166sourceFromFoldable :: Int -> Int -> Stream m Int
167sourceFromFoldable value n = Stream.fromFoldable [n..n+value]
169#ifndef USE_PRELUDE
170{-# INLINE benchIOSink #-}
171benchIOSink
172 :: (NFData b)
173 => Int -> String -> (Stream IO Int -> IO b) -> Benchmark
174benchIOSink value name f =
175 bench name $ nfIO $ randomRIO (1,1) >>= f . sourceUnfoldrM value
176#endif
178-- | Takes a source, and uses it with a default drain/fold method.
179{-# INLINE benchIOSrc #-}
180benchIOSrc
181 :: String
182 -> (Int -> Stream IO a)
183 -> Benchmark
184benchIOSrc name f =
185 bench name $ nfIO $ randomRIO (1,1) >>= drain . f
187{-# NOINLINE benchIO #-}
188benchIO :: (NFData b) => String -> (Int -> IO b) -> Benchmark
189benchIO name f = bench name $ nfIO $ randomRIO (1,1) >>= f
191#ifndef USE_PRELUDE
192{-# INLINE concatStreamsWith #-}
193concatStreamsWith
194 :: (Stream IO Int -> Stream IO Int -> Stream IO Int)
195 -> Int
196 -> Int
197 -> Int
198 -> IO ()
199concatStreamsWith op outer inner n =
200 drain $ Stream.concatMapWith op
201 (sourceUnfoldrM inner)
202 (sourceUnfoldrM outer n)
204{-# INLINE concatPairsWith #-}
205concatPairsWith
206 :: (Stream IO Int -> Stream IO Int -> Stream IO Int)
207 -> Int
208 -> Int
209 -> Int
210 -> IO ()
211concatPairsWith op outer inner n =
212 drain $ Stream.concatPairsWith op
213 (sourceUnfoldrM inner)
214 (sourceUnfoldrM outer n)
216{-# INLINE sourceConcatMapId #-}
217sourceConcatMapId :: (Monad m)
218 => Int -> Int -> Stream m (Stream m Int)
219sourceConcatMapId value n =
220 Stream.fromFoldable $ fmap (Stream.fromEffect . return) [n..n+value]
221#endif
223{-# INLINE apDiscardFst #-}
224apDiscardFst :: MonadAsync m =>
225 Int -> Int -> m ()
226apDiscardFst linearCount start = drain $
227 sourceUnfoldrM nestedCount2 start
228 *> sourceUnfoldrM nestedCount2 start
230 where
232 nestedCount2 = round (fromIntegral linearCount**(1/2::Double))
234{-# INLINE apDiscardSnd #-}
235apDiscardSnd :: MonadAsync m => Int -> Int -> m ()
236apDiscardSnd linearCount start = drain $
237 sourceUnfoldrM nestedCount2 start
238 <* sourceUnfoldrM nestedCount2 start
240 where
242 nestedCount2 = round (fromIntegral linearCount**(1/2::Double))
244{-# INLINE apLiftA2 #-}
245apLiftA2 :: MonadAsync m => Int -> Int -> m ()
246apLiftA2 linearCount start = drain $
247 liftA2 (+) (sourceUnfoldrM nestedCount2 start)
248 (sourceUnfoldrM nestedCount2 start)
250 where
252 nestedCount2 = round (fromIntegral linearCount**(1/2::Double))
254{-# INLINE toNullAp #-}
255toNullAp :: MonadAsync m => Int -> Int -> m ()
256toNullAp linearCount start = drain $
257 (+) <$> sourceUnfoldrM nestedCount2 start
258 <*> sourceUnfoldrM nestedCount2 start
260 where
262 nestedCount2 = round (fromIntegral linearCount**(1/2::Double))
264{-# INLINE monadThen #-}
265monadThen :: MonadAsync m => Int -> Int -> m ()
266monadThen linearCount start = drain $ do
267 sourceUnfoldrM nestedCount2 start >>
268 sourceUnfoldrM nestedCount2 start
270 where
272 nestedCount2 = round (fromIntegral linearCount**(1/2::Double))
274{-# INLINE toNullM #-}
275toNullM :: MonadAsync m => Int -> Int -> m ()
276toNullM linearCount start = drain $ do
277 x <- sourceUnfoldrM nestedCount2 start
278 y <- sourceUnfoldrM nestedCount2 start
279 return $ x + y
281 where
283 nestedCount2 = round (fromIntegral linearCount**(1/2::Double))
285{-# INLINE toNullM3 #-}
286toNullM3 :: MonadAsync m => Int -> Int -> m ()
287toNullM3 linearCount start = drain $ do
288 x <- sourceUnfoldrM nestedCount3 start
289 y <- sourceUnfoldrM nestedCount3 start
290 z <- sourceUnfoldrM nestedCount3 start
291 return $ x + y + z
292 where
293 nestedCount3 = round (fromIntegral linearCount**(1/3::Double))
295{-# INLINE filterAllOutM #-}
296filterAllOutM :: MonadAsync m => Int -> Int -> m ()
297filterAllOutM linearCount start = drain $ do
298 x <- sourceUnfoldrM nestedCount2 start
299 y <- sourceUnfoldrM nestedCount2 start
300 let s = x + y
301 if s < 0
302 then return s
303 else Stream.nil
304 where
305 nestedCount2 = round (fromIntegral linearCount**(1/2::Double))
307{-# INLINE filterAllInM #-}
308filterAllInM :: MonadAsync m => Int -> Int -> m ()
309filterAllInM linearCount start = drain $ do
310 x <- sourceUnfoldrM nestedCount2 start
311 y <- sourceUnfoldrM nestedCount2 start
312 let s = x + y
313 if s > 0
314 then return s
315 else Stream.nil
316 where
317 nestedCount2 = round (fromIntegral linearCount**(1/2::Double))
319{-# INLINE filterSome #-}
320filterSome :: MonadAsync m => Int -> Int -> m ()
321filterSome linearCount start = drain $ do
322 x <- sourceUnfoldrM nestedCount2 start
323 y <- sourceUnfoldrM nestedCount2 start
324 let s = x + y
325 if s > 1100000
326 then return s
327 else Stream.nil
328 where
329 nestedCount2 = round (fromIntegral linearCount**(1/2::Double))
331{-# INLINE breakAfterSome #-}
332breakAfterSome :: Int -> Int -> IO ()
333breakAfterSome linearCount start = do
334 (_ :: Either ErrorCall ()) <- try $ drain $ do
335 x <- sourceUnfoldrM nestedCount2 start
336 y <- sourceUnfoldrM nestedCount2 start
337 let s = x + y
338 if s > 1100000
339 then error "break"
340 else return s
341 return ()
342 where
343 nestedCount2 = round (fromIntegral linearCount**(1/2::Double))
345{-# INLINE toListM #-}
346toListM :: MonadAsync m => Int -> Int -> m [Int]
347toListM linearCount start = Stream.fold Fold.toList $ do
348 x <- sourceUnfoldrM nestedCount2 start
349 y <- sourceUnfoldrM nestedCount2 start
350 return $ x + y
351 where
352 nestedCount2 = round (fromIntegral linearCount**(1/2::Double))
354-- Taking a specified number of elements is very expensive in logict so we have
355-- a test to measure the same.
356{-# INLINE toListSome #-}
357toListSome :: MonadAsync m => Int -> Int -> m [Int]
358toListSome linearCount start =
359 Stream.fold Fold.toList $ Stream.take 10000 $ do
360 x <- sourceUnfoldrM nestedCount2 start
361 y <- sourceUnfoldrM nestedCount2 start
362 return $ x + y
363 where
364 nestedCount2 = round (fromIntegral linearCount**(1/2::Double))
366#ifndef USE_PRELUDE
367{-# INLINE composeN #-}
368composeN ::
369 (Monad m)
370 => Int
371 -> (Stream m Int -> Stream m Int)
372 -> Stream m Int
373 -> m ()
374composeN n f =
375 case n of
376 1 -> drain . f
377 2 -> drain . f . f
378 3 -> drain . f . f . f
379 4 -> drain . f . f . f . f
380 _ -> undefined
381#endif
383{-# INLINE mapN #-}
384mapN ::
385 Monad m
386 => Int
387 -> Stream m Int
388 -> m ()
389mapN n = composeN n $ fmap (+ 1)
391{-# INLINE mapM #-}
392mapM ::
393 MonadAsync m
394 => Int
395 -> Stream m Int
396 -> m ()
397mapM n = composeN n $ Stream.mapM return
399#ifndef USE_PRELUDE
400foldl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> m b
401foldl' f z = Stream.fold (Fold.foldl' f z)
403scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
404scanl' f z = Stream.scan (Fold.foldl' f z)
405#endif
407{-# INLINE transformMapM #-}
408transformMapM ::
409 (Monad m)
410 => Int
411 -> Stream m Int
412 -> m ()
413transformMapM n = composeN n $ Stream.transform (Pipe.mapM return)
415{-# INLINE transformComposeMapM #-}
416transformComposeMapM ::
417 (Monad m)
418 => Int
419 -> Stream m Int
420 -> m ()
421transformComposeMapM n =
422 composeN n $
423 Stream.transform
424 (Pipe.mapM (\x -> return (x + 1)) `Pipe.compose`
425 Pipe.mapM (\x -> return (x + 2)))
427{-# INLINE transformTeeMapM #-}
428transformTeeMapM ::
429 (Monad m)
430 => Int
431 -> Stream m Int
432 -> m ()
433transformTeeMapM n =
434 composeN n $
435 Stream.transform
436 (Pipe.mapM (\x -> return (x + 1)) `Pipe.tee`
437 Pipe.mapM (\x -> return (x + 2)))
439{-# INLINE transformZipMapM #-}
440transformZipMapM ::
441 (Monad m)
442 => Int
443 -> Stream m Int
444 -> m ()
445transformZipMapM n =
446 composeN n $
447 Stream.transform
448 (Pipe.zipWith
449 (+)
450 (Pipe.mapM (\x -> return (x + 1)))
451 (Pipe.mapM (\x -> return (x + 2))))