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