git.iko.soy/streamly
on cedc81c help sign in
streamly/core/src/Streamly/Internal/Data/Stream/Transformer.hs
182 lines 5.8 kB history raw
1{-# LANGUAGE CPP #-}
2-- |
3-- Module : Streamly.Internal.Data.Stream.Transformer
4-- Copyright : (c) 2018 Composewell Technologies
5-- License : BSD-3-Clause
6-- Maintainer : streamly@composewell.com
7-- Stability : experimental
8-- Portability : GHC
9--
10-- Transform the underlying monad of a stream using a monad transfomer.
12module Streamly.Internal.Data.Stream.Transformer
13 (
14 foldlT
15 , foldrT
17 -- * Transform Inner Monad
18 , liftInner
19 , runReaderT
20 , usingReaderT
21 , evalStateT
22 , runStateT
23 , usingStateT
24 )
25where
27#include "inline.hs"
29import Control.Monad.Trans.Class (MonadTrans(lift))
30import Control.Monad.Trans.Reader (ReaderT)
31import Control.Monad.Trans.State.Strict (StateT)
32import GHC.Types (SPEC(..))
33import Streamly.Internal.Data.SVar.Type (defState, adaptState)
35import qualified Control.Monad.Trans.Reader as Reader
36import qualified Control.Monad.Trans.State.Strict as State
38import Streamly.Internal.Data.Stream.Type
40#include "DocTestDataStream.hs"
42-- | Lazy left fold to a transformer monad.
43--
44{-# INLINE_NORMAL foldlT #-}
45foldlT :: (Monad m, Monad (s m), MonadTrans s)
46 => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
47foldlT fstep begin (Stream step state) = go SPEC begin state
48 where
49 go !_ acc st = do
50 r <- lift $ step defState st
51 case r of
52 Yield x s -> go SPEC (fstep acc x) s
53 Skip s -> go SPEC acc s
54 Stop -> acc
56-- | Right fold to a transformer monad. This is the most general right fold
57-- function. 'foldrS' is a special case of 'foldrT', however 'foldrS'
58-- implementation can be more efficient:
59--
60-- >>> foldrS = Stream.foldrT
61--
62-- >>> step f x xs = lift $ f x (runIdentityT xs)
63-- >>> foldrM f z s = runIdentityT $ Stream.foldrT (step f) (lift z) s
64--
65-- 'foldrT' can be used to translate streamly streams to other transformer
66-- monads e.g. to a different streaming type.
67--
68-- /Pre-release/
69{-# INLINE_NORMAL foldrT #-}
70foldrT :: (Monad m, Monad (t m), MonadTrans t)
71 => (a -> t m b -> t m b) -> t m b -> Stream m a -> t m b
72foldrT f final (Stream step state) = go SPEC state
73 where
74 {-# INLINE_LATE go #-}
75 go !_ st = do
76 r <- lift $ step defState st
77 case r of
78 Yield x s -> f x (go SPEC s)
79 Skip s -> go SPEC s
80 Stop -> final
82-------------------------------------------------------------------------------
83-- Transform Inner Monad
84-------------------------------------------------------------------------------
86-- | Lift the inner monad @m@ of @Stream m a@ to @t m@ where @t@ is a monad
87-- transformer.
88--
89{-# INLINE_NORMAL liftInner #-}
90liftInner :: (Monad m, MonadTrans t, Monad (t m))
91 => Stream m a -> Stream (t m) a
92liftInner (Stream step state) = Stream step' state
93 where
94 {-# INLINE_LATE step' #-}
95 step' gst st = do
96 r <- lift $ step (adaptState gst) st
97 return $ case r of
98 Yield x s -> Yield x s
99 Skip s -> Skip s
100 Stop -> Stop
102------------------------------------------------------------------------------
103-- Sharing read only state in a stream
104------------------------------------------------------------------------------
106-- | Evaluate the inner monad of a stream as 'ReaderT'.
107--
108{-# INLINE_NORMAL runReaderT #-}
109runReaderT :: Monad m => m s -> Stream (ReaderT s m) a -> Stream m a
110runReaderT env (Stream step state) = Stream step' (state, env)
111 where
112 {-# INLINE_LATE step' #-}
113 step' gst (st, action) = do
114 sv <- action
115 r <- Reader.runReaderT (step (adaptState gst) st) sv
116 return $ case r of
117 Yield x s -> Yield x (s, return sv)
118 Skip s -> Skip (s, return sv)
119 Stop -> Stop
121-- | Run a stream transformation using a given environment.
122--
123{-# INLINE usingReaderT #-}
124usingReaderT
125 :: Monad m
126 => m r
127 -> (Stream (ReaderT r m) a -> Stream (ReaderT r m) a)
128 -> Stream m a
129 -> Stream m a
130usingReaderT r f xs = runReaderT r $ f $ liftInner xs
132------------------------------------------------------------------------------
133-- Sharing read write state in a stream
134------------------------------------------------------------------------------
136-- | Evaluate the inner monad of a stream as 'StateT'.
137--
138-- >>> evalStateT s = fmap snd . Stream.runStateT s
139--
140{-# INLINE_NORMAL evalStateT #-}
141evalStateT :: Monad m => m s -> Stream (StateT s m) a -> Stream m a
142evalStateT initial (Stream step state) = Stream step' (state, initial)
143 where
144 {-# INLINE_LATE step' #-}
145 step' gst (st, action) = do
146 sv <- action
147 (r, !sv') <- State.runStateT (step (adaptState gst) st) sv
148 return $ case r of
149 Yield x s -> Yield x (s, return sv')
150 Skip s -> Skip (s, return sv')
151 Stop -> Stop
153-- | Evaluate the inner monad of a stream as 'StateT' and emit the resulting
154-- state and value pair after each step.
155--
156{-# INLINE_NORMAL runStateT #-}
157runStateT :: Monad m => m s -> Stream (StateT s m) a -> Stream m (s, a)
158runStateT initial (Stream step state) = Stream step' (state, initial)
159 where
160 {-# INLINE_LATE step' #-}
161 step' gst (st, action) = do
162 sv <- action
163 (r, !sv') <- State.runStateT (step (adaptState gst) st) sv
164 return $ case r of
165 Yield x s -> Yield (sv', x) (s, return sv')
166 Skip s -> Skip (s, return sv')
167 Stop -> Stop
169-- | Run a stateful (StateT) stream transformation using a given state.
170--
171-- >>> usingStateT s f = Stream.evalStateT s . f . Stream.liftInner
172--
173-- See also: 'scan'
174--
175{-# INLINE usingStateT #-}
176usingStateT
177 :: Monad m
178 => m s
179 -> (Stream (StateT s m) a -> Stream (StateT s m) a)
180 -> Stream m a
181 -> Stream m a
182usingStateT s f = evalStateT s . f . liftInner