on 43dda21 help sign in
1265 lines 39 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
205 , postscanlM'
206 , postscan
207 , scan
208 , scanMany
209 , foldMany
210 -- pipe
212 -- ** Either Wrapped Input
213 , either
215 -- ** Filtering
216 , takeWhileM
217 , takeWhile
218 , take
219 , filter
220 , filterM
221 , drop
222 , dropWhile
223 , dropWhileM
225 -- ** Zipping
226 , zipWithM
227 , zipWith
229 -- ** Cross product
230 , crossWithM
231 , crossWith
232 , cross
233 , apply
235 -- ** Nesting
236 , ConcatState (..)
237 , many
238 , concatMapM
239 , bind
241 -- ** Resource Management
242 , gbracket_
243 , gbracket
244 , before
245 , after
246 , after_
247 , finally
248 , finally_
249 , bracket
250 , bracket_
252 -- ** Exceptions
253 , onException
254 , handle
255 )
256where
258import Control.Exception (Exception, mask_)
259import Control.Monad.Catch (MonadCatch)
260import Data.Functor (($>))
261import GHC.Types (SPEC(..))
262import Streamly.Internal.Control.Concurrent (MonadRunInIO, MonadAsync, withRunInIO)
263import Streamly.Internal.Data.Fold.Type (Fold(..))
264import Streamly.Internal.Data.IOFinalizer
265 (newIOFinalizer, runIOFinalizer, clearingIOFinalizer)
266import Streamly.Internal.Data.Stream.Serial (SerialT(..))
267import Streamly.Internal.Data.Stream.StreamD.Type (Stream(..), Step(..))
268import Streamly.Internal.Data.SVar.Type (defState)
270import qualified Control.Monad.Catch as MC
271import qualified Data.Tuple as Tuple
272import qualified Streamly.Internal.Data.Fold.Type as FL
273import qualified Streamly.Internal.Data.Stream.StreamD.Type as D
274import qualified Streamly.Internal.Data.Stream.StreamK.Type as K
276import Streamly.Internal.Data.Unfold.Enumeration
277import Streamly.Internal.Data.Unfold.Type
278import Prelude
279 hiding (map, mapM, takeWhile, take, filter, const, zipWith
280 , drop, dropWhile, either)
282-- $setup
283-- >>> import qualified Streamly.Data.Fold as Fold
284-- >>> import qualified Streamly.Internal.Data.Unfold as Unfold
285-- >>> import qualified Streamly.Prelude as Stream
288-- | Convert an 'Unfold' into an unfold accepting a tuple as an argument,
289-- using the argument of the original fold as the second element of tuple and
290-- discarding the first element of the tuple.
291--
292-- @
293-- discardFirst = Unfold.lmap snd
294-- @
295--
296-- /Pre-release/
297--
298{-# INLINE_NORMAL discardFirst #-}
299discardFirst :: Unfold m a b -> Unfold m (c, a) b
300discardFirst = lmap snd
302-- | Convert an 'Unfold' into an unfold accepting a tuple as an argument,
303-- using the argument of the original fold as the first element of tuple and
304-- discarding the second element of the tuple.
305--
306-- @
307-- discardSecond = Unfold.lmap fst
308-- @
309--
310-- /Pre-release/
311--
312{-# INLINE_NORMAL discardSecond #-}
313discardSecond :: Unfold m a b -> Unfold m (a, c) b
314discardSecond = lmap fst
316-- | Convert an 'Unfold' that accepts a tuple as an argument into an unfold
317-- that accepts a tuple with elements swapped.
318--
319-- @
320-- swap = Unfold.lmap Tuple.swap
321-- @
322--
323-- /Pre-release/
324--
325{-# INLINE_NORMAL swap #-}
326swap :: Unfold m (a, c) b -> Unfold m (c, a) b
327swap = lmap Tuple.swap
329-------------------------------------------------------------------------------
330-- Output operations
331-------------------------------------------------------------------------------
333-- XXX Do we need this combinator or the stream based idiom is enough?
335-- | Compose an 'Unfold' and a 'Fold'. Given an @Unfold m a b@ and a
336-- @Fold m b c@, returns a monadic action @a -> m c@ representing the
337-- application of the fold on the unfolded stream.
338--
339-- >>> Unfold.fold Fold.sum Unfold.fromList [1..100]
340-- 5050
341--
342-- >>> fold f u = Stream.fold f . Stream.unfold u
343--
344-- /Pre-release/
345--
346{-# INLINE_NORMAL fold #-}
347fold :: Monad m => Fold m b c -> Unfold m a b -> a -> m c
348fold (Fold fstep initial extract) (Unfold ustep inject) a = do
349 res <- initial
350 case res of
351 FL.Partial x -> inject a >>= go SPEC x
352 FL.Done b -> return b
354 where
356 {-# INLINE_LATE go #-}
357 go !_ !fs st = do
358 r <- ustep st
359 case r of
360 Yield x s -> do
361 res <- fstep fs x
362 case res of
363 FL.Partial fs1 -> go SPEC fs1 s
364 FL.Done c -> return c
365 Skip s -> go SPEC fs s
366 Stop -> extract fs
368-- {-# ANN type FoldMany Fuse #-}
369data FoldMany s fs b a
370 = FoldManyStart s
371 | FoldManyFirst fs s
372 | FoldManyLoop s fs
373 | FoldManyYield b (FoldMany s fs b a)
374 | FoldManyDone
376-- | Apply a fold multiple times on the output of an unfold.
377--
378-- /Pre-release/
379{-# INLINE_NORMAL foldMany #-}
380foldMany :: Monad m => Fold m b c -> Unfold m a b -> Unfold m a c
381foldMany (Fold fstep initial extract) (Unfold ustep inject1) =
382 Unfold step inject
384 where
386 inject x = do
387 r <- inject1 x
388 return (FoldManyStart r)
390 {-# INLINE consume #-}
391 consume x s fs = do
392 res <- fstep fs x
393 return
394 $ Skip
395 $ case res of
396 FL.Done b -> FoldManyYield b (FoldManyStart s)
397 FL.Partial ps -> FoldManyLoop s ps
399 {-# INLINE_LATE step #-}
400 step (FoldManyStart st) = do
401 r <- initial
402 return
403 $ Skip
404 $ case r of
405 FL.Done b -> FoldManyYield b (FoldManyStart st)
406 FL.Partial fs -> FoldManyFirst fs st
407 step (FoldManyFirst fs st) = do
408 r <- ustep st
409 case r of
410 Yield x s -> consume x s fs
411 Skip s -> return $ Skip (FoldManyFirst fs s)
412 Stop -> return Stop
413 step (FoldManyLoop st fs) = do
414 r <- ustep st
415 case r of
416 Yield x s -> consume x s fs
417 Skip s -> return $ Skip (FoldManyLoop s fs)
418 Stop -> do
419 b <- extract fs
420 return $ Skip (FoldManyYield b FoldManyDone)
421 step (FoldManyYield b next) = return $ Yield b next
422 step FoldManyDone = return Stop
424-- | Apply a monadic function to each element of the stream and replace it
425-- with the output of the resulting action.
426--
427-- /Since: 0.8.0/
428--
429{-# INLINE_NORMAL mapM #-}
430mapM :: Monad m => (b -> m c) -> Unfold m a b -> Unfold m a c
431mapM f (Unfold ustep uinject) = Unfold step uinject
432 where
433 {-# INLINE_LATE step #-}
434 step st = do
435 r <- ustep st
436 case r of
437 Yield x s -> f x >>= \a -> return $ Yield a s
438 Skip s -> return $ Skip s
439 Stop -> return Stop
441{-# INLINE_NORMAL mapMWithInput #-}
442mapMWithInput :: Monad m => (a -> b -> m c) -> Unfold m a b -> Unfold m a c
443mapMWithInput f (Unfold ustep uinject) = Unfold step inject
444 where
445 inject a = do
446 r <- uinject a
447 return (a, r)
449 {-# INLINE_LATE step #-}
450 step (inp, st) = do
451 r <- ustep st
452 case r of
453 Yield x s -> f inp x >>= \a -> return $ Yield a (inp, s)
454 Skip s -> return $ Skip (inp, s)
455 Stop -> return Stop
457-------------------------------------------------------------------------------
458-- Either
459-------------------------------------------------------------------------------
461-- | Make an unfold operate on values wrapped in an @Either a a@ type. 'Right
462-- a' translates to 'Right b' and 'Left a' translates to 'Left b'.
463--
464-- /Internal/
465{-# INLINE_NORMAL either #-}
466either :: Applicative m => Unfold m a b -> Unfold m (Either a a) (Either b b)
467either (Unfold step1 inject1) = Unfold step inject
469 where
471 inject (Left a) = (, Left) <$> inject1 a
472 inject (Right a) = (, Right) <$> inject1 a
474 {-# INLINE_LATE step #-}
475 step (st, f) = do
476 (\case
477 Yield x s -> Yield (f x) (s, f)
478 Skip s -> Skip (s, f)
479 Stop -> Stop) <$> step1 st
481-- | Scan the output of an 'Unfold' to change it in a stateful manner.
482--
483-- /Pre-release/
484{-# INLINE_NORMAL postscan #-}
485postscan :: Monad m => Fold m b c -> Unfold m a b -> Unfold m a c
486postscan (Fold stepF initial extract) (Unfold stepU injectU) =
487 Unfold step inject
489 where
491 inject a = do
492 r <- initial
493 case r of
494 FL.Partial fs -> Just . (fs,) <$> injectU a
495 FL.Done _ -> return Nothing
497 {-# INLINE_LATE step #-}
498 step (Just (fs, us)) = do
499 ru <- stepU us
500 case ru of
501 Yield x s -> do
502 rf <- stepF fs x
503 case rf of
504 FL.Done v -> return $ Yield v Nothing
505 FL.Partial fs1 -> do
506 v <- extract fs1
507 return $ Yield v (Just (fs1, s))
508 Skip s -> return $ Skip (Just (fs, s))
509 Stop -> return Stop
511 step Nothing = return Stop
513data ScanState s f = ScanInit s | ScanDo s !f | ScanDone
515{-# INLINE_NORMAL scanWith #-}
516scanWith :: Monad m => Bool -> Fold m b c -> Unfold m a b -> Unfold m a c
517scanWith restart (Fold fstep initial extract) (Unfold stepU injectU) =
518 Unfold step inject
520 where
522 inject a = ScanInit <$> injectU a
524 {-# INLINE runStep #-}
525 runStep us action = do
526 r <- action
527 case r of
528 FL.Partial fs -> do
529 !b <- extract fs
530 return $ Yield b (ScanDo us fs)
531 FL.Done b ->
532 let next = if restart then ScanInit us else ScanDone
533 in return $ Yield b next
535 {-# INLINE_LATE step #-}
536 step (ScanInit us) = runStep us initial
537 step (ScanDo us fs) = do
538 res <- stepU us
539 case res of
540 Yield x s -> runStep s (fstep fs x)
541 Skip s -> return $ Skip $ ScanDo s fs
542 Stop -> return Stop
543 step ScanDone = return Stop
545-- | Scan the output of an 'Unfold' to change it in a stateful manner.
546-- Once fold is done it will restart from its initial state.
547--
548-- >>> u = Unfold.scanMany (Fold.take 2 Fold.sum) Unfold.fromList
549-- >>> Unfold.fold Fold.toList u [1,2,3,4,5]
550-- [0,1,3,0,3,7,0,5]
551--
552-- /Pre-release/
553{-# INLINE_NORMAL scanMany #-}
554scanMany :: Monad m => Fold m b c -> Unfold m a b -> Unfold m a c
555scanMany = scanWith True
557-- | Scan the output of an 'Unfold' to change it in a stateful manner.
558-- Once fold is done it will stop.
559--
560-- >>> u = Unfold.scan (Fold.take 2 Fold.sum) Unfold.fromList
561-- >>> Unfold.fold Fold.toList u [1,2,3,4,5]
562-- [0,1,3]
563--
564-- /Pre-release/
565{-# INLINE_NORMAL scan #-}
566scan :: Monad m => Fold m b c -> Unfold m a b -> Unfold m a c
567scan = scanWith False
569-- | Scan the output of an 'Unfold' to change it in a stateful manner.
570--
571-- /Pre-release/
572{-# INLINE_NORMAL postscanlM' #-}
573postscanlM' :: Monad m => (b -> a -> m b) -> m b -> Unfold m c a -> Unfold m c b
574postscanlM' f z = postscan (FL.foldlM' f z)
576-------------------------------------------------------------------------------
577-- Convert streams into unfolds
578-------------------------------------------------------------------------------
580{-# INLINE_NORMAL fromStreamD #-}
581fromStreamD :: Applicative m => Unfold m (Stream m a) a
582fromStreamD = Unfold step pure
584 where
586 {-# INLINE_LATE step #-}
587 step (UnStream step1 state1) =
588 (\case
589 Yield x s -> Yield x (Stream step1 s)
590 Skip s -> Skip (Stream step1 s)
591 Stop -> Stop) <$> step1 defState state1
593{-# INLINE_NORMAL fromStreamK #-}
594fromStreamK :: Applicative m => Unfold m (K.Stream m a) a
595fromStreamK = Unfold step pure
597 where
599 {-# INLINE_LATE step #-}
600 step stream = do
601 (\case
602 Just (x, xs) -> Yield x xs
603 Nothing -> Stop) <$> K.uncons stream
605-- XXX Using Unfold.fromStreamD seems to be faster (using cross product test
606-- case) than using fromStream even if it is implemented using fromStreamD.
607-- Check if StreamK to StreamD rewrite rules are working correctly when
608-- implementing fromStream using fromStreamD.
609--
610-- | Convert a stream into an 'Unfold'. Note that a stream converted to an
611-- 'Unfold' may not be as efficient as an 'Unfold' in some situations.
612--
613-- /Since: 0.8.0/
614--
615{-# INLINE_NORMAL fromStream #-}
616fromStream :: Applicative m => Unfold m (SerialT m a) a
617fromStream = lmap getSerialT fromStreamK
619-------------------------------------------------------------------------------
620-- Unfolds
621-------------------------------------------------------------------------------
623-- | Lift a monadic function into an unfold generating a nil stream with a side
624-- effect.
625--
626{-# INLINE nilM #-}
627nilM :: Applicative m => (a -> m c) -> Unfold m a b
628nilM f = Unfold step pure
630 where
632 {-# INLINE_LATE step #-}
633 step x = f x $> Stop
635-- | Prepend a monadic single element generator function to an 'Unfold'. The
636-- same seed is used in the action as well as the unfold.
637--
638-- /Pre-release/
639{-# INLINE_NORMAL consM #-}
640consM :: Applicative m => (a -> m b) -> Unfold m a b -> Unfold m a b
641consM action unf = Unfold step inject
643 where
645 inject = pure . Left
647 {-# INLINE_LATE step #-}
648 step (Left a) = (`Yield` Right (D.unfold unf a)) <$> action a
649 step (Right (UnStream step1 st)) = do
650 (\case
651 Yield x s -> Yield x (Right (Stream step1 s))
652 Skip s -> Skip (Right (Stream step1 s))
653 Stop -> Stop) <$> step1 defState st
655-- XXX Check if "unfold (fromList [1..10])" fuses, if it doesn't we can use
656-- rewrite rules to rewrite list enumerations to unfold enumerations.
657--
658-- | Convert a list of pure values to a 'Stream'
659--
660-- /Since: 0.8.0/
661--
662{-# INLINE_LATE fromList #-}
663fromList :: Applicative m => Unfold m [a] a
664fromList = Unfold step pure
666 where
668 {-# INLINE_LATE step #-}
669 step (x:xs) = pure $ Yield x xs
670 step [] = pure Stop
672-- | Convert a list of monadic values to a 'Stream'
673--
674-- /Since: 0.8.0/
675--
676{-# INLINE_LATE fromListM #-}
677fromListM :: Applicative m => Unfold m [m a] a
678fromListM = Unfold step pure
680 where
682 {-# INLINE_LATE step #-}
683 step (x:xs) = (`Yield` xs) <$> x
684 step [] = pure Stop
686------------------------------------------------------------------------------
687-- Specialized Generation
688------------------------------------------------------------------------------
690-- | Generates a stream replicating the seed @n@ times.
691--
692-- /Since: 0.8.0/
693--
694{-# INLINE replicateM #-}
695replicateM :: Applicative m => Int -> Unfold m (m a) a
696replicateM n = Unfold step inject
698 where
700 inject action = pure (action, n)
702 {-# INLINE_LATE step #-}
703 step (action, i) =
704 if i <= 0
705 then pure Stop
706 else (\x -> Yield x (action, i - 1)) <$> action
708-- | Generates an infinite stream repeating the seed.
709--
710-- /Since: 0.8.0/
711--
712{-# INLINE repeatM #-}
713repeatM :: Applicative m => Unfold m (m a) a
714repeatM = Unfold step pure
716 where
718 {-# INLINE_LATE step #-}
719 step action = (`Yield` action) <$> action
721-- | Generates an infinite stream starting with the given seed and applying the
722-- given function repeatedly.
723--
724-- /Since: 0.8.0/
725--
726{-# INLINE iterateM #-}
727iterateM :: Applicative m => (a -> m a) -> Unfold m (m a) a
728iterateM f = Unfold step id
730 where
732 {-# INLINE_LATE step #-}
733 step x = Yield x <$> f x
735-- | @fromIndicesM gen@ generates an infinite stream of values using @gen@
736-- starting from the seed.
737--
738-- @
739-- fromIndicesM f = Unfold.mapM f $ Unfold.enumerateFrom 0
740-- @
741--
742-- /Pre-release/
743--
744{-# INLINE_NORMAL fromIndicesM #-}
745fromIndicesM :: Applicative m => (Int -> m a) -> Unfold m Int a
746fromIndicesM gen = Unfold step pure
748 where
750 {-# INLINE_LATE step #-}
751 step i = (`Yield` (i + 1)) <$> gen i
753-------------------------------------------------------------------------------
754-- Filtering
755-------------------------------------------------------------------------------
757-- |
758-- >>> u = Unfold.take 2 Unfold.fromList
759-- >>> Unfold.fold Fold.toList u [1..100]
760-- [1,2]
761--
762-- /Since: 0.8.0/
763--
764{-# INLINE_NORMAL take #-}
765take :: Applicative m => Int -> Unfold m a b -> Unfold m a b
766take n (Unfold step1 inject1) = Unfold step inject
768 where
770 inject x = (, 0) <$> inject1 x
772 {-# INLINE_LATE step #-}
773 step (st, i) | i < n = do
774 (\case
775 Yield x s -> Yield x (s, i + 1)
776 Skip s -> Skip (s, i)
777 Stop -> Stop) <$> step1 st
778 step (_, _) = pure Stop
780-- | Same as 'filter' but with a monadic predicate.
781--
782-- /Since: 0.8.0/
783--
784{-# INLINE_NORMAL filterM #-}
785filterM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
786filterM f (Unfold step1 inject1) = Unfold step inject1
787 where
788 {-# INLINE_LATE step #-}
789 step st = do
790 r <- step1 st
791 case r of
792 Yield x s -> do
793 b <- f x
794 return $ if b then Yield x s else Skip s
795 Skip s -> return $ Skip s
796 Stop -> return Stop
798-- | Include only those elements that pass a predicate.
799--
800-- /Since: 0.8.0/
801--
802{-# INLINE filter #-}
803filter :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
804filter f = filterM (return . f)
806-- | @drop n unf@ drops @n@ elements from the stream generated by @unf@.
807--
808-- /Since: 0.8.0/
809--
810{-# INLINE_NORMAL drop #-}
811drop :: Applicative m => Int -> Unfold m a b -> Unfold m a b
812drop n (Unfold step inject) = Unfold step' inject'
814 where
816 inject' a = (, n) <$> inject a
818 {-# INLINE_LATE step' #-}
819 step' (st, i)
820 | i > 0 = do
821 (\case
822 Yield _ s -> Skip (s, i - 1)
823 Skip s -> Skip (s, i)
824 Stop -> Stop) <$> step st
825 | otherwise = do
826 (\case
827 Yield x s -> Yield x (s, 0)
828 Skip s -> Skip (s, 0)
829 Stop -> Stop) <$> step st
831-- | @dropWhileM f unf@ drops elements from the stream generated by @unf@ while
832-- the condition holds true. The condition function @f@ is /monadic/ in nature.
833--
834-- /Since: 0.8.0/
835--
836{-# INLINE_NORMAL dropWhileM #-}
837dropWhileM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
838dropWhileM f (Unfold step inject) = Unfold step' inject'
840 where
842 inject' a = do
843 b <- inject a
844 return $ Left b
846 {-# INLINE_LATE step' #-}
847 step' (Left st) = do
848 r <- step st
849 case r of
850 Yield x s -> do
851 b <- f x
852 return
853 $ if b
854 then Skip (Left s)
855 else Yield x (Right s)
856 Skip s -> return $ Skip (Left s)
857 Stop -> return Stop
858 step' (Right st) = do
859 r <- step st
860 return
861 $ case r of
862 Yield x s -> Yield x (Right s)
863 Skip s -> Skip (Right s)
864 Stop -> Stop
866-- | Similar to 'dropWhileM' but with a pure condition function.
867--
868-- /Since: 0.8.0/
869--
870{-# INLINE dropWhile #-}
871dropWhile :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
872dropWhile f = dropWhileM (return . f)
874------------------------------------------------------------------------------
875-- Exceptions
876------------------------------------------------------------------------------
878-- | Like 'gbracket' but with following differences:
879--
880-- * alloc action @a -> m c@ runs with async exceptions enabled
881-- * cleanup action @c -> m d@ won't run if the stream is garbage collected
882-- after partial evaluation.
883-- * does not require a 'MonadAsync' constraint.
884--
885-- /Inhibits stream fusion/
886--
887-- /Pre-release/
888--
889{-# INLINE_NORMAL gbracket_ #-}
890gbracket_
891 :: Monad m
892 => (a -> m c) -- ^ before
893 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
894 -> (c -> m d) -- ^ after, on normal stop
895 -> Unfold m (c, e) b -- ^ on exception
896 -> Unfold m c b -- ^ unfold to run
897 -> Unfold m a b
898gbracket_ bef exc aft (Unfold estep einject) (Unfold step1 inject1) =
899 Unfold step inject
901 where
903 inject x = do
904 r <- bef x
905 s <- inject1 r
906 return $ Right (s, r)
908 {-# INLINE_LATE step #-}
909 step (Right (st, v)) = do
910 res <- exc $ step1 st
911 case res of
912 Right r -> case r of
913 Yield x s -> return $ Yield x (Right (s, v))
914 Skip s -> return $ Skip (Right (s, v))
915 Stop -> aft v >> return Stop
916 -- XXX Do not handle async exceptions, just rethrow them.
917 Left e -> do
918 r <- einject (v, e)
919 return $ Skip (Left r)
920 step (Left st) = do
921 res <- estep st
922 return $ case res of
923 Yield x s -> Yield x (Left s)
924 Skip s -> Skip (Left s)
925 Stop -> Stop
927-- | Run the alloc action @a -> m c@ with async exceptions disabled but keeping
928-- blocking operations interruptible (see 'Control.Exception.mask'). Use the
929-- output @c@ as input to @Unfold m c b@ to generate an output stream. When
930-- unfolding use the supplied @try@ operation @forall s. m s -> m (Either e s)@
931-- to catch synchronous exceptions. If an exception occurs run the exception
932-- handling unfold @Unfold m (c, e) b@.
933--
934-- The cleanup action @c -> m d@, runs whenever the stream ends normally, due
935-- to a sync or async exception or if it gets garbage collected after a partial
936-- lazy evaluation. See 'bracket' for the semantics of the cleanup action.
937--
938-- 'gbracket' can express all other exception handling combinators.
939--
940-- /Inhibits stream fusion/
941--
942-- /Pre-release/
943{-# INLINE_NORMAL gbracket #-}
944gbracket
945 :: MonadRunInIO m
946 => (a -> m c) -- ^ before
947 -> (c -> m d) -- ^ after, on normal stop, or GC
948 -> Unfold m (c, e) b -- ^ on exception
949 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
950 -> Unfold m c b -- ^ unfold to run
951 -> Unfold m a b
952gbracket bef aft (Unfold estep einject) ftry (Unfold step1 inject1) =
953 Unfold step inject
955 where
957 inject x = do
958 -- Mask asynchronous exceptions to make the execution of 'bef' and
959 -- the registration of 'aft' atomic. See comment in 'D.gbracketIO'.
960 (r, ref) <- withRunInIO $ \run -> mask_ $ run $ do
961 r <- bef x
962 ref <- newIOFinalizer (aft r)
963 return (r, ref)
964 s <- inject1 r
965 return $ Right (s, r, ref)
967 {-# INLINE_LATE step #-}
968 step (Right (st, v, ref)) = do
969 res <- ftry $ step1 st
970 case res of
971 Right r -> case r of
972 Yield x s -> return $ Yield x (Right (s, v, ref))
973 Skip s -> return $ Skip (Right (s, v, ref))
974 Stop -> do
975 runIOFinalizer ref
976 return Stop
977 -- XXX Do not handle async exceptions, just rethrow them.
978 Left e -> do
979 -- Clearing of finalizer and running of exception handler must
980 -- be atomic wrt async exceptions. Otherwise if we have cleared
981 -- the finalizer and have not run the exception handler then we
982 -- may leak the resource.
983 r <- clearingIOFinalizer ref (einject (v, e))
984 return $ Skip (Left r)
985 step (Left st) = do
986 res <- estep st
987 return $ case res of
988 Yield x s -> Yield x (Left s)
989 Skip s -> Skip (Left s)
990 Stop -> Stop
992-- | Run a side effect @a -> m c@ on the input @a@ before unfolding it using
993-- @Unfold m a b@.
994--
995-- > before f = lmapM (\a -> f a >> return a)
996--
997-- /Pre-release/
998{-# INLINE_NORMAL before #-}
999before :: (a -> m c) -> Unfold m a b -> Unfold m a b
1000before action (Unfold step inject) = Unfold step (action >> inject)
1002-- The custom implementation of "after_" is slightly faster (5-7%) than
1003-- "_after". This is just to document and make sure that we can always use
1004-- gbracket to implement after_ The same applies to other combinators as well.
1006{-# INLINE_NORMAL _after #-}
1007_after :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
1008_after aft = gbracket_ return (fmap Right) aft undefined
1010-- | Like 'after' with following differences:
1012-- * action @a -> m c@ won't run if the stream is garbage collected
1013-- after partial evaluation.
1014-- * Monad @m@ does not require any other constraints.
1016-- /Pre-release/
1017{-# INLINE_NORMAL after_ #-}
1018after_ :: Monad m => (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 return (s, x)
1027 {-# INLINE_LATE step #-}
1028 step (st, v) = do
1029 res <- step1 st
1030 case res of
1031 Yield x s -> return $ Yield x (s, v)
1032 Skip s -> return $ Skip (s, v)
1033 Stop -> action v >> return Stop
1035-- | Unfold the input @a@ using @Unfold m a b@, run an action on @a@ whenever
1036-- the unfold stops normally, or if it is garbage collected after a partial
1037-- lazy evaluation.
1039-- The semantics of the action @a -> m c@ are similar to the cleanup action
1040-- semantics in 'bracket'.
1042-- /See also 'after_'/
1044-- /Pre-release/
1045{-# INLINE_NORMAL after #-}
1046after :: MonadRunInIO m
1047 => (a -> m c) -> Unfold m a b -> Unfold m a b
1048after action (Unfold step1 inject1) = Unfold step inject
1050 where
1052 inject x = do
1053 s <- inject1 x
1054 ref <- newIOFinalizer (action x)
1055 return (s, ref)
1057 {-# INLINE_LATE step #-}
1058 step (st, ref) = do
1059 res <- step1 st
1060 case res of
1061 Yield x s -> return $ Yield x (s, ref)
1062 Skip s -> return $ Skip (s, ref)
1063 Stop -> do
1064 runIOFinalizer ref
1065 return Stop
1067{-# INLINE_NORMAL _onException #-}
1068_onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
1069_onException action =
1070 gbracket_ return MC.try
1071 (\_ -> return ())
1072 (nilM (\(a, e :: MC.SomeException) -> action a >> MC.throwM e))
1074-- | Unfold the input @a@ using @Unfold m a b@, run the action @a -> m c@ on
1075-- @a@ if the unfold aborts due to an exception.
1077-- /Inhibits stream fusion/
1079-- /Pre-release/
1080{-# INLINE_NORMAL onException #-}
1081onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
1082onException action (Unfold step1 inject1) = Unfold step inject
1084 where
1086 inject x = do
1087 s <- inject1 x
1088 return (s, x)
1090 {-# INLINE_LATE step #-}
1091 step (st, v) = do
1092 res <- step1 st `MC.onException` action v
1093 return $ case res of
1094 Yield x s -> Yield x (s, v)
1095 Skip s -> Skip (s, v)
1096 Stop -> Stop
1098{-# INLINE_NORMAL _finally #-}
1099_finally :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
1100_finally action =
1101 gbracket_ return MC.try action
1102 (nilM (\(a, e :: MC.SomeException) -> action a >> MC.throwM e))
1104-- | Like 'finally' with following differences:
1106-- * action @a -> m c@ won't run if the stream is garbage collected
1107-- after partial evaluation.
1108-- * does not require a 'MonadAsync' constraint.
1110-- /Inhibits stream fusion/
1112-- /Pre-release/
1113{-# INLINE_NORMAL finally_ #-}
1114finally_ :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
1115finally_ action (Unfold step1 inject1) = Unfold step inject
1117 where
1119 inject x = do
1120 s <- inject1 x
1121 return (s, x)
1123 {-# INLINE_LATE step #-}
1124 step (st, v) = do
1125 res <- step1 st `MC.onException` action v
1126 case res of
1127 Yield x s -> return $ Yield x (s, v)
1128 Skip s -> return $ Skip (s, v)
1129 Stop -> action v >> return Stop
1131-- | Unfold the input @a@ using @Unfold m a b@, run an action on @a@ whenever
1132-- the unfold stops normally, aborts due to an exception or if it is garbage
1133-- collected after a partial lazy evaluation.
1135-- The semantics of the action @a -> m c@ are similar to the cleanup action
1136-- semantics in 'bracket'.
1138-- @
1139-- finally release = bracket return release
1140-- @
1142-- /See also 'finally_'/
1144-- /Inhibits stream fusion/
1146-- /Pre-release/
1147{-# INLINE_NORMAL finally #-}
1148finally :: (MonadAsync m, MonadCatch m)
1149 => (a -> m c) -> Unfold m a b -> Unfold m a b
1150finally action (Unfold step1 inject1) = Unfold step inject
1152 where
1154 inject x = do
1155 s <- inject1 x
1156 ref <- newIOFinalizer (action x)
1157 return (s, ref)
1159 {-# INLINE_LATE step #-}
1160 step (st, ref) = do
1161 res <- step1 st `MC.onException` runIOFinalizer ref
1162 case res of
1163 Yield x s -> return $ Yield x (s, ref)
1164 Skip s -> return $ Skip (s, ref)
1165 Stop -> do
1166 runIOFinalizer ref
1167 return Stop
1169{-# INLINE_NORMAL _bracket #-}
1170_bracket :: MonadCatch m
1171 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1172_bracket bef aft =
1173 gbracket_ bef MC.try aft (nilM (\(a, e :: MC.SomeException) -> aft a >>
1174 MC.throwM e))
1176-- | Like 'bracket' but with following differences:
1178-- * alloc action @a -> m c@ runs with async exceptions enabled
1179-- * cleanup action @c -> m d@ won't run if the stream is garbage collected
1180-- after partial evaluation.
1181-- * does not require a 'MonadAsync' constraint.
1183-- /Inhibits stream fusion/
1185-- /Pre-release/
1186{-# INLINE_NORMAL bracket_ #-}
1187bracket_ :: MonadCatch m
1188 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1189bracket_ bef aft (Unfold step1 inject1) = Unfold step inject
1191 where
1193 inject x = do
1194 r <- bef x
1195 s <- inject1 r
1196 return (s, r)
1198 {-# INLINE_LATE step #-}
1199 step (st, v) = do
1200 res <- step1 st `MC.onException` aft v
1201 case res of
1202 Yield x s -> return $ Yield x (s, v)
1203 Skip s -> return $ Skip (s, v)
1204 Stop -> aft v >> return Stop
1206-- | Run the alloc action @a -> m c@ with async exceptions disabled but keeping
1207-- blocking operations interruptible (see 'Control.Exception.mask'). Use the
1208-- output @c@ as input to @Unfold m c b@ to generate an output stream.
1210-- @c@ is usually a resource under the state of monad @m@, e.g. a file
1211-- handle, that requires a cleanup after use. The cleanup action @c -> m d@,
1212-- runs whenever the stream ends normally, due to a sync or async exception or
1213-- if it gets garbage collected after a partial lazy evaluation.
1215-- 'bracket' only guarantees that the cleanup action runs, and it runs with
1216-- async exceptions enabled. The action must ensure that it can successfully
1217-- cleanup the resource in the face of sync or async exceptions.
1219-- When the stream ends normally or on a sync exception, cleanup action runs
1220-- immediately in the current thread context, whereas in other cases it runs in
1221-- the GC context, therefore, cleanup may be delayed until the GC gets to run.
1223-- /See also: 'bracket_', 'gbracket'/
1225-- /Inhibits stream fusion/
1227-- /Pre-release/
1228{-# INLINE_NORMAL bracket #-}
1229bracket :: (MonadAsync m, MonadCatch m)
1230 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1231bracket bef aft (Unfold step1 inject1) = Unfold step inject
1233 where
1235 inject x = do
1236 -- Mask asynchronous exceptions to make the execution of 'bef' and
1237 -- the registration of 'aft' atomic. See comment in 'D.gbracketIO'.
1238 (r, ref) <- withRunInIO $ \run -> mask_ $ run $ do
1239 r <- bef x
1240 ref <- newIOFinalizer (aft r)
1241 return (r, ref)
1242 s <- inject1 r
1243 return (s, ref)
1245 {-# INLINE_LATE step #-}
1246 step (st, ref) = do
1247 res <- step1 st `MC.onException` runIOFinalizer ref
1248 case res of
1249 Yield x s -> return $ Yield x (s, ref)
1250 Skip s -> return $ Skip (s, ref)
1251 Stop -> do
1252 runIOFinalizer ref
1253 return Stop
1255-- | When unfolding @Unfold m a b@ if an exception @e@ occurs, unfold @e@ using
1256-- @Unfold m e b@.
1258-- /Inhibits stream fusion/
1260-- /Pre-release/
1261{-# INLINE_NORMAL handle #-}
1262handle :: (MonadCatch m, Exception e)
1263 => Unfold m e b -> Unfold m a b -> Unfold m a b
1264handle exc =
1265 gbracket_ return MC.try (\_ -> return ()) (discardFirst exc)