git.iko.soy/streamly
on f67df51 help sign in
streamly/src/Streamly/Internal/Data/Unfold.hs
1076 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 , fromStreamK
81 , fromStreamD
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 Fusion.Plugin.Types (Fuse(..))
143import Streamly.Internal.Data.Stream.StreamD.Type (Stream(..), Step(..))
144#if __GLASGOW_HASKELL__ < 800
145import Streamly.Internal.Data.Stream.StreamD.Type (pattern Stream)
146#endif
147import Streamly.Internal.Data.Unfold.Types (Unfold(..))
148import Streamly.Internal.Data.Fold.Types (Fold(..))
149import Streamly.Internal.Data.SVar (defState, MonadAsync)
150import Control.Monad.Catch (MonadCatch)
152import qualified Prelude
153import qualified Control.Monad.Catch as MC
154import qualified Data.Tuple as Tuple
155import qualified Streamly.Internal.Data.Stream.StreamK as K
156import qualified Streamly.Internal.Data.Stream.StreamD as D
158-------------------------------------------------------------------------------
159-- Input operations
160-------------------------------------------------------------------------------
162-- | Map a function on the input argument of the 'Unfold'.
163--
164-- @
165-- lmap f = concat (singleton f)
166-- @
167--
168-- /Internal/
169{-# INLINE_NORMAL lmap #-}
170lmap :: (a -> c) -> Unfold m c b -> Unfold m a b
171lmap f (Unfold ustep uinject) = Unfold ustep (uinject . f)
173-- | Map an action on the input argument of the 'Unfold'.
174--
175-- @
176-- lmapM f = concat (singletonM f)
177-- @
178--
179-- /Internal/
180{-# INLINE_NORMAL lmapM #-}
181lmapM :: Monad m => (a -> m c) -> Unfold m c b -> Unfold m a b
182lmapM f (Unfold ustep uinject) = Unfold ustep (\x -> f x >>= uinject)
184-- XXX change the signature to the following?
185-- supply :: a -> Unfold m a b -> Unfold m Void b
186--
187-- | Supply the seed to an unfold closing the input end of the unfold.
188--
189-- /Internal/
190--
191{-# INLINE_NORMAL supply #-}
192supply :: Unfold m a b -> a -> Unfold m Void b
193supply unf a = lmap (Prelude.const a) unf
195-- XXX change the signature to the following?
196-- supplyFirst :: a -> Unfold m (a, b) c -> Unfold m b c
197--
198-- | Supply the first component of the tuple to an unfold that accepts a tuple
199-- as a seed resulting in a fold that accepts the second component of the tuple
200-- as a seed.
201--
202-- /Internal/
203--
204{-# INLINE_NORMAL supplyFirst #-}
205supplyFirst :: Unfold m (a, b) c -> a -> Unfold m b c
206supplyFirst unf a = lmap (a, ) unf
208-- XXX change the signature to the following?
209-- supplySecond :: b -> Unfold m (a, b) c -> Unfold m a c
210--
211-- | Supply the second component of the tuple to an unfold that accepts a tuple
212-- as a seed resulting in a fold that accepts the first component of the tuple
213-- as a seed.
214--
215-- /Internal/
216--
217{-# INLINE_NORMAL supplySecond #-}
218supplySecond :: Unfold m (a, b) c -> b -> Unfold m a c
219supplySecond unf b = lmap (, b) unf
221-- | Convert an 'Unfold' into an unfold accepting a tuple as an argument,
222-- using the argument of the original fold as the second element of tuple and
223-- discarding the first element of the tuple.
224--
225-- /Internal/
226--
227{-# INLINE_NORMAL discardFirst #-}
228discardFirst :: Unfold m a b -> Unfold m (c, a) b
229discardFirst = lmap snd
231-- | Convert an 'Unfold' into an unfold accepting a tuple as an argument,
232-- using the argument of the original fold as the first element of tuple and
233-- discarding the second element of the tuple.
234--
235-- /Internal/
236--
237{-# INLINE_NORMAL discardSecond #-}
238discardSecond :: Unfold m a b -> Unfold m (a, c) b
239discardSecond = lmap fst
241-- | Convert an 'Unfold' that accepts a tuple as an argument into an unfold
242-- that accepts a tuple with elements swapped.
243--
244-- /Internal/
245--
246{-# INLINE_NORMAL swap #-}
247swap :: Unfold m (a, c) b -> Unfold m (c, a) b
248swap = lmap Tuple.swap
250-------------------------------------------------------------------------------
251-- Output operations
252-------------------------------------------------------------------------------
254-- | Compose an 'Unfold' and a 'Fold'. Given an @Unfold m a b@ and a
255-- @Fold m b c@, returns a monadic action @a -> m c@ representing the
256-- application of the fold on the unfolded stream.
257--
258-- /Internal/
259--
260{-# INLINE_NORMAL fold #-}
261fold :: Monad m => Unfold m a b -> Fold m b c -> a -> m c
262fold (Unfold ustep inject) (Fold fstep initial extract) a =
263 initial >>= \x -> inject a >>= go SPEC x
264 where
265 -- XXX !acc?
266 {-# INLINE_LATE go #-}
267 go !_ acc st = acc `seq` do
268 r <- ustep st
269 case r of
270 Yield x s -> do
271 acc' <- fstep acc x
272 go SPEC acc' s
273 Skip s -> go SPEC acc s
274 Stop -> extract acc
276{-# INLINE_NORMAL map #-}
277map :: Monad m => (b -> c) -> Unfold m a b -> Unfold m a c
278map f (Unfold ustep uinject) = Unfold step uinject
279 where
280 {-# INLINE_LATE step #-}
281 step st = do
282 r <- ustep st
283 return $ case r of
284 Yield x s -> Yield (f x) s
285 Skip s -> Skip s
286 Stop -> Stop
288{-# INLINE_NORMAL mapM #-}
289mapM :: Monad m => (b -> m c) -> Unfold m a b -> Unfold m a c
290mapM f (Unfold ustep uinject) = Unfold step uinject
291 where
292 {-# INLINE_LATE step #-}
293 step st = do
294 r <- ustep st
295 case r of
296 Yield x s -> f x >>= \a -> return $ Yield a s
297 Skip s -> return $ Skip s
298 Stop -> return $ Stop
300{-# INLINE_NORMAL mapMWithInput #-}
301mapMWithInput :: Monad m => (a -> b -> m c) -> Unfold m a b -> Unfold m a c
302mapMWithInput f (Unfold ustep uinject) = Unfold step inject
303 where
304 inject a = do
305 r <- uinject a
306 return (a, r)
308 {-# INLINE_LATE step #-}
309 step (inp, st) = do
310 r <- ustep st
311 case r of
312 Yield x s -> f inp x >>= \a -> return $ Yield a (inp, s)
313 Skip s -> return $ Skip (inp, s)
314 Stop -> return $ Stop
316-------------------------------------------------------------------------------
317-- Convert streams into unfolds
318-------------------------------------------------------------------------------
320-- this is horribly slow
321{-# INLINE_NORMAL fromStreamD #-}
322fromStreamD :: Monad m => Unfold m (Stream m a) a
323fromStreamD = Unfold step return
324 where
326 {-# INLINE_LATE step #-}
327 step (Stream step1 state1) = do
328 r <- step1 defState state1
329 return $ case r of
330 Yield x s -> Yield x (Stream step1 s)
331 Skip s -> Skip (Stream step1 s)
332 Stop -> Stop
334{-# INLINE_NORMAL fromStreamK #-}
335fromStreamK :: Monad m => Unfold m (K.Stream m a) a
336fromStreamK = Unfold step return
338 where
340 {-# INLINE_LATE step #-}
341 step stream = do
342 r <- K.uncons stream
343 return $ case r of
344 Just (x, xs) -> Yield x xs
345 Nothing -> Stop
347-- | Convert a stream into an 'Unfold'. Note that a stream converted to an
348-- 'Unfold' may not be as efficient as an 'Unfold' in some situations.
349--
350-- /Internal/
351{-# INLINE_NORMAL fromStream #-}
352fromStream :: (K.IsStream t, Monad m) => Unfold m (t m a) a
353fromStream = lmap K.toStream fromStreamK
355-------------------------------------------------------------------------------
356-- Unfolds
357-------------------------------------------------------------------------------
359-- | Lift a monadic function into an unfold generating a nil stream with a side
360-- effect.
361--
362{-# INLINE nilM #-}
363nilM :: Monad m => (a -> m c) -> Unfold m a b
364nilM f = Unfold step return
365 where
366 {-# INLINE_LATE step #-}
367 step x = f x >> return Stop
369-- | Prepend a monadic single element generator function to an 'Unfold'.
370--
371-- /Internal/
372{-# INLINE_NORMAL consM #-}
373consM :: Monad m => (a -> m b) -> Unfold m a b -> Unfold m a b
374consM action unf = Unfold step inject
376 where
378 inject = return . Left
380 {-# INLINE_LATE step #-}
381 step (Left a) = do
382 action a >>= \r -> return $ Yield r (Right (D.unfold unf a))
383 step (Right (UnStream step1 st)) = do
384 res <- step1 defState st
385 case res of
386 Yield x s -> return $ Yield x (Right (Stream step1 s))
387 Skip s -> return $ Skip (Right (Stream step1 s))
388 Stop -> return Stop
390-- | Lift a monadic effect into an unfold generating a singleton stream.
391--
392{-# INLINE effect #-}
393effect :: Monad m => m b -> Unfold m Void b
394effect eff = Unfold step inject
395 where
396 inject _ = return True
397 {-# INLINE_LATE step #-}
398 step True = eff >>= \r -> return $ Yield r False
399 step False = return Stop
401-- XXX change it to yieldM or change yieldM in Prelude to singletonM
402--
403-- | Lift a monadic function into an unfold generating a singleton stream.
404--
405{-# INLINE singletonM #-}
406singletonM :: Monad m => (a -> m b) -> Unfold m a b
407singletonM f = Unfold step inject
408 where
409 inject x = return $ Just x
410 {-# INLINE_LATE step #-}
411 step (Just x) = f x >>= \r -> return $ Yield r Nothing
412 step Nothing = return Stop
414-- | Lift a pure function into an unfold generating a singleton stream.
415--
416{-# INLINE singleton #-}
417singleton :: Monad m => (a -> b) -> Unfold m a b
418singleton f = singletonM $ return . f
420-- | Identity unfold. Generates a singleton stream with the seed as the only
421-- element in the stream.
422--
423-- > identity = singletonM return
424--
425{-# INLINE identity #-}
426identity :: Monad m => Unfold m a a
427identity = singletonM return
429const :: Monad m => m b -> Unfold m a b
430const m = Unfold step inject
431 where
432 inject _ = return ()
433 step () = m >>= \r -> return $ Yield r ()
435-- | Generates a stream replicating the seed @n@ times.
436--
437{-# INLINE replicateM #-}
438replicateM :: Monad m => Int -> Unfold m a a
439replicateM n = Unfold step inject
440 where
441 inject x = return (x, n)
442 {-# INLINE_LATE step #-}
443 step (x, i) = return $
444 if i <= 0
445 then Stop
446 else Yield x (x, (i - 1))
448-- | Generates an infinite stream repeating the seed.
449--
450{-# INLINE repeatM #-}
451repeatM :: Monad m => Unfold m a a
452repeatM = Unfold step return
453 where
454 {-# INLINE_LATE step #-}
455 step x = return $ Yield x x
457-- | Convert a list of pure values to a 'Stream'
458{-# INLINE_LATE fromList #-}
459fromList :: Monad m => Unfold m [a] a
460fromList = Unfold step inject
461 where
462 inject x = return x
463 {-# INLINE_LATE step #-}
464 step (x:xs) = return $ Yield x xs
465 step [] = return Stop
467-- | Convert a list of monadic values to a 'Stream'
468{-# INLINE_LATE fromListM #-}
469fromListM :: Monad m => Unfold m [m a] a
470fromListM = Unfold step inject
471 where
472 inject x = return x
473 {-# INLINE_LATE step #-}
474 step (x:xs) = x >>= \r -> return $ Yield r xs
475 step [] = return Stop
477-------------------------------------------------------------------------------
478-- Filtering
479-------------------------------------------------------------------------------
481{-# INLINE_NORMAL take #-}
482take :: Monad m => Int -> Unfold m a b -> Unfold m a b
483take n (Unfold step1 inject1) = Unfold step inject
484 where
485 inject x = do
486 s <- inject1 x
487 return (s, 0)
488 {-# INLINE_LATE step #-}
489 step (st, i) | i < n = do
490 r <- step1 st
491 return $ case r of
492 Yield x s -> Yield x (s, i + 1)
493 Skip s -> Skip (s, i)
494 Stop -> Stop
495 step (_, _) = return Stop
497{-# INLINE_NORMAL takeWhileM #-}
498takeWhileM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
499takeWhileM f (Unfold step1 inject1) = Unfold step inject1
500 where
501 {-# INLINE_LATE step #-}
502 step st = do
503 r <- step1 st
504 case r of
505 Yield x s -> do
506 b <- f x
507 return $ if b then Yield x s else Stop
508 Skip s -> return $ Skip s
509 Stop -> return Stop
511{-# INLINE takeWhile #-}
512takeWhile :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
513takeWhile f = takeWhileM (return . f)
515{-# INLINE_NORMAL filterM #-}
516filterM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
517filterM f (Unfold step1 inject1) = Unfold step inject1
518 where
519 {-# INLINE_LATE step #-}
520 step st = do
521 r <- step1 st
522 case r of
523 Yield x s -> do
524 b <- f x
525 return $ if b then Yield x s else Skip s
526 Skip s -> return $ Skip s
527 Stop -> return Stop
529{-# INLINE filter #-}
530filter :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
531filter f = filterM (return . f)
533-------------------------------------------------------------------------------
534-- Enumeration
535-------------------------------------------------------------------------------
537-- | Can be used to enumerate unbounded integrals. This does not check for
538-- overflow or underflow for bounded integrals.
539{-# INLINE_NORMAL enumerateFromStepIntegral #-}
540enumerateFromStepIntegral :: (Integral a, Monad m) => Unfold m (a, a) a
541enumerateFromStepIntegral = Unfold step inject
542 where
543 inject (from, stride) = from `seq` stride `seq` return (from, stride)
544 {-# INLINE_LATE step #-}
545 step !(x, stride) = return $ Yield x $! (x + stride, stride)
547-- We are assuming that "to" is constrained by the type to be within
548-- max/min bounds.
549{-# INLINE enumerateFromToIntegral #-}
550enumerateFromToIntegral :: (Monad m, Integral a) => a -> Unfold m a a
551enumerateFromToIntegral to =
552 takeWhile (<= to) $ supplySecond enumerateFromStepIntegral 1
554{-# INLINE enumerateFromIntegral #-}
555enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => Unfold m a a
556enumerateFromIntegral = enumerateFromToIntegral maxBound
558-------------------------------------------------------------------------------
559-- Zipping
560-------------------------------------------------------------------------------
562{-# INLINE_NORMAL zipWithM #-}
563zipWithM :: Monad m
564 => (a -> b -> m c) -> Unfold m x a -> Unfold m y b -> Unfold m (x, y) c
565zipWithM f (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject
567 where
569 inject (x, y) = do
570 s1 <- inject1 x
571 s2 <- inject2 y
572 return (s1, s2, Nothing)
574 {-# INLINE_LATE step #-}
575 step (s1, s2, Nothing) = do
576 r <- step1 s1
577 return $
578 case r of
579 Yield x s -> Skip (s, s2, Just x)
580 Skip s -> Skip (s, s2, Nothing)
581 Stop -> Stop
583 step (s1, s2, Just x) = do
584 r <- step2 s2
585 case r of
586 Yield y s -> do
587 z <- f x y
588 return $ Yield z (s1, s, Nothing)
589 Skip s -> return $ Skip (s1, s, Just x)
590 Stop -> return Stop
592-- | Divide the input into two unfolds and then zip the outputs to a single
593-- stream.
594--
595-- @
596-- S.mapM_ print
597-- $ S.concatUnfold (UF.zipWith (,) UF.identity (UF.singleton sqrt))
598-- $ S.map (\x -> (x,x))
599-- $ S.fromList [1..10]
600-- @
601--
602-- /Internal/
603--
604{-# INLINE zipWith #-}
605zipWith :: Monad m
606 => (a -> b -> c) -> Unfold m x a -> Unfold m y b -> Unfold m (x, y) c
607zipWith f = zipWithM (\a b -> return (f a b))
609-- | Distribute the input to two unfolds and then zip the outputs to a single
610-- stream.
611--
612-- @
613-- S.mapM_ print $ S.concatUnfold (UF.teeZipWith (,) UF.identity (UF.singleton sqrt)) $ S.fromList [1..10]
614-- @
615--
616-- /Internal/
617--
618{-# INLINE_NORMAL teeZipWith #-}
619teeZipWith :: Monad m
620 => (a -> b -> c) -> Unfold m x a -> Unfold m x b -> Unfold m x c
621teeZipWith f unf1 unf2 = lmap (\x -> (x,x)) $ zipWith f unf1 unf2
623-------------------------------------------------------------------------------
624-- Nested
625-------------------------------------------------------------------------------
627{-# ANN type ConcatState Fuse #-}
628data ConcatState s1 s2 = ConcatOuter s1 | ConcatInner s1 s2
630-- | Apply the second unfold to each output element of the first unfold and
631-- flatten the output in a single stream.
632--
633-- /Internal/
634--
635{-# INLINE_NORMAL concat #-}
636concat :: Monad m => Unfold m a b -> Unfold m b c -> Unfold m a c
637concat (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject
638 where
639 inject x = do
640 s <- inject1 x
641 return $ ConcatOuter s
643 {-# INLINE_LATE step #-}
644 step (ConcatOuter st) = do
645 r <- step1 st
646 case r of
647 Yield x s -> do
648 innerSt <- inject2 x
649 return $ Skip (ConcatInner s innerSt)
650 Skip s -> return $ Skip (ConcatOuter s)
651 Stop -> return Stop
653 step (ConcatInner ost ist) = do
654 r <- step2 ist
655 return $ case r of
656 Yield x s -> Yield x (ConcatInner ost s)
657 Skip s -> Skip (ConcatInner ost s)
658 Stop -> Skip (ConcatOuter ost)
660data OuterProductState s1 s2 sy x y =
661 OuterProductOuter s1 y | OuterProductInner s1 sy s2 x
663-- | Create an outer product (vector product or cartesian product) of the
664-- output streams of two unfolds.
665--
666{-# INLINE_NORMAL outerProduct #-}
667outerProduct :: Monad m
668 => Unfold m a b -> Unfold m c d -> Unfold m (a, c) (b, d)
669outerProduct (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject
670 where
671 inject (x, y) = do
672 s1 <- inject1 x
673 return $ OuterProductOuter s1 y
675 {-# INLINE_LATE step #-}
676 step (OuterProductOuter st1 sy) = do
677 r <- step1 st1
678 case r of
679 Yield x s -> do
680 s2 <- inject2 sy
681 return $ Skip (OuterProductInner s sy s2 x)
682 Skip s -> return $ Skip (OuterProductOuter s sy)
683 Stop -> return Stop
685 step (OuterProductInner ost sy ist x) = do
686 r <- step2 ist
687 return $ case r of
688 Yield y s -> Yield (x, y) (OuterProductInner ost sy s x)
689 Skip s -> Skip (OuterProductInner ost sy s x)
690 Stop -> Skip (OuterProductOuter ost sy)
692-- XXX This can be used to implement a Monad instance for "Unfold m ()".
694data ConcatMapState s1 s2 = ConcatMapOuter s1 | ConcatMapInner s1 s2
696-- | Map an unfold generating action to each element of an unfold and
697-- flattern the results into a single stream.
698--
699{-# INLINE_NORMAL concatMapM #-}
700concatMapM :: Monad m
701 => (b -> m (Unfold m () c)) -> Unfold m a b -> Unfold m a c
702concatMapM f (Unfold step1 inject1) = Unfold step inject
703 where
704 inject x = do
705 s <- inject1 x
706 return $ ConcatMapOuter s
708 {-# INLINE_LATE step #-}
709 step (ConcatMapOuter st) = do
710 r <- step1 st
711 case r of
712 Yield x s -> do
713 Unfold step2 inject2 <- f x
714 innerSt <- inject2 ()
715 return $ Skip (ConcatMapInner s (Stream (\_ ss -> step2 ss)
716 innerSt))
717 Skip s -> return $ Skip (ConcatMapOuter s)
718 Stop -> return Stop
720 step (ConcatMapInner ost (UnStream istep ist)) = do
721 r <- istep defState ist
722 return $ case r of
723 Yield x s -> Yield x (ConcatMapInner ost (Stream istep s))
724 Skip s -> Skip (ConcatMapInner ost (Stream istep s))
725 Stop -> Skip (ConcatMapOuter ost)
727------------------------------------------------------------------------------
728-- Exceptions
729------------------------------------------------------------------------------
731-- | The most general bracketing and exception combinator. All other
732-- combinators can be expressed in terms of this combinator. This can also be
733-- used for cases which are not covered by the standard combinators.
734--
735-- /Internal/
736--
737{-# INLINE_NORMAL gbracket #-}
738gbracket
739 :: Monad m
740 => (a -> m c) -- ^ before
741 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
742 -> (c -> m d) -- ^ after, on normal stop
743 -> Unfold m (c, e) b -- ^ on exception
744 -> Unfold m c b -- ^ unfold to run
745 -> Unfold m a b
746gbracket bef exc aft (Unfold estep einject) (Unfold step1 inject1) =
747 Unfold step inject
749 where
751 inject x = do
752 r <- bef x
753 s <- inject1 r
754 return $ Right (s, r)
756 {-# INLINE_LATE step #-}
757 step (Right (st, v)) = do
758 res <- exc $ step1 st
759 case res of
760 Right r -> case r of
761 Yield x s -> return $ Yield x (Right (s, v))
762 Skip s -> return $ Skip (Right (s, v))
763 Stop -> aft v >> return Stop
764 Left e -> do
765 r <- einject (v, e)
766 return $ Skip (Left r)
767 step (Left st) = do
768 res <- estep st
769 case res of
770 Yield x s -> return $ Yield x (Left s)
771 Skip s -> return $ Skip (Left s)
772 Stop -> return Stop
774-- | The most general bracketing and exception combinator. All other
775-- combinators can be expressed in terms of this combinator. This can also be
776-- used for cases which are not covered by the standard combinators.
777--
778-- /Internal/
779--
780{-# INLINE_NORMAL gbracketIO #-}
781gbracketIO
782 :: (MonadIO m, MonadBaseControl IO m)
783 => (a -> m c) -- ^ before
784 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
785 -> (c -> m d) -- ^ after, on normal stop, or GC
786 -> Unfold m (c, e) b -- ^ on exception
787 -> Unfold m c b -- ^ unfold to run
788 -> Unfold m a b
789gbracketIO bef exc aft (Unfold estep einject) (Unfold step1 inject1) =
790 Unfold step inject
792 where
794 inject x = do
795 -- Mask asynchronous exceptions to make the execution of 'bef' and
796 -- the registration of 'aft' atomic. See comment in 'D.gbracketIO'.
797 (r, ref) <- liftBaseOp_ mask_ $ do
798 r <- bef x
799 ref <- D.newFinalizedIORef (aft r)
800 return (r, ref)
801 s <- inject1 r
802 return $ Right (s, r, ref)
804 {-# INLINE_LATE step #-}
805 step (Right (st, v, ref)) = do
806 res <- exc $ step1 st
807 case res of
808 Right r -> case r of
809 Yield x s -> return $ Yield x (Right (s, v, ref))
810 Skip s -> return $ Skip (Right (s, v, ref))
811 Stop -> do
812 D.runIORefFinalizer ref
813 return Stop
814 Left e -> do
815 D.clearIORefFinalizer ref
816 r <- einject (v, e)
817 return $ Skip (Left r)
818 step (Left st) = do
819 res <- estep st
820 case res of
821 Yield x s -> return $ Yield x (Left s)
822 Skip s -> return $ Skip (Left s)
823 Stop -> return Stop
825-- The custom implementation of "before" is slightly faster (5-7%) than
826-- "_before". This is just to document and make sure that we can always use
827-- gbracket to implement before. The same applies to other combinators as well.
828--
829{-# INLINE_NORMAL _before #-}
830_before :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
831_before action unf = gbracket (\x -> action x >> return x) (fmap Right)
832 (\_ -> return ()) undefined unf
834-- | Run a side effect before the unfold yields its first element.
835--
836-- /Internal/
837{-# INLINE_NORMAL before #-}
838before :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
839before action (Unfold step1 inject1) = Unfold step inject
841 where
843 inject x = do
844 _ <- action x
845 st <- inject1 x
846 return st
848 {-# INLINE_LATE step #-}
849 step st = do
850 res <- step1 st
851 case res of
852 Yield x s -> return $ Yield x s
853 Skip s -> return $ Skip s
854 Stop -> return Stop
856{-# INLINE_NORMAL _after #-}
857_after :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
858_after aft = gbracket return (fmap Right) aft undefined
860-- | Run a side effect whenever the unfold stops normally.
861--
862-- Prefer afterIO over this as the @after@ action in this combinator is not
863-- executed if the unfold is partially evaluated lazily and then garbage
864-- collected.
865--
866-- /Internal/
867{-# INLINE_NORMAL after #-}
868after :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
869after action (Unfold step1 inject1) = Unfold step inject
871 where
873 inject x = do
874 s <- inject1 x
875 return (s, x)
877 {-# INLINE_LATE step #-}
878 step (st, v) = do
879 res <- step1 st
880 case res of
881 Yield x s -> return $ Yield x (s, v)
882 Skip s -> return $ Skip (s, v)
883 Stop -> action v >> return Stop
885-- | Run a side effect whenever the unfold stops normally
886-- or is garbage collected after a partial lazy evaluation.
887--
888-- /Internal/
889{-# INLINE_NORMAL afterIO #-}
890afterIO :: (MonadIO m, MonadBaseControl IO m)
891 => (a -> m c) -> Unfold m a b -> Unfold m a b
892afterIO action (Unfold step1 inject1) = Unfold step inject
894 where
896 inject x = do
897 s <- inject1 x
898 ref <- D.newFinalizedIORef (action x)
899 return (s, ref)
901 {-# INLINE_LATE step #-}
902 step (st, ref) = do
903 res <- step1 st
904 case res of
905 Yield x s -> return $ Yield x (s, ref)
906 Skip s -> return $ Skip (s, ref)
907 Stop -> do
908 D.runIORefFinalizer ref
909 return Stop
911{-# INLINE_NORMAL _onException #-}
912_onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
913_onException action unf =
914 gbracket return MC.try
915 (\_ -> return ())
916 (nilM (\(a, (e :: MC.SomeException)) -> action a >> MC.throwM e)) unf
918-- | Run a side effect whenever the unfold aborts due to an exception.
919--
920-- /Internal/
921{-# INLINE_NORMAL onException #-}
922onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
923onException action (Unfold step1 inject1) = Unfold step inject
925 where
927 inject x = do
928 s <- inject1 x
929 return (s, x)
931 {-# INLINE_LATE step #-}
932 step (st, v) = do
933 res <- step1 st `MC.onException` action v
934 case res of
935 Yield x s -> return $ Yield x (s, v)
936 Skip s -> return $ Skip (s, v)
937 Stop -> return Stop
939{-# INLINE_NORMAL _finally #-}
940_finally :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
941_finally action unf =
942 gbracket return MC.try action
943 (nilM (\(a, (e :: MC.SomeException)) -> action a >> MC.throwM e)) unf
945-- | Run a side effect whenever the unfold stops normally or aborts due to an
946-- exception.
947--
948-- Prefer finallyIO over this as the @after@ action in this combinator is not
949-- executed if the unfold is partially evaluated lazily and then garbage
950-- collected.
951--
952-- /Internal/
953{-# INLINE_NORMAL finally #-}
954finally :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
955finally action (Unfold step1 inject1) = Unfold step inject
957 where
959 inject x = do
960 s <- inject1 x
961 return (s, x)
963 {-# INLINE_LATE step #-}
964 step (st, v) = do
965 res <- step1 st `MC.onException` action v
966 case res of
967 Yield x s -> return $ Yield x (s, v)
968 Skip s -> return $ Skip (s, v)
969 Stop -> action v >> return Stop
971-- | Run a side effect whenever the unfold stops normally, aborts due to an
972-- exception or if it is garbage collected after a partial lazy evaluation.
973--
974-- /Internal/
975{-# INLINE_NORMAL finallyIO #-}
976finallyIO :: (MonadAsync m, MonadCatch m)
977 => (a -> m c) -> Unfold m a b -> Unfold m a b
978finallyIO action (Unfold step1 inject1) = Unfold step inject
980 where
982 inject x = do
983 s <- inject1 x
984 ref <- D.newFinalizedIORef (action x)
985 return (s, ref)
987 {-# INLINE_LATE step #-}
988 step (st, ref) = do
989 res <- step1 st `MC.onException` D.runIORefFinalizer ref
990 case res of
991 Yield x s -> return $ Yield x (s, ref)
992 Skip s -> return $ Skip (s, ref)
993 Stop -> do
994 D.runIORefFinalizer ref
995 return Stop
997{-# INLINE_NORMAL _bracket #-}
998_bracket :: MonadCatch m
999 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1000_bracket bef aft unf =
1001 gbracket bef MC.try aft (nilM (\(a, (e :: MC.SomeException)) -> aft a >>
1002 MC.throwM e)) unf
1004-- | @bracket before after between@ runs the @before@ action and then unfolds
1005-- its output using the @between@ unfold. When the @between@ unfold is done or
1006-- if an exception occurs then the @after@ action is run with the output of
1007-- @before@ as argument.
1009-- Prefer bracketIO over this as the @after@ action in this combinator is not
1010-- executed if the unfold is partially evaluated lazily and then garbage
1011-- collected.
1013-- /Internal/
1014{-# INLINE_NORMAL bracket #-}
1015bracket :: MonadCatch m
1016 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1017bracket bef aft (Unfold step1 inject1) = Unfold step inject
1019 where
1021 inject x = do
1022 r <- bef x
1023 s <- inject1 r
1024 return (s, r)
1026 {-# INLINE_LATE step #-}
1027 step (st, v) = do
1028 res <- step1 st `MC.onException` aft v
1029 case res of
1030 Yield x s -> return $ Yield x (s, v)
1031 Skip s -> return $ Skip (s, v)
1032 Stop -> aft v >> return Stop
1034-- | @bracket before after between@ runs the @before@ action and then unfolds
1035-- its output using the @between@ unfold. When the @between@ unfold is done or
1036-- if an exception occurs then the @after@ action is run with the output of
1037-- @before@ as argument. The after action is also executed if the unfold is
1038-- paritally evaluated and then garbage collected.
1040-- /Internal/
1041{-# INLINE_NORMAL bracketIO #-}
1042bracketIO :: (MonadAsync m, MonadCatch m)
1043 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1044bracketIO bef aft (Unfold step1 inject1) = Unfold step inject
1046 where
1048 inject x = do
1049 -- Mask asynchronous exceptions to make the execution of 'bef' and
1050 -- the registration of 'aft' atomic. See comment in 'D.gbracketIO'.
1051 (r, ref) <- liftBaseOp_ mask_ $ do
1052 r <- bef x
1053 ref <- D.newFinalizedIORef (aft r)
1054 return (r, ref)
1055 s <- inject1 r
1056 return (s, ref)
1058 {-# INLINE_LATE step #-}
1059 step (st, ref) = do
1060 res <- step1 st `MC.onException` D.runIORefFinalizer ref
1061 case res of
1062 Yield x s -> return $ Yield x (s, ref)
1063 Skip s -> return $ Skip (s, ref)
1064 Stop -> do
1065 D.runIORefFinalizer ref
1066 return Stop
1068-- | When unfolding if an exception occurs, unfold the exception using the
1069-- exception unfold supplied as the first argument to 'handle'.
1071-- /Internal/
1072{-# INLINE_NORMAL handle #-}
1073handle :: (MonadCatch m, Exception e)
1074 => Unfold m e b -> Unfold m a b -> Unfold m a b
1075handle exc unf =
1076 gbracket return MC.try (\_ -> return ()) (discardFirst exc) unf