git.iko.soy/streamly
on f39b982 help sign in
streamly/src/Streamly/Internal/Data/Unfold.hs
956 lines 31 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 , singleton
86 , identity
87 , const
88 , replicateM
89 , repeatM
90 , fromList
91 , fromListM
92 , enumerateFromStepIntegral
93 , enumerateFromToIntegral
94 , enumerateFromIntegral
96 -- * Transformations
97 , map
98 , mapM
99 , mapMWithInput
101 -- * Filtering
102 , takeWhileM
103 , takeWhile
104 , take
105 , filter
106 , filterM
108 -- * Nesting
109 , concat
110 , concatMapM
111 , outerProduct
113 -- * Exceptions
114 , gbracket
115 , gbracketIO
116 , before
117 , after
118 , afterIO
119 , onException
120 , finally
121 , finallyIO
122 , bracket
123 , bracketIO
124 , handle
125 )
126where
128import Control.Exception (Exception)
129import Control.Monad.IO.Class (MonadIO(..))
130import Control.Monad.Trans.Control (MonadBaseControl)
131import Data.Void (Void)
132import GHC.Types (SPEC(..))
133import Prelude hiding (concat, map, mapM, takeWhile, take, filter, const)
135import Streamly.Internal.Data.Stream.StreamD.Type (Stream(..), Step(..))
136#if __GLASGOW_HASKELL__ < 800
137import Streamly.Internal.Data.Stream.StreamD.Type (pattern Stream)
138#endif
139import Streamly.Internal.Data.Unfold.Types (Unfold(..))
140import Streamly.Internal.Data.Fold.Types (Fold(..))
141import Streamly.Internal.Data.SVar (defState, MonadAsync)
142import Control.Monad.Catch (MonadCatch)
144import qualified Prelude
145import qualified Control.Monad.Catch as MC
146import qualified Data.Tuple as Tuple
147import qualified Streamly.Internal.Data.Stream.StreamK as K
148import qualified Streamly.Internal.Data.Stream.StreamD as D
150-------------------------------------------------------------------------------
151-- Input operations
152-------------------------------------------------------------------------------
154-- | Map a function on the input argument of the 'Unfold'.
155--
156-- /Internal/
157{-# INLINE_NORMAL lmap #-}
158lmap :: (a -> c) -> Unfold m c b -> Unfold m a b
159lmap f (Unfold ustep uinject) = Unfold ustep (uinject . f)
161-- | Map an action on the input argument of the 'Unfold'.
162--
163-- /Internal/
164{-# INLINE_NORMAL lmapM #-}
165lmapM :: Monad m => (a -> m c) -> Unfold m c b -> Unfold m a b
166lmapM f (Unfold ustep uinject) = Unfold ustep (\x -> f x >>= uinject)
168-- | Supply the seed to an unfold closing the input end of the unfold.
169--
170-- /Internal/
171--
172{-# INLINE_NORMAL supply #-}
173supply :: Unfold m a b -> a -> Unfold m Void b
174supply unf a = lmap (Prelude.const a) unf
176-- | Supply the first component of the tuple to an unfold that accepts a tuple
177-- as a seed resulting in a fold that accepts the second component of the tuple
178-- as a seed.
179--
180-- /Internal/
181--
182{-# INLINE_NORMAL supplyFirst #-}
183supplyFirst :: Unfold m (a, b) c -> a -> Unfold m b c
184supplyFirst unf a = lmap (a, ) unf
186-- | Supply the second component of the tuple to an unfold that accepts a tuple
187-- as a seed resulting in a fold that accepts the first component of the tuple
188-- as a seed.
189--
190-- /Internal/
191--
192{-# INLINE_NORMAL supplySecond #-}
193supplySecond :: Unfold m (a, b) c -> b -> Unfold m a c
194supplySecond unf b = lmap (, b) unf
196-- | Convert an 'Unfold' into an unfold accepting a tuple as an argument,
197-- using the argument of the original fold as the second element of tuple and
198-- discarding the first element of the tuple.
199--
200-- /Internal/
201--
202{-# INLINE_NORMAL discardFirst #-}
203discardFirst :: Unfold m a b -> Unfold m (c, a) b
204discardFirst = lmap snd
206-- | Convert an 'Unfold' into an unfold accepting a tuple as an argument,
207-- using the argument of the original fold as the first element of tuple and
208-- discarding the second element of the tuple.
209--
210-- /Internal/
211--
212{-# INLINE_NORMAL discardSecond #-}
213discardSecond :: Unfold m a b -> Unfold m (a, c) b
214discardSecond = lmap fst
216-- | Convert an 'Unfold' that accepts a tuple as an argument into an unfold
217-- that accepts a tuple with elements swapped.
218--
219-- /Internal/
220--
221{-# INLINE_NORMAL swap #-}
222swap :: Unfold m (a, c) b -> Unfold m (c, a) b
223swap = lmap Tuple.swap
225-------------------------------------------------------------------------------
226-- Output operations
227-------------------------------------------------------------------------------
229-- | Compose an 'Unfold' and a 'Fold'. Given an @Unfold m a b@ and a
230-- @Fold m b c@, returns a monadic action @a -> m c@ representing the
231-- application of the fold on the unfolded stream.
232--
233-- /Internal/
234--
235{-# INLINE_NORMAL fold #-}
236fold :: Monad m => Unfold m a b -> Fold m b c -> a -> m c
237fold (Unfold ustep inject) (Fold fstep initial extract) a =
238 initial >>= \x -> inject a >>= go SPEC x
239 where
240 -- XXX !acc?
241 {-# INLINE_LATE go #-}
242 go !_ acc st = acc `seq` do
243 r <- ustep st
244 case r of
245 Yield x s -> do
246 acc' <- fstep acc x
247 go SPEC acc' s
248 Skip s -> go SPEC acc s
249 Stop -> extract acc
251{-# INLINE_NORMAL map #-}
252map :: Monad m => (b -> c) -> Unfold m a b -> Unfold m a c
253map f (Unfold ustep uinject) = Unfold step uinject
254 where
255 {-# INLINE_LATE step #-}
256 step st = do
257 r <- ustep st
258 return $ case r of
259 Yield x s -> Yield (f x) s
260 Skip s -> Skip s
261 Stop -> Stop
263{-# INLINE_NORMAL mapM #-}
264mapM :: Monad m => (b -> m c) -> Unfold m a b -> Unfold m a c
265mapM f (Unfold ustep uinject) = Unfold step uinject
266 where
267 {-# INLINE_LATE step #-}
268 step st = do
269 r <- ustep st
270 case r of
271 Yield x s -> f x >>= \a -> return $ Yield a s
272 Skip s -> return $ Skip s
273 Stop -> return $ Stop
275{-# INLINE_NORMAL mapMWithInput #-}
276mapMWithInput :: Monad m => (a -> b -> m c) -> Unfold m a b -> Unfold m a c
277mapMWithInput f (Unfold ustep uinject) = Unfold step inject
278 where
279 inject a = do
280 r <- uinject a
281 return (a, r)
283 {-# INLINE_LATE step #-}
284 step (inp, st) = do
285 r <- ustep st
286 case r of
287 Yield x s -> f inp x >>= \a -> return $ Yield a (inp, s)
288 Skip s -> return $ Skip (inp, s)
289 Stop -> return $ Stop
291-------------------------------------------------------------------------------
292-- Convert streams into unfolds
293-------------------------------------------------------------------------------
295{-# INLINE_LATE streamStep #-}
296streamStep :: Monad m => Stream m a -> m (Step (Stream m a) a)
297streamStep (Stream step1 state) = do
298 r <- step1 defState state
299 return $ case r of
300 Yield x s -> Yield x (Stream step1 s)
301 Skip s -> Skip (Stream step1 s)
302 Stop -> Stop
304-- | Convert a stream into an 'Unfold'. Note that a stream converted to an
305-- 'Unfold' may not be as efficient as an 'Unfold' in some situations.
306--
307-- /Internal/
308fromStream :: (K.IsStream t, Monad m) => t m b -> Unfold m Void b
309fromStream str = Unfold streamStep (\_ -> return $ D.toStreamD str)
311-- | Convert a single argument stream generator function into an
312-- 'Unfold'. Note that a stream converted to an 'Unfold' may not be as
313-- efficient as an 'Unfold' in some situations.
314--
315-- /Internal/
316fromStream1 :: (K.IsStream t, Monad m) => (a -> t m b) -> Unfold m a b
317fromStream1 f = Unfold streamStep (return . D.toStreamD . f)
319-- | Convert a two argument stream generator function into an 'Unfold'. Note
320-- that a stream converted to an 'Unfold' may not be as efficient as an
321-- 'Unfold' in some situations.
322--
323-- /Internal/
324fromStream2 :: (K.IsStream t, Monad m)
325 => (a -> b -> t m c) -> Unfold m (a, b) c
326fromStream2 f = Unfold streamStep (\(a, b) -> return $ D.toStreamD $ f a b)
328-------------------------------------------------------------------------------
329-- Unfolds
330-------------------------------------------------------------------------------
332-- | Lift a monadic function into an unfold generating a nil stream with a side
333-- effect.
334--
335{-# INLINE nilM #-}
336nilM :: Monad m => (a -> m c) -> Unfold m a b
337nilM f = Unfold step return
338 where
339 {-# INLINE_LATE step #-}
340 step x = f x >> return Stop
342-- | Prepend a monadic single element generator function to an 'Unfold'.
343--
344-- /Internal/
345{-# INLINE_NORMAL consM #-}
346consM :: Monad m => (a -> m b) -> Unfold m a b -> Unfold m a b
347consM action unf = Unfold step inject
349 where
351 inject = return . Left
353 {-# INLINE_LATE step #-}
354 step (Left a) = do
355 action a >>= \r -> return $ Yield r (Right (D.unfold unf a))
356 step (Right (UnStream step1 st)) = do
357 res <- step1 defState st
358 case res of
359 Yield x s -> return $ Yield x (Right (Stream step1 s))
360 Skip s -> return $ Skip (Right (Stream step1 s))
361 Stop -> return Stop
363-- | Lift a monadic effect into an unfold generating a singleton stream.
364--
365{-# INLINE effect #-}
366effect :: Monad m => m b -> Unfold m Void b
367effect eff = Unfold step inject
368 where
369 inject _ = return True
370 {-# INLINE_LATE step #-}
371 step True = eff >>= \r -> return $ Yield r False
372 step False = return Stop
374-- | Lift a monadic function into an unfold generating a singleton stream.
375--
376{-# INLINE singleton #-}
377singleton :: Monad m => (a -> m b) -> Unfold m a b
378singleton f = Unfold step inject
379 where
380 inject x = return $ Just x
381 {-# INLINE_LATE step #-}
382 step (Just x) = f x >>= \r -> return $ Yield r Nothing
383 step Nothing = return Stop
385-- | Identity unfold. Generates a singleton stream with the seed as the only
386-- element in the stream.
387--
388-- > identity = singleton return
389--
390{-# INLINE identity #-}
391identity :: Monad m => Unfold m a a
392identity = singleton return
394const :: Monad m => m b -> Unfold m a b
395const m = Unfold step inject
396 where
397 inject _ = return ()
398 step () = m >>= \r -> return $ Yield r ()
400-- | Generates a stream replicating the seed @n@ times.
401--
402{-# INLINE replicateM #-}
403replicateM :: Monad m => Int -> Unfold m a a
404replicateM n = Unfold step inject
405 where
406 inject x = return (x, n)
407 {-# INLINE_LATE step #-}
408 step (x, i) = return $
409 if i <= 0
410 then Stop
411 else Yield x (x, (i - 1))
413-- | Generates an infinite stream repeating the seed.
414--
415{-# INLINE repeatM #-}
416repeatM :: Monad m => Unfold m a a
417repeatM = Unfold step return
418 where
419 {-# INLINE_LATE step #-}
420 step x = return $ Yield x x
422-- | Convert a list of pure values to a 'Stream'
423{-# INLINE_LATE fromList #-}
424fromList :: Monad m => Unfold m [a] a
425fromList = Unfold step inject
426 where
427 inject x = return x
428 {-# INLINE_LATE step #-}
429 step (x:xs) = return $ Yield x xs
430 step [] = return Stop
432-- | Convert a list of monadic values to a 'Stream'
433{-# INLINE_LATE fromListM #-}
434fromListM :: Monad m => Unfold m [m a] a
435fromListM = Unfold step inject
436 where
437 inject x = return x
438 {-# INLINE_LATE step #-}
439 step (x:xs) = x >>= \r -> return $ Yield r xs
440 step [] = return Stop
442-------------------------------------------------------------------------------
443-- Filtering
444-------------------------------------------------------------------------------
446{-# INLINE_NORMAL take #-}
447take :: Monad m => Int -> Unfold m a b -> Unfold m a b
448take n (Unfold step1 inject1) = Unfold step inject
449 where
450 inject x = do
451 s <- inject1 x
452 return (s, 0)
453 {-# INLINE_LATE step #-}
454 step (st, i) | i < n = do
455 r <- step1 st
456 return $ case r of
457 Yield x s -> Yield x (s, i + 1)
458 Skip s -> Skip (s, i)
459 Stop -> Stop
460 step (_, _) = return Stop
462{-# INLINE_NORMAL takeWhileM #-}
463takeWhileM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
464takeWhileM f (Unfold step1 inject1) = Unfold step inject1
465 where
466 {-# INLINE_LATE step #-}
467 step st = do
468 r <- step1 st
469 case r of
470 Yield x s -> do
471 b <- f x
472 return $ if b then Yield x s else Stop
473 Skip s -> return $ Skip s
474 Stop -> return Stop
476{-# INLINE takeWhile #-}
477takeWhile :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
478takeWhile f = takeWhileM (return . f)
480{-# INLINE_NORMAL filterM #-}
481filterM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
482filterM f (Unfold step1 inject1) = Unfold step inject1
483 where
484 {-# INLINE_LATE step #-}
485 step st = do
486 r <- step1 st
487 case r of
488 Yield x s -> do
489 b <- f x
490 return $ if b then Yield x s else Skip s
491 Skip s -> return $ Skip s
492 Stop -> return Stop
494{-# INLINE filter #-}
495filter :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
496filter f = filterM (return . f)
498-------------------------------------------------------------------------------
499-- Enumeration
500-------------------------------------------------------------------------------
502-- | Can be used to enumerate unbounded integrals. This does not check for
503-- overflow or underflow for bounded integrals.
504{-# INLINE_NORMAL enumerateFromStepIntegral #-}
505enumerateFromStepIntegral :: (Integral a, Monad m) => Unfold m (a, a) a
506enumerateFromStepIntegral = Unfold step inject
507 where
508 inject (from, stride) = from `seq` stride `seq` return (from, stride)
509 {-# INLINE_LATE step #-}
510 step !(x, stride) = return $ Yield x $! (x + stride, stride)
512-- We are assuming that "to" is constrained by the type to be within
513-- max/min bounds.
514{-# INLINE enumerateFromToIntegral #-}
515enumerateFromToIntegral :: (Monad m, Integral a) => a -> Unfold m a a
516enumerateFromToIntegral to =
517 takeWhile (<= to) $ supplySecond enumerateFromStepIntegral 1
519{-# INLINE enumerateFromIntegral #-}
520enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => Unfold m a a
521enumerateFromIntegral = enumerateFromToIntegral maxBound
523-------------------------------------------------------------------------------
524-- Nested
525-------------------------------------------------------------------------------
527data ConcatState s1 s2 = ConcatOuter s1 | ConcatInner s1 s2
529{-# INLINE_NORMAL concat #-}
530concat :: Monad m => Unfold m a b -> Unfold m b c -> Unfold m a c
531concat (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject
532 where
533 inject x = do
534 s <- inject1 x
535 return $ ConcatOuter s
537 {-# INLINE_LATE step #-}
538 step (ConcatOuter st) = do
539 r <- step1 st
540 case r of
541 Yield x s -> do
542 innerSt <- inject2 x
543 return $ Skip (ConcatInner s innerSt)
544 Skip s -> return $ Skip (ConcatOuter s)
545 Stop -> return Stop
547 step (ConcatInner ost ist) = do
548 r <- step2 ist
549 return $ case r of
550 Yield x s -> Yield x (ConcatInner ost s)
551 Skip s -> Skip (ConcatInner ost s)
552 Stop -> Skip (ConcatOuter ost)
554data OuterProductState s1 s2 sy x y =
555 OuterProductOuter s1 y | OuterProductInner s1 sy s2 x
557{-# INLINE_NORMAL outerProduct #-}
558outerProduct :: Monad m
559 => Unfold m a b -> Unfold m c d -> Unfold m (a, c) (b, d)
560outerProduct (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject
561 where
562 inject (x, y) = do
563 s1 <- inject1 x
564 return $ OuterProductOuter s1 y
566 {-# INLINE_LATE step #-}
567 step (OuterProductOuter st1 sy) = do
568 r <- step1 st1
569 case r of
570 Yield x s -> do
571 s2 <- inject2 sy
572 return $ Skip (OuterProductInner s sy s2 x)
573 Skip s -> return $ Skip (OuterProductOuter s sy)
574 Stop -> return Stop
576 step (OuterProductInner ost sy ist x) = do
577 r <- step2 ist
578 return $ case r of
579 Yield y s -> Yield (x, y) (OuterProductInner ost sy s x)
580 Skip s -> Skip (OuterProductInner ost sy s x)
581 Stop -> Skip (OuterProductOuter ost sy)
583-- XXX This can be used to implement a Monad instance for "Unfold m ()".
585data ConcatMapState s1 s2 = ConcatMapOuter s1 | ConcatMapInner s1 s2
587{-# INLINE_NORMAL concatMapM #-}
588concatMapM :: Monad m
589 => (b -> m (Unfold m () c)) -> Unfold m a b -> Unfold m a c
590concatMapM f (Unfold step1 inject1) = Unfold step inject
591 where
592 inject x = do
593 s <- inject1 x
594 return $ ConcatMapOuter s
596 {-# INLINE_LATE step #-}
597 step (ConcatMapOuter st) = do
598 r <- step1 st
599 case r of
600 Yield x s -> do
601 Unfold step2 inject2 <- f x
602 innerSt <- inject2 ()
603 return $ Skip (ConcatMapInner s (Stream (\_ ss -> step2 ss)
604 innerSt))
605 Skip s -> return $ Skip (ConcatMapOuter s)
606 Stop -> return Stop
608 step (ConcatMapInner ost (UnStream istep ist)) = do
609 r <- istep defState ist
610 return $ case r of
611 Yield x s -> Yield x (ConcatMapInner ost (Stream istep s))
612 Skip s -> Skip (ConcatMapInner ost (Stream istep s))
613 Stop -> Skip (ConcatMapOuter ost)
615------------------------------------------------------------------------------
616-- Exceptions
617------------------------------------------------------------------------------
619-- | The most general bracketing and exception combinator. All other
620-- combinators can be expressed in terms of this combinator. This can also be
621-- used for cases which are not covered by the standard combinators.
622--
623-- /Internal/
624--
625{-# INLINE_NORMAL gbracket #-}
626gbracket
627 :: Monad m
628 => (a -> m c) -- ^ before
629 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
630 -> (c -> m d) -- ^ after, on normal stop
631 -> Unfold m (c, e) b -- ^ on exception
632 -> Unfold m c b -- ^ unfold to run
633 -> Unfold m a b
634gbracket bef exc aft (Unfold estep einject) (Unfold step1 inject1) =
635 Unfold step inject
637 where
639 inject x = do
640 r <- bef x
641 s <- inject1 r
642 return $ Right (s, r)
644 {-# INLINE_LATE step #-}
645 step (Right (st, v)) = do
646 res <- exc $ step1 st
647 case res of
648 Right r -> case r of
649 Yield x s -> return $ Yield x (Right (s, v))
650 Skip s -> return $ Skip (Right (s, v))
651 Stop -> aft v >> return Stop
652 Left e -> do
653 r <- einject (v, e)
654 return $ Skip (Left r)
655 step (Left st) = do
656 res <- estep st
657 case res of
658 Yield x s -> return $ Yield x (Left s)
659 Skip s -> return $ Skip (Left s)
660 Stop -> return Stop
662-- | The most general bracketing and exception combinator. All other
663-- combinators can be expressed in terms of this combinator. This can also be
664-- used for cases which are not covered by the standard combinators.
665--
666-- /Internal/
667--
668{-# INLINE_NORMAL gbracketIO #-}
669gbracketIO
670 :: (MonadIO m, MonadBaseControl IO m)
671 => (a -> m c) -- ^ before
672 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
673 -> (c -> m d) -- ^ after, on normal stop, or GC
674 -> Unfold m (c, e) b -- ^ on exception
675 -> Unfold m c b -- ^ unfold to run
676 -> Unfold m a b
677gbracketIO bef exc aft (Unfold estep einject) (Unfold step1 inject1) =
678 Unfold step inject
680 where
682 inject x = do
683 r <- bef x
684 ref <- D.newFinalizedIORef (aft r)
685 s <- inject1 r
686 return $ Right (s, r, ref)
688 {-# INLINE_LATE step #-}
689 step (Right (st, v, ref)) = do
690 res <- exc $ step1 st
691 case res of
692 Right r -> case r of
693 Yield x s -> return $ Yield x (Right (s, v, ref))
694 Skip s -> return $ Skip (Right (s, v, ref))
695 Stop -> do
696 D.runIORefFinalizer ref
697 return Stop
698 Left e -> do
699 D.clearIORefFinalizer ref
700 r <- einject (v, e)
701 return $ Skip (Left r)
702 step (Left st) = do
703 res <- estep st
704 case res of
705 Yield x s -> return $ Yield x (Left s)
706 Skip s -> return $ Skip (Left s)
707 Stop -> return Stop
709-- The custom implementation of "before" is slightly faster (5-7%) than
710-- "_before". This is just to document and make sure that we can always use
711-- gbracket to implement before. The same applies to other combinators as well.
712--
713{-# INLINE_NORMAL _before #-}
714_before :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
715_before action unf = gbracket (\x -> action x >> return x) (fmap Right)
716 (\_ -> return ()) undefined unf
718-- | Run a side effect before the unfold yields its first element.
719--
720-- /Internal/
721{-# INLINE_NORMAL before #-}
722before :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
723before action (Unfold step1 inject1) = Unfold step inject
725 where
727 inject x = do
728 _ <- action x
729 st <- inject1 x
730 return st
732 {-# INLINE_LATE step #-}
733 step st = do
734 res <- step1 st
735 case res of
736 Yield x s -> return $ Yield x s
737 Skip s -> return $ Skip s
738 Stop -> return Stop
740{-# INLINE_NORMAL _after #-}
741_after :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
742_after aft = gbracket return (fmap Right) aft undefined
744-- | Run a side effect whenever the unfold stops normally.
745--
746-- Prefer afterIO over this as the @after@ action in this combinator is not
747-- executed if the unfold is partially evaluated lazily and then garbage
748-- collected.
749--
750-- /Internal/
751{-# INLINE_NORMAL after #-}
752after :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
753after action (Unfold step1 inject1) = Unfold step inject
755 where
757 inject x = do
758 s <- inject1 x
759 return (s, x)
761 {-# INLINE_LATE step #-}
762 step (st, v) = do
763 res <- step1 st
764 case res of
765 Yield x s -> return $ Yield x (s, v)
766 Skip s -> return $ Skip (s, v)
767 Stop -> action v >> return Stop
769-- | Run a side effect whenever the unfold stops normally
770-- or is garbage collected after a partial lazy evaluation.
771--
772-- /Internal/
773{-# INLINE_NORMAL afterIO #-}
774afterIO :: (MonadIO m, MonadBaseControl IO m)
775 => (a -> m c) -> Unfold m a b -> Unfold m a b
776afterIO action (Unfold step1 inject1) = Unfold step inject
778 where
780 inject x = do
781 s <- inject1 x
782 ref <- D.newFinalizedIORef (action x)
783 return (s, ref)
785 {-# INLINE_LATE step #-}
786 step (st, ref) = do
787 res <- step1 st
788 case res of
789 Yield x s -> return $ Yield x (s, ref)
790 Skip s -> return $ Skip (s, ref)
791 Stop -> do
792 D.runIORefFinalizer ref
793 return Stop
795{-# INLINE_NORMAL _onException #-}
796_onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
797_onException action unf =
798 gbracket return MC.try
799 (\_ -> return ())
800 (nilM (\(a, (e :: MC.SomeException)) -> action a >> MC.throwM e)) unf
802-- | Run a side effect whenever the unfold aborts due to an exception.
803--
804-- /Internal/
805{-# INLINE_NORMAL onException #-}
806onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
807onException action (Unfold step1 inject1) = Unfold step inject
809 where
811 inject x = do
812 s <- inject1 x
813 return (s, x)
815 {-# INLINE_LATE step #-}
816 step (st, v) = do
817 res <- step1 st `MC.onException` action v
818 case res of
819 Yield x s -> return $ Yield x (s, v)
820 Skip s -> return $ Skip (s, v)
821 Stop -> return Stop
823{-# INLINE_NORMAL _finally #-}
824_finally :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
825_finally action unf =
826 gbracket return MC.try action
827 (nilM (\(a, (e :: MC.SomeException)) -> action a >> MC.throwM e)) unf
829-- | Run a side effect whenever the unfold stops normally or aborts due to an
830-- exception.
831--
832-- Prefer finallyIO over this as the @after@ action in this combinator is not
833-- executed if the unfold is partially evaluated lazily and then garbage
834-- collected.
835--
836-- /Internal/
837{-# INLINE_NORMAL finally #-}
838finally :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
839finally action (Unfold step1 inject1) = Unfold step inject
841 where
843 inject x = do
844 s <- inject1 x
845 return (s, x)
847 {-# INLINE_LATE step #-}
848 step (st, v) = do
849 res <- step1 st `MC.onException` action v
850 case res of
851 Yield x s -> return $ Yield x (s, v)
852 Skip s -> return $ Skip (s, v)
853 Stop -> action v >> return Stop
855-- | Run a side effect whenever the unfold stops normally, aborts due to an
856-- exception or if it is garbage collected after a partial lazy evaluation.
857--
858-- /Internal/
859{-# INLINE_NORMAL finallyIO #-}
860finallyIO :: (MonadAsync m, MonadCatch m)
861 => (a -> m c) -> Unfold m a b -> Unfold m a b
862finallyIO action (Unfold step1 inject1) = Unfold step inject
864 where
866 inject x = do
867 s <- inject1 x
868 ref <- D.newFinalizedIORef (action x)
869 return (s, ref)
871 {-# INLINE_LATE step #-}
872 step (st, ref) = do
873 res <- step1 st `MC.onException` D.runIORefFinalizer ref
874 case res of
875 Yield x s -> return $ Yield x (s, ref)
876 Skip s -> return $ Skip (s, ref)
877 Stop -> do
878 D.runIORefFinalizer ref
879 return Stop
881{-# INLINE_NORMAL _bracket #-}
882_bracket :: MonadCatch m
883 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
884_bracket bef aft unf =
885 gbracket bef MC.try aft (nilM (\(a, (e :: MC.SomeException)) -> aft a >>
886 MC.throwM e)) unf
888-- | @bracket before after between@ runs the @before@ action and then unfolds
889-- its output using the @between@ unfold. When the @between@ unfold is done or
890-- if an exception occurs then the @after@ action is run with the output of
891-- @before@ as argument.
892--
893-- Prefer bracketIO over this as the @after@ action in this combinator is not
894-- executed if the unfold is partially evaluated lazily and then garbage
895-- collected.
896--
897-- /Internal/
898{-# INLINE_NORMAL bracket #-}
899bracket :: MonadCatch m
900 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
901bracket bef aft (Unfold step1 inject1) = Unfold step inject
903 where
905 inject x = do
906 r <- bef x
907 s <- inject1 r
908 return (s, r)
910 {-# INLINE_LATE step #-}
911 step (st, v) = do
912 res <- step1 st `MC.onException` aft v
913 case res of
914 Yield x s -> return $ Yield x (s, v)
915 Skip s -> return $ Skip (s, v)
916 Stop -> aft v >> return Stop
918-- | @bracket before after between@ runs the @before@ action and then unfolds
919-- its output using the @between@ unfold. When the @between@ unfold is done or
920-- if an exception occurs then the @after@ action is run with the output of
921-- @before@ as argument. The after action is also executed if the unfold is
922-- paritally evaluated and then garbage collected.
923--
924-- /Internal/
925{-# INLINE_NORMAL bracketIO #-}
926bracketIO :: (MonadAsync m, MonadCatch m)
927 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
928bracketIO bef aft (Unfold step1 inject1) = Unfold step inject
930 where
932 inject x = do
933 r <- bef x
934 s <- inject1 r
935 ref <- D.newFinalizedIORef (aft r)
936 return (s, ref)
938 {-# INLINE_LATE step #-}
939 step (st, ref) = do
940 res <- step1 st `MC.onException` D.runIORefFinalizer ref
941 case res of
942 Yield x s -> return $ Yield x (s, ref)
943 Skip s -> return $ Skip (s, ref)
944 Stop -> do
945 D.runIORefFinalizer ref
946 return Stop
948-- | When unfolding if an exception occurs, unfold the exception using the
949-- exception unfold supplied as the first argument to 'handle'.
950--
951-- /Internal/
952{-# INLINE_NORMAL handle #-}
953handle :: (MonadCatch m, Exception e)
954 => Unfold m e b -> Unfold m a b -> Unfold m a b
955handle exc unf =
956 gbracket return MC.try (\_ -> return ()) (discardFirst exc) unf