on da36a01 help sign in
1236 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.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 -- ** From Memory
181 , fromPtr
183 -- ** From Stream
184 , fromStreamK
185 , fromStreamD
186 , fromStream
188 -- * Combinators
189 -- ** Mapping on Input
190 , lmap
191 , lmapM
192 , both
193 , first
194 , second
195 , discardFirst
196 , discardSecond
197 , swap
198 -- coapply
199 -- comonad
201 -- * Folding
202 , fold
204 -- XXX Add "WithInput" versions of all these, map2, postscan2, takeWhile2,
205 -- filter2 etc. Alternatively, we can use the default operations with
206 -- input, but that might make the common case more inconvenient.
208 -- ** Mapping on Output
209 , map
210 , map2
211 , mapM
212 , mapM2
214 , postscanlM'
215 , postscan
216 , scan
217 , scanMany
218 , foldMany
219 -- pipe
221 -- ** Either Wrapped Input
222 , either
224 -- ** Filtering
225 , takeWhileM
226 , takeWhile
227 , take
228 , filter
229 , filterM
230 , drop
231 , dropWhile
232 , dropWhileM
234 -- ** Zipping
235 , zipWithM
236 , zipWith
238 -- ** Cross product
239 , crossWithM
240 , crossWith
241 , cross
242 , apply
244 -- ** Nesting
245 , ConcatState (..)
246 , many
247 , many2
248 , concatMapM
249 , bind
251 -- ** Resource Management
252 , gbracket_
253 , gbracket
254 , before
255 , after
256 , after_
257 , finally
258 , finally_
259 , bracket
260 , bracket_
262 -- ** Exceptions
263 , onException
264 , handle
265 )
266where
268#include "ArrayMacros.h"
270import Control.Exception (Exception, mask_)
271import Control.Monad.Catch (MonadCatch)
272import Data.Functor (($>))
273import GHC.Types (SPEC(..))
274import Streamly.Internal.Control.Concurrent (MonadRunInIO, MonadAsync, withRunInIO)
275import Streamly.Internal.Data.Fold.Type (Fold(..))
276import Streamly.Internal.Data.IOFinalizer
277 (newIOFinalizer, runIOFinalizer, clearingIOFinalizer)
278import Streamly.Internal.Data.Stream.StreamD.Type (Stream(..), Step(..))
279import Streamly.Internal.Data.SVar.Type (defState)
281import qualified Control.Monad.Catch as MC
282import qualified Data.Tuple as Tuple
283import qualified Streamly.Internal.Data.Fold.Type as FL
284import qualified Streamly.Internal.Data.Stream.StreamD.Type as D
285import qualified Streamly.Internal.Data.Stream.StreamK.Type as K
286import qualified Streamly.Internal.Data.Stream.Type as Stream
288import Streamly.Internal.Data.Unfold.Enumeration
289import Streamly.Internal.Data.Unfold.Type
290import Prelude
291 hiding (map, mapM, takeWhile, take, filter, const, zipWith
292 , drop, dropWhile, either)
293import Control.Monad.IO.Class (MonadIO (liftIO))
294import Foreign (Storable, peek, sizeOf)
295import Foreign.Ptr
297-- $setup
298-- >>> import qualified Streamly.Data.Fold as Fold
299-- >>> import qualified Streamly.Internal.Data.Unfold as Unfold
300-- >>> import qualified Streamly.Data.Stream as Stream
303-- | Convert an 'Unfold' into an unfold accepting a tuple as an argument,
304-- using the argument of the original fold as the second element of tuple and
305-- discarding the first element of the tuple.
306--
307-- @
308-- discardFirst = Unfold.lmap snd
309-- @
310--
311-- /Pre-release/
312--
313{-# INLINE_NORMAL discardFirst #-}
314discardFirst :: Unfold m a b -> Unfold m (c, a) b
315discardFirst = lmap snd
317-- | Convert an 'Unfold' into an unfold accepting a tuple as an argument,
318-- using the argument of the original fold as the first element of tuple and
319-- discarding the second element of the tuple.
320--
321-- @
322-- discardSecond = Unfold.lmap fst
323-- @
324--
325-- /Pre-release/
326--
327{-# INLINE_NORMAL discardSecond #-}
328discardSecond :: Unfold m a b -> Unfold m (a, c) b
329discardSecond = lmap fst
331-- | Convert an 'Unfold' that accepts a tuple as an argument into an unfold
332-- that accepts a tuple with elements swapped.
333--
334-- @
335-- swap = Unfold.lmap Tuple.swap
336-- @
337--
338-- /Pre-release/
339--
340{-# INLINE_NORMAL swap #-}
341swap :: Unfold m (a, c) b -> Unfold m (c, a) b
342swap = lmap Tuple.swap
344-------------------------------------------------------------------------------
345-- Output operations
346-------------------------------------------------------------------------------
348-- XXX Do we need this combinator or the stream based idiom is enough?
350-- | Compose an 'Unfold' and a 'Fold'. Given an @Unfold m a b@ and a
351-- @Fold m b c@, returns a monadic action @a -> m c@ representing the
352-- application of the fold on the unfolded stream.
353--
354-- >>> Unfold.fold Fold.sum Unfold.fromList [1..100]
355-- 5050
356--
357-- >>> fold f u = Stream.fold f . Stream.unfold u
358--
359-- /Pre-release/
360--
361{-# INLINE_NORMAL fold #-}
362fold :: Monad m => Fold m b c -> Unfold m a b -> a -> m c
363fold (Fold fstep initial extract) (Unfold ustep inject) a = do
364 res <- initial
365 case res of
366 FL.Partial x -> inject a >>= go SPEC x
367 FL.Done b -> return b
369 where
371 {-# INLINE_LATE go #-}
372 go !_ !fs st = do
373 r <- ustep st
374 case r of
375 Yield x s -> do
376 res <- fstep fs x
377 case res of
378 FL.Partial fs1 -> go SPEC fs1 s
379 FL.Done c -> return c
380 Skip s -> go SPEC fs s
381 Stop -> extract fs
383-- {-# ANN type FoldMany Fuse #-}
384data FoldMany s fs b a
385 = FoldManyStart s
386 | FoldManyFirst fs s
387 | FoldManyLoop s fs
388 | FoldManyYield b (FoldMany s fs b a)
389 | FoldManyDone
391-- | Apply a fold multiple times on the output of an unfold.
392--
393-- /Pre-release/
394{-# INLINE_NORMAL foldMany #-}
395foldMany :: Monad m => Fold m b c -> Unfold m a b -> Unfold m a c
396foldMany (Fold fstep initial extract) (Unfold ustep inject1) =
397 Unfold step inject
399 where
401 inject x = do
402 r <- inject1 x
403 return (FoldManyStart r)
405 {-# INLINE consume #-}
406 consume x s fs = do
407 res <- fstep fs x
408 return
409 $ Skip
410 $ case res of
411 FL.Done b -> FoldManyYield b (FoldManyStart s)
412 FL.Partial ps -> FoldManyLoop s ps
414 {-# INLINE_LATE step #-}
415 step (FoldManyStart st) = do
416 r <- initial
417 return
418 $ Skip
419 $ case r of
420 FL.Done b -> FoldManyYield b (FoldManyStart st)
421 FL.Partial fs -> FoldManyFirst fs st
422 step (FoldManyFirst fs st) = do
423 r <- ustep st
424 case r of
425 Yield x s -> consume x s fs
426 Skip s -> return $ Skip (FoldManyFirst fs s)
427 Stop -> return Stop
428 step (FoldManyLoop st fs) = do
429 r <- ustep st
430 case r of
431 Yield x s -> consume x s fs
432 Skip s -> return $ Skip (FoldManyLoop s fs)
433 Stop -> do
434 b <- extract fs
435 return $ Skip (FoldManyYield b FoldManyDone)
436 step (FoldManyYield b next) = return $ Yield b next
437 step FoldManyDone = return Stop
439-------------------------------------------------------------------------------
440-- Either
441-------------------------------------------------------------------------------
443-- | Make an unfold operate on values wrapped in an @Either a a@ type. 'Right
444-- a' translates to 'Right b' and 'Left a' translates to 'Left b'.
445--
446-- /Internal/
447{-# INLINE_NORMAL either #-}
448either :: Applicative m => Unfold m a b -> Unfold m (Either a a) (Either b b)
449either (Unfold step1 inject1) = Unfold step inject
451 where
453 inject (Left a) = (, Left) <$> inject1 a
454 inject (Right a) = (, Right) <$> inject1 a
456 {-# INLINE_LATE step #-}
457 step (st, f) = do
458 (\case
459 Yield x s -> Yield (f x) (s, f)
460 Skip s -> Skip (s, f)
461 Stop -> Stop) <$> step1 st
463-- postscan2 :: Monad m => Refold m a b c -> Unfold m a b -> Unfold m a c
465-- | Scan the output of an 'Unfold' to change it in a stateful manner.
466--
467-- /Pre-release/
468{-# INLINE_NORMAL postscan #-}
469postscan :: Monad m => Fold m b c -> Unfold m a b -> Unfold m a c
470postscan (Fold stepF initial extract) (Unfold stepU injectU) =
471 Unfold step inject
473 where
475 inject a = do
476 r <- initial
477 case r of
478 FL.Partial fs -> Just . (fs,) <$> injectU a
479 FL.Done _ -> return Nothing
481 {-# INLINE_LATE step #-}
482 step (Just (fs, us)) = do
483 ru <- stepU us
484 case ru of
485 Yield x s -> do
486 rf <- stepF fs x
487 case rf of
488 FL.Done v -> return $ Yield v Nothing
489 FL.Partial fs1 -> do
490 v <- extract fs1
491 return $ Yield v (Just (fs1, s))
492 Skip s -> return $ Skip (Just (fs, s))
493 Stop -> return Stop
495 step Nothing = return Stop
497data ScanState s f = ScanInit s | ScanDo s !f | ScanDone
499{-# INLINE_NORMAL scanWith #-}
500scanWith :: Monad m => Bool -> Fold m b c -> Unfold m a b -> Unfold m a c
501scanWith restart (Fold fstep initial extract) (Unfold stepU injectU) =
502 Unfold step inject
504 where
506 inject a = ScanInit <$> injectU a
508 {-# INLINE runStep #-}
509 runStep us action = do
510 r <- action
511 case r of
512 FL.Partial fs -> do
513 !b <- extract fs
514 return $ Yield b (ScanDo us fs)
515 FL.Done b ->
516 let next = if restart then ScanInit us else ScanDone
517 in return $ Yield b next
519 {-# INLINE_LATE step #-}
520 step (ScanInit us) = runStep us initial
521 step (ScanDo us fs) = do
522 res <- stepU us
523 case res of
524 Yield x s -> runStep s (fstep fs x)
525 Skip s -> return $ Skip $ ScanDo s fs
526 Stop -> return Stop
527 step ScanDone = return Stop
529-- | Scan the output of an 'Unfold' to change it in a stateful manner.
530-- Once fold is done it will restart from its initial state.
531--
532-- >>> u = Unfold.scanMany (Fold.take 2 Fold.sum) Unfold.fromList
533-- >>> Unfold.fold Fold.toList u [1,2,3,4,5]
534-- [0,1,3,0,3,7,0,5]
535--
536-- /Pre-release/
537{-# INLINE_NORMAL scanMany #-}
538scanMany :: Monad m => Fold m b c -> Unfold m a b -> Unfold m a c
539scanMany = scanWith True
541-- scan2 :: Monad m => Refold m a b c -> Unfold m a b -> Unfold m a c
543-- | Scan the output of an 'Unfold' to change it in a stateful manner.
544-- Once fold is done it will stop.
545--
546-- >>> u = Unfold.scan (Fold.take 2 Fold.sum) Unfold.fromList
547-- >>> Unfold.fold Fold.toList u [1,2,3,4,5]
548-- [0,1,3]
549--
550-- /Pre-release/
551{-# INLINE_NORMAL scan #-}
552scan :: Monad m => Fold m b c -> Unfold m a b -> Unfold m a c
553scan = scanWith False
555-- | Scan the output of an 'Unfold' to change it in a stateful manner.
556--
557-- /Pre-release/
558{-# INLINE_NORMAL postscanlM' #-}
559postscanlM' :: Monad m => (b -> a -> m b) -> m b -> Unfold m c a -> Unfold m c b
560postscanlM' f z = postscan (FL.foldlM' f z)
562-------------------------------------------------------------------------------
563-- Convert streams into unfolds
564-------------------------------------------------------------------------------
566{-# INLINE_NORMAL fromStreamD #-}
567fromStreamD :: Applicative m => Unfold m (Stream m a) a
568fromStreamD = Unfold step pure
570 where
572 {-# INLINE_LATE step #-}
573 step (UnStream step1 state1) =
574 (\case
575 Yield x s -> Yield x (Stream step1 s)
576 Skip s -> Skip (Stream step1 s)
577 Stop -> Stop) <$> step1 defState state1
579{-# INLINE_NORMAL fromStreamK #-}
580fromStreamK :: Applicative m => Unfold m (K.Stream m a) a
581fromStreamK = Unfold step pure
583 where
585 {-# INLINE_LATE step #-}
586 step stream = do
587 (\case
588 Just (x, xs) -> Yield x xs
589 Nothing -> Stop) <$> K.uncons stream
591{-# INLINE fromStream #-}
592fromStream :: Applicative m => Unfold m (Stream.Stream m a) a
593fromStream = lmap Stream.toStreamK fromStreamK
595-------------------------------------------------------------------------------
596-- Unfolds
597-------------------------------------------------------------------------------
599-- | Lift a monadic function into an unfold generating a nil stream with a side
600-- effect.
601--
602{-# INLINE nilM #-}
603nilM :: Applicative m => (a -> m c) -> Unfold m a b
604nilM f = Unfold step pure
606 where
608 {-# INLINE_LATE step #-}
609 step x = f x $> Stop
611-- | Prepend a monadic single element generator function to an 'Unfold'. The
612-- same seed is used in the action as well as the unfold.
613--
614-- /Pre-release/
615{-# INLINE_NORMAL consM #-}
616consM :: Applicative m => (a -> m b) -> Unfold m a b -> Unfold m a b
617consM action unf = Unfold step inject
619 where
621 inject = pure . Left
623 {-# INLINE_LATE step #-}
624 step (Left a) = (`Yield` Right (D.unfold unf a)) <$> action a
625 step (Right (UnStream step1 st)) = do
626 (\case
627 Yield x s -> Yield x (Right (Stream step1 s))
628 Skip s -> Skip (Right (Stream step1 s))
629 Stop -> Stop) <$> step1 defState st
631-- | Convert a list of monadic values to a 'Stream'
632--
633-- /Since: 0.8.0/
634--
635{-# INLINE_LATE fromListM #-}
636fromListM :: Applicative m => Unfold m [m a] a
637fromListM = Unfold step pure
639 where
641 {-# INLINE_LATE step #-}
642 step (x:xs) = (`Yield` xs) <$> x
643 step [] = pure Stop
645{-# INLINE fromPtr #-}
646fromPtr :: forall m a. (MonadIO m, Storable a) => Unfold m (Ptr a) a
647fromPtr = Unfold step return
649 where
651 {-# INLINE_LATE step #-}
652 step p = do
653 x <- liftIO $ peek p
654 return $ Yield x (PTR_NEXT(p, a))
656------------------------------------------------------------------------------
657-- Specialized Generation
658------------------------------------------------------------------------------
660-- | Given a seed @(n, action)@, generates a stream replicating the @action@ @n@
661-- times.
662--
663-- /Since: 0.9.0/
664--
665{-# INLINE replicateM #-}
666replicateM :: Applicative m => Unfold m (Int, m a) a
667replicateM = Unfold step inject
669 where
671 inject seed = pure seed
673 {-# INLINE_LATE step #-}
674 step (i, action) =
675 if i <= 0
676 then pure Stop
677 else (\x -> Yield x (i - 1, action)) <$> action
679-- | Generates an infinite stream repeating the seed.
680--
681-- /Since: 0.8.0/
682--
683{-# INLINE repeatM #-}
684repeatM :: Applicative m => Unfold m (m a) a
685repeatM = Unfold step pure
687 where
689 {-# INLINE_LATE step #-}
690 step action = (`Yield` action) <$> action
692-- | Generates an infinite stream starting with the given seed and applying the
693-- given function repeatedly.
694--
695-- /Since: 0.8.0/
696--
697{-# INLINE iterateM #-}
698iterateM :: Applicative m => (a -> m a) -> Unfold m (m a) a
699iterateM f = Unfold step id
701 where
703 {-# INLINE_LATE step #-}
704 step x = Yield x <$> f x
706-- | @fromIndicesM gen@ generates an infinite stream of values using @gen@
707-- starting from the seed.
708--
709-- @
710-- fromIndicesM f = Unfold.mapM f $ Unfold.enumerateFrom 0
711-- @
712--
713-- /Pre-release/
714--
715{-# INLINE_NORMAL fromIndicesM #-}
716fromIndicesM :: Applicative m => (Int -> m a) -> Unfold m Int a
717fromIndicesM gen = Unfold step pure
719 where
721 {-# INLINE_LATE step #-}
722 step i = (`Yield` (i + 1)) <$> gen i
724-------------------------------------------------------------------------------
725-- Filtering
726-------------------------------------------------------------------------------
728-- |
729-- >>> u = Unfold.take 2 Unfold.fromList
730-- >>> Unfold.fold Fold.toList u [1..100]
731-- [1,2]
732--
733-- /Since: 0.8.0/
734--
735{-# INLINE_NORMAL take #-}
736take :: Applicative m => Int -> Unfold m a b -> Unfold m a b
737take n (Unfold step1 inject1) = Unfold step inject
739 where
741 inject x = (, 0) <$> inject1 x
743 {-# INLINE_LATE step #-}
744 step (st, i) | i < n = do
745 (\case
746 Yield x s -> Yield x (s, i + 1)
747 Skip s -> Skip (s, i)
748 Stop -> Stop) <$> step1 st
749 step (_, _) = pure Stop
751-- | Same as 'filter' but with a monadic predicate.
752--
753-- /Since: 0.8.0/
754--
755{-# INLINE_NORMAL filterM #-}
756filterM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
757filterM f (Unfold step1 inject1) = Unfold step inject1
758 where
759 {-# INLINE_LATE step #-}
760 step st = do
761 r <- step1 st
762 case r of
763 Yield x s -> do
764 b <- f x
765 return $ if b then Yield x s else Skip s
766 Skip s -> return $ Skip s
767 Stop -> return Stop
769-- | Include only those elements that pass a predicate.
770--
771-- /Since: 0.8.0/
772--
773{-# INLINE filter #-}
774filter :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
775filter f = filterM (return . f)
777-- | @drop n unf@ drops @n@ elements from the stream generated by @unf@.
778--
779-- /Since: 0.8.0/
780--
781{-# INLINE_NORMAL drop #-}
782drop :: Applicative m => Int -> Unfold m a b -> Unfold m a b
783drop n (Unfold step inject) = Unfold step' inject'
785 where
787 inject' a = (, n) <$> inject a
789 {-# INLINE_LATE step' #-}
790 step' (st, i)
791 | i > 0 = do
792 (\case
793 Yield _ s -> Skip (s, i - 1)
794 Skip s -> Skip (s, i)
795 Stop -> Stop) <$> step st
796 | otherwise = do
797 (\case
798 Yield x s -> Yield x (s, 0)
799 Skip s -> Skip (s, 0)
800 Stop -> Stop) <$> step st
802-- | @dropWhileM f unf@ drops elements from the stream generated by @unf@ while
803-- the condition holds true. The condition function @f@ is /monadic/ in nature.
804--
805-- /Since: 0.8.0/
806--
807{-# INLINE_NORMAL dropWhileM #-}
808dropWhileM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
809dropWhileM f (Unfold step inject) = Unfold step' inject'
811 where
813 inject' a = do
814 b <- inject a
815 return $ Left b
817 {-# INLINE_LATE step' #-}
818 step' (Left st) = do
819 r <- step st
820 case r of
821 Yield x s -> do
822 b <- f x
823 return
824 $ if b
825 then Skip (Left s)
826 else Yield x (Right s)
827 Skip s -> return $ Skip (Left s)
828 Stop -> return Stop
829 step' (Right st) = do
830 r <- step st
831 return
832 $ case r of
833 Yield x s -> Yield x (Right s)
834 Skip s -> Skip (Right s)
835 Stop -> Stop
837-- | Similar to 'dropWhileM' but with a pure condition function.
838--
839-- /Since: 0.8.0/
840--
841{-# INLINE dropWhile #-}
842dropWhile :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
843dropWhile f = dropWhileM (return . f)
845------------------------------------------------------------------------------
846-- Exceptions
847------------------------------------------------------------------------------
849-- | Like 'gbracket' but with following differences:
850--
851-- * alloc action @a -> m c@ runs with async exceptions enabled
852-- * cleanup action @c -> m d@ won't run if the stream is garbage collected
853-- after partial evaluation.
854-- * does not require a 'MonadAsync' constraint.
855--
856-- /Inhibits stream fusion/
857--
858-- /Pre-release/
859--
860{-# INLINE_NORMAL gbracket_ #-}
861gbracket_
862 :: Monad m
863 => (a -> m c) -- ^ before
864 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
865 -> (c -> m d) -- ^ after, on normal stop
866 -> Unfold m (c, e) b -- ^ on exception
867 -> Unfold m c b -- ^ unfold to run
868 -> Unfold m a b
869gbracket_ bef exc aft (Unfold estep einject) (Unfold step1 inject1) =
870 Unfold step inject
872 where
874 inject x = do
875 r <- bef x
876 s <- inject1 r
877 return $ Right (s, r)
879 {-# INLINE_LATE step #-}
880 step (Right (st, v)) = do
881 res <- exc $ step1 st
882 case res of
883 Right r -> case r of
884 Yield x s -> return $ Yield x (Right (s, v))
885 Skip s -> return $ Skip (Right (s, v))
886 Stop -> aft v >> return Stop
887 -- XXX Do not handle async exceptions, just rethrow them.
888 Left e -> do
889 r <- einject (v, e)
890 return $ Skip (Left r)
891 step (Left st) = do
892 res <- estep st
893 return $ case res of
894 Yield x s -> Yield x (Left s)
895 Skip s -> Skip (Left s)
896 Stop -> Stop
898-- | Run the alloc action @a -> m c@ with async exceptions disabled but keeping
899-- blocking operations interruptible (see 'Control.Exception.mask'). Use the
900-- output @c@ as input to @Unfold m c b@ to generate an output stream. When
901-- unfolding use the supplied @try@ operation @forall s. m s -> m (Either e s)@
902-- to catch synchronous exceptions. If an exception occurs run the exception
903-- handling unfold @Unfold m (c, e) b@.
904--
905-- The cleanup action @c -> m d@, runs whenever the stream ends normally, due
906-- to a sync or async exception or if it gets garbage collected after a partial
907-- lazy evaluation. See 'bracket' for the semantics of the cleanup action.
908--
909-- 'gbracket' can express all other exception handling combinators.
910--
911-- /Inhibits stream fusion/
912--
913-- /Pre-release/
914{-# INLINE_NORMAL gbracket #-}
915gbracket
916 :: MonadRunInIO m
917 => (a -> m c) -- ^ before
918 -> (c -> m d) -- ^ after, on normal stop, or GC
919 -> Unfold m (c, e) b -- ^ on exception
920 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
921 -> Unfold m c b -- ^ unfold to run
922 -> Unfold m a b
923gbracket bef aft (Unfold estep einject) ftry (Unfold step1 inject1) =
924 Unfold step inject
926 where
928 inject x = do
929 -- Mask asynchronous exceptions to make the execution of 'bef' and
930 -- the registration of 'aft' atomic. See comment in 'D.gbracketIO'.
931 (r, ref) <- withRunInIO $ \run -> mask_ $ run $ do
932 r <- bef x
933 ref <- newIOFinalizer (aft r)
934 return (r, ref)
935 s <- inject1 r
936 return $ Right (s, r, ref)
938 {-# INLINE_LATE step #-}
939 step (Right (st, v, ref)) = do
940 res <- ftry $ step1 st
941 case res of
942 Right r -> case r of
943 Yield x s -> return $ Yield x (Right (s, v, ref))
944 Skip s -> return $ Skip (Right (s, v, ref))
945 Stop -> do
946 runIOFinalizer ref
947 return Stop
948 -- XXX Do not handle async exceptions, just rethrow them.
949 Left e -> do
950 -- Clearing of finalizer and running of exception handler must
951 -- be atomic wrt async exceptions. Otherwise if we have cleared
952 -- the finalizer and have not run the exception handler then we
953 -- may leak the resource.
954 r <- clearingIOFinalizer ref (einject (v, e))
955 return $ Skip (Left r)
956 step (Left st) = do
957 res <- estep st
958 return $ case res of
959 Yield x s -> Yield x (Left s)
960 Skip s -> Skip (Left s)
961 Stop -> Stop
963-- | Run a side effect @a -> m c@ on the input @a@ before unfolding it using
964-- @Unfold m a b@.
965--
966-- > before f = lmapM (\a -> f a >> return a)
967--
968-- /Pre-release/
969{-# INLINE_NORMAL before #-}
970before :: (a -> m c) -> Unfold m a b -> Unfold m a b
971before action (Unfold step inject) = Unfold step (action >> inject)
973-- The custom implementation of "after_" is slightly faster (5-7%) than
974-- "_after". This is just to document and make sure that we can always use
975-- gbracket to implement after_ The same applies to other combinators as well.
976--
977{-# INLINE_NORMAL _after #-}
978_after :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
979_after aft = gbracket_ return (fmap Right) aft undefined
981-- | Like 'after' with following differences:
982--
983-- * action @a -> m c@ won't run if the stream is garbage collected
984-- after partial evaluation.
985-- * Monad @m@ does not require any other constraints.
986--
987-- /Pre-release/
988{-# INLINE_NORMAL after_ #-}
989after_ :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
990after_ action (Unfold step1 inject1) = Unfold step inject
992 where
994 inject x = do
995 s <- inject1 x
996 return (s, x)
998 {-# INLINE_LATE step #-}
999 step (st, v) = do
1000 res <- step1 st
1001 case res of
1002 Yield x s -> return $ Yield x (s, v)
1003 Skip s -> return $ Skip (s, v)
1004 Stop -> action v >> return Stop
1006-- | Unfold the input @a@ using @Unfold m a b@, run an action on @a@ whenever
1007-- the unfold stops normally, or if it is garbage collected after a partial
1008-- lazy evaluation.
1010-- The semantics of the action @a -> m c@ are similar to the cleanup action
1011-- semantics in 'bracket'.
1013-- /See also 'after_'/
1015-- /Pre-release/
1016{-# INLINE_NORMAL after #-}
1017after :: MonadRunInIO m
1018 => (a -> m c) -> Unfold m a b -> Unfold m a b
1019after action (Unfold step1 inject1) = Unfold step inject
1021 where
1023 inject x = do
1024 s <- inject1 x
1025 ref <- newIOFinalizer (action x)
1026 return (s, ref)
1028 {-# INLINE_LATE step #-}
1029 step (st, ref) = do
1030 res <- step1 st
1031 case res of
1032 Yield x s -> return $ Yield x (s, ref)
1033 Skip s -> return $ Skip (s, ref)
1034 Stop -> do
1035 runIOFinalizer ref
1036 return Stop
1038{-# INLINE_NORMAL _onException #-}
1039_onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
1040_onException action =
1041 gbracket_ return MC.try
1042 (\_ -> return ())
1043 (nilM (\(a, e :: MC.SomeException) -> action a >> MC.throwM e))
1045-- | Unfold the input @a@ using @Unfold m a b@, run the action @a -> m c@ on
1046-- @a@ if the unfold aborts due to an exception.
1048-- /Inhibits stream fusion/
1050-- /Pre-release/
1051{-# INLINE_NORMAL onException #-}
1052onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
1053onException action (Unfold step1 inject1) = Unfold step inject
1055 where
1057 inject x = do
1058 s <- inject1 x
1059 return (s, x)
1061 {-# INLINE_LATE step #-}
1062 step (st, v) = do
1063 res <- step1 st `MC.onException` action v
1064 return $ case res of
1065 Yield x s -> Yield x (s, v)
1066 Skip s -> Skip (s, v)
1067 Stop -> Stop
1069{-# INLINE_NORMAL _finally #-}
1070_finally :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
1071_finally action =
1072 gbracket_ return MC.try action
1073 (nilM (\(a, e :: MC.SomeException) -> action a >> MC.throwM e))
1075-- | Like 'finally' with following differences:
1077-- * action @a -> m c@ won't run if the stream is garbage collected
1078-- after partial evaluation.
1079-- * does not require a 'MonadAsync' constraint.
1081-- /Inhibits stream fusion/
1083-- /Pre-release/
1084{-# INLINE_NORMAL finally_ #-}
1085finally_ :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
1086finally_ action (Unfold step1 inject1) = Unfold step inject
1088 where
1090 inject x = do
1091 s <- inject1 x
1092 return (s, x)
1094 {-# INLINE_LATE step #-}
1095 step (st, v) = do
1096 res <- step1 st `MC.onException` action v
1097 case res of
1098 Yield x s -> return $ Yield x (s, v)
1099 Skip s -> return $ Skip (s, v)
1100 Stop -> action v >> return Stop
1102-- | Unfold the input @a@ using @Unfold m a b@, run an action on @a@ whenever
1103-- the unfold stops normally, aborts due to an exception or if it is garbage
1104-- collected after a partial lazy evaluation.
1106-- The semantics of the action @a -> m c@ are similar to the cleanup action
1107-- semantics in 'bracket'.
1109-- @
1110-- finally release = bracket return release
1111-- @
1113-- /See also 'finally_'/
1115-- /Inhibits stream fusion/
1117-- /Pre-release/
1118{-# INLINE_NORMAL finally #-}
1119finally :: (MonadAsync m, MonadCatch m)
1120 => (a -> m c) -> Unfold m a b -> Unfold m a b
1121finally action (Unfold step1 inject1) = Unfold step inject
1123 where
1125 inject x = do
1126 s <- inject1 x
1127 ref <- newIOFinalizer (action x)
1128 return (s, ref)
1130 {-# INLINE_LATE step #-}
1131 step (st, ref) = do
1132 res <- step1 st `MC.onException` runIOFinalizer ref
1133 case res of
1134 Yield x s -> return $ Yield x (s, ref)
1135 Skip s -> return $ Skip (s, ref)
1136 Stop -> do
1137 runIOFinalizer ref
1138 return Stop
1140{-# INLINE_NORMAL _bracket #-}
1141_bracket :: MonadCatch m
1142 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1143_bracket bef aft =
1144 gbracket_ bef MC.try aft (nilM (\(a, e :: MC.SomeException) -> aft a >>
1145 MC.throwM e))
1147-- | Like 'bracket' but with following differences:
1149-- * alloc action @a -> m c@ runs with async exceptions enabled
1150-- * cleanup action @c -> m d@ won't run if the stream is garbage collected
1151-- after partial evaluation.
1152-- * does not require a 'MonadAsync' constraint.
1154-- /Inhibits stream fusion/
1156-- /Pre-release/
1157{-# INLINE_NORMAL bracket_ #-}
1158bracket_ :: MonadCatch m
1159 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1160bracket_ bef aft (Unfold step1 inject1) = Unfold step inject
1162 where
1164 inject x = do
1165 r <- bef x
1166 s <- inject1 r
1167 return (s, r)
1169 {-# INLINE_LATE step #-}
1170 step (st, v) = do
1171 res <- step1 st `MC.onException` aft v
1172 case res of
1173 Yield x s -> return $ Yield x (s, v)
1174 Skip s -> return $ Skip (s, v)
1175 Stop -> aft v >> return Stop
1177-- | Run the alloc action @a -> m c@ with async exceptions disabled but keeping
1178-- blocking operations interruptible (see 'Control.Exception.mask'). Use the
1179-- output @c@ as input to @Unfold m c b@ to generate an output stream.
1181-- @c@ is usually a resource under the state of monad @m@, e.g. a file
1182-- handle, that requires a cleanup after use. The cleanup action @c -> m d@,
1183-- runs whenever the stream ends normally, due to a sync or async exception or
1184-- if it gets garbage collected after a partial lazy evaluation.
1186-- 'bracket' only guarantees that the cleanup action runs, and it runs with
1187-- async exceptions enabled. The action must ensure that it can successfully
1188-- cleanup the resource in the face of sync or async exceptions.
1190-- When the stream ends normally or on a sync exception, cleanup action runs
1191-- immediately in the current thread context, whereas in other cases it runs in
1192-- the GC context, therefore, cleanup may be delayed until the GC gets to run.
1194-- /See also: 'bracket_', 'gbracket'/
1196-- /Inhibits stream fusion/
1198-- /Pre-release/
1199{-# INLINE_NORMAL bracket #-}
1200bracket :: (MonadAsync m, MonadCatch m)
1201 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1202bracket bef aft (Unfold step1 inject1) = Unfold step inject
1204 where
1206 inject x = do
1207 -- Mask asynchronous exceptions to make the execution of 'bef' and
1208 -- the registration of 'aft' atomic. See comment in 'D.gbracketIO'.
1209 (r, ref) <- withRunInIO $ \run -> mask_ $ run $ do
1210 r <- bef x
1211 ref <- newIOFinalizer (aft r)
1212 return (r, ref)
1213 s <- inject1 r
1214 return (s, ref)
1216 {-# INLINE_LATE step #-}
1217 step (st, ref) = do
1218 res <- step1 st `MC.onException` runIOFinalizer ref
1219 case res of
1220 Yield x s -> return $ Yield x (s, ref)
1221 Skip s -> return $ Skip (s, ref)
1222 Stop -> do
1223 runIOFinalizer ref
1224 return Stop
1226-- | When unfolding @Unfold m a b@ if an exception @e@ occurs, unfold @e@ using
1227-- @Unfold m e b@.
1229-- /Inhibits stream fusion/
1231-- /Pre-release/
1232{-# INLINE_NORMAL handle #-}
1233handle :: (MonadCatch m, Exception e)
1234 => Unfold m e b -> Unfold m a b -> Unfold m a b
1235handle exc =
1236 gbracket_ return MC.try (\_ -> return ()) (discardFirst exc)