git.iko.soy/streamly
on 5cbcd5b help sign in
streamly/src/Streamly/Streams/StreamD.hs
669 lines 20 kB history raw
1{-# LANGUAGE BangPatterns #-}
2{-# LANGUAGE CPP #-}
3{-# LANGUAGE ConstraintKinds #-}
4{-# LANGUAGE ExistentialQuantification #-}
5{-# LANGUAGE FlexibleContexts #-}
6{-# LANGUAGE FlexibleInstances #-}
7{-# LANGUAGE MultiParamTypeClasses #-}
8{-# LANGUAGE RankNTypes #-}
10#include "inline.h"
12-- |
13-- Module : Streamly.Streams.StreamD
14-- Copyright : (c) 2018 Harendra Kumar
15--
16-- License : BSD3
17-- Maintainer : harendra.kumar@gmail.com
18-- Stability : experimental
19-- Portability : GHC
20--
21-- Direct style re-implementation of CPS style stream in StreamK module. The
22-- symbol or suffix 'D' in this module denotes the "Direct" style. GHC is able
23-- to INLINE and fuse direct style better, providing better performance than
24-- CPS implementation.
25--
26-- @
27-- import qualified Streamly.Streams.StreamD as D
28-- @
30-- Some of functions in this file have been adapted from the vector
31-- library, https://hackage.haskell.org/package/vector.
33module Streamly.Streams.StreamD
34 (
35 -- * The stream type
36 Step (..)
37 , Stream (..)
39 -- * Construction
40 , nil
41 , cons
43 -- * Deconstruction
44 , uncons
46 -- * Generation
47 -- ** Unfolds
48 , unfoldr
49 , unfoldrM
51 -- ** Specialized Generation
52 -- | Generate a monadic stream from a seed.
53 , repeat
54 , enumFromStepN
56 -- ** Conversions
57 -- | Transform an input structure into a stream.
58 -- | Direct style stream does not support @fromFoldable@.
59 , yield
60 , yieldM
61 , fromList
62 , fromListM
63 , fromStreamK
65 -- * Elimination
66 -- ** General Folds
67 , foldr
68 , foldrM
69 , foldl'
70 , foldlM'
72 -- ** Specialized Folds
73 , runStream
74 , null
75 , head
76 , tail
77 , last
78 , elem
79 , notElem
80 , all
81 , any
82 , maximum
83 , minimum
85 -- ** Map and Fold
86 , mapM_
88 -- ** Conversions
89 -- | Transform a stream into another type.
90 , toList
91 , toStreamK
93 -- * Transformation
94 -- ** By folding (scans)
95 , scanlM'
97 -- * Filtering
98 , filter
99 , filterM
100 , take
101 , takeWhile
102 , takeWhileM
103 , drop
104 , dropWhile
105 , dropWhileM
107 -- * Mapping
108 , map
109 , mapM
111 -- ** Map and Filter
112 , mapMaybe
113 , mapMaybeM
115 -- * Zipping
116 , zipWith
117 , zipWithM
118 )
119where
121import Data.Maybe (fromJust, isJust)
122import GHC.Types ( SPEC(..) )
123import Prelude
124 hiding (map, mapM, mapM_, repeat, foldr, last, take, filter,
125 takeWhile, drop, dropWhile, all, any, maximum, minimum, elem,
126 notElem, null, head, tail, zipWith)
128import Streamly.SVar (MonadAsync, State(..), defState, rstState)
129import qualified Streamly.Streams.StreamK as K
131------------------------------------------------------------------------------
132-- The direct style stream type
133------------------------------------------------------------------------------
135-- | A stream is a succession of 'Step's. A 'Yield' produces a single value and
136-- the next state of the stream. 'Stop' indicates there are no more values in
137-- the stream.
138data Step s a = Yield a s | Stop
140instance Functor (Step s) where
141 {-# INLINE fmap #-}
142 fmap f (Yield x s) = Yield (f x) s
143 fmap _ Stop = Stop
145-- gst = global state
146-- | A stream consists of a step function that generates the next step given a
147-- current state, and the current state.
148data Stream m a = forall s. Stream (State K.Stream m a -> s -> m (Step s a)) s
150------------------------------------------------------------------------------
151-- Construction
152------------------------------------------------------------------------------
154-- | An empty 'Stream'.
155{-# INLINE_NORMAL nil #-}
156nil :: Monad m => Stream m a
157nil = Stream (\_ _ -> return Stop) ()
159-- | Can fuse but has O(n^2) complexity.
160cons :: Monad m => a -> Stream m a -> Stream m a
161cons x (Stream step state) = Stream step1 Nothing
162 where
163 step1 _ Nothing = return $ Yield x (Just state)
164 step1 gst (Just st) = do
165 r <- step (rstState gst) st
166 case r of
167 Yield a s -> return $ Yield a (Just s)
168 Stop -> return Stop
170-------------------------------------------------------------------------------
171-- Deconstruction
172-------------------------------------------------------------------------------
174-- Does not fuse, has the same performance as the StreamK version.
175{-# INLINE_NORMAL uncons #-}
176uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
177uncons (Stream step state) = go state
178 where
179 go st = do
180 r <- step defState st
181 return $ case r of
182 Yield x s -> Just (x, Stream step s)
183 Stop -> Nothing
185------------------------------------------------------------------------------
186-- Generation by unfold
187------------------------------------------------------------------------------
189{-# INLINE_NORMAL unfoldrM #-}
190unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
191unfoldrM next state = Stream step state
192 where
193 {-# INLINE_LATE step #-}
194 step _ st = do
195 r <- next st
196 return $ case r of
197 Just (x, s) -> Yield x s
198 Nothing -> Stop
200{-# INLINE_LATE unfoldr #-}
201unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
202unfoldr f = unfoldrM (return . f)
204------------------------------------------------------------------------------
205-- Specialized Generation
206------------------------------------------------------------------------------
208repeat :: Monad m => a -> Stream m a
209repeat x = Stream (\_ _ -> return $ Yield x ()) ()
211{-# INLINE_NORMAL enumFromStepN #-}
212enumFromStepN :: (Num a, Monad m) => a -> a -> Int -> Stream m a
213enumFromStepN from stride n =
214 from `seq` stride `seq` n `seq` Stream step (from, n)
215 where
216 {-# INLINE_LATE step #-}
217 step _ (x, i) | i > 0 = return $ Yield x (x + stride, i - 1)
218 | otherwise = return Stop
220-------------------------------------------------------------------------------
221-- Generation by Conversion
222-------------------------------------------------------------------------------
224-- | Create a singleton 'Stream' from a pure value.
225{-# INLINE_NORMAL yield #-}
226yield :: Monad m => a -> Stream m a
227yield x = Stream (\_ s -> return $ step undefined s) True
228 where
229 {-# INLINE_LATE step #-}
230 step _ True = Yield x False
231 step _ False = Stop
233-- | Create a singleton 'Stream' from a monadic action.
234{-# INLINE_NORMAL yieldM #-}
235yieldM :: Monad m => m a -> Stream m a
236yieldM m = Stream step True
237 where
238 {-# INLINE_LATE step #-}
239 step _ True = m >>= \x -> return $ Yield x False
240 step _ False = return Stop
242-- XXX we need the MonadAsync constraint because of a rewrite rule.
243-- | Convert a list of monadic actions to a 'Stream'
244{-# INLINE_LATE fromListM #-}
245fromListM :: MonadAsync m => [m a] -> Stream m a
246fromListM = Stream step
247 where
248 {-# INLINE_LATE step #-}
249 step _ (m:ms) = m >>= \x -> return $ Yield x ms
250 step _ [] = return Stop
252-- | Convert a list of pure values to a 'Stream'
253{-# INLINE_LATE fromList #-}
254fromList :: Monad m => [a] -> Stream m a
255fromList = Stream step
256 where
257 {-# INLINE_LATE step #-}
258 step _ (x:xs) = return $ Yield x xs
259 step _ [] = return Stop
261-- XXX pass the state to streamD
262{-# INLINE_LATE fromStreamK #-}
263fromStreamK :: Monad m => K.Stream m a -> Stream m a
264fromStreamK = Stream step
265 where
266 step gst m1 =
267 let stop = return Stop
268 single a = return $ Yield a K.nil
269 yieldk a r = return $ Yield a r
270 in K.unStream m1 gst stop single yieldk
272------------------------------------------------------------------------------
273-- Elimination by Folds
274------------------------------------------------------------------------------
276{-# INLINE_NORMAL foldrM #-}
277foldrM :: Monad m => (a -> b -> m b) -> b -> Stream m a -> m b
278foldrM f z (Stream step state) = go SPEC state
279 where
280 go !_ st = do
281 r <- step defState st
282 case r of
283 Yield x s -> go SPEC s >>= f x
284 Stop -> return z
286{-# INLINE_NORMAL foldr #-}
287foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b
288foldr f = foldrM (\a b -> return (f a b))
290{-# INLINE_NORMAL foldlM' #-}
291foldlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> m b
292foldlM' fstep begin (Stream step state) = go SPEC begin state
293 where
294 go !_ acc st = acc `seq` do
295 r <- step defState st
296 case r of
297 Yield x s -> do
298 acc' <- fstep acc x
299 go SPEC acc' s
300 Stop -> return acc
302{-# INLINE foldl' #-}
303foldl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> m b
304foldl' fstep = foldlM' (\b a -> return (fstep b a))
306------------------------------------------------------------------------------
307-- Specialized Folds
308------------------------------------------------------------------------------
310-- | Run a streaming composition, discard the results.
311{-# INLINE_LATE runStream #-}
312runStream :: Monad m => Stream m a -> m ()
313runStream (Stream step state) = go SPEC state
314 where
315 go !_ st = do
316 r <- step defState st
317 case r of
318 Yield _ s -> go SPEC s
319 Stop -> return ()
321{-# INLINE_NORMAL null #-}
322null :: Monad m => Stream m a -> m Bool
323null (Stream step state) = go state
324 where
325 go st = do
326 r <- step defState st
327 case r of
328 Yield _ _ -> return False
329 Stop -> return True
331-- XXX SPEC?
332{-# INLINE_NORMAL head #-}
333head :: Monad m => Stream m a -> m (Maybe a)
334head (Stream step state) = go state
335 where
336 go st = do
337 r <- step defState st
338 case r of
339 Yield x _ -> return (Just x)
340 Stop -> return Nothing
342-- Does not fuse, has the same performance as the StreamK version.
343{-# INLINE_NORMAL tail #-}
344tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
345tail (Stream step state) = go state
346 where
347 go st = do
348 r <- step defState st
349 case r of
350 Yield _ s -> return (Just $ Stream step s)
351 Stop -> return Nothing
353-- XXX will it fuse? need custom impl?
354{-# INLINE_NORMAL last #-}
355last :: Monad m => Stream m a -> m (Maybe a)
356last = foldl' (\_ y -> Just y) Nothing
358{-# INLINE_NORMAL elem #-}
359elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
360elem e (Stream step state) = go state
361 where
362 go st = do
363 r <- step defState st
364 case r of
365 Yield x s ->
366 if x == e
367 then return True
368 else go s
369 Stop -> return False
371{-# INLINE_NORMAL notElem #-}
372notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
373notElem e (Stream step state) = go state
374 where
375 go st = do
376 r <- step defState st
377 case r of
378 Yield x s ->
379 if x == e
380 then return False
381 else go s
382 Stop -> return True
384{-# INLINE_NORMAL all #-}
385all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
386all p (Stream step state) = go state
387 where
388 go st = do
389 r <- step defState st
390 case r of
391 Yield x s ->
392 if p x
393 then go s
394 else return False
395 Stop -> return True
397{-# INLINE_NORMAL any #-}
398any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
399any p (Stream step state) = go state
400 where
401 go st = do
402 r <- step defState st
403 case r of
404 Yield x s ->
405 if p x
406 then return True
407 else go s
408 Stop -> return False
410{-# INLINE_NORMAL maximum #-}
411maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
412maximum (Stream step state) = go Nothing state
413 where
414 go Nothing st = do
415 r <- step defState st
416 case r of
417 Yield x s -> go (Just x) s
418 Stop -> return Nothing
419 go (Just acc) st = do
420 r <- step defState st
421 case r of
422 Yield x s ->
423 if acc <= x
424 then go (Just x) s
425 else go (Just acc) s
426 Stop -> return (Just acc)
428{-# INLINE_NORMAL minimum #-}
429minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
430minimum (Stream step state) = go Nothing state
431 where
432 go Nothing st = do
433 r <- step defState st
434 case r of
435 Yield x s -> go (Just x) s
436 Stop -> return Nothing
437 go (Just acc) st = do
438 r <- step defState st
439 case r of
440 Yield x s ->
441 if acc <= x
442 then go (Just acc) s
443 else go (Just x) s
444 Stop -> return (Just acc)
446------------------------------------------------------------------------------
447-- Map and Fold
448------------------------------------------------------------------------------
450-- | Execute a monadic action for each element of the 'Stream'
451{-# INLINE_NORMAL mapM_ #-}
452mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
453mapM_ m = runStream . mapM m
455------------------------------------------------------------------------------
456-- Converting folds
457------------------------------------------------------------------------------
459{-# INLINE toList #-}
460toList :: Monad m => Stream m a -> m [a]
461toList = foldr (:) []
463-- Convert a direct stream to and from CPS encoded stream
464{-# INLINE_LATE toStreamK #-}
465toStreamK :: Monad m => Stream m a -> K.Stream m a
466toStreamK (Stream step state) = go state
467 where
468 go st = K.Stream $ \gst stp _ yld -> do
469 r <- step gst st
470 case r of
471 Yield x s -> yld x (go s)
472 Stop -> stp
474#ifndef DISABLE_FUSION
475{-# RULES "fromStreamK/toStreamK fusion"
476 forall s. toStreamK (fromStreamK s) = s #-}
477{-# RULES "toStreamK/fromStreamK fusion"
478 forall s. fromStreamK (toStreamK s) = s #-}
479#endif
481------------------------------------------------------------------------------
482-- Transformation by Folding (Scans)
483------------------------------------------------------------------------------
485{-# INLINE_NORMAL postscanlM' #-}
486postscanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
487postscanlM' fstep begin (Stream step state) =
488 begin `seq` Stream step' (state, begin)
489 where
490 {-# INLINE_LATE step' #-}
491 step' gst (st, acc) = acc `seq` do
492 r <- step (rstState gst) st
493 case r of
494 Yield x s -> do
495 y <- fstep acc x
496 y `seq` return (Yield y (s, y))
497 Stop -> return Stop
499{-# INLINE scanlM' #-}
500scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
501scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)
503-------------------------------------------------------------------------------
504-- Filtering
505-------------------------------------------------------------------------------
507{-# INLINE_NORMAL take #-}
508take :: Monad m => Int -> Stream m a -> Stream m a
509take n (Stream step state) = n `seq` Stream step' (state, 0)
510 where
511 {-# INLINE_LATE step' #-}
512 step' gst (st, i) | i < n = do
513 r <- step (rstState gst) st
514 return $ case r of
515 Yield x s -> Yield x (s, i + 1)
516 Stop -> Stop
517 step' _ (_, _) = return Stop
519{-# INLINE_NORMAL takeWhileM #-}
520takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
521takeWhileM f (Stream step state) = Stream step' state
522 where
523 {-# INLINE_LATE step' #-}
524 step' gst st = do
525 r <- step (rstState gst) st
526 case r of
527 Yield x s -> do
528 b <- f x
529 return $ if b then Yield x s else Stop
530 Stop -> return Stop
532{-# INLINE takeWhile #-}
533takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
534takeWhile f = takeWhileM (return . f)
536{-# INLINE_NORMAL drop #-}
537drop :: Monad m => Int -> Stream m a -> Stream m a
538drop n (Stream step state) = Stream step' (state, n)
539 where
540 {-# INLINE_LATE step' #-}
541 step' gst (st, i) = do
542 r <- step (rstState gst) st
543 case r of
544 Yield _ s | i > 0 -> step' gst (s, i - 1)
545 Yield x s -> return $ Yield x (s, 0)
546 Stop -> return Stop
548data DropWhileState s a
549 = DropWhileDrop s
550 | DropWhileYield a s
551 | DropWhileNext s
553{-# INLINE_NORMAL dropWhileM #-}
554dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
555dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
556 where
557 {-# INLINE_LATE step' #-}
558 step' gst (DropWhileDrop st) = do
559 r <- step (rstState gst) st
560 case r of
561 Yield x s -> do
562 b <- f x
563 if b
564 then step' gst (DropWhileDrop s)
565 else step' gst (DropWhileYield x s)
566 Stop -> return Stop
568 step' gst (DropWhileNext st) = do
569 r <- step (rstState gst) st
570 case r of
571 Yield x s -> step' gst (DropWhileYield x s)
572 Stop -> return Stop
574 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
576{-# INLINE dropWhile #-}
577dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
578dropWhile f = dropWhileM (return . f)
580{-# INLINE_NORMAL filterM #-}
581filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
582filterM f (Stream step state) = Stream step' state
583 where
584 {-# INLINE_LATE step' #-}
585 step' gst st = do
586 r <- step (rstState gst) st
587 case r of
588 Yield x s -> do
589 b <- f x
590 if b
591 then return $ Yield x s
592 else step' gst s
593 Stop -> return Stop
595{-# INLINE filter #-}
596filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
597filter f = filterM (return . f)
599------------------------------------------------------------------------------
600-- Transformation by Mapping
601------------------------------------------------------------------------------
603-- | Map a monadic function over a 'Stream'
604{-# INLINE_NORMAL mapM #-}
605mapM :: Monad m => (a -> m b) -> Stream m a -> Stream m b
606mapM f (Stream step state) = Stream step' state
607 where
608 {-# INLINE_LATE step' #-}
609 step' gst st = do
610 r <- step (rstState gst) st
611 case r of
612 Yield x s -> f x >>= \a -> return $ Yield a s
613 Stop -> return Stop
615{-# INLINE map #-}
616map :: Monad m => (a -> b) -> Stream m a -> Stream m b
617map f = mapM (return . f)
619------------------------------------------------------------------------------
620-- Transformation by Map and Filter
621------------------------------------------------------------------------------
623-- XXX Will this always fuse properly?
624{-# INLINE_NORMAL mapMaybe #-}
625mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
626mapMaybe f = fmap fromJust . filter isJust . map f
628{-# INLINE_NORMAL mapMaybeM #-}
629mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
630mapMaybeM f = fmap fromJust . filter isJust . mapM f
632------------------------------------------------------------------------------
633-- Instances
634------------------------------------------------------------------------------
636{-# INLINE_NORMAL zipWithM #-}
637zipWithM :: Monad m
638 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
639zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
640 where
641 {-# INLINE_LATE step #-}
642 step gst (sa, sb, Nothing) = do
643 r <- stepa (rstState gst) sa
644 case r of
645 Yield x sa' -> step gst (sa', sb, Just x)
646 Stop -> return Stop
648 step gst (sa, sb, Just x) = do
649 r <- stepb (rstState gst) sb
650 case r of
651 Yield y sb' -> do
652 z <- f x y
653 return $ Yield z (sa, sb', Nothing)
654 Stop -> return Stop
656{-# RULES "zipWithM xs xs"
657 forall f xs. zipWithM f xs xs = mapM (\x -> f x x) xs #-}
659{-# INLINE zipWith #-}
660zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
661zipWith f = zipWithM (\a b -> return (f a b))
663------------------------------------------------------------------------------
664-- Instances
665------------------------------------------------------------------------------
667instance Monad m => Functor (Stream m) where
668 {-# INLINE fmap #-}
669 fmap = map