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