git.iko.soy/streamly
on bc5617f help sign in
streamly/src/Streamly/Internal/Data/Stream/StreamDK.hs
156 lines 4.1 kB history raw
1#include "inline.hs"
3-- |
4-- Module : Streamly.Internal.Data.Stream.StreamDK
5-- Copyright : (c) 2019 Composewell Technologies
6-- License : BSD3
7-- Maintainer : streamly@composewell.com
8-- Stability : experimental
9-- Portability : GHC
10--
12module Streamly.Internal.Data.Stream.StreamDK
13 (
14 -- * Stream Type
16 Stream
17 , Step (..)
19 -- * Construction
20 , nil
21 , cons
22 , consM
23 , unfoldr
24 , unfoldrM
25 , replicateM
27 -- * Folding
28 , uncons
29 , foldrS
31 -- * Specific Folds
32 , drain
33 )
34where
36import Streamly.Internal.Data.Stream.StreamDK.Type (Stream(..), Step(..))
38-------------------------------------------------------------------------------
39-- Construction
40-------------------------------------------------------------------------------
42nil :: Monad m => Stream m a
43nil = Stream $ return Stop
45{-# INLINE_NORMAL cons #-}
46cons :: Monad m => a -> Stream m a -> Stream m a
47cons x xs = Stream $ return $ Yield x xs
49consM :: Monad m => m a -> Stream m a -> Stream m a
50consM eff xs = Stream $ eff >>= \x -> return $ Yield x xs
52unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
53unfoldrM next state = Stream (step' state)
54 where
55 step' st = do
56 r <- next st
57 return $ case r of
58 Just (x, s) -> Yield x (Stream (step' s))
59 Nothing -> Stop
60{-
61unfoldrM next s0 = buildM $ \yld stp ->
62 let go s = do
63 r <- next s
64 case r of
65 Just (a, b) -> yld a (go b)
66 Nothing -> stp
67 in go s0
68-}
70{-# INLINE unfoldr #-}
71unfoldr :: Monad m => (b -> Maybe (a, b)) -> b -> Stream m a
72unfoldr next s0 = build $ \yld stp ->
73 let go s =
74 case next s of
75 Just (a, b) -> yld a (go b)
76 Nothing -> stp
77 in go s0
79replicateM :: Monad m => Int -> a -> Stream m a
80replicateM n x = Stream (step n)
81 where
82 step i = return $
83 if i <= 0
84 then Stop
85 else Yield x (Stream (step (i - 1)))
87-------------------------------------------------------------------------------
88-- Folding
89-------------------------------------------------------------------------------
91uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
92uncons (Stream step) = do
93 r <- step
94 return $ case r of
95 Yield x xs -> Just (x, xs)
96 Stop -> Nothing
98-- | Lazy right associative fold to a stream.
99{-# INLINE_NORMAL foldrS #-}
100foldrS :: Monad m
101 => (a -> Stream m b -> Stream m b)
102 -> Stream m b
103 -> Stream m a
104 -> Stream m b
105foldrS f streamb = go
106 where
107 go (Stream stepa) = Stream $ do
108 r <- stepa
109 case r of
110 Yield x xs -> let Stream step = f x (go xs) in step
111 Stop -> let Stream step = streamb in step
113{-# INLINE_LATE foldrM #-}
114foldrM :: Monad m => (a -> m b -> m b) -> m b -> Stream m a -> m b
115foldrM fstep acc = go
116 where
117 go (Stream step) = do
118 r <- step
119 case r of
120 Yield x xs -> fstep x (go xs)
121 Stop -> acc
123{-# INLINE_NORMAL build #-}
124build :: Monad m
125 => forall a. (forall b. (a -> b -> b) -> b -> b) -> Stream m a
126build g = g cons nil
128{-# RULES
129"foldrM/build" forall k z (g :: forall b. (a -> b -> b) -> b -> b).
130 foldrM k z (build g) = g k z #-}
132{-
133-- To fuse foldrM with unfoldrM we need the type m1 to be polymorphic such that
134-- it is either Monad m or Stream m. So that we can use cons/nil as well as
135-- monadic construction function as its arguments.
136--
137{-# INLINE_NORMAL buildM #-}
138buildM :: Monad m
139 => forall a. (forall b. (a -> m1 b -> m1 b) -> m1 b -> m1 b) -> Stream m a
140buildM g = g cons nil
141-}
143-------------------------------------------------------------------------------
144-- Specific folds
145-------------------------------------------------------------------------------
147{-# INLINE drain #-}
148drain :: Monad m => Stream m a -> m ()
149drain = foldrM (\_ xs -> xs) (return ())
150{-
151drain (Stream step) = do
152 r <- step
153 case r of
154 Yield _ next -> drain next
155 Stop -> return ()
156 -}