git.iko.soy/streamly
on 5054813 help sign in
streamly/src/Streamly/Internal/Data/Unfold.hs
958 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.IORef (readIORef, writeIORef)
132import Data.Void (Void)
133import GHC.Types (SPEC(..))
134import Prelude hiding (concat, map, mapM, takeWhile, take, filter, const)
136import Streamly.Internal.Data.Stream.StreamD.Type (Stream(..), Step(..))
137#if __GLASGOW_HASKELL__ < 800
138import Streamly.Internal.Data.Stream.StreamD.Type (pattern Stream)
139#endif
140import Streamly.Internal.Data.Unfold.Types (Unfold(..))
141import Streamly.Internal.Data.Fold.Types (Fold(..))
142import Streamly.Internal.Data.SVar (defState, MonadAsync)
143import Control.Monad.Catch (MonadCatch)
145import qualified Prelude
146import qualified Control.Monad.Catch as MC
147import qualified Data.Tuple as Tuple
148import qualified Streamly.Internal.Data.Stream.StreamK as K
149import qualified Streamly.Internal.Data.Stream.StreamD as D
151-------------------------------------------------------------------------------
152-- Input operations
153-------------------------------------------------------------------------------
155-- | Map a function on the input argument of the 'Unfold'.
156--
157-- /Internal/
158{-# INLINE_NORMAL lmap #-}
159lmap :: (a -> c) -> Unfold m c b -> Unfold m a b
160lmap f (Unfold ustep uinject) = Unfold ustep (uinject . f)
162-- | Map an action on the input argument of the 'Unfold'.
163--
164-- /Internal/
165{-# INLINE_NORMAL lmapM #-}
166lmapM :: Monad m => (a -> m c) -> Unfold m c b -> Unfold m a b
167lmapM f (Unfold ustep uinject) = Unfold ustep (\x -> f x >>= uinject)
169-- | Supply the seed to an unfold closing the input end of the unfold.
170--
171-- /Internal/
172--
173{-# INLINE_NORMAL supply #-}
174supply :: Unfold m a b -> a -> Unfold m Void b
175supply unf a = lmap (Prelude.const a) unf
177-- | Supply the first component of the tuple to an unfold that accepts a tuple
178-- as a seed resulting in a fold that accepts the second component of the tuple
179-- as a seed.
180--
181-- /Internal/
182--
183{-# INLINE_NORMAL supplyFirst #-}
184supplyFirst :: Unfold m (a, b) c -> a -> Unfold m b c
185supplyFirst unf a = lmap (a, ) unf
187-- | Supply the second component of the tuple to an unfold that accepts a tuple
188-- as a seed resulting in a fold that accepts the first component of the tuple
189-- as a seed.
190--
191-- /Internal/
192--
193{-# INLINE_NORMAL supplySecond #-}
194supplySecond :: Unfold m (a, b) c -> b -> Unfold m a c
195supplySecond unf b = lmap (, b) unf
197-- | Convert an 'Unfold' into an unfold accepting a tuple as an argument,
198-- using the argument of the original fold as the second element of tuple and
199-- discarding the first element of the tuple.
200--
201-- /Internal/
202--
203{-# INLINE_NORMAL discardFirst #-}
204discardFirst :: Unfold m a b -> Unfold m (c, a) b
205discardFirst = lmap snd
207-- | Convert an 'Unfold' into an unfold accepting a tuple as an argument,
208-- using the argument of the original fold as the first element of tuple and
209-- discarding the second element of the tuple.
210--
211-- /Internal/
212--
213{-# INLINE_NORMAL discardSecond #-}
214discardSecond :: Unfold m a b -> Unfold m (a, c) b
215discardSecond = lmap fst
217-- | Convert an 'Unfold' that accepts a tuple as an argument into an unfold
218-- that accepts a tuple with elements swapped.
219--
220-- /Internal/
221--
222{-# INLINE_NORMAL swap #-}
223swap :: Unfold m (a, c) b -> Unfold m (c, a) b
224swap = lmap Tuple.swap
226-------------------------------------------------------------------------------
227-- Output operations
228-------------------------------------------------------------------------------
230-- | Compose an 'Unfold' and a 'Fold'. Given an @Unfold m a b@ and a
231-- @Fold m b c@, returns a monadic action @a -> m c@ representing the
232-- application of the fold on the unfolded stream.
233--
234-- /Internal/
235--
236{-# INLINE_NORMAL fold #-}
237fold :: Monad m => Unfold m a b -> Fold m b c -> a -> m c
238fold (Unfold ustep inject) (Fold fstep initial extract) a =
239 initial >>= \x -> inject a >>= go SPEC x
240 where
241 -- XXX !acc?
242 {-# INLINE_LATE go #-}
243 go !_ acc st = acc `seq` do
244 r <- ustep st
245 case r of
246 Yield x s -> do
247 acc' <- fstep acc x
248 go SPEC acc' s
249 Skip s -> go SPEC acc s
250 Stop -> extract acc
252{-# INLINE_NORMAL map #-}
253map :: Monad m => (b -> c) -> Unfold m a b -> Unfold m a c
254map f (Unfold ustep uinject) = Unfold step uinject
255 where
256 {-# INLINE_LATE step #-}
257 step st = do
258 r <- ustep st
259 return $ case r of
260 Yield x s -> Yield (f x) s
261 Skip s -> Skip s
262 Stop -> Stop
264{-# INLINE_NORMAL mapM #-}
265mapM :: Monad m => (b -> m c) -> Unfold m a b -> Unfold m a c
266mapM f (Unfold ustep uinject) = Unfold step uinject
267 where
268 {-# INLINE_LATE step #-}
269 step st = do
270 r <- ustep st
271 case r of
272 Yield x s -> f x >>= \a -> return $ Yield a s
273 Skip s -> return $ Skip s
274 Stop -> return $ Stop
276{-# INLINE_NORMAL mapMWithInput #-}
277mapMWithInput :: Monad m => (a -> b -> m c) -> Unfold m a b -> Unfold m a c
278mapMWithInput f (Unfold ustep uinject) = Unfold step inject
279 where
280 inject a = do
281 r <- uinject a
282 return (a, r)
284 {-# INLINE_LATE step #-}
285 step (inp, st) = do
286 r <- ustep st
287 case r of
288 Yield x s -> f inp x >>= \a -> return $ Yield a (inp, s)
289 Skip s -> return $ Skip (inp, s)
290 Stop -> return $ Stop
292-------------------------------------------------------------------------------
293-- Convert streams into unfolds
294-------------------------------------------------------------------------------
296{-# INLINE_LATE streamStep #-}
297streamStep :: Monad m => Stream m a -> m (Step (Stream m a) a)
298streamStep (Stream step1 state) = do
299 r <- step1 defState state
300 return $ case r of
301 Yield x s -> Yield x (Stream step1 s)
302 Skip s -> Skip (Stream step1 s)
303 Stop -> Stop
305-- | Convert a stream into an 'Unfold'. Note that a stream converted to an
306-- 'Unfold' may not be as efficient as an 'Unfold' in some situations.
307--
308-- /Internal/
309fromStream :: (K.IsStream t, Monad m) => t m b -> Unfold m Void b
310fromStream str = Unfold streamStep (\_ -> return $ D.toStreamD str)
312-- | Convert a single argument stream generator function into an
313-- 'Unfold'. Note that a stream converted to an 'Unfold' may not be as
314-- efficient as an 'Unfold' in some situations.
315--
316-- /Internal/
317fromStream1 :: (K.IsStream t, Monad m) => (a -> t m b) -> Unfold m a b
318fromStream1 f = Unfold streamStep (return . D.toStreamD . f)
320-- | Convert a two argument stream generator function into an 'Unfold'. Note
321-- that a stream converted to an 'Unfold' may not be as efficient as an
322-- 'Unfold' in some situations.
323--
324-- /Internal/
325fromStream2 :: (K.IsStream t, Monad m)
326 => (a -> b -> t m c) -> Unfold m (a, b) c
327fromStream2 f = Unfold streamStep (\(a, b) -> return $ D.toStreamD $ f a b)
329-------------------------------------------------------------------------------
330-- Unfolds
331-------------------------------------------------------------------------------
333-- | Lift a monadic function into an unfold generating a nil stream with a side
334-- effect.
335--
336{-# INLINE nilM #-}
337nilM :: Monad m => (a -> m c) -> Unfold m a b
338nilM f = Unfold step return
339 where
340 {-# INLINE_LATE step #-}
341 step x = f x >> return Stop
343-- | Prepend a monadic single element generator function to an 'Unfold'.
344--
345-- /Internal/
346{-# INLINE_NORMAL consM #-}
347consM :: Monad m => (a -> m b) -> Unfold m a b -> Unfold m a b
348consM action unf = Unfold step inject
350 where
352 inject = return . Left
354 {-# INLINE_LATE step #-}
355 step (Left a) = do
356 action a >>= \r -> return $ Yield r (Right (D.unfold unf a))
357 step (Right (UnStream step1 st)) = do
358 res <- step1 defState st
359 case res of
360 Yield x s -> return $ Yield x (Right (Stream step1 s))
361 Skip s -> return $ Skip (Right (Stream step1 s))
362 Stop -> return Stop
364-- | Lift a monadic effect into an unfold generating a singleton stream.
365--
366{-# INLINE effect #-}
367effect :: Monad m => m b -> Unfold m Void b
368effect eff = Unfold step inject
369 where
370 inject _ = return True
371 {-# INLINE_LATE step #-}
372 step True = eff >>= \r -> return $ Yield r False
373 step False = return Stop
375-- | Lift a monadic function into an unfold generating a singleton stream.
376--
377{-# INLINE singleton #-}
378singleton :: Monad m => (a -> m b) -> Unfold m a b
379singleton f = Unfold step inject
380 where
381 inject x = return $ Just x
382 {-# INLINE_LATE step #-}
383 step (Just x) = f x >>= \r -> return $ Yield r Nothing
384 step Nothing = return Stop
386-- | Identity unfold. Generates a singleton stream with the seed as the only
387-- element in the stream.
388--
389-- > identity = singleton return
390--
391{-# INLINE identity #-}
392identity :: Monad m => Unfold m a a
393identity = singleton return
395const :: Monad m => m b -> Unfold m a b
396const m = Unfold step inject
397 where
398 inject _ = return ()
399 step () = m >>= \r -> return $ Yield r ()
401-- | Generates a stream replicating the seed @n@ times.
402--
403{-# INLINE replicateM #-}
404replicateM :: Monad m => Int -> Unfold m a a
405replicateM n = Unfold step inject
406 where
407 inject x = return (x, n)
408 {-# INLINE_LATE step #-}
409 step (x, i) = return $
410 if i <= 0
411 then Stop
412 else Yield x (x, (i - 1))
414-- | Generates an infinite stream repeating the seed.
415--
416{-# INLINE repeatM #-}
417repeatM :: Monad m => Unfold m a a
418repeatM = Unfold step return
419 where
420 {-# INLINE_LATE step #-}
421 step x = return $ Yield x x
423-- | Convert a list of pure values to a 'Stream'
424{-# INLINE_LATE fromList #-}
425fromList :: Monad m => Unfold m [a] a
426fromList = Unfold step inject
427 where
428 inject x = return x
429 {-# INLINE_LATE step #-}
430 step (x:xs) = return $ Yield x xs
431 step [] = return Stop
433-- | Convert a list of monadic values to a 'Stream'
434{-# INLINE_LATE fromListM #-}
435fromListM :: Monad m => Unfold m [m a] a
436fromListM = Unfold step inject
437 where
438 inject x = return x
439 {-# INLINE_LATE step #-}
440 step (x:xs) = x >>= \r -> return $ Yield r xs
441 step [] = return Stop
443-------------------------------------------------------------------------------
444-- Filtering
445-------------------------------------------------------------------------------
447{-# INLINE_NORMAL take #-}
448take :: Monad m => Int -> Unfold m a b -> Unfold m a b
449take n (Unfold step1 inject1) = Unfold step inject
450 where
451 inject x = do
452 s <- inject1 x
453 return (s, 0)
454 {-# INLINE_LATE step #-}
455 step (st, i) | i < n = do
456 r <- step1 st
457 return $ case r of
458 Yield x s -> Yield x (s, i + 1)
459 Skip s -> Skip (s, i)
460 Stop -> Stop
461 step (_, _) = return Stop
463{-# INLINE_NORMAL takeWhileM #-}
464takeWhileM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
465takeWhileM f (Unfold step1 inject1) = Unfold step inject1
466 where
467 {-# INLINE_LATE step #-}
468 step st = do
469 r <- step1 st
470 case r of
471 Yield x s -> do
472 b <- f x
473 return $ if b then Yield x s else Stop
474 Skip s -> return $ Skip s
475 Stop -> return Stop
477{-# INLINE takeWhile #-}
478takeWhile :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
479takeWhile f = takeWhileM (return . f)
481{-# INLINE_NORMAL filterM #-}
482filterM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
483filterM f (Unfold step1 inject1) = Unfold step inject1
484 where
485 {-# INLINE_LATE step #-}
486 step st = do
487 r <- step1 st
488 case r of
489 Yield x s -> do
490 b <- f x
491 return $ if b then Yield x s else Skip s
492 Skip s -> return $ Skip s
493 Stop -> return Stop
495{-# INLINE filter #-}
496filter :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
497filter f = filterM (return . f)
499-------------------------------------------------------------------------------
500-- Enumeration
501-------------------------------------------------------------------------------
503-- | Can be used to enumerate unbounded integrals. This does not check for
504-- overflow or underflow for bounded integrals.
505{-# INLINE_NORMAL enumerateFromStepIntegral #-}
506enumerateFromStepIntegral :: (Integral a, Monad m) => Unfold m (a, a) a
507enumerateFromStepIntegral = Unfold step inject
508 where
509 inject (from, stride) = from `seq` stride `seq` return (from, stride)
510 {-# INLINE_LATE step #-}
511 step !(x, stride) = return $ Yield x $! (x + stride, stride)
513-- We are assuming that "to" is constrained by the type to be within
514-- max/min bounds.
515{-# INLINE enumerateFromToIntegral #-}
516enumerateFromToIntegral :: (Monad m, Integral a) => a -> Unfold m a a
517enumerateFromToIntegral to =
518 takeWhile (<= to) $ supplySecond enumerateFromStepIntegral 1
520{-# INLINE enumerateFromIntegral #-}
521enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => Unfold m a a
522enumerateFromIntegral = enumerateFromToIntegral maxBound
524-------------------------------------------------------------------------------
525-- Nested
526-------------------------------------------------------------------------------
528data ConcatState s1 s2 = ConcatOuter s1 | ConcatInner s1 s2
530{-# INLINE_NORMAL concat #-}
531concat :: Monad m => Unfold m a b -> Unfold m b c -> Unfold m a c
532concat (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject
533 where
534 inject x = do
535 s <- inject1 x
536 return $ ConcatOuter s
538 {-# INLINE_LATE step #-}
539 step (ConcatOuter st) = do
540 r <- step1 st
541 case r of
542 Yield x s -> do
543 innerSt <- inject2 x
544 return $ Skip (ConcatInner s innerSt)
545 Skip s -> return $ Skip (ConcatOuter s)
546 Stop -> return Stop
548 step (ConcatInner ost ist) = do
549 r <- step2 ist
550 return $ case r of
551 Yield x s -> Yield x (ConcatInner ost s)
552 Skip s -> Skip (ConcatInner ost s)
553 Stop -> Skip (ConcatOuter ost)
555data OuterProductState s1 s2 sy x y =
556 OuterProductOuter s1 y | OuterProductInner s1 sy s2 x
558{-# INLINE_NORMAL outerProduct #-}
559outerProduct :: Monad m
560 => Unfold m a b -> Unfold m c d -> Unfold m (a, c) (b, d)
561outerProduct (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject
562 where
563 inject (x, y) = do
564 s1 <- inject1 x
565 return $ OuterProductOuter s1 y
567 {-# INLINE_LATE step #-}
568 step (OuterProductOuter st1 sy) = do
569 r <- step1 st1
570 case r of
571 Yield x s -> do
572 s2 <- inject2 sy
573 return $ Skip (OuterProductInner s sy s2 x)
574 Skip s -> return $ Skip (OuterProductOuter s sy)
575 Stop -> return Stop
577 step (OuterProductInner ost sy ist x) = do
578 r <- step2 ist
579 return $ case r of
580 Yield y s -> Yield (x, y) (OuterProductInner ost sy s x)
581 Skip s -> Skip (OuterProductInner ost sy s x)
582 Stop -> Skip (OuterProductOuter ost sy)
584-- XXX This can be used to implement a Monad instance for "Unfold m ()".
586data ConcatMapState s1 s2 = ConcatMapOuter s1 | ConcatMapInner s1 s2
588{-# INLINE_NORMAL concatMapM #-}
589concatMapM :: Monad m
590 => (b -> m (Unfold m () c)) -> Unfold m a b -> Unfold m a c
591concatMapM f (Unfold step1 inject1) = Unfold step inject
592 where
593 inject x = do
594 s <- inject1 x
595 return $ ConcatMapOuter s
597 {-# INLINE_LATE step #-}
598 step (ConcatMapOuter st) = do
599 r <- step1 st
600 case r of
601 Yield x s -> do
602 Unfold step2 inject2 <- f x
603 innerSt <- inject2 ()
604 return $ Skip (ConcatMapInner s (Stream (\_ ss -> step2 ss)
605 innerSt))
606 Skip s -> return $ Skip (ConcatMapOuter s)
607 Stop -> return Stop
609 step (ConcatMapInner ost (UnStream istep ist)) = do
610 r <- istep defState ist
611 return $ case r of
612 Yield x s -> Yield x (ConcatMapInner ost (Stream istep s))
613 Skip s -> Skip (ConcatMapInner ost (Stream istep s))
614 Stop -> Skip (ConcatMapOuter ost)
616------------------------------------------------------------------------------
617-- Exceptions
618------------------------------------------------------------------------------
620-- | The most general bracketing and exception combinator. All other
621-- combinators can be expressed in terms of this combinator. This can also be
622-- used for cases which are not covered by the standard combinators.
623--
624-- /Internal/
625--
626{-# INLINE_NORMAL gbracket #-}
627gbracket
628 :: Monad m
629 => (a -> m c) -- ^ before
630 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
631 -> (c -> m d) -- ^ after, on normal stop
632 -> Unfold m (c, e) b -- ^ on exception
633 -> Unfold m c b -- ^ unfold to run
634 -> Unfold m a b
635gbracket bef exc aft (Unfold estep einject) (Unfold step1 inject1) =
636 Unfold step inject
638 where
640 inject x = do
641 r <- bef x
642 s <- inject1 r
643 return $ Right (s, r)
645 {-# INLINE_LATE step #-}
646 step (Right (st, v)) = do
647 res <- exc $ step1 st
648 case res of
649 Right r -> case r of
650 Yield x s -> return $ Yield x (Right (s, v))
651 Skip s -> return $ Skip (Right (s, v))
652 Stop -> aft v >> return Stop
653 Left e -> do
654 r <- einject (v, e)
655 return $ Skip (Left r)
656 step (Left st) = do
657 res <- estep st
658 case res of
659 Yield x s -> return $ Yield x (Left s)
660 Skip s -> return $ Skip (Left s)
661 Stop -> return Stop
663-- | The most general bracketing and exception combinator. All other
664-- combinators can be expressed in terms of this combinator. This can also be
665-- used for cases which are not covered by the standard combinators.
666--
667-- /Internal/
668--
669{-# INLINE_NORMAL gbracketIO #-}
670gbracketIO
671 :: (MonadIO m, MonadBaseControl IO m)
672 => (a -> m c) -- ^ before
673 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
674 -> (c -> m d) -- ^ after, on normal stop, or GC
675 -> Unfold m (c, e) b -- ^ on exception
676 -> Unfold m c b -- ^ unfold to run
677 -> Unfold m a b
678gbracketIO bef exc aft (Unfold estep einject) (Unfold step1 inject1) =
679 Unfold step inject
681 where
683 inject x = do
684 r <- bef x
685 ref <- D.newFinalizedIORef (aft r)
686 s <- inject1 r
687 return $ Right (s, r, ref)
689 {-# INLINE_LATE step #-}
690 step (Right (st, v, ref)) = do
691 res <- exc $ step1 st
692 case res of
693 Right r -> case r of
694 Yield x s -> return $ Yield x (Right (s, v, ref))
695 Skip s -> return $ Skip (Right (s, v, ref))
696 Stop -> do
697 liftIO $ writeIORef ref Nothing
698 aft v >> return Stop
699 Left e -> do
700 liftIO (writeIORef ref Nothing)
701 r <- einject (v, e)
702 return $ Skip (Left r)
703 step (Left st) = do
704 res <- estep st
705 case res of
706 Yield x s -> return $ Yield x (Left s)
707 Skip s -> return $ Skip (Left s)
708 Stop -> return Stop
710-- The custom implementation of "before" is slightly faster (5-7%) than
711-- "_before". This is just to document and make sure that we can always use
712-- gbracket to implement before. The same applies to other combinators as well.
713--
714{-# INLINE_NORMAL _before #-}
715_before :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
716_before action unf = gbracket (\x -> action x >> return x) (fmap Right)
717 (\_ -> return ()) undefined unf
719-- | Run a side effect before the unfold yields its first element.
720--
721-- /Internal/
722{-# INLINE_NORMAL before #-}
723before :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
724before action (Unfold step1 inject1) = Unfold step inject
726 where
728 inject x = do
729 _ <- action x
730 st <- inject1 x
731 return st
733 {-# INLINE_LATE step #-}
734 step st = do
735 res <- step1 st
736 case res of
737 Yield x s -> return $ Yield x s
738 Skip s -> return $ Skip s
739 Stop -> return Stop
741{-# INLINE_NORMAL _after #-}
742_after :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
743_after aft = gbracket return (fmap Right) aft undefined
745-- | Run a side effect whenever the unfold stops normally.
746--
747-- Prefer afterIO over this as the @after@ action in this combinator is not
748-- executed if the unfold is partially evaluated lazily and then garbage
749-- collected.
750--
751-- /Internal/
752{-# INLINE_NORMAL after #-}
753after :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
754after action (Unfold step1 inject1) = Unfold step inject
756 where
758 inject x = do
759 s <- inject1 x
760 return (s, x)
762 {-# INLINE_LATE step #-}
763 step (st, v) = do
764 res <- step1 st
765 case res of
766 Yield x s -> return $ Yield x (s, v)
767 Skip s -> return $ Skip (s, v)
768 Stop -> action v >> return Stop
770-- | Run a side effect whenever the unfold stops normally
771-- or is garbage collected after a partial lazy evaluation.
772--
773-- /Internal/
774{-# INLINE_NORMAL afterIO #-}
775afterIO :: (MonadIO m, MonadBaseControl IO m)
776 => (a -> m c) -> Unfold m a b -> Unfold m a b
777afterIO action (Unfold step1 inject1) = Unfold step inject
779 where
781 inject x = do
782 s <- inject1 x
783 ref <- D.newFinalizedIORef (action x)
784 return (s, ref)
786 {-# INLINE_LATE step #-}
787 step (st, ref) = do
788 res <- step1 st
789 case res of
790 Yield x s -> return $ Yield x (s, ref)
791 Skip s -> return $ Skip (s, ref)
792 Stop -> liftIO $ do
793 Just f <- readIORef ref
794 writeIORef ref Nothing
795 f >> return Stop
797{-# INLINE_NORMAL _onException #-}
798_onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
799_onException action unf =
800 gbracket return MC.try
801 (\_ -> return ())
802 (nilM (\(a, (e :: MC.SomeException)) -> action a >> MC.throwM e)) unf
804-- | Run a side effect whenever the unfold aborts due to an exception.
805--
806-- /Internal/
807{-# INLINE_NORMAL onException #-}
808onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
809onException action (Unfold step1 inject1) = Unfold step inject
811 where
813 inject x = do
814 s <- inject1 x
815 return (s, x)
817 {-# INLINE_LATE step #-}
818 step (st, v) = do
819 res <- step1 st `MC.onException` action v
820 case res of
821 Yield x s -> return $ Yield x (s, v)
822 Skip s -> return $ Skip (s, v)
823 Stop -> return Stop
825{-# INLINE_NORMAL _finally #-}
826_finally :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
827_finally action unf =
828 gbracket return MC.try action
829 (nilM (\(a, (e :: MC.SomeException)) -> action a >> MC.throwM e)) unf
831-- | Run a side effect whenever the unfold stops normally or aborts due to an
832-- exception.
833--
834-- Prefer finallyIO over this as the @after@ action in this combinator is not
835-- executed if the unfold is partially evaluated lazily and then garbage
836-- collected.
837--
838-- /Internal/
839{-# INLINE_NORMAL finally #-}
840finally :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
841finally action (Unfold step1 inject1) = Unfold step inject
843 where
845 inject x = do
846 s <- inject1 x
847 return (s, x)
849 {-# INLINE_LATE step #-}
850 step (st, v) = do
851 res <- step1 st `MC.onException` action v
852 case res of
853 Yield x s -> return $ Yield x (s, v)
854 Skip s -> return $ Skip (s, v)
855 Stop -> action v >> return Stop
857-- | Run a side effect whenever the unfold stops normally, aborts due to an
858-- exception or if it is garbage collected after a partial lazy evaluation.
859--
860-- /Internal/
861{-# INLINE_NORMAL finallyIO #-}
862finallyIO :: (MonadAsync m, MonadCatch m)
863 => (a -> m c) -> Unfold m a b -> Unfold m a b
864finallyIO action (Unfold step1 inject1) = Unfold step inject
866 where
868 inject x = do
869 s <- inject1 x
870 ref <- D.newFinalizedIORef (action x)
871 return (s, x, ref)
873 {-# INLINE_LATE step #-}
874 step (st, v, ref) = do
875 res <- step1 st `MC.onException` action v
876 case res of
877 Yield x s -> return $ Yield x (s, v, ref)
878 Skip s -> return $ Skip (s, v, ref)
879 Stop -> do
880 liftIO (writeIORef ref Nothing)
881 action v >> return Stop
883{-# INLINE_NORMAL _bracket #-}
884_bracket :: MonadCatch m
885 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
886_bracket bef aft unf =
887 gbracket bef MC.try aft (nilM (\(a, (e :: MC.SomeException)) -> aft a >>
888 MC.throwM e)) unf
890-- | @bracket before after between@ runs the @before@ action and then unfolds
891-- its output using the @between@ unfold. When the @between@ unfold is done or
892-- if an exception occurs then the @after@ action is run with the output of
893-- @before@ as argument.
894--
895-- Prefer bracketIO over this as the @after@ action in this combinator is not
896-- executed if the unfold is partially evaluated lazily and then garbage
897-- collected.
898--
899-- /Internal/
900{-# INLINE_NORMAL bracket #-}
901bracket :: MonadCatch m
902 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
903bracket bef aft (Unfold step1 inject1) = Unfold step inject
905 where
907 inject x = do
908 r <- bef x
909 s <- inject1 r
910 return (s, r)
912 {-# INLINE_LATE step #-}
913 step (st, v) = do
914 res <- step1 st `MC.onException` aft v
915 case res of
916 Yield x s -> return $ Yield x (s, v)
917 Skip s -> return $ Skip (s, v)
918 Stop -> aft v >> return Stop
920-- | @bracket before after between@ runs the @before@ action and then unfolds
921-- its output using the @between@ unfold. When the @between@ unfold is done or
922-- if an exception occurs then the @after@ action is run with the output of
923-- @before@ as argument. The after action is also executed if the unfold is
924-- paritally evaluated and then garbage collected.
925--
926-- /Internal/
927{-# INLINE_NORMAL bracketIO #-}
928bracketIO :: (MonadAsync m, MonadCatch m)
929 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
930bracketIO bef aft (Unfold step1 inject1) = Unfold step inject
932 where
934 inject x = do
935 r <- bef x
936 s <- inject1 r
937 ref <- D.newFinalizedIORef (aft r)
938 return (s, r, ref)
940 {-# INLINE_LATE step #-}
941 step (st, v, ref) = do
942 res <- step1 st `MC.onException` aft v
943 case res of
944 Yield x s -> return $ Yield x (s, v, ref)
945 Skip s -> return $ Skip (s, v, ref)
946 Stop -> do
947 liftIO (writeIORef ref Nothing)
948 aft v >> return Stop
950-- | When unfolding if an exception occurs, unfold the exception using the
951-- exception unfold supplied as the first argument to 'handle'.
952--
953-- /Internal/
954{-# INLINE_NORMAL handle #-}
955handle :: (MonadCatch m, Exception e)
956 => Unfold m e b -> Unfold m a b -> Unfold m a b
957handle exc unf =
958 gbracket return MC.try (\_ -> return ()) (discardFirst exc) unf