git.iko.soy/streamly
on 628355f help sign in
streamly/src/Streamly/Streams/StreamD.hs
701 lines 22 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.hs"
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'
96 , scanl'
98 -- * Filtering
99 , filter
100 , filterM
101 , take
102 , takeWhile
103 , takeWhileM
104 , drop
105 , dropWhile
106 , dropWhileM
108 -- * Mapping
109 , map
110 , mapM
112 -- ** Map and Filter
113 , mapMaybe
114 , mapMaybeM
116 -- * Zipping
117 , zipWith
118 , zipWithM
119 )
120where
122import Data.Maybe (fromJust, isJust)
123import GHC.Types ( SPEC(..) )
124import Prelude
125 hiding (map, mapM, mapM_, repeat, foldr, last, take, filter,
126 takeWhile, drop, dropWhile, all, any, maximum, minimum, elem,
127 notElem, null, head, tail, zipWith)
129import Streamly.SVar (MonadAsync, State(..), defState, rstState)
130import qualified Streamly.Streams.StreamK as K
132------------------------------------------------------------------------------
133-- The direct style stream type
134------------------------------------------------------------------------------
136-- | A stream is a succession of 'Step's. A 'Yield' produces a single value and
137-- the next state of the stream. 'Stop' indicates there are no more values in
138-- the stream.
139data Step s a = Yield a s | Skip s | Stop
141instance Functor (Step s) where
142 {-# INLINE fmap #-}
143 fmap f (Yield x s) = Yield (f x) s
144 fmap _ (Skip s) = Skip s
145 fmap _ Stop = Stop
147-- gst = global state
148-- | A stream consists of a step function that generates the next step given a
149-- current state, and the current state.
150data Stream m a = forall s. Stream (State K.Stream m a -> s -> m (Step s a)) s
152------------------------------------------------------------------------------
153-- Construction
154------------------------------------------------------------------------------
156-- | An empty 'Stream'.
157{-# INLINE_NORMAL nil #-}
158nil :: Monad m => Stream m a
159nil = Stream (\_ _ -> return Stop) ()
161-- | Can fuse but has O(n^2) complexity.
162{-# INLINE_NORMAL cons #-}
163cons :: Monad m => a -> Stream m a -> Stream m a
164cons x (Stream step state) = Stream step1 Nothing
165 where
166 {-# INLINE_LATE step1 #-}
167 step1 _ Nothing = return $ Yield x (Just state)
168 step1 gst (Just st) = do
169 r <- step (rstState gst) st
170 return $
171 case r of
172 Yield a s -> Yield a (Just s)
173 Skip s -> Skip (Just s)
174 Stop -> Stop
176-------------------------------------------------------------------------------
177-- Deconstruction
178-------------------------------------------------------------------------------
180-- Does not fuse, has the same performance as the StreamK version.
181{-# INLINE_NORMAL uncons #-}
182uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
183uncons (Stream step state) = go state
184 where
185 go st = do
186 r <- step defState st
187 case r of
188 Yield x s -> return $ Just (x, Stream step s)
189 Skip s -> go s
190 Stop -> return Nothing
192------------------------------------------------------------------------------
193-- Generation by unfold
194------------------------------------------------------------------------------
196{-# INLINE_NORMAL unfoldrM #-}
197unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
198unfoldrM next state = Stream step state
199 where
200 {-# INLINE_LATE step #-}
201 step _ st = do
202 r <- next st
203 return $ case r of
204 Just (x, s) -> Yield x s
205 Nothing -> Stop
207{-# INLINE_LATE unfoldr #-}
208unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
209unfoldr f = unfoldrM (return . f)
211------------------------------------------------------------------------------
212-- Specialized Generation
213------------------------------------------------------------------------------
215repeat :: Monad m => a -> Stream m a
216repeat x = Stream (\_ _ -> return $ Yield x ()) ()
218{-# INLINE_NORMAL enumFromStepN #-}
219enumFromStepN :: (Num a, Monad m) => a -> a -> Int -> Stream m a
220enumFromStepN from stride n =
221 from `seq` stride `seq` n `seq` Stream step (from, n)
222 where
223 {-# INLINE_LATE step #-}
224 step _ (x, i) | i > 0 = return $ Yield x (x + stride, i - 1)
225 | otherwise = return Stop
227-------------------------------------------------------------------------------
228-- Generation by Conversion
229-------------------------------------------------------------------------------
231-- | Create a singleton 'Stream' from a pure value.
232{-# INLINE_NORMAL yield #-}
233yield :: Monad m => a -> Stream m a
234yield x = Stream (\_ s -> return $ step undefined s) True
235 where
236 {-# INLINE_LATE step #-}
237 step _ True = Yield x False
238 step _ False = Stop
240-- | Create a singleton 'Stream' from a monadic action.
241{-# INLINE_NORMAL yieldM #-}
242yieldM :: Monad m => m a -> Stream m a
243yieldM m = Stream step True
244 where
245 {-# INLINE_LATE step #-}
246 step _ True = m >>= \x -> return $ Yield x False
247 step _ False = return Stop
249-- XXX we need the MonadAsync constraint because of a rewrite rule.
250-- | Convert a list of monadic actions to a 'Stream'
251{-# INLINE_LATE fromListM #-}
252fromListM :: MonadAsync m => [m a] -> Stream m a
253fromListM = Stream step
254 where
255 {-# INLINE_LATE step #-}
256 step _ (m:ms) = m >>= \x -> return $ Yield x ms
257 step _ [] = return Stop
259-- | Convert a list of pure values to a 'Stream'
260{-# INLINE_LATE fromList #-}
261fromList :: Monad m => [a] -> Stream m a
262fromList = Stream step
263 where
264 {-# INLINE_LATE step #-}
265 step _ (x:xs) = return $ Yield x xs
266 step _ [] = return Stop
268-- XXX pass the state to streamD
269{-# INLINE_LATE fromStreamK #-}
270fromStreamK :: Monad m => K.Stream m a -> Stream m a
271fromStreamK = Stream step
272 where
273 step gst m1 =
274 let stop = return Stop
275 single a = return $ Yield a K.nil
276 yieldk a r = return $ Yield a r
277 in K.unStream m1 gst stop single yieldk
279------------------------------------------------------------------------------
280-- Elimination by Folds
281------------------------------------------------------------------------------
283{-# INLINE_NORMAL foldrM #-}
284foldrM :: Monad m => (a -> b -> m b) -> b -> Stream m a -> m b
285foldrM f z (Stream step state) = go SPEC state
286 where
287 go !_ st = do
288 r <- step defState st
289 case r of
290 Yield x s -> go SPEC s >>= f x
291 Skip s -> go SPEC s
292 Stop -> return z
294{-# INLINE_NORMAL foldr #-}
295foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b
296foldr f = foldrM (\a b -> return (f a b))
298{-# INLINE_NORMAL foldlM' #-}
299foldlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> m b
300foldlM' fstep begin (Stream step state) = go SPEC begin state
301 where
302 go !_ acc st = acc `seq` do
303 r <- step defState st
304 case r of
305 Yield x s -> do
306 acc' <- fstep acc x
307 go SPEC acc' s
308 Skip s -> go SPEC acc s
309 Stop -> return acc
311{-# INLINE foldl' #-}
312foldl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> m b
313foldl' fstep = foldlM' (\b a -> return (fstep b a))
315------------------------------------------------------------------------------
316-- Specialized Folds
317------------------------------------------------------------------------------
319-- | Run a streaming composition, discard the results.
320{-# INLINE_LATE runStream #-}
321runStream :: Monad m => Stream m a -> m ()
322runStream (Stream step state) = go SPEC state
323 where
324 go !_ st = do
325 r <- step defState st
326 case r of
327 Yield _ s -> go SPEC s
328 Skip s -> go SPEC s
329 Stop -> return ()
331{-# INLINE_NORMAL null #-}
332null :: Monad m => Stream m a -> m Bool
333null (Stream step state) = go state
334 where
335 go st = do
336 r <- step defState st
337 case r of
338 Yield _ _ -> return False
339 Skip s -> go s
340 Stop -> return True
342-- XXX SPEC?
343{-# INLINE_NORMAL head #-}
344head :: Monad m => Stream m a -> m (Maybe a)
345head (Stream step state) = go state
346 where
347 go st = do
348 r <- step defState st
349 case r of
350 Yield x _ -> return (Just x)
351 Skip s -> go s
352 Stop -> return Nothing
354-- Does not fuse, has the same performance as the StreamK version.
355{-# INLINE_NORMAL tail #-}
356tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
357tail (Stream step state) = go state
358 where
359 go st = do
360 r <- step defState st
361 case r of
362 Yield _ s -> return (Just $ Stream step s)
363 Skip s -> go s
364 Stop -> return Nothing
366-- XXX will it fuse? need custom impl?
367{-# INLINE_NORMAL last #-}
368last :: Monad m => Stream m a -> m (Maybe a)
369last = foldl' (\_ y -> Just y) Nothing
371{-# INLINE_NORMAL elem #-}
372elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
373elem 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 | x == e -> return True
380 | otherwise -> go s
381 Skip s -> go s
382 Stop -> return False
384{-# INLINE_NORMAL notElem #-}
385notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
386notElem e s = fmap not (elem e s)
388{-# INLINE_NORMAL all #-}
389all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
390all p (Stream step state) = go state
391 where
392 go st = do
393 r <- step defState st
394 case r of
395 Yield x s
396 | p x -> go s
397 | otherwise -> return False
398 Skip s -> go s
399 Stop -> return True
401{-# INLINE_NORMAL any #-}
402any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
403any p (Stream step state) = go state
404 where
405 go st = do
406 r <- step defState st
407 case r of
408 Yield x s
409 | p x -> return True
410 | otherwise -> go s
411 Skip s -> go s
412 Stop -> return False
414{-# INLINE_NORMAL maximum #-}
415maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
416maximum (Stream step state) = go Nothing state
417 where
418 go Nothing st = do
419 r <- step defState st
420 case r of
421 Yield x s -> go (Just x) s
422 Skip s -> go Nothing s
423 Stop -> return Nothing
424 go (Just acc) st = do
425 r <- step defState st
426 case r of
427 Yield x s
428 | acc <= x -> go (Just x) s
429 | otherwise -> go (Just acc) s
430 Skip s -> go (Just acc) s
431 Stop -> return (Just acc)
433{-# INLINE_NORMAL minimum #-}
434minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
435minimum (Stream step state) = go Nothing state
436 where
437 go Nothing st = do
438 r <- step defState st
439 case r of
440 Yield x s -> go (Just x) s
441 Skip s -> go Nothing s
442 Stop -> return Nothing
443 go (Just acc) st = do
444 r <- step defState st
445 case r of
446 Yield x s
447 | acc <= x -> go (Just acc) s
448 | otherwise -> go (Just x) s
449 Skip s -> go (Just acc) s
450 Stop -> return (Just acc)
452------------------------------------------------------------------------------
453-- Map and Fold
454------------------------------------------------------------------------------
456-- | Execute a monadic action for each element of the 'Stream'
457{-# INLINE_NORMAL mapM_ #-}
458mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
459mapM_ m = runStream . mapM m
461------------------------------------------------------------------------------
462-- Converting folds
463------------------------------------------------------------------------------
465{-# INLINE toList #-}
466toList :: Monad m => Stream m a -> m [a]
467toList = foldr (:) []
469-- Convert a direct stream to and from CPS encoded stream
470{-# INLINE_LATE toStreamK #-}
471toStreamK :: Monad m => Stream m a -> K.Stream m a
472toStreamK (Stream step state) = go state
473 where
474 go st = K.Stream $ \gst stp sng yld -> do
475 r <- step gst st
476 case r of
477 Yield x s -> yld x (go s)
478 Skip s -> K.unStream (go s) gst stp sng yld
479 Stop -> stp
481#ifndef DISABLE_FUSION
482{-# RULES "fromStreamK/toStreamK fusion"
483 forall s. toStreamK (fromStreamK s) = s #-}
484{-# RULES "toStreamK/fromStreamK fusion"
485 forall s. fromStreamK (toStreamK s) = s #-}
486#endif
488------------------------------------------------------------------------------
489-- Transformation by Folding (Scans)
490------------------------------------------------------------------------------
492{-# INLINE_NORMAL postscanlM' #-}
493postscanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
494postscanlM' fstep begin (Stream step state) =
495 begin `seq` Stream step' (state, begin)
496 where
497 {-# INLINE_LATE step' #-}
498 step' gst (st, acc) = acc `seq` do
499 r <- step (rstState gst) st
500 case r of
501 Yield x s -> do
502 y <- fstep acc x
503 y `seq` return (Yield y (s, y))
504 Skip s -> return $ Skip (s, acc)
505 Stop -> return Stop
507{-# INLINE_LATE scanlM' #-}
508scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
509scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)
511{-# INLINE scanl' #-}
512scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
513scanl' f = scanlM' (\a b -> return (f a b))
515-------------------------------------------------------------------------------
516-- Filtering
517-------------------------------------------------------------------------------
519{-# INLINE_NORMAL take #-}
520take :: Monad m => Int -> Stream m a -> Stream m a
521take n (Stream step state) = n `seq` Stream step' (state, 0)
522 where
523 {-# INLINE_LATE step' #-}
524 step' gst (st, i) | i < n = do
525 r <- step (rstState gst) st
526 return $ case r of
527 Yield x s -> Yield x (s, i + 1)
528 Skip s -> Skip (s, i)
529 Stop -> Stop
530 step' _ (_, _) = return Stop
532{-# INLINE_NORMAL takeWhileM #-}
533takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
534takeWhileM f (Stream step state) = Stream step' state
535 where
536 {-# INLINE_LATE step' #-}
537 step' gst st = do
538 r <- step (rstState gst) st
539 case r of
540 Yield x s -> do
541 b <- f x
542 return $ if b then Yield x s else Stop
543 Skip s -> return $ Skip s
544 Stop -> return Stop
546{-# INLINE takeWhile #-}
547takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
548takeWhile f = takeWhileM (return . f)
550{-# INLINE_NORMAL drop #-}
551drop :: Monad m => Int -> Stream m a -> Stream m a
552drop n (Stream step state) = Stream step' (state, Just n)
553 where
554 {-# INLINE_LATE step' #-}
555 step' gst (st, Just i)
556 | i > 0 = do
557 r <- step (rstState gst) st
558 return $
559 case r of
560 Yield _ s -> Skip (s, Just (i - 1))
561 Skip s -> Skip (s, Just i)
562 Stop -> Stop
563 | otherwise = return $ Skip (st, Nothing)
565 step' gst (st, Nothing) = do
566 r <- step (rstState gst) st
567 return $
568 case r of
569 Yield x s -> Yield x (s, Nothing)
570 Skip s -> Skip (s, Nothing)
571 Stop -> Stop
573data DropWhileState s a
574 = DropWhileDrop s
575 | DropWhileYield a s
576 | DropWhileNext s
578{-# INLINE_NORMAL dropWhileM #-}
579dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
580dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
581 where
582 {-# INLINE_LATE step' #-}
583 step' gst (DropWhileDrop st) = do
584 r <- step (rstState gst) st
585 case r of
586 Yield x s -> do
587 b <- f x
588 if b
589 then return $ Skip (DropWhileDrop s)
590 else return $ Skip (DropWhileYield x s)
591 Skip s -> return $ Skip (DropWhileDrop s)
592 Stop -> return Stop
594 step' gst (DropWhileNext st) = do
595 r <- step (rstState gst) st
596 case r of
597 Yield x s -> return $ Skip (DropWhileYield x s)
598 Skip s -> return $ Skip (DropWhileNext s)
599 Stop -> return Stop
601 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
603{-# INLINE dropWhile #-}
604dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
605dropWhile f = dropWhileM (return . f)
607{-# INLINE_NORMAL filterM #-}
608filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
609filterM 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 -> do
616 b <- f x
617 return $ if b
618 then Yield x s
619 else Skip s
620 Skip s -> return $ Skip s
621 Stop -> return Stop
623{-# INLINE filter #-}
624filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
625filter f = filterM (return . f)
627------------------------------------------------------------------------------
628-- Transformation by Mapping
629------------------------------------------------------------------------------
631-- | Map a monadic function over a 'Stream'
632{-# INLINE_NORMAL mapM #-}
633mapM :: Monad m => (a -> m b) -> Stream m a -> Stream m b
634mapM f (Stream step state) = Stream step' state
635 where
636 {-# INLINE_LATE step' #-}
637 step' gst st = do
638 r <- step (rstState gst) st
639 case r of
640 Yield x s -> f x >>= \a -> return $ Yield a s
641 Skip s -> return $ Skip s
642 Stop -> return Stop
644{-# INLINE map #-}
645map :: Monad m => (a -> b) -> Stream m a -> Stream m b
646map f = mapM (return . f)
648------------------------------------------------------------------------------
649-- Transformation by Map and Filter
650------------------------------------------------------------------------------
652-- XXX Will this always fuse properly?
653{-# INLINE_NORMAL mapMaybe #-}
654mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
655mapMaybe f = fmap fromJust . filter isJust . map f
657{-# INLINE_NORMAL mapMaybeM #-}
658mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
659mapMaybeM f = fmap fromJust . filter isJust . mapM f
661------------------------------------------------------------------------------
662-- Instances
663------------------------------------------------------------------------------
665{-# INLINE_NORMAL zipWithM #-}
666zipWithM :: Monad m
667 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
668zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
669 where
670 {-# INLINE_LATE step #-}
671 step gst (sa, sb, Nothing) = do
672 r <- stepa (rstState gst) sa
673 return $
674 case r of
675 Yield x sa' -> Skip (sa', sb, Just x)
676 Skip sa' -> Skip (sa', sb, Nothing)
677 Stop -> Stop
679 step gst (sa, sb, Just x) = do
680 r <- stepb (rstState gst) sb
681 case r of
682 Yield y sb' -> do
683 z <- f x y
684 return $ Yield z (sa, sb', Nothing)
685 Skip sb' -> return $ Skip (sa, sb', Just x)
686 Stop -> return Stop
688{-# RULES "zipWithM xs xs"
689 forall f xs. zipWithM f xs xs = mapM (\x -> f x x) xs #-}
691{-# INLINE zipWith #-}
692zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
693zipWith f = zipWithM (\a b -> return (f a b))
695------------------------------------------------------------------------------
696-- Instances
697------------------------------------------------------------------------------
699instance Monad m => Functor (Stream m) where
700 {-# INLINE fmap #-}
701 fmap = map