git.iko.soy/streamly
on cc8cd01 help sign in
streamly/src/Streamly/Internal/Data/Unfold.hs
1072 lines 34 kB history raw
1{-# LANGUAGE BangPatterns #-}
2{-# LANGUAGE CPP #-}
3{-# LANGUAGE ExistentialQuantification #-}
4{-# LANGUAGE FlexibleContexts #-}
5{-# LANGUAGE PatternSynonyms #-}
6{-# LANGUAGE RankNTypes #-}
7{-# LANGUAGE RecordWildCards #-}
8{-# LANGUAGE ScopedTypeVariables #-}
9{-# LANGUAGE TupleSections #-}
11#include "inline.hs"
13-- |
14-- Module : Streamly.Internal.Data.Unfold
15-- Copyright : (c) 2019 Composewell Technologies
16-- License : BSD3
17-- Maintainer : streamly@composewell.com
18-- Stability : experimental
19-- Portability : GHC
20--
21-- Streams forcing a closed control flow loop can be categorized under
22-- two types, unfolds and folds, both of these are duals of each other.
23--
24-- Unfold streams are really generators of a sequence of elements, we can also
25-- call them pull style streams. These are lazy producers of streams. On each
26-- evaluation the producer generates the next element. A consumer can
27-- therefore pull elements from the stream whenever it wants to. A stream
28-- consumer can multiplex pull streams by pulling elements from the chosen
29-- streams, therefore, pull streams allow merging or multiplexing. On the
30-- other hand, with this representation we cannot split or demultiplex a
31-- stream. So really these are stream sources that can be generated from a
32-- seed and can be merged or zipped into a single stream.
33--
34-- The dual of Unfolds are Folds. Folds can also be called as push style
35-- streams or reducers. These are strict consumers of streams. We keep pushing
36-- elements to a fold and we can extract the result at any point. A driver can
37-- choose which fold to push to and can also push the same element to multiple
38-- folds. Therefore, folds allow splitting or demultiplexing a stream. On the
39-- other hand, we cannot merge streams using this representation. So really
40-- these are stream consumers that reduce the stream to a single value, these
41-- consumers can be composed such that a stream can be split over multiple
42-- consumers.
43--
44-- Performance:
45--
46-- Composing a tree or graph of computations with unfolds can be much more
47-- efficient compared to composing with the Monad instance. The reason is that
48-- unfolds allow the compiler to statically know the state and optimize it
49-- using stream fusion whereas it is not possible with the monad bind because
50-- the state is determined dynamically.
52-- Open control flow style streams can also have two representations. StreamK
53-- is a producer style representation. We can also have a consumer style
54-- representation. We can use that for composable folds in StreamK
55-- representation.
56--
57module Streamly.Internal.Data.Unfold
58 (
59 -- * Unfold Type
60 Unfold
62 -- * Operations on Input
63 , lmap
64 , lmapM
65 , supply
66 , supplyFirst
67 , supplySecond
68 , discardFirst
69 , discardSecond
70 , swap
71 -- coapply
72 -- comonad
74 -- * Operations on Output
75 , fold
76 -- pipe
78 -- * Unfolds
79 , fromStream
80 , fromStream1
81 , fromStream2
82 , nilM
83 , consM
84 , effect
85 , singletonM
86 , singleton
87 , identity
88 , const
89 , replicateM
90 , repeatM
91 , fromList
92 , fromListM
93 , enumerateFromStepIntegral
94 , enumerateFromToIntegral
95 , enumerateFromIntegral
97 -- * Transformations
98 , map
99 , mapM
100 , mapMWithInput
102 -- * Filtering
103 , takeWhileM
104 , takeWhile
105 , take
106 , filter
107 , filterM
109 -- * Zipping
110 , zipWithM
111 , zipWith
112 , teeZipWith
114 -- * Nesting
115 , concat
116 , concatMapM
117 , outerProduct
119 -- * Exceptions
120 , gbracket
121 , gbracketIO
122 , before
123 , after
124 , afterIO
125 , onException
126 , finally
127 , finallyIO
128 , bracket
129 , bracketIO
130 , handle
131 )
132where
134import Control.Exception (Exception, mask_)
135import Control.Monad.IO.Class (MonadIO(..))
136import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_)
137import Data.Void (Void)
138import GHC.Types (SPEC(..))
139import Prelude
140 hiding (concat, map, mapM, takeWhile, take, filter, const, zipWith)
142import Streamly.Internal.Data.Stream.StreamD.Type (Stream(..), Step(..))
143#if __GLASGOW_HASKELL__ < 800
144import Streamly.Internal.Data.Stream.StreamD.Type (pattern Stream)
145#endif
146import Streamly.Internal.Data.Unfold.Types (Unfold(..))
147import Streamly.Internal.Data.Fold.Types (Fold(..))
148import Streamly.Internal.Data.SVar (defState, MonadAsync)
149import Control.Monad.Catch (MonadCatch)
151import qualified Prelude
152import qualified Control.Monad.Catch as MC
153import qualified Data.Tuple as Tuple
154import qualified Streamly.Internal.Data.Stream.StreamK as K
155import qualified Streamly.Internal.Data.Stream.StreamD as D
157-------------------------------------------------------------------------------
158-- Input operations
159-------------------------------------------------------------------------------
161-- | Map a function on the input argument of the 'Unfold'.
162--
163-- @
164-- lmap f = concat (singleton f)
165-- @
166--
167-- /Internal/
168{-# INLINE_NORMAL lmap #-}
169lmap :: (a -> c) -> Unfold m c b -> Unfold m a b
170lmap f (Unfold ustep uinject) = Unfold ustep (uinject . f)
172-- | Map an action on the input argument of the 'Unfold'.
173--
174-- @
175-- lmapM f = concat (singletonM f)
176-- @
177--
178-- /Internal/
179{-# INLINE_NORMAL lmapM #-}
180lmapM :: Monad m => (a -> m c) -> Unfold m c b -> Unfold m a b
181lmapM f (Unfold ustep uinject) = Unfold ustep (\x -> f x >>= uinject)
183-- XXX change the signature to the following?
184-- supply :: a -> Unfold m a b -> Unfold m Void b
185--
186-- | Supply the seed to an unfold closing the input end of the unfold.
187--
188-- /Internal/
189--
190{-# INLINE_NORMAL supply #-}
191supply :: Unfold m a b -> a -> Unfold m Void b
192supply unf a = lmap (Prelude.const a) unf
194-- XXX change the signature to the following?
195-- supplyFirst :: a -> Unfold m (a, b) c -> Unfold m b c
196--
197-- | Supply the first component of the tuple to an unfold that accepts a tuple
198-- as a seed resulting in a fold that accepts the second component of the tuple
199-- as a seed.
200--
201-- /Internal/
202--
203{-# INLINE_NORMAL supplyFirst #-}
204supplyFirst :: Unfold m (a, b) c -> a -> Unfold m b c
205supplyFirst unf a = lmap (a, ) unf
207-- XXX change the signature to the following?
208-- supplySecond :: b -> Unfold m (a, b) c -> Unfold m a c
209--
210-- | Supply the second component of the tuple to an unfold that accepts a tuple
211-- as a seed resulting in a fold that accepts the first component of the tuple
212-- as a seed.
213--
214-- /Internal/
215--
216{-# INLINE_NORMAL supplySecond #-}
217supplySecond :: Unfold m (a, b) c -> b -> Unfold m a c
218supplySecond unf b = lmap (, b) unf
220-- | Convert an 'Unfold' into an unfold accepting a tuple as an argument,
221-- using the argument of the original fold as the second element of tuple and
222-- discarding the first element of the tuple.
223--
224-- /Internal/
225--
226{-# INLINE_NORMAL discardFirst #-}
227discardFirst :: Unfold m a b -> Unfold m (c, a) b
228discardFirst = lmap snd
230-- | Convert an 'Unfold' into an unfold accepting a tuple as an argument,
231-- using the argument of the original fold as the first element of tuple and
232-- discarding the second element of the tuple.
233--
234-- /Internal/
235--
236{-# INLINE_NORMAL discardSecond #-}
237discardSecond :: Unfold m a b -> Unfold m (a, c) b
238discardSecond = lmap fst
240-- | Convert an 'Unfold' that accepts a tuple as an argument into an unfold
241-- that accepts a tuple with elements swapped.
242--
243-- /Internal/
244--
245{-# INLINE_NORMAL swap #-}
246swap :: Unfold m (a, c) b -> Unfold m (c, a) b
247swap = lmap Tuple.swap
249-------------------------------------------------------------------------------
250-- Output operations
251-------------------------------------------------------------------------------
253-- | Compose an 'Unfold' and a 'Fold'. Given an @Unfold m a b@ and a
254-- @Fold m b c@, returns a monadic action @a -> m c@ representing the
255-- application of the fold on the unfolded stream.
256--
257-- /Internal/
258--
259{-# INLINE_NORMAL fold #-}
260fold :: Monad m => Unfold m a b -> Fold m b c -> a -> m c
261fold (Unfold ustep inject) (Fold fstep initial extract) a =
262 initial >>= \x -> inject a >>= go SPEC x
263 where
264 -- XXX !acc?
265 {-# INLINE_LATE go #-}
266 go !_ acc st = acc `seq` do
267 r <- ustep st
268 case r of
269 Yield x s -> do
270 acc' <- fstep acc x
271 go SPEC acc' s
272 Skip s -> go SPEC acc s
273 Stop -> extract acc
275{-# INLINE_NORMAL map #-}
276map :: Monad m => (b -> c) -> Unfold m a b -> Unfold m a c
277map f (Unfold ustep uinject) = Unfold step uinject
278 where
279 {-# INLINE_LATE step #-}
280 step st = do
281 r <- ustep st
282 return $ case r of
283 Yield x s -> Yield (f x) s
284 Skip s -> Skip s
285 Stop -> Stop
287{-# INLINE_NORMAL mapM #-}
288mapM :: Monad m => (b -> m c) -> Unfold m a b -> Unfold m a c
289mapM f (Unfold ustep uinject) = Unfold step uinject
290 where
291 {-# INLINE_LATE step #-}
292 step st = do
293 r <- ustep st
294 case r of
295 Yield x s -> f x >>= \a -> return $ Yield a s
296 Skip s -> return $ Skip s
297 Stop -> return $ Stop
299{-# INLINE_NORMAL mapMWithInput #-}
300mapMWithInput :: Monad m => (a -> b -> m c) -> Unfold m a b -> Unfold m a c
301mapMWithInput f (Unfold ustep uinject) = Unfold step inject
302 where
303 inject a = do
304 r <- uinject a
305 return (a, r)
307 {-# INLINE_LATE step #-}
308 step (inp, st) = do
309 r <- ustep st
310 case r of
311 Yield x s -> f inp x >>= \a -> return $ Yield a (inp, s)
312 Skip s -> return $ Skip (inp, s)
313 Stop -> return $ Stop
315-------------------------------------------------------------------------------
316-- Convert streams into unfolds
317-------------------------------------------------------------------------------
319{-# INLINE_LATE streamStep #-}
320streamStep :: Monad m => Stream m a -> m (Step (Stream m a) a)
321streamStep (Stream step1 state) = do
322 r <- step1 defState state
323 return $ case r of
324 Yield x s -> Yield x (Stream step1 s)
325 Skip s -> Skip (Stream step1 s)
326 Stop -> Stop
328-- | Convert a stream into an 'Unfold'. Note that a stream converted to an
329-- 'Unfold' may not be as efficient as an 'Unfold' in some situations.
330--
331-- /Internal/
332fromStream :: (K.IsStream t, Monad m) => t m b -> Unfold m Void b
333fromStream str = Unfold streamStep (\_ -> return $ D.toStreamD str)
335-- | Convert a single argument stream generator function into an
336-- 'Unfold'. Note that a stream converted to an 'Unfold' may not be as
337-- efficient as an 'Unfold' in some situations.
338--
339-- /Internal/
340fromStream1 :: (K.IsStream t, Monad m) => (a -> t m b) -> Unfold m a b
341fromStream1 f = Unfold streamStep (return . D.toStreamD . f)
343-- | Convert a two argument stream generator function into an 'Unfold'. Note
344-- that a stream converted to an 'Unfold' may not be as efficient as an
345-- 'Unfold' in some situations.
346--
347-- /Internal/
348fromStream2 :: (K.IsStream t, Monad m)
349 => (a -> b -> t m c) -> Unfold m (a, b) c
350fromStream2 f = Unfold streamStep (\(a, b) -> return $ D.toStreamD $ f a b)
352-------------------------------------------------------------------------------
353-- Unfolds
354-------------------------------------------------------------------------------
356-- | Lift a monadic function into an unfold generating a nil stream with a side
357-- effect.
358--
359{-# INLINE nilM #-}
360nilM :: Monad m => (a -> m c) -> Unfold m a b
361nilM f = Unfold step return
362 where
363 {-# INLINE_LATE step #-}
364 step x = f x >> return Stop
366-- | Prepend a monadic single element generator function to an 'Unfold'.
367--
368-- /Internal/
369{-# INLINE_NORMAL consM #-}
370consM :: Monad m => (a -> m b) -> Unfold m a b -> Unfold m a b
371consM action unf = Unfold step inject
373 where
375 inject = return . Left
377 {-# INLINE_LATE step #-}
378 step (Left a) = do
379 action a >>= \r -> return $ Yield r (Right (D.unfold unf a))
380 step (Right (UnStream step1 st)) = do
381 res <- step1 defState st
382 case res of
383 Yield x s -> return $ Yield x (Right (Stream step1 s))
384 Skip s -> return $ Skip (Right (Stream step1 s))
385 Stop -> return Stop
387-- | Lift a monadic effect into an unfold generating a singleton stream.
388--
389{-# INLINE effect #-}
390effect :: Monad m => m b -> Unfold m Void b
391effect eff = Unfold step inject
392 where
393 inject _ = return True
394 {-# INLINE_LATE step #-}
395 step True = eff >>= \r -> return $ Yield r False
396 step False = return Stop
398-- XXX change it to yieldM or change yieldM in Prelude to singletonM
399--
400-- | Lift a monadic function into an unfold generating a singleton stream.
401--
402{-# INLINE singletonM #-}
403singletonM :: Monad m => (a -> m b) -> Unfold m a b
404singletonM f = Unfold step inject
405 where
406 inject x = return $ Just x
407 {-# INLINE_LATE step #-}
408 step (Just x) = f x >>= \r -> return $ Yield r Nothing
409 step Nothing = return Stop
411-- | Lift a pure function into an unfold generating a singleton stream.
412--
413{-# INLINE singleton #-}
414singleton :: Monad m => (a -> b) -> Unfold m a b
415singleton f = singletonM $ return . f
417-- | Identity unfold. Generates a singleton stream with the seed as the only
418-- element in the stream.
419--
420-- > identity = singletonM return
421--
422{-# INLINE identity #-}
423identity :: Monad m => Unfold m a a
424identity = singletonM return
426const :: Monad m => m b -> Unfold m a b
427const m = Unfold step inject
428 where
429 inject _ = return ()
430 step () = m >>= \r -> return $ Yield r ()
432-- | Generates a stream replicating the seed @n@ times.
433--
434{-# INLINE replicateM #-}
435replicateM :: Monad m => Int -> Unfold m a a
436replicateM n = Unfold step inject
437 where
438 inject x = return (x, n)
439 {-# INLINE_LATE step #-}
440 step (x, i) = return $
441 if i <= 0
442 then Stop
443 else Yield x (x, (i - 1))
445-- | Generates an infinite stream repeating the seed.
446--
447{-# INLINE repeatM #-}
448repeatM :: Monad m => Unfold m a a
449repeatM = Unfold step return
450 where
451 {-# INLINE_LATE step #-}
452 step x = return $ Yield x x
454-- | Convert a list of pure values to a 'Stream'
455{-# INLINE_LATE fromList #-}
456fromList :: Monad m => Unfold m [a] a
457fromList = Unfold step inject
458 where
459 inject x = return x
460 {-# INLINE_LATE step #-}
461 step (x:xs) = return $ Yield x xs
462 step [] = return Stop
464-- | Convert a list of monadic values to a 'Stream'
465{-# INLINE_LATE fromListM #-}
466fromListM :: Monad m => Unfold m [m a] a
467fromListM = Unfold step inject
468 where
469 inject x = return x
470 {-# INLINE_LATE step #-}
471 step (x:xs) = x >>= \r -> return $ Yield r xs
472 step [] = return Stop
474-------------------------------------------------------------------------------
475-- Filtering
476-------------------------------------------------------------------------------
478{-# INLINE_NORMAL take #-}
479take :: Monad m => Int -> Unfold m a b -> Unfold m a b
480take n (Unfold step1 inject1) = Unfold step inject
481 where
482 inject x = do
483 s <- inject1 x
484 return (s, 0)
485 {-# INLINE_LATE step #-}
486 step (st, i) | i < n = do
487 r <- step1 st
488 return $ case r of
489 Yield x s -> Yield x (s, i + 1)
490 Skip s -> Skip (s, i)
491 Stop -> Stop
492 step (_, _) = return Stop
494{-# INLINE_NORMAL takeWhileM #-}
495takeWhileM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
496takeWhileM f (Unfold step1 inject1) = Unfold step inject1
497 where
498 {-# INLINE_LATE step #-}
499 step st = do
500 r <- step1 st
501 case r of
502 Yield x s -> do
503 b <- f x
504 return $ if b then Yield x s else Stop
505 Skip s -> return $ Skip s
506 Stop -> return Stop
508{-# INLINE takeWhile #-}
509takeWhile :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
510takeWhile f = takeWhileM (return . f)
512{-# INLINE_NORMAL filterM #-}
513filterM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
514filterM f (Unfold step1 inject1) = Unfold step inject1
515 where
516 {-# INLINE_LATE step #-}
517 step st = do
518 r <- step1 st
519 case r of
520 Yield x s -> do
521 b <- f x
522 return $ if b then Yield x s else Skip s
523 Skip s -> return $ Skip s
524 Stop -> return Stop
526{-# INLINE filter #-}
527filter :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
528filter f = filterM (return . f)
530-------------------------------------------------------------------------------
531-- Enumeration
532-------------------------------------------------------------------------------
534-- | Can be used to enumerate unbounded integrals. This does not check for
535-- overflow or underflow for bounded integrals.
536{-# INLINE_NORMAL enumerateFromStepIntegral #-}
537enumerateFromStepIntegral :: (Integral a, Monad m) => Unfold m (a, a) a
538enumerateFromStepIntegral = Unfold step inject
539 where
540 inject (from, stride) = from `seq` stride `seq` return (from, stride)
541 {-# INLINE_LATE step #-}
542 step !(x, stride) = return $ Yield x $! (x + stride, stride)
544-- We are assuming that "to" is constrained by the type to be within
545-- max/min bounds.
546{-# INLINE enumerateFromToIntegral #-}
547enumerateFromToIntegral :: (Monad m, Integral a) => a -> Unfold m a a
548enumerateFromToIntegral to =
549 takeWhile (<= to) $ supplySecond enumerateFromStepIntegral 1
551{-# INLINE enumerateFromIntegral #-}
552enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => Unfold m a a
553enumerateFromIntegral = enumerateFromToIntegral maxBound
555-------------------------------------------------------------------------------
556-- Zipping
557-------------------------------------------------------------------------------
559{-# INLINE_NORMAL zipWithM #-}
560zipWithM :: Monad m
561 => (a -> b -> m c) -> Unfold m x a -> Unfold m y b -> Unfold m (x, y) c
562zipWithM f (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject
564 where
566 inject (x, y) = do
567 s1 <- inject1 x
568 s2 <- inject2 y
569 return (s1, s2, Nothing)
571 {-# INLINE_LATE step #-}
572 step (s1, s2, Nothing) = do
573 r <- step1 s1
574 return $
575 case r of
576 Yield x s -> Skip (s, s2, Just x)
577 Skip s -> Skip (s, s2, Nothing)
578 Stop -> Stop
580 step (s1, s2, Just x) = do
581 r <- step2 s2
582 case r of
583 Yield y s -> do
584 z <- f x y
585 return $ Yield z (s1, s, Nothing)
586 Skip s -> return $ Skip (s1, s, Just x)
587 Stop -> return Stop
589-- | Divide the input into two unfolds and then zip the outputs to a single
590-- stream.
591--
592-- @
593-- S.mapM_ print
594-- $ S.concatUnfold (UF.zipWith (,) UF.identity (UF.singleton sqrt))
595-- $ S.map (\x -> (x,x))
596-- $ S.fromList [1..10]
597-- @
598--
599-- /Internal/
600--
601{-# INLINE zipWith #-}
602zipWith :: Monad m
603 => (a -> b -> c) -> Unfold m x a -> Unfold m y b -> Unfold m (x, y) c
604zipWith f = zipWithM (\a b -> return (f a b))
606-- | Distribute the input to two unfolds and then zip the outputs to a single
607-- stream.
608--
609-- @
610-- S.mapM_ print $ S.concatUnfold (UF.teeZipWith (,) UF.identity (UF.singleton sqrt)) $ S.fromList [1..10]
611-- @
612--
613-- /Internal/
614--
615{-# INLINE_NORMAL teeZipWith #-}
616teeZipWith :: Monad m
617 => (a -> b -> c) -> Unfold m x a -> Unfold m x b -> Unfold m x c
618teeZipWith f unf1 unf2 = lmap (\x -> (x,x)) $ zipWith f unf1 unf2
620-------------------------------------------------------------------------------
621-- Nested
622-------------------------------------------------------------------------------
624data ConcatState s1 s2 = ConcatOuter s1 | ConcatInner s1 s2
626-- | Apply the second unfold to each output element of the first unfold and
627-- flatten the output in a single stream.
628--
629-- /Internal/
630--
631{-# INLINE_NORMAL concat #-}
632concat :: Monad m => Unfold m a b -> Unfold m b c -> Unfold m a c
633concat (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject
634 where
635 inject x = do
636 s <- inject1 x
637 return $ ConcatOuter s
639 {-# INLINE_LATE step #-}
640 step (ConcatOuter st) = do
641 r <- step1 st
642 case r of
643 Yield x s -> do
644 innerSt <- inject2 x
645 return $ Skip (ConcatInner s innerSt)
646 Skip s -> return $ Skip (ConcatOuter s)
647 Stop -> return Stop
649 step (ConcatInner ost ist) = do
650 r <- step2 ist
651 return $ case r of
652 Yield x s -> Yield x (ConcatInner ost s)
653 Skip s -> Skip (ConcatInner ost s)
654 Stop -> Skip (ConcatOuter ost)
656data OuterProductState s1 s2 sy x y =
657 OuterProductOuter s1 y | OuterProductInner s1 sy s2 x
659-- | Create an outer product (vector product or cartesian product) of the
660-- output streams of two unfolds.
661--
662{-# INLINE_NORMAL outerProduct #-}
663outerProduct :: Monad m
664 => Unfold m a b -> Unfold m c d -> Unfold m (a, c) (b, d)
665outerProduct (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject
666 where
667 inject (x, y) = do
668 s1 <- inject1 x
669 return $ OuterProductOuter s1 y
671 {-# INLINE_LATE step #-}
672 step (OuterProductOuter st1 sy) = do
673 r <- step1 st1
674 case r of
675 Yield x s -> do
676 s2 <- inject2 sy
677 return $ Skip (OuterProductInner s sy s2 x)
678 Skip s -> return $ Skip (OuterProductOuter s sy)
679 Stop -> return Stop
681 step (OuterProductInner ost sy ist x) = do
682 r <- step2 ist
683 return $ case r of
684 Yield y s -> Yield (x, y) (OuterProductInner ost sy s x)
685 Skip s -> Skip (OuterProductInner ost sy s x)
686 Stop -> Skip (OuterProductOuter ost sy)
688-- XXX This can be used to implement a Monad instance for "Unfold m ()".
690data ConcatMapState s1 s2 = ConcatMapOuter s1 | ConcatMapInner s1 s2
692-- | Map an unfold generating action to each element of an unfold and
693-- flattern the results into a single stream.
694--
695{-# INLINE_NORMAL concatMapM #-}
696concatMapM :: Monad m
697 => (b -> m (Unfold m () c)) -> Unfold m a b -> Unfold m a c
698concatMapM f (Unfold step1 inject1) = Unfold step inject
699 where
700 inject x = do
701 s <- inject1 x
702 return $ ConcatMapOuter s
704 {-# INLINE_LATE step #-}
705 step (ConcatMapOuter st) = do
706 r <- step1 st
707 case r of
708 Yield x s -> do
709 Unfold step2 inject2 <- f x
710 innerSt <- inject2 ()
711 return $ Skip (ConcatMapInner s (Stream (\_ ss -> step2 ss)
712 innerSt))
713 Skip s -> return $ Skip (ConcatMapOuter s)
714 Stop -> return Stop
716 step (ConcatMapInner ost (UnStream istep ist)) = do
717 r <- istep defState ist
718 return $ case r of
719 Yield x s -> Yield x (ConcatMapInner ost (Stream istep s))
720 Skip s -> Skip (ConcatMapInner ost (Stream istep s))
721 Stop -> Skip (ConcatMapOuter ost)
723------------------------------------------------------------------------------
724-- Exceptions
725------------------------------------------------------------------------------
727-- | The most general bracketing and exception combinator. All other
728-- combinators can be expressed in terms of this combinator. This can also be
729-- used for cases which are not covered by the standard combinators.
730--
731-- /Internal/
732--
733{-# INLINE_NORMAL gbracket #-}
734gbracket
735 :: Monad m
736 => (a -> m c) -- ^ before
737 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
738 -> (c -> m d) -- ^ after, on normal stop
739 -> Unfold m (c, e) b -- ^ on exception
740 -> Unfold m c b -- ^ unfold to run
741 -> Unfold m a b
742gbracket bef exc aft (Unfold estep einject) (Unfold step1 inject1) =
743 Unfold step inject
745 where
747 inject x = do
748 r <- bef x
749 s <- inject1 r
750 return $ Right (s, r)
752 {-# INLINE_LATE step #-}
753 step (Right (st, v)) = do
754 res <- exc $ step1 st
755 case res of
756 Right r -> case r of
757 Yield x s -> return $ Yield x (Right (s, v))
758 Skip s -> return $ Skip (Right (s, v))
759 Stop -> aft v >> return Stop
760 Left e -> do
761 r <- einject (v, e)
762 return $ Skip (Left r)
763 step (Left st) = do
764 res <- estep st
765 case res of
766 Yield x s -> return $ Yield x (Left s)
767 Skip s -> return $ Skip (Left s)
768 Stop -> return Stop
770-- | The most general bracketing and exception combinator. All other
771-- combinators can be expressed in terms of this combinator. This can also be
772-- used for cases which are not covered by the standard combinators.
773--
774-- /Internal/
775--
776{-# INLINE_NORMAL gbracketIO #-}
777gbracketIO
778 :: (MonadIO m, MonadBaseControl IO m)
779 => (a -> m c) -- ^ before
780 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
781 -> (c -> m d) -- ^ after, on normal stop, or GC
782 -> Unfold m (c, e) b -- ^ on exception
783 -> Unfold m c b -- ^ unfold to run
784 -> Unfold m a b
785gbracketIO bef exc aft (Unfold estep einject) (Unfold step1 inject1) =
786 Unfold step inject
788 where
790 inject x = do
791 -- Mask asynchronous exceptions to make the execution of 'bef' and
792 -- the registration of 'aft' atomic. See comment in 'D.gbracketIO'.
793 (r, ref) <- liftBaseOp_ mask_ $ do
794 r <- bef x
795 ref <- D.newFinalizedIORef (aft r)
796 return (r, ref)
797 s <- inject1 r
798 return $ Right (s, r, ref)
800 {-# INLINE_LATE step #-}
801 step (Right (st, v, ref)) = do
802 res <- exc $ step1 st
803 case res of
804 Right r -> case r of
805 Yield x s -> return $ Yield x (Right (s, v, ref))
806 Skip s -> return $ Skip (Right (s, v, ref))
807 Stop -> do
808 D.runIORefFinalizer ref
809 return Stop
810 Left e -> do
811 D.clearIORefFinalizer ref
812 r <- einject (v, e)
813 return $ Skip (Left r)
814 step (Left st) = do
815 res <- estep st
816 case res of
817 Yield x s -> return $ Yield x (Left s)
818 Skip s -> return $ Skip (Left s)
819 Stop -> return Stop
821-- The custom implementation of "before" is slightly faster (5-7%) than
822-- "_before". This is just to document and make sure that we can always use
823-- gbracket to implement before. The same applies to other combinators as well.
824--
825{-# INLINE_NORMAL _before #-}
826_before :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
827_before action unf = gbracket (\x -> action x >> return x) (fmap Right)
828 (\_ -> return ()) undefined unf
830-- | Run a side effect before the unfold yields its first element.
831--
832-- /Internal/
833{-# INLINE_NORMAL before #-}
834before :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
835before action (Unfold step1 inject1) = Unfold step inject
837 where
839 inject x = do
840 _ <- action x
841 st <- inject1 x
842 return st
844 {-# INLINE_LATE step #-}
845 step st = do
846 res <- step1 st
847 case res of
848 Yield x s -> return $ Yield x s
849 Skip s -> return $ Skip s
850 Stop -> return Stop
852{-# INLINE_NORMAL _after #-}
853_after :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
854_after aft = gbracket return (fmap Right) aft undefined
856-- | Run a side effect whenever the unfold stops normally.
857--
858-- Prefer afterIO over this as the @after@ action in this combinator is not
859-- executed if the unfold is partially evaluated lazily and then garbage
860-- collected.
861--
862-- /Internal/
863{-# INLINE_NORMAL after #-}
864after :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
865after action (Unfold step1 inject1) = Unfold step inject
867 where
869 inject x = do
870 s <- inject1 x
871 return (s, x)
873 {-# INLINE_LATE step #-}
874 step (st, v) = do
875 res <- step1 st
876 case res of
877 Yield x s -> return $ Yield x (s, v)
878 Skip s -> return $ Skip (s, v)
879 Stop -> action v >> return Stop
881-- | Run a side effect whenever the unfold stops normally
882-- or is garbage collected after a partial lazy evaluation.
883--
884-- /Internal/
885{-# INLINE_NORMAL afterIO #-}
886afterIO :: (MonadIO m, MonadBaseControl IO m)
887 => (a -> m c) -> Unfold m a b -> Unfold m a b
888afterIO action (Unfold step1 inject1) = Unfold step inject
890 where
892 inject x = do
893 s <- inject1 x
894 ref <- D.newFinalizedIORef (action x)
895 return (s, ref)
897 {-# INLINE_LATE step #-}
898 step (st, ref) = do
899 res <- step1 st
900 case res of
901 Yield x s -> return $ Yield x (s, ref)
902 Skip s -> return $ Skip (s, ref)
903 Stop -> do
904 D.runIORefFinalizer ref
905 return Stop
907{-# INLINE_NORMAL _onException #-}
908_onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
909_onException action unf =
910 gbracket return MC.try
911 (\_ -> return ())
912 (nilM (\(a, (e :: MC.SomeException)) -> action a >> MC.throwM e)) unf
914-- | Run a side effect whenever the unfold aborts due to an exception.
915--
916-- /Internal/
917{-# INLINE_NORMAL onException #-}
918onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
919onException action (Unfold step1 inject1) = Unfold step inject
921 where
923 inject x = do
924 s <- inject1 x
925 return (s, x)
927 {-# INLINE_LATE step #-}
928 step (st, v) = do
929 res <- step1 st `MC.onException` action v
930 case res of
931 Yield x s -> return $ Yield x (s, v)
932 Skip s -> return $ Skip (s, v)
933 Stop -> return Stop
935{-# INLINE_NORMAL _finally #-}
936_finally :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
937_finally action unf =
938 gbracket return MC.try action
939 (nilM (\(a, (e :: MC.SomeException)) -> action a >> MC.throwM e)) unf
941-- | Run a side effect whenever the unfold stops normally or aborts due to an
942-- exception.
943--
944-- Prefer finallyIO over this as the @after@ action in this combinator is not
945-- executed if the unfold is partially evaluated lazily and then garbage
946-- collected.
947--
948-- /Internal/
949{-# INLINE_NORMAL finally #-}
950finally :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
951finally action (Unfold step1 inject1) = Unfold step inject
953 where
955 inject x = do
956 s <- inject1 x
957 return (s, x)
959 {-# INLINE_LATE step #-}
960 step (st, v) = do
961 res <- step1 st `MC.onException` action v
962 case res of
963 Yield x s -> return $ Yield x (s, v)
964 Skip s -> return $ Skip (s, v)
965 Stop -> action v >> return Stop
967-- | Run a side effect whenever the unfold stops normally, aborts due to an
968-- exception or if it is garbage collected after a partial lazy evaluation.
969--
970-- /Internal/
971{-# INLINE_NORMAL finallyIO #-}
972finallyIO :: (MonadAsync m, MonadCatch m)
973 => (a -> m c) -> Unfold m a b -> Unfold m a b
974finallyIO action (Unfold step1 inject1) = Unfold step inject
976 where
978 inject x = do
979 s <- inject1 x
980 ref <- D.newFinalizedIORef (action x)
981 return (s, ref)
983 {-# INLINE_LATE step #-}
984 step (st, ref) = do
985 res <- step1 st `MC.onException` D.runIORefFinalizer ref
986 case res of
987 Yield x s -> return $ Yield x (s, ref)
988 Skip s -> return $ Skip (s, ref)
989 Stop -> do
990 D.runIORefFinalizer ref
991 return Stop
993{-# INLINE_NORMAL _bracket #-}
994_bracket :: MonadCatch m
995 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
996_bracket bef aft unf =
997 gbracket bef MC.try aft (nilM (\(a, (e :: MC.SomeException)) -> aft a >>
998 MC.throwM e)) unf
1000-- | @bracket before after between@ runs the @before@ action and then unfolds
1001-- its output using the @between@ unfold. When the @between@ unfold is done or
1002-- if an exception occurs then the @after@ action is run with the output of
1003-- @before@ as argument.
1005-- Prefer bracketIO over this as the @after@ action in this combinator is not
1006-- executed if the unfold is partially evaluated lazily and then garbage
1007-- collected.
1009-- /Internal/
1010{-# INLINE_NORMAL bracket #-}
1011bracket :: MonadCatch m
1012 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1013bracket bef aft (Unfold step1 inject1) = Unfold step inject
1015 where
1017 inject x = do
1018 r <- bef x
1019 s <- inject1 r
1020 return (s, r)
1022 {-# INLINE_LATE step #-}
1023 step (st, v) = do
1024 res <- step1 st `MC.onException` aft v
1025 case res of
1026 Yield x s -> return $ Yield x (s, v)
1027 Skip s -> return $ Skip (s, v)
1028 Stop -> aft v >> return Stop
1030-- | @bracket before after between@ runs the @before@ action and then unfolds
1031-- its output using the @between@ unfold. When the @between@ unfold is done or
1032-- if an exception occurs then the @after@ action is run with the output of
1033-- @before@ as argument. The after action is also executed if the unfold is
1034-- paritally evaluated and then garbage collected.
1036-- /Internal/
1037{-# INLINE_NORMAL bracketIO #-}
1038bracketIO :: (MonadAsync m, MonadCatch m)
1039 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1040bracketIO bef aft (Unfold step1 inject1) = Unfold step inject
1042 where
1044 inject x = do
1045 -- Mask asynchronous exceptions to make the execution of 'bef' and
1046 -- the registration of 'aft' atomic. See comment in 'D.gbracketIO'.
1047 (r, ref) <- liftBaseOp_ mask_ $ do
1048 r <- bef x
1049 ref <- D.newFinalizedIORef (aft r)
1050 return (r, ref)
1051 s <- inject1 r
1052 return (s, ref)
1054 {-# INLINE_LATE step #-}
1055 step (st, ref) = do
1056 res <- step1 st `MC.onException` D.runIORefFinalizer ref
1057 case res of
1058 Yield x s -> return $ Yield x (s, ref)
1059 Skip s -> return $ Skip (s, ref)
1060 Stop -> do
1061 D.runIORefFinalizer ref
1062 return Stop
1064-- | When unfolding if an exception occurs, unfold the exception using the
1065-- exception unfold supplied as the first argument to 'handle'.
1067-- /Internal/
1068{-# INLINE_NORMAL handle #-}
1069handle :: (MonadCatch m, Exception e)
1070 => Unfold m e b -> Unfold m a b -> Unfold m a b
1071handle exc unf =
1072 gbracket return MC.try (\_ -> return ()) (discardFirst exc) unf