git.iko.soy/streamly
on c3dcabc help sign in
streamly/src/Streamly/Internal/Data/Unfold.hs
1098 lines 35 kB history raw
1#include "inline.hs"
3-- |
4-- Module : Streamly.Internal.Data.Unfold
5-- Copyright : (c) 2019 Composewell Technologies
6-- License : BSD3
7-- Maintainer : streamly@composewell.com
8-- Stability : experimental
9-- Portability : GHC
10--
11-- Streams forcing a closed control flow loop can be categorized under
12-- two types, unfolds and folds, both of these are duals of each other.
13--
14-- Unfold streams are really generators of a sequence of elements, we can also
15-- call them pull style streams. These are lazy producers of streams. On each
16-- evaluation the producer generates the next element. A consumer can
17-- therefore pull elements from the stream whenever it wants to. A stream
18-- consumer can multiplex pull streams by pulling elements from the chosen
19-- streams, therefore, pull streams allow merging or multiplexing. On the
20-- other hand, with this representation we cannot split or demultiplex a
21-- stream. So really these are stream sources that can be generated from a
22-- seed and can be merged or zipped into a single stream.
23--
24-- The dual of Unfolds are Folds. Folds can also be called as push style
25-- streams or reducers. These are strict consumers of streams. We keep pushing
26-- elements to a fold and we can extract the result at any point. A driver can
27-- choose which fold to push to and can also push the same element to multiple
28-- folds. Therefore, folds allow splitting or demultiplexing a stream. On the
29-- other hand, we cannot merge streams using this representation. So really
30-- these are stream consumers that reduce the stream to a single value, these
31-- consumers can be composed such that a stream can be split over multiple
32-- consumers.
33--
34-- Performance:
35--
36-- Composing a tree or graph of computations with unfolds can be much more
37-- efficient compared to composing with the Monad instance. The reason is that
38-- unfolds allow the compiler to statically know the state and optimize it
39-- using stream fusion whereas it is not possible with the monad bind because
40-- the state is determined dynamically.
42-- Open control flow style streams can also have two representations. StreamK
43-- is a producer style representation. We can also have a consumer style
44-- representation. We can use that for composable folds in StreamK
45-- representation.
46--
47module Streamly.Internal.Data.Unfold
48 (
49 -- * Unfold Type
50 Unfold
52 -- * Operations on Input
53 , lmap
54 , lmapM
55 , supply
56 , supplyFirst
57 , supplySecond
58 , discardFirst
59 , discardSecond
60 , swap
61 -- coapply
62 -- comonad
64 -- * Operations on Output
65 , fold
66 -- pipe
68 -- * Unfolds
69 , fromStream
70 , fromStreamK
71 , fromStreamD
72 , nilM
73 , consM
74 , effect
75 , singletonM
76 , singleton
77 , identity
78 , const
79 , replicateM
80 , repeatM
81 , fromList
82 , fromListM
83 , enumerateFromStepIntegral
84 , enumerateFromToIntegral
85 , enumerateFromIntegral
87 -- * Transformations
88 , map
89 , mapM
90 , mapMWithInput
92 -- * Filtering
93 , takeWhileM
94 , takeWhile
95 , take
96 , filter
97 , filterM
99 -- * Zipping
100 , zipWithM
101 , zipWith
102 , teeZipWith
104 -- * Nesting
105 , ConcatState (..)
106 , concat
107 , concatMapM
108 , outerProduct
109 , ap
110 , apDiscardFst
111 , apDiscardSnd
113 -- * Exceptions
114 , gbracket_
115 , gbracket
116 , before
117 , after_
118 , after
119 , onException
120 , finally_
121 , finally
122 , bracket_
123 , bracket
124 , handle
125 )
126where
128import Control.Exception (Exception, mask_)
129import Control.Monad ((>=>))
130import Control.Monad.IO.Class (MonadIO(..))
131import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_)
132import Data.Void (Void)
133import GHC.Types (SPEC(..))
134import Prelude
135 hiding (concat, map, mapM, takeWhile, take, filter, const, zipWith)
137import Fusion.Plugin.Types (Fuse(..))
138import Streamly.Internal.Data.Stream.StreamD.Type (Stream(..), Step(..))
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-- @
157-- lmap f = concat (singleton f)
158-- @
159--
160-- /Internal/
161{-# INLINE_NORMAL lmap #-}
162lmap :: (a -> c) -> Unfold m c b -> Unfold m a b
163lmap f (Unfold ustep uinject) = Unfold ustep (uinject . f)
165-- | Map an action on the input argument of the 'Unfold'.
166--
167-- @
168-- lmapM f = concat (singletonM f)
169-- @
170--
171-- /Internal/
172{-# INLINE_NORMAL lmapM #-}
173lmapM :: Monad m => (a -> m c) -> Unfold m c b -> Unfold m a b
174lmapM f (Unfold ustep uinject) = Unfold ustep (f >=> uinject)
176-- XXX change the signature to the following?
177-- supply :: a -> Unfold m a b -> Unfold m Void b
178--
179-- | Supply the seed to an unfold closing the input end of the unfold.
180--
181-- /Internal/
182--
183{-# INLINE_NORMAL supply #-}
184supply :: Unfold m a b -> a -> Unfold m Void b
185supply unf a = lmap (Prelude.const a) unf
187-- XXX change the signature to the following?
188-- supplyFirst :: a -> Unfold m (a, b) c -> Unfold m b c
189--
190-- | Supply the first component of the tuple to an unfold that accepts a tuple
191-- as a seed resulting in a fold that accepts the second component of the tuple
192-- as a seed.
193--
194-- /Internal/
195--
196{-# INLINE_NORMAL supplyFirst #-}
197supplyFirst :: Unfold m (a, b) c -> a -> Unfold m b c
198supplyFirst unf a = lmap (a, ) unf
200-- XXX change the signature to the following?
201-- supplySecond :: b -> Unfold m (a, b) c -> Unfold m a c
202--
203-- | Supply the second component of the tuple to an unfold that accepts a tuple
204-- as a seed resulting in a fold that accepts the first component of the tuple
205-- as a seed.
206--
207-- /Internal/
208--
209{-# INLINE_NORMAL supplySecond #-}
210supplySecond :: Unfold m (a, b) c -> b -> Unfold m a c
211supplySecond unf b = lmap (, b) unf
213-- | Convert an 'Unfold' into an unfold accepting a tuple as an argument,
214-- using the argument of the original fold as the second element of tuple and
215-- discarding the first element of the tuple.
216--
217-- /Internal/
218--
219{-# INLINE_NORMAL discardFirst #-}
220discardFirst :: Unfold m a b -> Unfold m (c, a) b
221discardFirst = lmap snd
223-- | Convert an 'Unfold' into an unfold accepting a tuple as an argument,
224-- using the argument of the original fold as the first element of tuple and
225-- discarding the second element of the tuple.
226--
227-- /Internal/
228--
229{-# INLINE_NORMAL discardSecond #-}
230discardSecond :: Unfold m a b -> Unfold m (a, c) b
231discardSecond = lmap fst
233-- | Convert an 'Unfold' that accepts a tuple as an argument into an unfold
234-- that accepts a tuple with elements swapped.
235--
236-- /Internal/
237--
238{-# INLINE_NORMAL swap #-}
239swap :: Unfold m (a, c) b -> Unfold m (c, a) b
240swap = lmap Tuple.swap
242-------------------------------------------------------------------------------
243-- Output operations
244-------------------------------------------------------------------------------
246-- | Compose an 'Unfold' and a 'Fold'. Given an @Unfold m a b@ and a
247-- @Fold m b c@, returns a monadic action @a -> m c@ representing the
248-- application of the fold on the unfolded stream.
249--
250-- /Internal/
251--
252{-# INLINE_NORMAL fold #-}
253fold :: Monad m => Unfold m a b -> Fold m b c -> a -> m c
254fold (Unfold ustep inject) (Fold fstep initial extract) a =
255 initial >>= \x -> inject a >>= go SPEC x
256 where
257 -- XXX !acc?
258 {-# INLINE_LATE go #-}
259 go !_ acc st = acc `seq` do
260 r <- ustep st
261 case r of
262 Yield x s -> do
263 acc' <- fstep acc x
264 go SPEC acc' s
265 Skip s -> go SPEC acc s
266 Stop -> extract acc
268{-# INLINE_NORMAL map #-}
269map :: Monad m => (b -> c) -> Unfold m a b -> Unfold m a c
270map f (Unfold ustep uinject) = Unfold step uinject
271 where
272 {-# INLINE_LATE step #-}
273 step st = do
274 r <- ustep st
275 return $ case r of
276 Yield x s -> Yield (f x) s
277 Skip s -> Skip s
278 Stop -> Stop
280{-# INLINE_NORMAL mapM #-}
281mapM :: Monad m => (b -> m c) -> Unfold m a b -> Unfold m a c
282mapM f (Unfold ustep uinject) = Unfold step uinject
283 where
284 {-# INLINE_LATE step #-}
285 step st = do
286 r <- ustep st
287 case r of
288 Yield x s -> f x >>= \a -> return $ Yield a s
289 Skip s -> return $ Skip s
290 Stop -> return Stop
292{-# INLINE_NORMAL mapMWithInput #-}
293mapMWithInput :: Monad m => (a -> b -> m c) -> Unfold m a b -> Unfold m a c
294mapMWithInput f (Unfold ustep uinject) = Unfold step inject
295 where
296 inject a = do
297 r <- uinject a
298 return (a, r)
300 {-# INLINE_LATE step #-}
301 step (inp, st) = do
302 r <- ustep st
303 case r of
304 Yield x s -> f inp x >>= \a -> return $ Yield a (inp, s)
305 Skip s -> return $ Skip (inp, s)
306 Stop -> return Stop
308-------------------------------------------------------------------------------
309-- Convert streams into unfolds
310-------------------------------------------------------------------------------
312-- this is horribly slow
313{-# INLINE_NORMAL fromStreamD #-}
314fromStreamD :: Monad m => Unfold m (Stream m a) a
315fromStreamD = Unfold step return
316 where
318 {-# INLINE_LATE step #-}
319 step (Stream step1 state1) = do
320 r <- step1 defState state1
321 return $ case r of
322 Yield x s -> Yield x (Stream step1 s)
323 Skip s -> Skip (Stream step1 s)
324 Stop -> Stop
326{-# INLINE_NORMAL fromStreamK #-}
327fromStreamK :: Monad m => Unfold m (K.Stream m a) a
328fromStreamK = Unfold step return
330 where
332 {-# INLINE_LATE step #-}
333 step stream = do
334 r <- K.uncons stream
335 return $ case r of
336 Just (x, xs) -> Yield x xs
337 Nothing -> Stop
339-- | Convert a stream into an 'Unfold'. Note that a stream converted to an
340-- 'Unfold' may not be as efficient as an 'Unfold' in some situations.
341--
342-- /Internal/
343{-# INLINE_NORMAL fromStream #-}
344fromStream :: (K.IsStream t, Monad m) => Unfold m (t m a) a
345fromStream = lmap K.toStream fromStreamK
347-------------------------------------------------------------------------------
348-- Unfolds
349-------------------------------------------------------------------------------
351-- | Lift a monadic function into an unfold generating a nil stream with a side
352-- effect.
353--
354{-# INLINE nilM #-}
355nilM :: Monad m => (a -> m c) -> Unfold m a b
356nilM f = Unfold step return
357 where
358 {-# INLINE_LATE step #-}
359 step x = f x >> return Stop
361-- | Prepend a monadic single element generator function to an 'Unfold'.
362--
363-- /Internal/
364{-# INLINE_NORMAL consM #-}
365consM :: Monad m => (a -> m b) -> Unfold m a b -> Unfold m a b
366consM action unf = Unfold step inject
368 where
370 inject = return . Left
372 {-# INLINE_LATE step #-}
373 step (Left a) =
374 action a >>= \r -> return $ Yield r (Right (D.unfold unf a))
375 step (Right (UnStream step1 st)) = do
376 res <- step1 defState st
377 case res of
378 Yield x s -> return $ Yield x (Right (Stream step1 s))
379 Skip s -> return $ Skip (Right (Stream step1 s))
380 Stop -> return Stop
382-- | Lift a monadic effect into an unfold generating a singleton stream.
383--
384{-# INLINE effect #-}
385effect :: Monad m => m b -> Unfold m Void b
386effect eff = Unfold step inject
387 where
388 inject _ = return True
389 {-# INLINE_LATE step #-}
390 step True = eff >>= \r -> return $ Yield r False
391 step False = return Stop
393-- XXX change it to yieldM or change yieldM in Prelude to singletonM
394--
395-- | Lift a monadic function into an unfold generating a singleton stream.
396--
397{-# INLINE singletonM #-}
398singletonM :: Monad m => (a -> m b) -> Unfold m a b
399singletonM f = Unfold step inject
400 where
401 inject x = return $ Just x
402 {-# INLINE_LATE step #-}
403 step (Just x) = f x >>= \r -> return $ Yield r Nothing
404 step Nothing = return Stop
406-- | Lift a pure function into an unfold generating a singleton stream.
407--
408{-# INLINE singleton #-}
409singleton :: Monad m => (a -> b) -> Unfold m a b
410singleton f = singletonM $ return . f
412-- | Identity unfold. Generates a singleton stream with the seed as the only
413-- element in the stream.
414--
415-- > identity = singletonM return
416--
417{-# INLINE identity #-}
418identity :: Monad m => Unfold m a a
419identity = singletonM return
421const :: Monad m => m b -> Unfold m a b
422const m = Unfold step inject
423 where
424 inject _ = return ()
425 step () = m >>= \r -> return $ Yield r ()
427-- | Generates a stream replicating the seed @n@ times.
428--
429{-# INLINE replicateM #-}
430replicateM :: Monad m => Int -> Unfold m a a
431replicateM n = Unfold step inject
432 where
433 inject x = return (x, n)
434 {-# INLINE_LATE step #-}
435 step (x, i) = return $
436 if i <= 0
437 then Stop
438 else Yield x (x, i - 1)
440-- | Generates an infinite stream repeating the seed.
441--
442{-# INLINE repeatM #-}
443repeatM :: Monad m => Unfold m a a
444repeatM = Unfold step return
445 where
446 {-# INLINE_LATE step #-}
447 step x = return $ Yield x x
449-- | Convert a list of pure values to a 'Stream'
450{-# INLINE_LATE fromList #-}
451fromList :: Monad m => Unfold m [a] a
452fromList = Unfold step inject
453 where
454 inject x = return x
455 {-# INLINE_LATE step #-}
456 step (x:xs) = return $ Yield x xs
457 step [] = return Stop
459-- | Convert a list of monadic values to a 'Stream'
460{-# INLINE_LATE fromListM #-}
461fromListM :: Monad m => Unfold m [m a] a
462fromListM = Unfold step inject
463 where
464 inject x = return x
465 {-# INLINE_LATE step #-}
466 step (x:xs) = x >>= \r -> return $ Yield r xs
467 step [] = return Stop
469-------------------------------------------------------------------------------
470-- Filtering
471-------------------------------------------------------------------------------
473{-# INLINE_NORMAL take #-}
474take :: Monad m => Int -> Unfold m a b -> Unfold m a b
475take n (Unfold step1 inject1) = Unfold step inject
476 where
477 inject x = do
478 s <- inject1 x
479 return (s, 0)
480 {-# INLINE_LATE step #-}
481 step (st, i) | i < n = do
482 r <- step1 st
483 return $ case r of
484 Yield x s -> Yield x (s, i + 1)
485 Skip s -> Skip (s, i)
486 Stop -> Stop
487 step (_, _) = return Stop
489{-# INLINE_NORMAL takeWhileM #-}
490takeWhileM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
491takeWhileM f (Unfold step1 inject1) = Unfold step inject1
492 where
493 {-# INLINE_LATE step #-}
494 step st = do
495 r <- step1 st
496 case r of
497 Yield x s -> do
498 b <- f x
499 return $ if b then Yield x s else Stop
500 Skip s -> return $ Skip s
501 Stop -> return Stop
503{-# INLINE takeWhile #-}
504takeWhile :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
505takeWhile f = takeWhileM (return . f)
507{-# INLINE_NORMAL filterM #-}
508filterM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
509filterM f (Unfold step1 inject1) = Unfold step inject1
510 where
511 {-# INLINE_LATE step #-}
512 step st = do
513 r <- step1 st
514 case r of
515 Yield x s -> do
516 b <- f x
517 return $ if b then Yield x s else Skip s
518 Skip s -> return $ Skip s
519 Stop -> return Stop
521{-# INLINE filter #-}
522filter :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
523filter f = filterM (return . f)
525-------------------------------------------------------------------------------
526-- Enumeration
527-------------------------------------------------------------------------------
529-- | Can be used to enumerate unbounded integrals. This does not check for
530-- overflow or underflow for bounded integrals.
531{-# INLINE_NORMAL enumerateFromStepIntegral #-}
532enumerateFromStepIntegral :: (Integral a, Monad m) => Unfold m (a, a) a
533enumerateFromStepIntegral = Unfold step inject
534 where
535 inject (from, stride) = from `seq` stride `seq` return (from, stride)
536 {-# INLINE_LATE step #-}
537 step (x, stride) = return $ Yield x $! (x + stride, stride)
539-- We are assuming that "to" is constrained by the type to be within
540-- max/min bounds.
541{-# INLINE enumerateFromToIntegral #-}
542enumerateFromToIntegral :: (Monad m, Integral a) => a -> Unfold m a a
543enumerateFromToIntegral to =
544 takeWhile (<= to) $ supplySecond enumerateFromStepIntegral 1
546{-# INLINE enumerateFromIntegral #-}
547enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => Unfold m a a
548enumerateFromIntegral = enumerateFromToIntegral maxBound
550-------------------------------------------------------------------------------
551-- Zipping
552-------------------------------------------------------------------------------
554{-# INLINE_NORMAL zipWithM #-}
555zipWithM :: Monad m
556 => (a -> b -> m c) -> Unfold m x a -> Unfold m y b -> Unfold m (x, y) c
557zipWithM f (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject
559 where
561 inject (x, y) = do
562 s1 <- inject1 x
563 s2 <- inject2 y
564 return (s1, s2, Nothing)
566 {-# INLINE_LATE step #-}
567 step (s1, s2, Nothing) = do
568 r <- step1 s1
569 return $
570 case r of
571 Yield x s -> Skip (s, s2, Just x)
572 Skip s -> Skip (s, s2, Nothing)
573 Stop -> Stop
575 step (s1, s2, Just x) = do
576 r <- step2 s2
577 case r of
578 Yield y s -> do
579 z <- f x y
580 return $ Yield z (s1, s, Nothing)
581 Skip s -> return $ Skip (s1, s, Just x)
582 Stop -> return Stop
584-- | Divide the input into two unfolds and then zip the outputs to a single
585-- stream.
586--
587-- @
588-- S.mapM_ print
589-- $ S.concatUnfold (UF.zipWith (,) UF.identity (UF.singleton sqrt))
590-- $ S.map (\x -> (x,x))
591-- $ S.fromList [1..10]
592-- @
593--
594-- /Internal/
595--
596{-# INLINE zipWith #-}
597zipWith :: Monad m
598 => (a -> b -> c) -> Unfold m x a -> Unfold m y b -> Unfold m (x, y) c
599zipWith f = zipWithM (\a b -> return (f a b))
601-- | Distribute the input to two unfolds and then zip the outputs to a single
602-- stream.
603--
604-- @
605-- S.mapM_ print $ S.concatUnfold (UF.teeZipWith (,) UF.identity (UF.singleton sqrt)) $ S.fromList [1..10]
606-- @
607--
608-- /Internal/
609--
610{-# INLINE_NORMAL teeZipWith #-}
611teeZipWith :: Monad m
612 => (a -> b -> c) -> Unfold m x a -> Unfold m x b -> Unfold m x c
613teeZipWith f unf1 unf2 = lmap (\x -> (x,x)) $ zipWith f unf1 unf2
615-------------------------------------------------------------------------------
616-- Nested
617-------------------------------------------------------------------------------
619{-# ANN type ConcatState Fuse #-}
620data ConcatState s1 s2 = ConcatOuter s1 | ConcatInner s1 s2
622-- | Apply the second unfold to each output element of the first unfold and
623-- flatten the output in a single stream.
624--
625-- /Internal/
626--
627{-# INLINE_NORMAL concat #-}
628concat :: Monad m => Unfold m a b -> Unfold m b c -> Unfold m a c
629concat (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject
630 where
631 inject x = do
632 s <- inject1 x
633 return $ ConcatOuter s
635 {-# INLINE_LATE step #-}
636 step (ConcatOuter st) = do
637 r <- step1 st
638 case r of
639 Yield x s -> do
640 innerSt <- inject2 x
641 return $ Skip (ConcatInner s innerSt)
642 Skip s -> return $ Skip (ConcatOuter s)
643 Stop -> return Stop
645 step (ConcatInner ost ist) = do
646 r <- step2 ist
647 return $ case r of
648 Yield x s -> Yield x (ConcatInner ost s)
649 Skip s -> Skip (ConcatInner ost s)
650 Stop -> Skip (ConcatOuter ost)
652data OuterProductState s1 s2 sy x y =
653 OuterProductOuter s1 y | OuterProductInner s1 sy s2 x
655-- | Create an outer product (vector product or cartesian product) of the
656-- output streams of two unfolds.
657--
658{-# INLINE_NORMAL outerProduct #-}
659outerProduct :: Monad m
660 => Unfold m a b -> Unfold m c d -> Unfold m (a, c) (b, d)
661outerProduct (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject
662 where
663 inject (x, y) = do
664 s1 <- inject1 x
665 return $ OuterProductOuter s1 y
667 {-# INLINE_LATE step #-}
668 step (OuterProductOuter st1 sy) = do
669 r <- step1 st1
670 case r of
671 Yield x s -> do
672 s2 <- inject2 sy
673 return $ Skip (OuterProductInner s sy s2 x)
674 Skip s -> return $ Skip (OuterProductOuter s sy)
675 Stop -> return Stop
677 step (OuterProductInner ost sy ist x) = do
678 r <- step2 ist
679 return $ case r of
680 Yield y s -> Yield (x, y) (OuterProductInner ost sy s x)
681 Skip s -> Skip (OuterProductInner ost sy s x)
682 Stop -> Skip (OuterProductOuter ost sy)
684-- Special cases of outer product
685-- | Outer product with a function application.
686--
687-- /Unimplemented/
688--
689{-# INLINE_NORMAL ap #-}
690ap :: -- Monad m =>
691 Unfold m a (b -> c) -> Unfold m d b -> Unfold m (a, d) c
692ap (Unfold _step1 _inject1) (Unfold _step2 _inject2) = undefined
694-- | Outer product discarding the first element.
695--
696-- /Unimplemented/
697--
698{-# INLINE_NORMAL apDiscardFst #-}
699apDiscardFst :: -- Monad m =>
700 Unfold m a b -> Unfold m c d -> Unfold m (a, c) d
701apDiscardFst (Unfold _step1 _inject1) (Unfold _step2 _inject2) = undefined
703-- | Outer product discarding the second element.
704--
705-- /Unimplemented/
706--
707{-# INLINE_NORMAL apDiscardSnd #-}
708apDiscardSnd :: -- Monad m =>
709 Unfold m a b -> Unfold m c d -> Unfold m (a, c) b
710apDiscardSnd (Unfold _step1 _inject1) (Unfold _step2 _inject2) = undefined
712-- XXX This can be used to implement a Monad instance for "Unfold m ()".
714data ConcatMapState s1 s2 = ConcatMapOuter s1 | ConcatMapInner s1 s2
716-- | Map an unfold generating action to each element of an unfold and
717-- flattern the results into a single stream.
718--
719{-# INLINE_NORMAL concatMapM #-}
720concatMapM :: Monad m
721 => (b -> m (Unfold m () c)) -> Unfold m a b -> Unfold m a c
722concatMapM f (Unfold step1 inject1) = Unfold step inject
723 where
724 inject x = do
725 s <- inject1 x
726 return $ ConcatMapOuter s
728 {-# INLINE_LATE step #-}
729 step (ConcatMapOuter st) = do
730 r <- step1 st
731 case r of
732 Yield x s -> do
733 Unfold step2 inject2 <- f x
734 innerSt <- inject2 ()
735 return $ Skip (ConcatMapInner s (Stream (\_ ss -> step2 ss)
736 innerSt))
737 Skip s -> return $ Skip (ConcatMapOuter s)
738 Stop -> return Stop
740 step (ConcatMapInner ost (UnStream istep ist)) = do
741 r <- istep defState ist
742 return $ case r of
743 Yield x s -> Yield x (ConcatMapInner ost (Stream istep s))
744 Skip s -> Skip (ConcatMapInner ost (Stream istep s))
745 Stop -> Skip (ConcatMapOuter ost)
747------------------------------------------------------------------------------
748-- Exceptions
749------------------------------------------------------------------------------
751-- | The most general bracketing and exception combinator. All other
752-- combinators can be expressed in terms of this combinator. This can also be
753-- used for cases which are not covered by the standard combinators.
754--
755-- /Internal/
756--
757{-# INLINE_NORMAL gbracket_ #-}
758gbracket_
759 :: Monad m
760 => (a -> m c) -- ^ before
761 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
762 -> (c -> m d) -- ^ after, on normal stop
763 -> Unfold m (c, e) b -- ^ on exception
764 -> Unfold m c b -- ^ unfold to run
765 -> Unfold m a b
766gbracket_ bef exc aft (Unfold estep einject) (Unfold step1 inject1) =
767 Unfold step inject
769 where
771 inject x = do
772 r <- bef x
773 s <- inject1 r
774 return $ Right (s, r)
776 {-# INLINE_LATE step #-}
777 step (Right (st, v)) = do
778 res <- exc $ step1 st
779 case res of
780 Right r -> case r of
781 Yield x s -> return $ Yield x (Right (s, v))
782 Skip s -> return $ Skip (Right (s, v))
783 Stop -> aft v >> return Stop
784 Left e -> do
785 r <- einject (v, e)
786 return $ Skip (Left r)
787 step (Left st) = do
788 res <- estep st
789 case res of
790 Yield x s -> return $ Yield x (Left s)
791 Skip s -> return $ Skip (Left s)
792 Stop -> return Stop
794-- | The most general bracketing and exception combinator. All other
795-- combinators can be expressed in terms of this combinator. This can also be
796-- used for cases which are not covered by the standard combinators.
797--
798-- /Internal/
799--
800{-# INLINE_NORMAL gbracket #-}
801gbracket
802 :: (MonadIO m, MonadBaseControl IO m)
803 => (a -> m c) -- ^ before
804 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
805 -> (c -> m d) -- ^ after, on normal stop, or GC
806 -> Unfold m (c, e) b -- ^ on exception
807 -> Unfold m c b -- ^ unfold to run
808 -> Unfold m a b
809gbracket bef exc aft (Unfold estep einject) (Unfold step1 inject1) =
810 Unfold step inject
812 where
814 inject x = do
815 -- Mask asynchronous exceptions to make the execution of 'bef' and
816 -- the registration of 'aft' atomic. See comment in 'D.gbracketIO'.
817 (r, ref) <- liftBaseOp_ mask_ $ do
818 r <- bef x
819 ref <- D.newFinalizedIORef (aft r)
820 return (r, ref)
821 s <- inject1 r
822 return $ Right (s, r, ref)
824 {-# INLINE_LATE step #-}
825 step (Right (st, v, ref)) = do
826 res <- exc $ step1 st
827 case res of
828 Right r -> case r of
829 Yield x s -> return $ Yield x (Right (s, v, ref))
830 Skip s -> return $ Skip (Right (s, v, ref))
831 Stop -> do
832 D.runIORefFinalizer ref
833 return Stop
834 Left e -> do
835 -- Clearing of finalizer and running of exception handler must
836 -- be atomic wrt async exceptions. Otherwise if we have cleared
837 -- the finalizer and have not run the exception handler then we
838 -- may leak the resource.
839 r <- D.withIORefFinalizer ref (einject (v, e))
840 return $ Skip (Left r)
841 step (Left st) = do
842 res <- estep st
843 case res of
844 Yield x s -> return $ Yield x (Left s)
845 Skip s -> return $ Skip (Left s)
846 Stop -> return Stop
848-- The custom implementation of "before" is slightly faster (5-7%) than
849-- "_before". This is just to document and make sure that we can always use
850-- gbracket to implement before. The same applies to other combinators as well.
851--
852{-# INLINE_NORMAL _before #-}
853_before :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
854_before action = gbracket_ (\x -> action x >> return x) (fmap Right)
855 (\_ -> return ()) undefined
857-- | Run a side effect before the unfold yields its first element.
858--
859-- /Internal/
860{-# INLINE_NORMAL before #-}
861before :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
862before action (Unfold step1 inject1) = Unfold step inject
864 where
866 inject x = do
867 _ <- action x
868 inject1 x
870 {-# INLINE_LATE step #-}
871 step st = do
872 res <- step1 st
873 case res of
874 Yield x s -> return $ Yield x s
875 Skip s -> return $ Skip s
876 Stop -> return Stop
878{-# INLINE_NORMAL _after #-}
879_after :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
880_after aft = gbracket_ return (fmap Right) aft undefined
882-- | Run a side effect whenever the unfold stops normally.
883--
884-- Prefer 'after' over this as the @after@ action in this combinator is not
885-- executed if the unfold is partially evaluated lazily and then garbage
886-- collected.
887--
888-- /Internal/
889{-# INLINE_NORMAL after_ #-}
890after_ :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
891after_ action (Unfold step1 inject1) = Unfold step inject
893 where
895 inject x = do
896 s <- inject1 x
897 return (s, x)
899 {-# INLINE_LATE step #-}
900 step (st, v) = do
901 res <- step1 st
902 case res of
903 Yield x s -> return $ Yield x (s, v)
904 Skip s -> return $ Skip (s, v)
905 Stop -> action v >> return Stop
907-- | Run a side effect whenever the unfold stops normally
908-- or is garbage collected after a partial lazy evaluation.
909--
910-- /Internal/
911{-# INLINE_NORMAL after #-}
912after :: (MonadIO m, MonadBaseControl IO m)
913 => (a -> m c) -> Unfold m a b -> Unfold m a b
914after action (Unfold step1 inject1) = Unfold step inject
916 where
918 inject x = do
919 s <- inject1 x
920 ref <- D.newFinalizedIORef (action x)
921 return (s, ref)
923 {-# INLINE_LATE step #-}
924 step (st, ref) = do
925 res <- step1 st
926 case res of
927 Yield x s -> return $ Yield x (s, ref)
928 Skip s -> return $ Skip (s, ref)
929 Stop -> do
930 D.runIORefFinalizer ref
931 return Stop
933{-# INLINE_NORMAL _onException #-}
934_onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
935_onException action =
936 gbracket_ return MC.try
937 (\_ -> return ())
938 (nilM (\(a, e :: MC.SomeException) -> action a >> MC.throwM e))
940-- | Run a side effect whenever the unfold aborts due to an exception.
941--
942-- /Internal/
943{-# INLINE_NORMAL onException #-}
944onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
945onException action (Unfold step1 inject1) = Unfold step inject
947 where
949 inject x = do
950 s <- inject1 x
951 return (s, x)
953 {-# INLINE_LATE step #-}
954 step (st, v) = do
955 res <- step1 st `MC.onException` action v
956 case res of
957 Yield x s -> return $ Yield x (s, v)
958 Skip s -> return $ Skip (s, v)
959 Stop -> return Stop
961{-# INLINE_NORMAL _finally #-}
962_finally :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
963_finally action =
964 gbracket_ return MC.try action
965 (nilM (\(a, e :: MC.SomeException) -> action a >> MC.throwM e))
967-- | Run a side effect whenever the unfold stops normally or aborts due to an
968-- exception.
969--
970-- Prefer 'finally' over this as the @after@ action in this combinator is not
971-- executed if the unfold is partially evaluated lazily and then garbage
972-- collected.
973--
974-- /Internal/
975{-# INLINE_NORMAL finally_ #-}
976finally_ :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
977finally_ action (Unfold step1 inject1) = Unfold step inject
979 where
981 inject x = do
982 s <- inject1 x
983 return (s, x)
985 {-# INLINE_LATE step #-}
986 step (st, v) = do
987 res <- step1 st `MC.onException` action v
988 case res of
989 Yield x s -> return $ Yield x (s, v)
990 Skip s -> return $ Skip (s, v)
991 Stop -> action v >> return Stop
993-- | Run a side effect whenever the unfold stops normally, aborts due to an
994-- exception or if it is garbage collected after a partial lazy evaluation.
995--
996-- /Internal/
997{-# INLINE_NORMAL finally #-}
998finally :: (MonadAsync m, MonadCatch m)
999 => (a -> m c) -> Unfold m a b -> Unfold m a b
1000finally action (Unfold step1 inject1) = Unfold step inject
1002 where
1004 inject x = do
1005 s <- inject1 x
1006 ref <- D.newFinalizedIORef (action x)
1007 return (s, ref)
1009 {-# INLINE_LATE step #-}
1010 step (st, ref) = do
1011 res <- step1 st `MC.onException` D.runIORefFinalizer ref
1012 case res of
1013 Yield x s -> return $ Yield x (s, ref)
1014 Skip s -> return $ Skip (s, ref)
1015 Stop -> do
1016 D.runIORefFinalizer ref
1017 return Stop
1019{-# INLINE_NORMAL _bracket #-}
1020_bracket :: MonadCatch m
1021 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1022_bracket bef aft =
1023 gbracket_ bef MC.try aft (nilM (\(a, e :: MC.SomeException) -> aft a >>
1024 MC.throwM e))
1026-- | @bracket_ before after between@ runs the @before@ action and then unfolds
1027-- its output using the @between@ unfold. When the @between@ unfold is done or
1028-- if an exception occurs then the @after@ action is run with the output of
1029-- @before@ as argument.
1031-- Prefer 'bracket' over this as the @after@ action in this combinator is not
1032-- executed if the unfold is partially evaluated lazily and then garbage
1033-- collected.
1035-- /Internal/
1036{-# INLINE_NORMAL bracket_ #-}
1037bracket_ :: MonadCatch m
1038 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1039bracket_ bef aft (Unfold step1 inject1) = Unfold step inject
1041 where
1043 inject x = do
1044 r <- bef x
1045 s <- inject1 r
1046 return (s, r)
1048 {-# INLINE_LATE step #-}
1049 step (st, v) = do
1050 res <- step1 st `MC.onException` aft v
1051 case res of
1052 Yield x s -> return $ Yield x (s, v)
1053 Skip s -> return $ Skip (s, v)
1054 Stop -> aft v >> return Stop
1056-- | @bracket before after between@ runs the @before@ action and then unfolds
1057-- its output using the @between@ unfold. When the @between@ unfold is done or
1058-- if an exception occurs then the @after@ action is run with the output of
1059-- @before@ as argument. The after action is also executed if the unfold is
1060-- paritally evaluated and then garbage collected.
1062-- /Internal/
1063{-# INLINE_NORMAL bracket #-}
1064bracket :: (MonadAsync m, MonadCatch m)
1065 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1066bracket bef aft (Unfold step1 inject1) = Unfold step inject
1068 where
1070 inject x = do
1071 -- Mask asynchronous exceptions to make the execution of 'bef' and
1072 -- the registration of 'aft' atomic. See comment in 'D.gbracketIO'.
1073 (r, ref) <- liftBaseOp_ mask_ $ do
1074 r <- bef x
1075 ref <- D.newFinalizedIORef (aft r)
1076 return (r, ref)
1077 s <- inject1 r
1078 return (s, ref)
1080 {-# INLINE_LATE step #-}
1081 step (st, ref) = do
1082 res <- step1 st `MC.onException` D.runIORefFinalizer ref
1083 case res of
1084 Yield x s -> return $ Yield x (s, ref)
1085 Skip s -> return $ Skip (s, ref)
1086 Stop -> do
1087 D.runIORefFinalizer ref
1088 return Stop
1090-- | When unfolding if an exception occurs, unfold the exception using the
1091-- exception unfold supplied as the first argument to 'handle'.
1093-- /Internal/
1094{-# INLINE_NORMAL handle #-}
1095handle :: (MonadCatch m, Exception e)
1096 => Unfold m e b -> Unfold m a b -> Unfold m a b
1097handle exc =
1098 gbracket_ return MC.try (\_ -> return ()) (discardFirst exc)