git.iko.soy/streamly
on 3feeec2 help sign in
streamly/src/Streamly/Internal/Data/Sink.hs
245 lines 7.3 kB history raw
1-- |
2-- Module : Streamly.Internal.Data.Sink
3-- Copyright : (c) 2019 Composewell Technologies
4-- License : BSD3
5-- Maintainer : streamly@composewell.com
6-- Stability : experimental
7-- Portability : GHC
8--
9--
10-- The 'Sink' type is a just a special case of 'Fold' and we can do without
11-- it. However, in some cases 'Sink' is a simpler type and may provide better
12-- performance than 'Fold' because it does not maintain any state. Folds can
13-- be used for both pure and monadic computations. Sinks are not applicable to
14-- pure computations.
16module Streamly.Internal.Data.Sink
17 (
18 Sink (..)
20 -- * Upgrading
21 , toFold
23 -- * Composing Sinks
24 -- ** Distribute
25 , tee
26 , distribute
28 -- ** Demultiplex
29 , demux
31 -- ** Unzip
32 , unzipM
33 , unzip
35 -- -- ** Group
36 -- , grouped
38 -- -- ** Nest
39 -- , concatFold
41 -- -- * Comonad
42 -- , duplicate
44 -- * Input Transformation
45 -- | These are contravariant operations i.e. they apply on the input of the
46 -- 'Sink', for this reason they are prefixed with 'l' for 'left'.
47 , lmap
48 , lmapM
49 , lfilter
50 , lfilterM
52 -- * Sinks
53 , drain
54 , drainM
55 -- , drainN
56 -- , drainWhile
57 )
58where
60import Control.Monad ((>=>), when, void)
61import Data.Map.Strict (Map)
62import Prelude
63 hiding (filter, drop, dropWhile, take, takeWhile, zipWith, foldr,
64 foldl, map, mapM, sequence, all, any, sum, product, elem,
65 notElem, maximum, minimum, head, last, tail, length, null,
66 reverse, iterate, init, and, or, lookup, foldr1, (!!),
67 scanl, scanl1, replicate, concatMap, mconcat, foldMap, unzip)
69import Streamly.Internal.Data.Fold.Type (Fold(..), Step(..))
70import Streamly.Internal.Data.Sink.Type (Sink(..))
72import qualified Data.Map.Strict as Map
74------------------------------------------------------------------------------
75-- Conversion
76------------------------------------------------------------------------------
78-- | Convert a 'Sink' to a 'Fold'. When you want to compose sinks and folds
79-- together, upgrade a sink to a fold before composing.
80toFold :: Monad m => Sink m a -> Fold m a ()
81toFold (Sink f) = Fold step begin done
82 where
83 begin = return $ Partial ()
84 step _ a = Partial <$> f a
85 done _ = return ()
87------------------------------------------------------------------------------
88-- Composing with sinks
89------------------------------------------------------------------------------
91-- | Distribute one copy each of the input to both the sinks.
92--
93-- @
94-- |-------Sink m a
95-- ---stream m a---|
96-- |-------Sink m a
97-- @
98-- @
99-- > let pr x = Sink.drainM (putStrLn . ((x ++ " ") ++) . show)
100-- > sink (Sink.tee (pr \"L") (pr \"R")) (S.enumerateFromTo 1 2)
101-- L 1
102-- R 1
103-- L 2
104-- R 2
105-- @
106--
107tee :: Monad m => Sink m a -> Sink m a -> Sink m a
108tee (Sink fL) (Sink fR) = Sink (\a -> fL a >> fR a)
110-- | Distribute copies of the input to all the sinks in a container.
111--
112-- @
113-- |-------Sink m a
114-- ---stream m a---|
115-- |-------Sink m a
116-- |
117-- ...
118-- @
119-- @
120-- > let pr x = Sink.drainM (putStrLn . ((x ++ " ") ++) . show)
121-- > sink (Sink.distribute [(pr \"L"), (pr \"R")]) (S.enumerateFromTo 1 2)
122-- L 1
123-- R 1
124-- L 2
125-- R 2
126-- @
127--
128-- This is the consumer side dual of the producer side 'sequence_' operation.
129{-# INLINE distribute #-}
130distribute :: Monad m => [Sink m a] -> Sink m a
131distribute ss = Sink (\a -> Prelude.mapM_ (\(Sink f) -> f a) ss)
133-- | Demultiplex to multiple consumers without collecting the results. Useful
134-- to run different effectful computations depending on the value of the stream
135-- elements, for example handling network packets of different types using
136-- different handlers.
137--
138-- @
139--
140-- |-------Sink m a
141-- -----stream m a-----Map-----|
142-- |-------Sink m a
143-- |
144-- ...
145-- @
146--
147-- @
148-- > let pr x = Sink.drainM (putStrLn . ((x ++ " ") ++) . show)
149-- > let table = Data.Map.fromList [(1, pr \"One"), (2, pr \"Two")]
150-- in Sink.sink (Sink.demux id table) (S.enumerateFromTo 1 100)
151-- One 1
152-- Two 2
153-- @
154{-
155demux :: (Monad m, Ord k) => (a -> k) -> Map k (Sink m a) -> Sink m a
156demux f kv = Sink step
158 where
160 step a =
161 -- XXX should we raise an exception in Nothing case?
162 -- Ideally we should enforce that it is a total map over k so that look
163 -- up never fails
164 case Map.lookup (f a) kv of
165 Nothing -> return ()
166 Just (Sink g) -> g a
167-}
169demux :: (Monad m, Ord k) => Map k (Sink m a) -> Sink m (a, k)
170demux kv = Sink step
172 where
174 step (a, k) =
175 -- XXX should we raise an exception in Nothing case?
176 -- Ideally we should enforce that it is a total map over k so that look
177 -- up never fails
178 case Map.lookup k kv of
179 Nothing -> return ()
180 Just (Sink g) -> g a
182-- | Split elements in the input stream into two parts using a monadic unzip
183-- function, direct each part to a different sink.
184--
185-- @
186--
187-- |-------Sink m b
188-- -----Stream m a----(b,c)--|
189-- |-------Sink m c
190-- @
191-- @
192-- > let pr x = Sink.drainM (putStrLn . ((x ++ " ") ++) . show)
193-- in Sink.sink (Sink.unzip return (pr \"L") (pr \"R")) (S.fromPure (1,2))
194-- L 1
195-- R 2
196-- @
197{-# INLINE unzipM #-}
198unzipM :: Monad m => (a -> m (b,c)) -> Sink m b -> Sink m c -> Sink m a
199unzipM f (Sink stepB) (Sink stepC) =
200 Sink (f >=> (\(b, c) -> stepB b >> stepC c))
202-- | Same as 'unzipM' but with a pure unzip function.
203{-# INLINE unzip #-}
204unzip :: Monad m => (a -> (b,c)) -> Sink m b -> Sink m c -> Sink m a
205unzip f = unzipM (return . f)
207------------------------------------------------------------------------------
208-- Input transformation
209------------------------------------------------------------------------------
211-- | Map a pure function on the input of a 'Sink'.
212{-# INLINABLE lmap #-}
213lmap :: (a -> b) -> Sink m b -> Sink m a
214lmap f (Sink step) = Sink (step . f)
216-- | Map a monadic function on the input of a 'Sink'.
217{-# INLINABLE lmapM #-}
218lmapM :: Monad m => (a -> m b) -> Sink m b -> Sink m a
219lmapM f (Sink step) = Sink (f >=> step)
221-- | Filter the input of a 'Sink' using a pure predicate function.
222{-# INLINABLE lfilter #-}
223lfilter :: Monad m => (a -> Bool) -> Sink m a -> Sink m a
224lfilter f (Sink step) = Sink (\a -> when (f a) $ step a)
226-- | Filter the input of a 'Sink' using a monadic predicate function.
227{-# INLINABLE lfilterM #-}
228lfilterM :: Monad m => (a -> m Bool) -> Sink m a -> Sink m a
229lfilterM f (Sink step) = Sink (\a -> f a >>= \use -> when use $ step a)
231------------------------------------------------------------------------------
232-- Sinks
233------------------------------------------------------------------------------
235-- | Drain all input, running the effects and discarding the results.
236drain :: Monad m => Sink m a
237drain = Sink (\_ -> return ())
239-- |
240-- > drainM f = lmapM f drain
241--
242-- Drain all input after passing it through a monadic function.
243{-# INLINABLE drainM #-}
244drainM :: Monad m => (a -> m b) -> Sink m a
245drainM f = Sink (void . f)