on aa2b166 help sign in
1206 lines 38 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-- To run the examples in this module:
12--
13-- >>> import qualified Streamly.Data.Fold as Fold
14-- >>> import qualified Streamly.Internal.Data.Unfold as Unfold
15--
16-- = Unfolds and Streams
17--
18-- An 'Unfold' type is the same as the direct style 'Stream' type except that
19-- it uses an inject function to determine the initial state of the stream
20-- based on an input. A stream is a special case of Unfold when the static
21-- input is unit or Void.
22--
23-- This allows an important optimization to occur in several cases, making the
24-- 'Unfold' a more efficient abstraction. Consider the 'concatMap' and
25-- 'unfoldMany' operations, the latter is more efficient. 'concatMap'
26-- generates a new stream object from each element in the stream by applying
27-- the supplied function to the element, the stream object includes the "step"
28-- function as well as the initial "state" of the stream. Since the stream is
29-- generated dynamically the compiler does not know the step function or the
30-- state type statically at compile time, therefore, it cannot inline it. On
31-- the other hand in case of 'unfoldMany' the compiler has visibility into
32-- the unfold's state generation function, therefore, the compiler knows all
33-- the types statically and it can inline the inject as well as the step
34-- functions, generating efficient code. Essentially, the stream is not opaque
35-- to the consumer in case of unfolds, the consumer knows how to generate the
36-- stream from a seed using a known "inject" and "step" functions.
37--
38-- A Stream is like a data object whereas unfold is like a function. Being
39-- function like, an Unfold is an instance of 'Category' and 'Arrow' type
40-- classes.
41--
42-- = Unfolds and Folds
43--
44-- Streams forcing a closed control flow loop can be categorized under
45-- two types, unfolds and folds, both of these are duals of each other.
46--
47-- Unfold streams are really generators of a sequence of elements, we can also
48-- call them pull style streams. These are lazy producers of streams. On each
49-- evaluation the producer generates the next element. A consumer can
50-- therefore pull elements from the stream whenever it wants to. A stream
51-- consumer can multiplex pull streams by pulling elements from the chosen
52-- streams, therefore, pull streams allow merging or multiplexing. On the
53-- other hand, with this representation we cannot split or demultiplex a
54-- stream. So really these are stream sources that can be generated from a
55-- seed and can be merged or zipped into a single stream.
56--
57-- The dual of Unfolds are Folds. Folds can also be called as push style
58-- streams or reducers. These are strict consumers of streams. We keep pushing
59-- elements to a fold and we can extract the result at any point. A driver can
60-- choose which fold to push to and can also push the same element to multiple
61-- folds. Therefore, folds allow splitting or demultiplexing a stream. On the
62-- other hand, we cannot merge streams using this representation. So really
63-- these are stream consumers that reduce the stream to a single value, these
64-- consumers can be composed such that a stream can be split over multiple
65-- consumers.
66--
67-- Performance:
68--
69-- Composing a tree or graph of computations with unfolds can be much more
70-- efficient compared to composing with the Monad instance. The reason is that
71-- unfolds allow the compiler to statically know the state and optimize it
72-- using stream fusion whereas it is not possible with the monad bind because
73-- the state is determined dynamically.
74--
75-- Reader:
76--
77-- An unfold acts as a reader (see 'Reader' monad). The input to an unfold acts
78-- as the read-only environment. The environment can be extracted using the
79-- 'identity' unfold (equivalent to 'ask') and transformed using 'lmap'.
81-- Open control flow style streams can also have two representations. StreamK
82-- is a producer style representation. We can also have a consumer style
83-- representation. We can use that for composable folds in StreamK
84-- representation.
85--
87-- = Performance Notes
88--
89-- 'Unfold' representation is more efficient than using streams when combining
90-- streams. 'Unfold' type allows multiple unfold actions to be composed into a
91-- single unfold function in an efficient manner by enabling the compiler to
92-- perform stream fusion optimization.
93-- @Unfold m a b@ can be considered roughly equivalent to an action @a -> t m
94-- b@ (where @t@ is a stream type). Instead of using an 'Unfold' one could just
95-- use a function of the shape @a -> t m b@. However, working with stream types
96-- like t'Streamly.SerialT' does not allow the compiler to perform stream fusion
97-- optimization when merging, appending or concatenating multiple streams.
98-- Even though stream based combinator have excellent performance, they are
99-- much less efficient when compared to combinators using 'Unfold'. For
100-- example, the 'Streamly.Prelude.concatMap' combinator which uses @a -> t m b@
101-- (where @t@ is a stream type) to generate streams is much less efficient
102-- compared to 'Streamly.Prelude.unfoldMany'.
103--
104-- On the other hand, transformation operations on stream types are as
105-- efficient as transformations on 'Unfold'.
106--
107-- We should note that in some cases working with stream types may be more
108-- convenient compared to working with the 'Unfold' type. However, if extra
109-- performance boost is important then 'Unfold' based composition should be
110-- preferred compared to stream based composition when merging or concatenating
111-- streams.
113module Streamly.Internal.Data.Unfold
114 (
115 -- * Unfold Type
116 Step(..)
117 , Unfold
119 -- * Unfolds
120 -- One to one correspondence with
121 -- "Streamly.Internal.Data.Stream.IsStream.Generate"
122 -- ** Basic Constructors
123 , mkUnfoldM
124 , mkUnfoldrM
125 , unfoldrM
126 , unfoldr
127 , functionM
128 , function
129 , identity
130 , nilM
131 , consM
133 -- ** From Values
134 , fromEffect
135 , fromPure
137 -- ** Generators
138 -- | Generate a monadic stream from a seed.
139 , repeatM
140 , replicateM
141 , fromIndicesM
142 , iterateM
144 -- ** Enumerations
145 , Enumerable (..)
147 -- ** Enumerate Num
148 , enumerateFromNum
149 , enumerateFromThenNum
150 , enumerateFromStepNum
152 -- ** Enumerating 'Bounded 'Integral' Types
153 , enumerateFromIntegralBounded
154 , enumerateFromThenIntegralBounded
155 , enumerateFromToIntegralBounded
156 , enumerateFromThenToIntegralBounded
158 -- ** Enumerating 'Unounded Integral' Types
159 , enumerateFromIntegral
160 , enumerateFromThenIntegral
161 , enumerateFromToIntegral
162 , enumerateFromThenToIntegral
164 -- ** Enumerating 'Small Integral' Types
165 , enumerateFromSmallBounded
166 , enumerateFromThenSmallBounded
167 , enumerateFromToSmall
168 , enumerateFromThenToSmall
170 -- ** Enumerating 'Fractional' Types
171 , enumerateFromFractional
172 , enumerateFromThenFractional
173 , enumerateFromToFractional
174 , enumerateFromThenToFractional
176 -- ** From Containers
177 , fromList
178 , fromListM
180 , fromStream
181 , fromStreamK
182 , fromStreamD
184 -- * Combinators
185 -- ** Mapping on Input
186 , lmap
187 , lmapM
188 , both
189 , first
190 , second
191 , discardFirst
192 , discardSecond
193 , swap
194 -- coapply
195 -- comonad
197 -- * Folding
198 , fold
200 -- ** Mapping on Output
201 , map
202 , mapM
203 , mapMWithInput
204 , postscanlM'
205 , postscan
206 , foldMany
207 -- pipe
209 -- ** Either Wrapped Input
210 , either
212 -- ** Filtering
213 , takeWhileM
214 , takeWhile
215 , take
216 , filter
217 , filterM
218 , drop
219 , dropWhile
220 , dropWhileM
222 -- ** Zipping
223 , zipWithM
224 , zipWith
226 -- ** Cross product
227 , crossWithM
228 , crossWith
229 , cross
230 , apply
232 -- ** Nesting
233 , ConcatState (..)
234 , many
235 , concatMapM
236 , bind
238 -- ** Resource Management
239 , gbracket_
240 , gbracket
241 , before
242 , after
243 , after_
244 , finally
245 , finally_
246 , bracket
247 , bracket_
249 -- ** Exceptions
250 , onException
251 , handle
252 )
253where
255import Control.Exception (Exception, mask_)
256import Control.Monad.Catch (MonadCatch)
257import Data.Functor (($>))
258import GHC.Types (SPEC(..))
259import Streamly.Internal.Control.Concurrent (MonadRunInIO, MonadAsync, withRunInIO)
260import Streamly.Internal.Data.Fold.Type (Fold(..))
261import Streamly.Internal.Data.IOFinalizer
262 (newIOFinalizer, runIOFinalizer, clearingIOFinalizer)
263import Streamly.Internal.Data.Stream.Serial (SerialT(..))
264import Streamly.Internal.Data.Stream.StreamD.Type (Stream(..), Step(..))
265import Streamly.Internal.Data.SVar.Type (defState)
267import qualified Control.Monad.Catch as MC
268import qualified Data.Tuple as Tuple
269import qualified Streamly.Internal.Data.Fold.Type as FL
270import qualified Streamly.Internal.Data.Stream.StreamD.Type as D
271import qualified Streamly.Internal.Data.Stream.StreamK.Type as K
273import Streamly.Internal.Data.Unfold.Enumeration
274import Streamly.Internal.Data.Unfold.Type
275import Prelude
276 hiding (map, mapM, takeWhile, take, filter, const, zipWith
277 , drop, dropWhile, either)
279-- $setup
280-- >>> import qualified Streamly.Data.Fold as Fold
281-- >>> import qualified Streamly.Internal.Data.Unfold as Unfold
282-- >>> import qualified Streamly.Prelude as Stream
285-- | Convert an 'Unfold' into an unfold accepting a tuple as an argument,
286-- using the argument of the original fold as the second element of tuple and
287-- discarding the first element of the tuple.
288--
289-- @
290-- discardFirst = Unfold.lmap snd
291-- @
292--
293-- /Pre-release/
294--
295{-# INLINE_NORMAL discardFirst #-}
296discardFirst :: Unfold m a b -> Unfold m (c, a) b
297discardFirst = lmap snd
299-- | Convert an 'Unfold' into an unfold accepting a tuple as an argument,
300-- using the argument of the original fold as the first element of tuple and
301-- discarding the second element of the tuple.
302--
303-- @
304-- discardSecond = Unfold.lmap fst
305-- @
306--
307-- /Pre-release/
308--
309{-# INLINE_NORMAL discardSecond #-}
310discardSecond :: Unfold m a b -> Unfold m (a, c) b
311discardSecond = lmap fst
313-- | Convert an 'Unfold' that accepts a tuple as an argument into an unfold
314-- that accepts a tuple with elements swapped.
315--
316-- @
317-- swap = Unfold.lmap Tuple.swap
318-- @
319--
320-- /Pre-release/
321--
322{-# INLINE_NORMAL swap #-}
323swap :: Unfold m (a, c) b -> Unfold m (c, a) b
324swap = lmap Tuple.swap
326-------------------------------------------------------------------------------
327-- Output operations
328-------------------------------------------------------------------------------
330-- XXX Do we need this combinator or the stream based idiom is enough?
332-- | Compose an 'Unfold' and a 'Fold'. Given an @Unfold m a b@ and a
333-- @Fold m b c@, returns a monadic action @a -> m c@ representing the
334-- application of the fold on the unfolded stream.
335--
336-- >>> Unfold.fold Fold.sum Unfold.fromList [1..100]
337-- 5050
338--
339-- >>> fold f u = Stream.fold f . Stream.unfold u
340--
341-- /Pre-release/
342--
343{-# INLINE_NORMAL fold #-}
344fold :: Monad m => Fold m b c -> Unfold m a b -> a -> m c
345fold (Fold fstep initial extract) (Unfold ustep inject) a = do
346 res <- initial
347 case res of
348 FL.Partial x -> inject a >>= go SPEC x
349 FL.Done b -> return b
351 where
353 {-# INLINE_LATE go #-}
354 go !_ !fs st = do
355 r <- ustep st
356 case r of
357 Yield x s -> do
358 res <- fstep fs x
359 case res of
360 FL.Partial fs1 -> go SPEC fs1 s
361 FL.Done c -> return c
362 Skip s -> go SPEC fs s
363 Stop -> extract fs
365-- {-# ANN type FoldMany Fuse #-}
366data FoldMany s fs b a
367 = FoldManyStart s
368 | FoldManyFirst fs s
369 | FoldManyLoop s fs
370 | FoldManyYield b (FoldMany s fs b a)
371 | FoldManyDone
373-- | Apply a fold multiple times on the output of an unfold.
374--
375-- /Pre-release/
376{-# INLINE_NORMAL foldMany #-}
377foldMany :: Monad m => Fold m b c -> Unfold m a b -> Unfold m a c
378foldMany (Fold fstep initial extract) (Unfold ustep inject1) =
379 Unfold step inject
381 where
383 inject x = do
384 r <- inject1 x
385 return (FoldManyStart r)
387 {-# INLINE consume #-}
388 consume x s fs = do
389 res <- fstep fs x
390 return
391 $ Skip
392 $ case res of
393 FL.Done b -> FoldManyYield b (FoldManyStart s)
394 FL.Partial ps -> FoldManyLoop s ps
396 {-# INLINE_LATE step #-}
397 step (FoldManyStart st) = do
398 r <- initial
399 return
400 $ Skip
401 $ case r of
402 FL.Done b -> FoldManyYield b (FoldManyStart st)
403 FL.Partial fs -> FoldManyFirst fs st
404 step (FoldManyFirst fs st) = do
405 r <- ustep st
406 case r of
407 Yield x s -> consume x s fs
408 Skip s -> return $ Skip (FoldManyFirst fs s)
409 Stop -> return Stop
410 step (FoldManyLoop st fs) = do
411 r <- ustep st
412 case r of
413 Yield x s -> consume x s fs
414 Skip s -> return $ Skip (FoldManyLoop s fs)
415 Stop -> do
416 b <- extract fs
417 return $ Skip (FoldManyYield b FoldManyDone)
418 step (FoldManyYield b next) = return $ Yield b next
419 step FoldManyDone = return Stop
421-- | Apply a monadic function to each element of the stream and replace it
422-- with the output of the resulting action.
423--
424-- /Since: 0.8.0/
425--
426{-# INLINE_NORMAL mapM #-}
427mapM :: Monad m => (b -> m c) -> Unfold m a b -> Unfold m a c
428mapM f (Unfold ustep uinject) = Unfold step uinject
429 where
430 {-# INLINE_LATE step #-}
431 step st = do
432 r <- ustep st
433 case r of
434 Yield x s -> f x >>= \a -> return $ Yield a s
435 Skip s -> return $ Skip s
436 Stop -> return Stop
438{-# INLINE_NORMAL mapMWithInput #-}
439mapMWithInput :: Monad m => (a -> b -> m c) -> Unfold m a b -> Unfold m a c
440mapMWithInput f (Unfold ustep uinject) = Unfold step inject
441 where
442 inject a = do
443 r <- uinject a
444 return (a, r)
446 {-# INLINE_LATE step #-}
447 step (inp, st) = do
448 r <- ustep st
449 case r of
450 Yield x s -> f inp x >>= \a -> return $ Yield a (inp, s)
451 Skip s -> return $ Skip (inp, s)
452 Stop -> return Stop
454-------------------------------------------------------------------------------
455-- Either
456-------------------------------------------------------------------------------
458-- | Make an unfold operate on values wrapped in an @Either a a@ type. 'Right
459-- a' translates to 'Right b' and 'Left a' translates to 'Left b'.
460--
461-- /Internal/
462{-# INLINE_NORMAL either #-}
463either :: Applicative m => Unfold m a b -> Unfold m (Either a a) (Either b b)
464either (Unfold step1 inject1) = Unfold step inject
466 where
468 inject (Left a) = (, Left) <$> inject1 a
469 inject (Right a) = (, Right) <$> inject1 a
471 {-# INLINE_LATE step #-}
472 step (st, f) = do
473 (\case
474 Yield x s -> Yield (f x) (s, f)
475 Skip s -> Skip (s, f)
476 Stop -> Stop) <$> step1 st
478-- | Scan the output of an 'Unfold' to change it in a stateful manner.
479--
480-- /Pre-release/
481{-# INLINE_NORMAL postscan #-}
482postscan :: Monad m => Fold m b c -> Unfold m a b -> Unfold m a c
483postscan (Fold stepF initial extract) (Unfold stepU injectU) =
484 Unfold step inject
486 where
488 inject a = do
489 r <- initial
490 case r of
491 FL.Partial fs -> Just . (fs,) <$> injectU a
492 FL.Done _ -> return Nothing
494 {-# INLINE_LATE step #-}
495 step (Just (fs, us)) = do
496 ru <- stepU us
497 case ru of
498 Yield x s -> do
499 rf <- stepF fs x
500 case rf of
501 FL.Done v -> return $ Yield v Nothing
502 FL.Partial fs1 -> do
503 v <- extract fs1
504 return $ Yield v (Just (fs1, s))
505 Skip s -> return $ Skip (Just (fs, s))
506 Stop -> return Stop
508 step Nothing = return Stop
510-- | Scan the output of an 'Unfold' to change it in a stateful manner.
511--
512-- /Pre-release/
513{-# INLINE_NORMAL postscanlM' #-}
514postscanlM' :: Monad m => (b -> a -> m b) -> m b -> Unfold m c a -> Unfold m c b
515postscanlM' f z = postscan (FL.foldlM' f z)
517-------------------------------------------------------------------------------
518-- Convert streams into unfolds
519-------------------------------------------------------------------------------
521{-# INLINE_NORMAL fromStreamD #-}
522fromStreamD :: Applicative m => Unfold m (Stream m a) a
523fromStreamD = Unfold step pure
525 where
527 {-# INLINE_LATE step #-}
528 step (UnStream step1 state1) =
529 (\case
530 Yield x s -> Yield x (Stream step1 s)
531 Skip s -> Skip (Stream step1 s)
532 Stop -> Stop) <$> step1 defState state1
534{-# INLINE_NORMAL fromStreamK #-}
535fromStreamK :: Applicative m => Unfold m (K.Stream m a) a
536fromStreamK = Unfold step pure
538 where
540 {-# INLINE_LATE step #-}
541 step stream = do
542 (\case
543 Just (x, xs) -> Yield x xs
544 Nothing -> Stop) <$> K.uncons stream
546-- XXX Using Unfold.fromStreamD seems to be faster (using cross product test
547-- case) than using fromStream even if it is implemented using fromStreamD.
548-- Check if StreamK to StreamD rewrite rules are working correctly when
549-- implementing fromStream using fromStreamD.
550--
551-- | Convert a stream into an 'Unfold'. Note that a stream converted to an
552-- 'Unfold' may not be as efficient as an 'Unfold' in some situations.
553--
554-- /Since: 0.8.0/
555--
556{-# INLINE_NORMAL fromStream #-}
557fromStream :: Applicative m => Unfold m (SerialT m a) a
558fromStream = lmap getSerialT fromStreamK
560-------------------------------------------------------------------------------
561-- Unfolds
562-------------------------------------------------------------------------------
564-- | Lift a monadic function into an unfold generating a nil stream with a side
565-- effect.
566--
567{-# INLINE nilM #-}
568nilM :: Applicative m => (a -> m c) -> Unfold m a b
569nilM f = Unfold step pure
571 where
573 {-# INLINE_LATE step #-}
574 step x = f x $> Stop
576-- | Prepend a monadic single element generator function to an 'Unfold'. The
577-- same seed is used in the action as well as the unfold.
578--
579-- /Pre-release/
580{-# INLINE_NORMAL consM #-}
581consM :: Applicative m => (a -> m b) -> Unfold m a b -> Unfold m a b
582consM action unf = Unfold step inject
584 where
586 inject = pure . Left
588 {-# INLINE_LATE step #-}
589 step (Left a) = (`Yield` Right (D.unfold unf a)) <$> action a
590 step (Right (UnStream step1 st)) = do
591 (\case
592 Yield x s -> Yield x (Right (Stream step1 s))
593 Skip s -> Skip (Right (Stream step1 s))
594 Stop -> Stop) <$> step1 defState st
596-- XXX Check if "unfold (fromList [1..10])" fuses, if it doesn't we can use
597-- rewrite rules to rewrite list enumerations to unfold enumerations.
598--
599-- | Convert a list of pure values to a 'Stream'
600--
601-- /Since: 0.8.0/
602--
603{-# INLINE_LATE fromList #-}
604fromList :: Applicative m => Unfold m [a] a
605fromList = Unfold step pure
607 where
609 {-# INLINE_LATE step #-}
610 step (x:xs) = pure $ Yield x xs
611 step [] = pure Stop
613-- | Convert a list of monadic values to a 'Stream'
614--
615-- /Since: 0.8.0/
616--
617{-# INLINE_LATE fromListM #-}
618fromListM :: Applicative m => Unfold m [m a] a
619fromListM = Unfold step pure
621 where
623 {-# INLINE_LATE step #-}
624 step (x:xs) = (`Yield` xs) <$> x
625 step [] = pure Stop
627------------------------------------------------------------------------------
628-- Specialized Generation
629------------------------------------------------------------------------------
631-- | Generates a stream replicating the seed @n@ times.
632--
633-- /Since: 0.8.0/
634--
635{-# INLINE replicateM #-}
636replicateM :: Applicative m => Int -> Unfold m (m a) a
637replicateM n = Unfold step inject
639 where
641 inject action = pure (action, n)
643 {-# INLINE_LATE step #-}
644 step (action, i) =
645 if i <= 0
646 then pure Stop
647 else (\x -> Yield x (action, i - 1)) <$> action
649-- | Generates an infinite stream repeating the seed.
650--
651-- /Since: 0.8.0/
652--
653{-# INLINE repeatM #-}
654repeatM :: Applicative m => Unfold m (m a) a
655repeatM = Unfold step pure
657 where
659 {-# INLINE_LATE step #-}
660 step action = (`Yield` action) <$> action
662-- | Generates an infinite stream starting with the given seed and applying the
663-- given function repeatedly.
664--
665-- /Since: 0.8.0/
666--
667{-# INLINE iterateM #-}
668iterateM :: Applicative m => (a -> m a) -> Unfold m (m a) a
669iterateM f = Unfold step id
671 where
673 {-# INLINE_LATE step #-}
674 step x = Yield x <$> f x
676-- | @fromIndicesM gen@ generates an infinite stream of values using @gen@
677-- starting from the seed.
678--
679-- @
680-- fromIndicesM f = Unfold.mapM f $ Unfold.enumerateFrom 0
681-- @
682--
683-- /Pre-release/
684--
685{-# INLINE_NORMAL fromIndicesM #-}
686fromIndicesM :: Applicative m => (Int -> m a) -> Unfold m Int a
687fromIndicesM gen = Unfold step pure
689 where
691 {-# INLINE_LATE step #-}
692 step i = (`Yield` (i + 1)) <$> gen i
694-------------------------------------------------------------------------------
695-- Filtering
696-------------------------------------------------------------------------------
698-- |
699-- >>> u = Unfold.take 2 Unfold.fromList
700-- >>> Unfold.fold Fold.toList u [1..100]
701-- [1,2]
702--
703-- /Since: 0.8.0/
704--
705{-# INLINE_NORMAL take #-}
706take :: Applicative m => Int -> Unfold m a b -> Unfold m a b
707take n (Unfold step1 inject1) = Unfold step inject
709 where
711 inject x = (, 0) <$> inject1 x
713 {-# INLINE_LATE step #-}
714 step (st, i) | i < n = do
715 (\case
716 Yield x s -> Yield x (s, i + 1)
717 Skip s -> Skip (s, i)
718 Stop -> Stop) <$> step1 st
719 step (_, _) = pure Stop
721-- | Same as 'filter' but with a monadic predicate.
722--
723-- /Since: 0.8.0/
724--
725{-# INLINE_NORMAL filterM #-}
726filterM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
727filterM f (Unfold step1 inject1) = Unfold step inject1
728 where
729 {-# INLINE_LATE step #-}
730 step st = do
731 r <- step1 st
732 case r of
733 Yield x s -> do
734 b <- f x
735 return $ if b then Yield x s else Skip s
736 Skip s -> return $ Skip s
737 Stop -> return Stop
739-- | Include only those elements that pass a predicate.
740--
741-- /Since: 0.8.0/
742--
743{-# INLINE filter #-}
744filter :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
745filter f = filterM (return . f)
747-- | @drop n unf@ drops @n@ elements from the stream generated by @unf@.
748--
749-- /Since: 0.8.0/
750--
751{-# INLINE_NORMAL drop #-}
752drop :: Applicative m => Int -> Unfold m a b -> Unfold m a b
753drop n (Unfold step inject) = Unfold step' inject'
755 where
757 inject' a = (, n) <$> inject a
759 {-# INLINE_LATE step' #-}
760 step' (st, i)
761 | i > 0 = do
762 (\case
763 Yield _ s -> Skip (s, i - 1)
764 Skip s -> Skip (s, i)
765 Stop -> Stop) <$> step st
766 | otherwise = do
767 (\case
768 Yield x s -> Yield x (s, 0)
769 Skip s -> Skip (s, 0)
770 Stop -> Stop) <$> step st
772-- | @dropWhileM f unf@ drops elements from the stream generated by @unf@ while
773-- the condition holds true. The condition function @f@ is /monadic/ in nature.
774--
775-- /Since: 0.8.0/
776--
777{-# INLINE_NORMAL dropWhileM #-}
778dropWhileM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
779dropWhileM f (Unfold step inject) = Unfold step' inject'
781 where
783 inject' a = do
784 b <- inject a
785 return $ Left b
787 {-# INLINE_LATE step' #-}
788 step' (Left st) = do
789 r <- step st
790 case r of
791 Yield x s -> do
792 b <- f x
793 return
794 $ if b
795 then Skip (Left s)
796 else Yield x (Right s)
797 Skip s -> return $ Skip (Left s)
798 Stop -> return Stop
799 step' (Right st) = do
800 r <- step st
801 return
802 $ case r of
803 Yield x s -> Yield x (Right s)
804 Skip s -> Skip (Right s)
805 Stop -> Stop
807-- | Similar to 'dropWhileM' but with a pure condition function.
808--
809-- /Since: 0.8.0/
810--
811{-# INLINE dropWhile #-}
812dropWhile :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
813dropWhile f = dropWhileM (return . f)
815------------------------------------------------------------------------------
816-- Exceptions
817------------------------------------------------------------------------------
819-- | Like 'gbracket' but with following differences:
820--
821-- * alloc action @a -> m c@ runs with async exceptions enabled
822-- * cleanup action @c -> m d@ won't run if the stream is garbage collected
823-- after partial evaluation.
824-- * does not require a 'MonadAsync' constraint.
825--
826-- /Inhibits stream fusion/
827--
828-- /Pre-release/
829--
830{-# INLINE_NORMAL gbracket_ #-}
831gbracket_
832 :: Monad m
833 => (a -> m c) -- ^ before
834 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
835 -> (c -> m d) -- ^ after, on normal stop
836 -> Unfold m (c, e) b -- ^ on exception
837 -> Unfold m c b -- ^ unfold to run
838 -> Unfold m a b
839gbracket_ bef exc aft (Unfold estep einject) (Unfold step1 inject1) =
840 Unfold step inject
842 where
844 inject x = do
845 r <- bef x
846 s <- inject1 r
847 return $ Right (s, r)
849 {-# INLINE_LATE step #-}
850 step (Right (st, v)) = do
851 res <- exc $ step1 st
852 case res of
853 Right r -> case r of
854 Yield x s -> return $ Yield x (Right (s, v))
855 Skip s -> return $ Skip (Right (s, v))
856 Stop -> aft v >> return Stop
857 -- XXX Do not handle async exceptions, just rethrow them.
858 Left e -> do
859 r <- einject (v, e)
860 return $ Skip (Left r)
861 step (Left st) = do
862 res <- estep st
863 return $ case res of
864 Yield x s -> Yield x (Left s)
865 Skip s -> Skip (Left s)
866 Stop -> Stop
868-- | Run the alloc action @a -> m c@ with async exceptions disabled but keeping
869-- blocking operations interruptible (see 'Control.Exception.mask'). Use the
870-- output @c@ as input to @Unfold m c b@ to generate an output stream. When
871-- unfolding use the supplied @try@ operation @forall s. m s -> m (Either e s)@
872-- to catch synchronous exceptions. If an exception occurs run the exception
873-- handling unfold @Unfold m (c, e) b@.
874--
875-- The cleanup action @c -> m d@, runs whenever the stream ends normally, due
876-- to a sync or async exception or if it gets garbage collected after a partial
877-- lazy evaluation. See 'bracket' for the semantics of the cleanup action.
878--
879-- 'gbracket' can express all other exception handling combinators.
880--
881-- /Inhibits stream fusion/
882--
883-- /Pre-release/
884{-# INLINE_NORMAL gbracket #-}
885gbracket
886 :: MonadRunInIO m
887 => (a -> m c) -- ^ before
888 -> (c -> m d) -- ^ after, on normal stop, or GC
889 -> Unfold m (c, e) b -- ^ on exception
890 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
891 -> Unfold m c b -- ^ unfold to run
892 -> Unfold m a b
893gbracket bef aft (Unfold estep einject) ftry (Unfold step1 inject1) =
894 Unfold step inject
896 where
898 inject x = do
899 -- Mask asynchronous exceptions to make the execution of 'bef' and
900 -- the registration of 'aft' atomic. See comment in 'D.gbracketIO'.
901 (r, ref) <- withRunInIO $ \run -> mask_ $ run $ do
902 r <- bef x
903 ref <- newIOFinalizer (aft r)
904 return (r, ref)
905 s <- inject1 r
906 return $ Right (s, r, ref)
908 {-# INLINE_LATE step #-}
909 step (Right (st, v, ref)) = do
910 res <- ftry $ step1 st
911 case res of
912 Right r -> case r of
913 Yield x s -> return $ Yield x (Right (s, v, ref))
914 Skip s -> return $ Skip (Right (s, v, ref))
915 Stop -> do
916 runIOFinalizer ref
917 return Stop
918 -- XXX Do not handle async exceptions, just rethrow them.
919 Left e -> do
920 -- Clearing of finalizer and running of exception handler must
921 -- be atomic wrt async exceptions. Otherwise if we have cleared
922 -- the finalizer and have not run the exception handler then we
923 -- may leak the resource.
924 r <- clearingIOFinalizer ref (einject (v, e))
925 return $ Skip (Left r)
926 step (Left st) = do
927 res <- estep st
928 return $ case res of
929 Yield x s -> Yield x (Left s)
930 Skip s -> Skip (Left s)
931 Stop -> Stop
933-- | Run a side effect @a -> m c@ on the input @a@ before unfolding it using
934-- @Unfold m a b@.
935--
936-- > before f = lmapM (\a -> f a >> return a)
937--
938-- /Pre-release/
939{-# INLINE_NORMAL before #-}
940before :: (a -> m c) -> Unfold m a b -> Unfold m a b
941before action (Unfold step inject) = Unfold step (action >> inject)
943-- The custom implementation of "after_" is slightly faster (5-7%) than
944-- "_after". This is just to document and make sure that we can always use
945-- gbracket to implement after_ The same applies to other combinators as well.
946--
947{-# INLINE_NORMAL _after #-}
948_after :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
949_after aft = gbracket_ return (fmap Right) aft undefined
951-- | Like 'after' with following differences:
952--
953-- * action @a -> m c@ won't run if the stream is garbage collected
954-- after partial evaluation.
955-- * Monad @m@ does not require any other constraints.
956--
957-- /Pre-release/
958{-# INLINE_NORMAL after_ #-}
959after_ :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
960after_ action (Unfold step1 inject1) = Unfold step inject
962 where
964 inject x = do
965 s <- inject1 x
966 return (s, x)
968 {-# INLINE_LATE step #-}
969 step (st, v) = do
970 res <- step1 st
971 case res of
972 Yield x s -> return $ Yield x (s, v)
973 Skip s -> return $ Skip (s, v)
974 Stop -> action v >> return Stop
976-- | Unfold the input @a@ using @Unfold m a b@, run an action on @a@ whenever
977-- the unfold stops normally, or if it is garbage collected after a partial
978-- lazy evaluation.
979--
980-- The semantics of the action @a -> m c@ are similar to the cleanup action
981-- semantics in 'bracket'.
982--
983-- /See also 'after_'/
984--
985-- /Pre-release/
986{-# INLINE_NORMAL after #-}
987after :: MonadRunInIO m
988 => (a -> m c) -> Unfold m a b -> Unfold m a b
989after action (Unfold step1 inject1) = Unfold step inject
991 where
993 inject x = do
994 s <- inject1 x
995 ref <- newIOFinalizer (action x)
996 return (s, ref)
998 {-# INLINE_LATE step #-}
999 step (st, ref) = do
1000 res <- step1 st
1001 case res of
1002 Yield x s -> return $ Yield x (s, ref)
1003 Skip s -> return $ Skip (s, ref)
1004 Stop -> do
1005 runIOFinalizer ref
1006 return Stop
1008{-# INLINE_NORMAL _onException #-}
1009_onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
1010_onException action =
1011 gbracket_ return MC.try
1012 (\_ -> return ())
1013 (nilM (\(a, e :: MC.SomeException) -> action a >> MC.throwM e))
1015-- | Unfold the input @a@ using @Unfold m a b@, run the action @a -> m c@ on
1016-- @a@ if the unfold aborts due to an exception.
1018-- /Inhibits stream fusion/
1020-- /Pre-release/
1021{-# INLINE_NORMAL onException #-}
1022onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
1023onException action (Unfold step1 inject1) = Unfold step inject
1025 where
1027 inject x = do
1028 s <- inject1 x
1029 return (s, x)
1031 {-# INLINE_LATE step #-}
1032 step (st, v) = do
1033 res <- step1 st `MC.onException` action v
1034 return $ case res of
1035 Yield x s -> Yield x (s, v)
1036 Skip s -> Skip (s, v)
1037 Stop -> Stop
1039{-# INLINE_NORMAL _finally #-}
1040_finally :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
1041_finally action =
1042 gbracket_ return MC.try action
1043 (nilM (\(a, e :: MC.SomeException) -> action a >> MC.throwM e))
1045-- | Like 'finally' with following differences:
1047-- * action @a -> m c@ won't run if the stream is garbage collected
1048-- after partial evaluation.
1049-- * does not require a 'MonadAsync' constraint.
1051-- /Inhibits stream fusion/
1053-- /Pre-release/
1054{-# INLINE_NORMAL finally_ #-}
1055finally_ :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
1056finally_ action (Unfold step1 inject1) = Unfold step inject
1058 where
1060 inject x = do
1061 s <- inject1 x
1062 return (s, x)
1064 {-# INLINE_LATE step #-}
1065 step (st, v) = do
1066 res <- step1 st `MC.onException` action v
1067 case res of
1068 Yield x s -> return $ Yield x (s, v)
1069 Skip s -> return $ Skip (s, v)
1070 Stop -> action v >> return Stop
1072-- | Unfold the input @a@ using @Unfold m a b@, run an action on @a@ whenever
1073-- the unfold stops normally, aborts due to an exception or if it is garbage
1074-- collected after a partial lazy evaluation.
1076-- The semantics of the action @a -> m c@ are similar to the cleanup action
1077-- semantics in 'bracket'.
1079-- @
1080-- finally release = bracket return release
1081-- @
1083-- /See also 'finally_'/
1085-- /Inhibits stream fusion/
1087-- /Pre-release/
1088{-# INLINE_NORMAL finally #-}
1089finally :: (MonadAsync m, MonadCatch m)
1090 => (a -> m c) -> Unfold m a b -> Unfold m a b
1091finally action (Unfold step1 inject1) = Unfold step inject
1093 where
1095 inject x = do
1096 s <- inject1 x
1097 ref <- newIOFinalizer (action x)
1098 return (s, ref)
1100 {-# INLINE_LATE step #-}
1101 step (st, ref) = do
1102 res <- step1 st `MC.onException` runIOFinalizer ref
1103 case res of
1104 Yield x s -> return $ Yield x (s, ref)
1105 Skip s -> return $ Skip (s, ref)
1106 Stop -> do
1107 runIOFinalizer ref
1108 return Stop
1110{-# INLINE_NORMAL _bracket #-}
1111_bracket :: MonadCatch m
1112 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1113_bracket bef aft =
1114 gbracket_ bef MC.try aft (nilM (\(a, e :: MC.SomeException) -> aft a >>
1115 MC.throwM e))
1117-- | Like 'bracket' but with following differences:
1119-- * alloc action @a -> m c@ runs with async exceptions enabled
1120-- * cleanup action @c -> m d@ won't run if the stream is garbage collected
1121-- after partial evaluation.
1122-- * does not require a 'MonadAsync' constraint.
1124-- /Inhibits stream fusion/
1126-- /Pre-release/
1127{-# INLINE_NORMAL bracket_ #-}
1128bracket_ :: MonadCatch m
1129 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1130bracket_ bef aft (Unfold step1 inject1) = Unfold step inject
1132 where
1134 inject x = do
1135 r <- bef x
1136 s <- inject1 r
1137 return (s, r)
1139 {-# INLINE_LATE step #-}
1140 step (st, v) = do
1141 res <- step1 st `MC.onException` aft v
1142 case res of
1143 Yield x s -> return $ Yield x (s, v)
1144 Skip s -> return $ Skip (s, v)
1145 Stop -> aft v >> return Stop
1147-- | Run the alloc action @a -> m c@ with async exceptions disabled but keeping
1148-- blocking operations interruptible (see 'Control.Exception.mask'). Use the
1149-- output @c@ as input to @Unfold m c b@ to generate an output stream.
1151-- @c@ is usually a resource under the state of monad @m@, e.g. a file
1152-- handle, that requires a cleanup after use. The cleanup action @c -> m d@,
1153-- runs whenever the stream ends normally, due to a sync or async exception or
1154-- if it gets garbage collected after a partial lazy evaluation.
1156-- 'bracket' only guarantees that the cleanup action runs, and it runs with
1157-- async exceptions enabled. The action must ensure that it can successfully
1158-- cleanup the resource in the face of sync or async exceptions.
1160-- When the stream ends normally or on a sync exception, cleanup action runs
1161-- immediately in the current thread context, whereas in other cases it runs in
1162-- the GC context, therefore, cleanup may be delayed until the GC gets to run.
1164-- /See also: 'bracket_', 'gbracket'/
1166-- /Inhibits stream fusion/
1168-- /Pre-release/
1169{-# INLINE_NORMAL bracket #-}
1170bracket :: (MonadAsync m, MonadCatch m)
1171 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1172bracket bef aft (Unfold step1 inject1) = Unfold step inject
1174 where
1176 inject x = do
1177 -- Mask asynchronous exceptions to make the execution of 'bef' and
1178 -- the registration of 'aft' atomic. See comment in 'D.gbracketIO'.
1179 (r, ref) <- withRunInIO $ \run -> mask_ $ run $ do
1180 r <- bef x
1181 ref <- newIOFinalizer (aft r)
1182 return (r, ref)
1183 s <- inject1 r
1184 return (s, ref)
1186 {-# INLINE_LATE step #-}
1187 step (st, ref) = do
1188 res <- step1 st `MC.onException` runIOFinalizer ref
1189 case res of
1190 Yield x s -> return $ Yield x (s, ref)
1191 Skip s -> return $ Skip (s, ref)
1192 Stop -> do
1193 runIOFinalizer ref
1194 return Stop
1196-- | When unfolding @Unfold m a b@ if an exception @e@ occurs, unfold @e@ using
1197-- @Unfold m e b@.
1199-- /Inhibits stream fusion/
1201-- /Pre-release/
1202{-# INLINE_NORMAL handle #-}
1203handle :: (MonadCatch m, Exception e)
1204 => Unfold m e b -> Unfold m a b -> Unfold m a b
1205handle exc =
1206 gbracket_ return MC.try (\_ -> return ()) (discardFirst exc)