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