on cbb57f2 help sign in
1337 lines 43 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.
75-- Open control flow style streams can also have two representations. StreamK
76-- is a producer style representation. We can also have a consumer style
77-- representation. We can use that for composable folds in StreamK
78-- representation.
79--
81-- = Performance Notes
82--
83-- 'Unfold' representation is more efficient than using streams when combining
84-- streams. 'Unfold' type allows multiple unfold actions to be composed into a
85-- single unfold function in an efficient manner by enabling the compiler to
86-- perform stream fusion optimization.
87-- @Unfold m a b@ can be considered roughly equivalent to an action @a -> t m
88-- b@ (where @t@ is a stream type). Instead of using an 'Unfold' one could just
89-- use a function of the shape @a -> t m b@. However, working with stream types
90-- like t'Streamly.SerialT' does not allow the compiler to perform stream fusion
91-- optimization when merging, appending or concatenating multiple streams.
92-- Even though stream based combinator have excellent performance, they are
93-- much less efficient when compared to combinators using 'Unfold'. For
94-- example, the 'Streamly.Prelude.concatMap' combinator which uses @a -> t m b@
95-- (where @t@ is a stream type) to generate streams is much less efficient
96-- compared to 'Streamly.Prelude.unfoldMany'.
97--
98-- On the other hand, transformation operations on stream types are as
99-- efficient as transformations on 'Unfold'.
100--
101-- We should note that in some cases working with stream types may be more
102-- convenient compared to working with the 'Unfold' type. However, if extra
103-- performance boost is important then 'Unfold' based composition should be
104-- preferred compared to stream based composition when merging or concatenating
105-- streams.
107module Streamly.Internal.Data.Unfold
108 (
109 -- * Unfold Type
110 Step(..)
111 , Unfold
113 -- * Folding
114 , fold
115 -- pipe
117 -- * Unfolds
118 -- One to one correspondence with
119 -- "Streamly.Internal.Data.Stream.IsStream.Generate"
120 -- ** Basic Constructors
121 , mkUnfoldM
122 , mkUnfoldrM
123 , unfoldrM
124 , unfoldr
125 , functionM
126 , function
127 , identity
128 , nilM
129 , consM
131 -- ** From Values
132 , fromEffect
133 , fromPure
135 -- ** Generators
136 -- | Generate a monadic stream from a seed.
137 , repeatM
138 , replicateM
139 , fromIndicesM
140 , iterateM
142 -- ** Enumerations
143 -- *** Enumerate Num
144 , enumerateFromStepNum
145 , numFrom
147 -- *** Enumerate Integral
148 , enumerateFromStepIntegral
149 , enumerateFromToIntegral
150 , enumerateFromIntegral
152 -- *** Enumerate Fractional
153 -- | Use 'Num' enumerations for fractional or floating point number
154 -- enumerations.
155 , enumerateFromToFractional
157 -- ** From Containers
158 , fromList
159 , fromListM
161 , fromStream
162 , fromStreamK
163 , fromStreamD
165 , fromSVar
166 , fromProducer
168 -- * Combinators
169 -- ** Mapping on Input
170 , lmap
171 , lmapM
172 , supply
173 , supplyFirst
174 , supplySecond
175 , discardFirst
176 , discardSecond
177 , swap
178 -- coapply
179 -- comonad
181 -- ** Mapping on Output
182 , map
183 , mapM
184 , mapMWithInput
186 -- ** Filtering
187 , takeWhileM
188 , takeWhile
189 , take
190 , filter
191 , filterM
192 , drop
193 , dropWhile
194 , dropWhileM
196 -- ** Zipping
197 , zipWithM
198 , zipWith
200 -- ** Cross product
201 , crossWithM
202 , crossWith
203 , cross
204 , apply
206 -- ** Nesting
207 , ConcatState (..)
208 , many
209 , concatMapM
210 , bind
212 -- ** Resource Management
213 , gbracket_
214 , gbracket
215 , before
216 , after
217 , after_
218 , finally
219 , finally_
220 , bracket
221 , bracket_
223 -- ** Exceptions
224 , onException
225 , handle
226 )
227where
229import Control.Exception (Exception, fromException, mask_)
230import Control.Monad ((>=>), when)
231import Control.Monad.Catch (MonadCatch)
232import Control.Monad.IO.Class (MonadIO(..))
233import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_)
234import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef)
235import Data.Maybe (isNothing)
236import Data.Void (Void)
237import GHC.Types (SPEC(..))
238import Streamly.Internal.Data.Fold.Type (Fold(..))
239import Streamly.Internal.Data.IOFinalizer
240 (newIOFinalizer, runIOFinalizer, clearingIOFinalizer)
241import Streamly.Internal.Data.Stream.StreamD.Type (Stream(..), Step(..))
242import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
243import System.Mem (performMajorGC)
245import qualified Prelude
246import qualified Control.Monad.Catch as MC
247import qualified Data.Tuple as Tuple
248import qualified Streamly.Internal.Data.Fold.Type as FL
249import qualified Streamly.Internal.Data.Stream.StreamD.Type as D
250import qualified Streamly.Internal.Data.Stream.StreamK as K (uncons)
251import qualified Streamly.Internal.Data.Stream.StreamK.Type as K
253import Streamly.Internal.Data.SVar
254import Streamly.Internal.Data.Unfold.Type
255import Prelude
256 hiding (map, mapM, takeWhile, take, filter, const, zipWith
257 , drop, dropWhile)
259-------------------------------------------------------------------------------
260-- Input operations
261-------------------------------------------------------------------------------
263-- | Map an action on the input argument of the 'Unfold'.
264--
265-- @
266-- lmapM f = Unfold.many (Unfold.functionM f)
267-- @
268--
269-- /Since: 0.8.0/
270{-# INLINE_NORMAL lmapM #-}
271lmapM :: Monad m => (a -> m c) -> Unfold m c b -> Unfold m a b
272lmapM f (Unfold ustep uinject) = Unfold ustep (f >=> uinject)
274-- | Supply the seed to an unfold closing the input end of the unfold.
275--
276-- @
277-- supply a = Unfold.lmap (Prelude.const a)
278-- @
279--
280-- /Pre-release/
281--
282{-# INLINE_NORMAL supply #-}
283supply :: a -> Unfold m a b -> Unfold m Void b
284supply a = lmap (Prelude.const a)
286-- | Supply the first component of the tuple to an unfold that accepts a tuple
287-- as a seed resulting in a fold that accepts the second component of the tuple
288-- as a seed.
289--
290-- @
291-- supplyFirst a = Unfold.lmap (a, )
292-- @
293--
294-- /Pre-release/
295--
296{-# INLINE_NORMAL supplyFirst #-}
297supplyFirst :: a -> Unfold m (a, b) c -> Unfold m b c
298supplyFirst a = lmap (a, )
300-- | Supply the second component of the tuple to an unfold that accepts a tuple
301-- as a seed resulting in a fold that accepts the first component of the tuple
302-- as a seed.
303--
304-- @
305-- supplySecond b = Unfold.lmap (, b)
306-- @
307--
308-- /Pre-release/
309--
310{-# INLINE_NORMAL supplySecond #-}
311supplySecond :: b -> Unfold m (a, b) c -> Unfold m a c
312supplySecond b = lmap (, b)
314-- | Convert an 'Unfold' into an unfold accepting a tuple as an argument,
315-- using the argument of the original fold as the second element of tuple and
316-- discarding the first element of the tuple.
317--
318-- @
319-- discardFirst = Unfold.lmap snd
320-- @
321--
322-- /Pre-release/
323--
324{-# INLINE_NORMAL discardFirst #-}
325discardFirst :: Unfold m a b -> Unfold m (c, a) b
326discardFirst = lmap snd
328-- | Convert an 'Unfold' into an unfold accepting a tuple as an argument,
329-- using the argument of the original fold as the first element of tuple and
330-- discarding the second element of the tuple.
331--
332-- @
333-- discardSecond = Unfold.lmap fst
334-- @
335--
336-- /Pre-release/
337--
338{-# INLINE_NORMAL discardSecond #-}
339discardSecond :: Unfold m a b -> Unfold m (a, c) b
340discardSecond = lmap fst
342-- | Convert an 'Unfold' that accepts a tuple as an argument into an unfold
343-- that accepts a tuple with elements swapped.
344--
345-- @
346-- swap = Unfold.lmap Tuple.swap
347-- @
348--
349-- /Pre-release/
350--
351{-# INLINE_NORMAL swap #-}
352swap :: Unfold m (a, c) b -> Unfold m (c, a) b
353swap = lmap Tuple.swap
355-------------------------------------------------------------------------------
356-- Output operations
357-------------------------------------------------------------------------------
359-- | Compose an 'Unfold' and a 'Fold'. Given an @Unfold m a b@ and a
360-- @Fold m b c@, returns a monadic action @a -> m c@ representing the
361-- application of the fold on the unfolded stream.
362--
363-- >>> Unfold.fold Fold.sum Unfold.fromList [1..100]
364-- 5050
365--
366-- /Pre-release/
367--
368{-# INLINE_NORMAL fold #-}
369fold :: Monad m => Fold m b c -> Unfold m a b -> a -> m c
370fold (Fold fstep initial extract) (Unfold ustep inject) a = do
371 res <- initial
372 case res of
373 FL.Partial x -> inject a >>= go SPEC x
374 FL.Done b -> return b
376 where
378 {-# INLINE_LATE go #-}
379 go !_ !fs st = do
380 r <- ustep st
381 case r of
382 Yield x s -> do
383 res <- fstep fs x
384 case res of
385 FL.Partial fs1 -> go SPEC fs1 s
386 FL.Done c -> return c
387 Skip s -> go SPEC fs s
388 Stop -> extract fs
390-- | Apply a monadic function to each element of the stream and replace it
391-- with the output of the resulting action.
392--
393-- /Since: 0.8.0/
394--
395{-# INLINE_NORMAL mapM #-}
396mapM :: Monad m => (b -> m c) -> Unfold m a b -> Unfold m a c
397mapM f (Unfold ustep uinject) = Unfold step uinject
398 where
399 {-# INLINE_LATE step #-}
400 step st = do
401 r <- ustep st
402 case r of
403 Yield x s -> f x >>= \a -> return $ Yield a s
404 Skip s -> return $ Skip s
405 Stop -> return Stop
407{-# INLINE_NORMAL mapMWithInput #-}
408mapMWithInput :: Monad m => (a -> b -> m c) -> Unfold m a b -> Unfold m a c
409mapMWithInput f (Unfold ustep uinject) = Unfold step inject
410 where
411 inject a = do
412 r <- uinject a
413 return (a, r)
415 {-# INLINE_LATE step #-}
416 step (inp, st) = do
417 r <- ustep st
418 case r of
419 Yield x s -> f inp x >>= \a -> return $ Yield a (inp, s)
420 Skip s -> return $ Skip (inp, s)
421 Stop -> return Stop
423-------------------------------------------------------------------------------
424-- Convert streams into unfolds
425-------------------------------------------------------------------------------
427{-# INLINE_NORMAL fromStreamD #-}
428fromStreamD :: Monad m => Unfold m (Stream m a) a
429fromStreamD = Unfold step return
430 where
432 {-# INLINE_LATE step #-}
433 step (UnStream step1 state1) = do
434 r <- step1 defState state1
435 return $ case r of
436 Yield x s -> Yield x (Stream step1 s)
437 Skip s -> Skip (Stream step1 s)
438 Stop -> Stop
440{-# INLINE_NORMAL fromStreamK #-}
441fromStreamK :: Monad m => Unfold m (K.Stream m a) a
442fromStreamK = Unfold step return
444 where
446 {-# INLINE_LATE step #-}
447 step stream = do
448 r <- K.uncons stream
449 return $ case r of
450 Just (x, xs) -> Yield x xs
451 Nothing -> Stop
453-- XXX Using Unfold.fromStreamD seems to be faster (using cross product test
454-- case) than using fromStream even if it is implemented using fromStreamD.
455--
456-- | Convert a stream into an 'Unfold'. Note that a stream converted to an
457-- 'Unfold' may not be as efficient as an 'Unfold' in some situations.
458--
459-- /Since: 0.8.0/
460--
461{-# INLINE_NORMAL fromStream #-}
462fromStream :: (K.IsStream t, Monad m) => Unfold m (t m a) a
463fromStream = lmap K.toStream fromStreamK
465-------------------------------------------------------------------------------
466-- Unfolds
467-------------------------------------------------------------------------------
469-- | Lift a monadic function into an unfold generating a nil stream with a side
470-- effect.
471--
472{-# INLINE nilM #-}
473nilM :: Monad m => (a -> m c) -> Unfold m a b
474nilM f = Unfold step return
475 where
476 {-# INLINE_LATE step #-}
477 step x = f x >> return Stop
479-- | Prepend a monadic single element generator function to an 'Unfold'. The
480-- same seed is used in the action as well as the unfold.
481--
482-- /Pre-release/
483{-# INLINE_NORMAL consM #-}
484consM :: Monad m => (a -> m b) -> Unfold m a b -> Unfold m a b
485consM action unf = Unfold step inject
487 where
489 inject = return . Left
491 {-# INLINE_LATE step #-}
492 step (Left a) =
493 action a >>= \r -> return $ Yield r (Right (D.unfold unf a))
494 step (Right (UnStream step1 st)) = do
495 res <- step1 defState st
496 case res of
497 Yield x s -> return $ Yield x (Right (Stream step1 s))
498 Skip s -> return $ Skip (Right (Stream step1 s))
499 Stop -> return Stop
501-- | Convert a list of pure values to a 'Stream'
502--
503-- /Since: 0.8.0/
504--
505{-# INLINE_LATE fromList #-}
506fromList :: Monad m => Unfold m [a] a
507fromList = Unfold step inject
508 where
509 inject = return
510 {-# INLINE_LATE step #-}
511 step (x:xs) = return $ Yield x xs
512 step [] = return Stop
514-- | Convert a list of monadic values to a 'Stream'
515--
516-- /Since: 0.8.0/
517--
518{-# INLINE_LATE fromListM #-}
519fromListM :: Monad m => Unfold m [m a] a
520fromListM = Unfold step inject
521 where
522 inject = return
523 {-# INLINE_LATE step #-}
524 step (x:xs) = x >>= \r -> return $ Yield r xs
525 step [] = return Stop
527------------------------------------------------------------------------------
528-- Specialized Generation
529------------------------------------------------------------------------------
531-- | Generates a stream replicating the seed @n@ times.
532--
533-- /Since: 0.8.0/
534--
535{-# INLINE replicateM #-}
536replicateM :: Monad m => Int -> Unfold m (m a) a
537replicateM n = Unfold step inject
539 where
541 inject x = return (x, n)
543 {-# INLINE_LATE step #-}
544 step (x, i) =
545 if i <= 0
546 then return Stop
547 else do
548 x1 <- x
549 return $ Yield x1 (x, i - 1)
551-- | Generates an infinite stream repeating the seed.
552--
553-- /Since: 0.8.0/
554--
555{-# INLINE repeatM #-}
556repeatM :: Monad m => Unfold m (m a) a
557repeatM = Unfold step return
558 where
559 {-# INLINE_LATE step #-}
560 step x = x >>= \x1 -> return $ Yield x1 x
562-- | Generates an infinite stream starting with the given seed and applying the
563-- given function repeatedly.
564--
565-- /Since: 0.8.0/
566--
567{-# INLINE iterateM #-}
568iterateM :: Monad m => (a -> m a) -> Unfold m (m a) a
569iterateM f = Unfold step id
570 where
571 {-# INLINE_LATE step #-}
572 step x = do
573 fx <- f x
574 return $ Yield x fx
576-- | @fromIndicesM gen@ generates an infinite stream of values using @gen@
577-- starting from the seed.
578--
579-- @
580-- fromIndicesM f = Unfold.mapM f $ Unfold.enumerateFrom 0
581-- @
582--
583-- /Pre-release/
584--
585{-# INLINE_NORMAL fromIndicesM #-}
586fromIndicesM :: Monad m => (Int -> m a) -> Unfold m Int a
587fromIndicesM gen = Unfold step return
588 where
589 {-# INLINE_LATE step #-}
590 step i = do
591 x <- gen i
592 return $ Yield x (i + 1)
594-------------------------------------------------------------------------------
595-- Filtering
596-------------------------------------------------------------------------------
598-- |
599-- >>> u = Unfold.take 2 Unfold.fromList
600-- >>> Unfold.fold Fold.toList u [1..100]
601-- [1,2]
602--
603-- /Since: 0.8.0/
604--
605{-# INLINE_NORMAL take #-}
606take :: Monad m => Int -> Unfold m a b -> Unfold m a b
607take n (Unfold step1 inject1) = Unfold step inject
608 where
609 inject x = do
610 s <- inject1 x
611 return (s, 0)
612 {-# INLINE_LATE step #-}
613 step (st, i) | i < n = do
614 r <- step1 st
615 return $ case r of
616 Yield x s -> Yield x (s, i + 1)
617 Skip s -> Skip (s, i)
618 Stop -> Stop
619 step (_, _) = return Stop
621-- | Same as 'takeWhile' but with a monadic predicate.
622--
623-- /Since: 0.8.0/
624--
625{-# INLINE_NORMAL takeWhileM #-}
626takeWhileM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
627takeWhileM f (Unfold step1 inject1) = Unfold step inject1
628 where
629 {-# INLINE_LATE step #-}
630 step st = do
631 r <- step1 st
632 case r of
633 Yield x s -> do
634 b <- f x
635 return $ if b then Yield x s else Stop
636 Skip s -> return $ Skip s
637 Stop -> return Stop
639-- | End the stream generated by the 'Unfold' as soon as the predicate fails
640-- on an element.
641--
642-- /Since: 0.8.0/
643--
644{-# INLINE takeWhile #-}
645takeWhile :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
646takeWhile f = takeWhileM (return . f)
648-- | Same as 'filter' but with a monadic predicate.
649--
650-- /Since: 0.8.0/
651--
652{-# INLINE_NORMAL filterM #-}
653filterM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
654filterM f (Unfold step1 inject1) = Unfold step inject1
655 where
656 {-# INLINE_LATE step #-}
657 step st = do
658 r <- step1 st
659 case r of
660 Yield x s -> do
661 b <- f x
662 return $ if b then Yield x s else Skip s
663 Skip s -> return $ Skip s
664 Stop -> return Stop
666-- | Include only those elements that pass a predicate.
667--
668-- /Since: 0.8.0/
669--
670{-# INLINE filter #-}
671filter :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
672filter f = filterM (return . f)
674-- | @drop n unf@ drops @n@ elements from the stream generated by @unf@.
675--
676-- /Since: 0.8.0/
677--
678{-# INLINE_NORMAL drop #-}
679drop :: Monad m => Int -> Unfold m a b -> Unfold m a b
680drop n (Unfold step inject) = Unfold step' inject'
682 where
684 inject' a = do
685 b <- inject a
686 return (b, n)
688 {-# INLINE_LATE step' #-}
689 step' (st, i)
690 | i > 0 = do
691 r <- step st
692 return
693 $ case r of
694 Yield _ s -> Skip (s, i - 1)
695 Skip s -> Skip (s, i)
696 Stop -> Stop
697 | otherwise = do
698 r <- step st
699 return
700 $ case r of
701 Yield x s -> Yield x (s, 0)
702 Skip s -> Skip (s, 0)
703 Stop -> Stop
705-- | @dropWhileM f unf@ drops elements from the stream generated by @unf@ while
706-- the condition holds true. The condition function @f@ is /monadic/ in nature.
707--
708-- /Since: 0.8.0/
709--
710{-# INLINE_NORMAL dropWhileM #-}
711dropWhileM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
712dropWhileM f (Unfold step inject) = Unfold step' inject'
714 where
716 inject' a = do
717 b <- inject a
718 return $ Left b
720 {-# INLINE_LATE step' #-}
721 step' (Left st) = do
722 r <- step st
723 case r of
724 Yield x s -> do
725 b <- f x
726 return
727 $ if b
728 then Skip (Left s)
729 else Yield x (Right s)
730 Skip s -> return $ Skip (Left s)
731 Stop -> return Stop
732 step' (Right st) = do
733 r <- step st
734 return
735 $ case r of
736 Yield x s -> Yield x (Right s)
737 Skip s -> Skip (Right s)
738 Stop -> Stop
740-- | Similar to 'dropWhileM' but with a pure condition function.
741--
742-- /Since: 0.8.0/
743--
744{-# INLINE dropWhile #-}
745dropWhile :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
746dropWhile f = dropWhileM (return . f)
748------------------------------------------------------------------------------
749-- Enumeration of Num
750------------------------------------------------------------------------------
752-- | Generate an infinite stream starting from a starting value with increments
753-- of the given stride. The implementation is numerically stable for floating
754-- point values.
755--
756-- Note 'enumerateFromStepIntegral' is faster for integrals.
757--
758-- /Pre-release/
759--
760{-# INLINE enumerateFromStepNum #-}
761enumerateFromStepNum :: (Monad m, Num a) => a -> Unfold m a a
762enumerateFromStepNum stride = Unfold step inject
764 where
766 inject !from = return (from, 0)
768 {-# INLINE_LATE step #-}
769 step (from, !i) = return $ (Yield $! (from + i * stride)) $! (from, i + 1)
771-- | @numFrom = enumerateFromStepNum 1@
772--
773-- /Pre-release/
774--
775{-# INLINE_NORMAL numFrom #-}
776numFrom :: (Monad m, Num a) => Unfold m a a
777numFrom = enumerateFromStepNum 1
779------------------------------------------------------------------------------
780-- Enumeration of Integrals
781------------------------------------------------------------------------------
783-- | Can be used to enumerate unbounded integrals. This does not check for
784-- overflow or underflow for bounded integrals.
785{-# INLINE_NORMAL enumerateFromStepIntegral #-}
786enumerateFromStepIntegral :: (Integral a, Monad m) => Unfold m (a, a) a
787enumerateFromStepIntegral = Unfold step inject
788 where
789 inject (from, stride) = from `seq` stride `seq` return (from, stride)
790 {-# INLINE_LATE step #-}
791 step (x, stride) = return $ Yield x $! (x + stride, stride)
793-- We are assuming that "to" is constrained by the type to be within
794-- max/min bounds.
795{-# INLINE enumerateFromToIntegral #-}
796enumerateFromToIntegral :: (Monad m, Integral a) => a -> Unfold m a a
797enumerateFromToIntegral to =
798 takeWhile (<= to) $ supplySecond 1 enumerateFromStepIntegral
800{-# INLINE enumerateFromIntegral #-}
801enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => Unfold m a a
802enumerateFromIntegral = enumerateFromToIntegral maxBound
804------------------------------------------------------------------------------
805-- Enumeration of Fractionals
806------------------------------------------------------------------------------
808-- | /Internal/
809--
810-- > enumerateFromToFractional to = takeWhile (<= to + 1 / 2) $ enumerateFromStepNum 1
811--
812{-# INLINE_NORMAL enumerateFromToFractional #-}
813enumerateFromToFractional :: (Monad m, Fractional a, Ord a) => a -> Unfold m a a
814enumerateFromToFractional to =
815 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum 1
817-------------------------------------------------------------------------------
818-- Generation from SVar
819-------------------------------------------------------------------------------
821data FromSVarState t m a =
822 FromSVarInit (SVar t m a)
823 | FromSVarRead (SVar t m a)
824 | FromSVarLoop (SVar t m a) [ChildEvent a]
825 | FromSVarDone (SVar t m a)
827-- | /Internal/
828--
829{-# INLINE_NORMAL fromSVar #-}
830fromSVar :: MonadAsync m => Unfold m (SVar t m a) a
831fromSVar = Unfold step (return . FromSVarInit)
832 where
834 {-# INLINE_LATE step #-}
835 step (FromSVarInit svar) = do
836 ref <- liftIO $ newIORef ()
837 _ <- liftIO $ mkWeakIORef ref hook
838 -- when this copy of svar gets garbage collected "ref" will get
839 -- garbage collected and our GC hook will be called.
840 let sv = svar{svarRef = Just ref}
841 return $ Skip (FromSVarRead sv)
843 where
845 {-# NOINLINE hook #-}
846 hook = do
847 when (svarInspectMode svar) $ do
848 r <- liftIO $ readIORef (svarStopTime (svarStats svar))
849 when (isNothing r) $
850 printSVar svar "SVar Garbage Collected"
851 cleanupSVar svar
852 -- If there are any SVars referenced by this SVar a GC will prompt
853 -- them to be cleaned up quickly.
854 when (svarInspectMode svar) performMajorGC
856 step (FromSVarRead sv) = do
857 list <- readOutputQ sv
858 -- Reversing the output is important to guarantee that we process the
859 -- outputs in the same order as they were generated by the constituent
860 -- streams.
861 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
863 step (FromSVarLoop sv []) = do
864 done <- postProcess sv
865 return $ Skip $ if done
866 then FromSVarDone sv
867 else FromSVarRead sv
869 step (FromSVarLoop sv (ev : es)) = do
870 case ev of
871 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
872 ChildStop tid e -> do
873 accountThread sv tid
874 case e of
875 Nothing -> do
876 stop <- shouldStop tid
877 if stop
878 then do
879 liftIO (cleanupSVar sv)
880 return $ Skip (FromSVarDone sv)
881 else return $ Skip (FromSVarLoop sv es)
882 Just ex ->
883 case fromException ex of
884 Just ThreadAbort ->
885 return $ Skip (FromSVarLoop sv es)
886 Nothing -> liftIO (cleanupSVar sv) >> MC.throwM ex
887 where
889 shouldStop tid =
890 case svarStopStyle sv of
891 StopNone -> return False
892 StopAny -> return True
893 StopBy -> do
894 sid <- liftIO $ readIORef (svarStopBy sv)
895 return $ tid == sid
897 step (FromSVarDone sv) = do
898 when (svarInspectMode sv) $ do
899 t <- liftIO $ getTime Monotonic
900 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
901 liftIO $ printSVar sv "SVar Done"
902 return Stop
904-------------------------------------------------------------------------------
905-- Process events received by a fold consumer from a stream producer
906-------------------------------------------------------------------------------
908-- XXX Have an error for "FromSVarInit" instead of undefined. The error can
909-- redirect the user to report the failure to the developers.
910-- | /Internal/
911--
912{-# INLINE_NORMAL fromProducer #-}
913fromProducer :: MonadAsync m => Unfold m (SVar t m a) a
914fromProducer = Unfold step (return . FromSVarRead)
915 where
917 {-# INLINE_LATE step #-}
918 step (FromSVarRead sv) = do
919 list <- readOutputQ sv
920 -- Reversing the output is important to guarantee that we process the
921 -- outputs in the same order as they were generated by the constituent
922 -- streams.
923 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
925 step (FromSVarLoop sv []) = return $ Skip $ FromSVarRead sv
926 step (FromSVarLoop sv (ev : es)) = do
927 case ev of
928 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
929 ChildStop tid e -> do
930 accountThread sv tid
931 case e of
932 Nothing -> do
933 sendStopToProducer sv
934 return $ Skip (FromSVarDone sv)
935 Just _ -> error "Bug: fromProducer: received exception"
937 step (FromSVarDone sv) = do
938 when (svarInspectMode sv) $ do
939 t <- liftIO $ getTime Monotonic
940 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
941 liftIO $ printSVar sv "SVar Done"
942 return Stop
944 step (FromSVarInit _) = undefined
946------------------------------------------------------------------------------
947-- Exceptions
948------------------------------------------------------------------------------
950-- | Like 'gbracket' but with following differences:
951--
952-- * alloc action @a -> m c@ runs with async exceptions enabled
953-- * cleanup action @c -> m d@ won't run if the stream is garbage collected
954-- after partial evaluation.
955-- * does not require a 'MonadAsync' constraint.
956--
957-- /Inhibits stream fusion/
958--
959-- /Pre-release/
960--
961{-# INLINE_NORMAL gbracket_ #-}
962gbracket_
963 :: Monad m
964 => (a -> m c) -- ^ before
965 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
966 -> (c -> m d) -- ^ after, on normal stop
967 -> Unfold m (c, e) b -- ^ on exception
968 -> Unfold m c b -- ^ unfold to run
969 -> Unfold m a b
970gbracket_ bef exc aft (Unfold estep einject) (Unfold step1 inject1) =
971 Unfold step inject
973 where
975 inject x = do
976 r <- bef x
977 s <- inject1 r
978 return $ Right (s, r)
980 {-# INLINE_LATE step #-}
981 step (Right (st, v)) = do
982 res <- exc $ step1 st
983 case res of
984 Right r -> case r of
985 Yield x s -> return $ Yield x (Right (s, v))
986 Skip s -> return $ Skip (Right (s, v))
987 Stop -> aft v >> return Stop
988 -- XXX Do not handle async exceptions, just rethrow them.
989 Left e -> do
990 r <- einject (v, e)
991 return $ Skip (Left r)
992 step (Left st) = do
993 res <- estep st
994 case res of
995 Yield x s -> return $ Yield x (Left s)
996 Skip s -> return $ Skip (Left s)
997 Stop -> return Stop
999-- | Run the alloc action @a -> m c@ with async exceptions disabled but keeping
1000-- blocking operations interruptible (see 'Control.Exception.mask'). Use the
1001-- output @c@ as input to @Unfold m c b@ to generate an output stream. When
1002-- unfolding use the supplied @try@ operation @forall s. m s -> m (Either e s)@
1003-- to catch synchronous exceptions. If an exception occurs run the exception
1004-- handling unfold @Unfold m (c, e) b@.
1006-- The cleanup action @c -> m d@, runs whenever the stream ends normally, due
1007-- to a sync or async exception or if it gets garbage collected after a partial
1008-- lazy evaluation. See 'bracket' for the semantics of the cleanup action.
1010-- 'gbracket' can express all other exception handling combinators.
1012-- /Inhibits stream fusion/
1014-- /Pre-release/
1015{-# INLINE_NORMAL gbracket #-}
1016gbracket
1017 :: (MonadIO m, MonadBaseControl IO m)
1018 => (a -> m c) -- ^ before
1019 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
1020 -> (c -> m d) -- ^ after, on normal stop, or GC
1021 -> Unfold m (c, e) b -- ^ on exception
1022 -> Unfold m c b -- ^ unfold to run
1023 -> Unfold m a b
1024gbracket bef exc aft (Unfold estep einject) (Unfold step1 inject1) =
1025 Unfold step inject
1027 where
1029 inject x = do
1030 -- Mask asynchronous exceptions to make the execution of 'bef' and
1031 -- the registration of 'aft' atomic. See comment in 'D.gbracketIO'.
1032 (r, ref) <- liftBaseOp_ mask_ $ do
1033 r <- bef x
1034 ref <- newIOFinalizer (aft r)
1035 return (r, ref)
1036 s <- inject1 r
1037 return $ Right (s, r, ref)
1039 {-# INLINE_LATE step #-}
1040 step (Right (st, v, ref)) = do
1041 res <- exc $ step1 st
1042 case res of
1043 Right r -> case r of
1044 Yield x s -> return $ Yield x (Right (s, v, ref))
1045 Skip s -> return $ Skip (Right (s, v, ref))
1046 Stop -> do
1047 runIOFinalizer ref
1048 return Stop
1049 -- XXX Do not handle async exceptions, just rethrow them.
1050 Left e -> do
1051 -- Clearing of finalizer and running of exception handler must
1052 -- be atomic wrt async exceptions. Otherwise if we have cleared
1053 -- the finalizer and have not run the exception handler then we
1054 -- may leak the resource.
1055 r <- clearingIOFinalizer ref (einject (v, e))
1056 return $ Skip (Left r)
1057 step (Left st) = do
1058 res <- estep st
1059 case res of
1060 Yield x s -> return $ Yield x (Left s)
1061 Skip s -> return $ Skip (Left s)
1062 Stop -> return Stop
1064-- | Run a side effect @a -> m c@ on the input @a@ before unfolding it using
1065-- @Unfold m a b@.
1067-- > before f = lmapM (\a -> f a >> return a)
1069-- /Pre-release/
1070{-# INLINE_NORMAL before #-}
1071before :: (a -> m c) -> Unfold m a b -> Unfold m a b
1072before action (Unfold step inject) = Unfold step (action >> inject)
1074-- The custom implementation of "after_" is slightly faster (5-7%) than
1075-- "_after". This is just to document and make sure that we can always use
1076-- gbracket to implement after_ The same applies to other combinators as well.
1078{-# INLINE_NORMAL _after #-}
1079_after :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
1080_after aft = gbracket_ return (fmap Right) aft undefined
1082-- | Like 'after' with following differences:
1084-- * action @a -> m c@ won't run if the stream is garbage collected
1085-- after partial evaluation.
1086-- * Monad @m@ does not require any other constraints.
1088-- /Pre-release/
1089{-# INLINE_NORMAL after_ #-}
1090after_ :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
1091after_ action (Unfold step1 inject1) = Unfold step inject
1093 where
1095 inject x = do
1096 s <- inject1 x
1097 return (s, x)
1099 {-# INLINE_LATE step #-}
1100 step (st, v) = do
1101 res <- step1 st
1102 case res of
1103 Yield x s -> return $ Yield x (s, v)
1104 Skip s -> return $ Skip (s, v)
1105 Stop -> action v >> return Stop
1107-- | Unfold the input @a@ using @Unfold m a b@, run an action on @a@ whenever
1108-- the unfold stops normally, or if it is garbage collected after a partial
1109-- lazy evaluation.
1111-- The semantics of the action @a -> m c@ are similar to the cleanup action
1112-- semantics in 'bracket'.
1114-- /See also 'after_'/
1116-- /Pre-release/
1117{-# INLINE_NORMAL after #-}
1118after :: (MonadIO m, MonadBaseControl IO m)
1119 => (a -> m c) -> Unfold m a b -> Unfold m a b
1120after action (Unfold step1 inject1) = Unfold step inject
1122 where
1124 inject x = do
1125 s <- inject1 x
1126 ref <- newIOFinalizer (action x)
1127 return (s, ref)
1129 {-# INLINE_LATE step #-}
1130 step (st, ref) = do
1131 res <- step1 st
1132 case res of
1133 Yield x s -> return $ Yield x (s, ref)
1134 Skip s -> return $ Skip (s, ref)
1135 Stop -> do
1136 runIOFinalizer ref
1137 return Stop
1139{-# INLINE_NORMAL _onException #-}
1140_onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
1141_onException action =
1142 gbracket_ return MC.try
1143 (\_ -> return ())
1144 (nilM (\(a, e :: MC.SomeException) -> action a >> MC.throwM e))
1146-- | Unfold the input @a@ using @Unfold m a b@, run the action @a -> m c@ on
1147-- @a@ if the unfold aborts due to an exception.
1149-- /Inhibits stream fusion/
1151-- /Pre-release/
1152{-# INLINE_NORMAL onException #-}
1153onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
1154onException action (Unfold step1 inject1) = Unfold step inject
1156 where
1158 inject x = do
1159 s <- inject1 x
1160 return (s, x)
1162 {-# INLINE_LATE step #-}
1163 step (st, v) = do
1164 res <- step1 st `MC.onException` action v
1165 case res of
1166 Yield x s -> return $ Yield x (s, v)
1167 Skip s -> return $ Skip (s, v)
1168 Stop -> return Stop
1170{-# INLINE_NORMAL _finally #-}
1171_finally :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
1172_finally action =
1173 gbracket_ return MC.try action
1174 (nilM (\(a, e :: MC.SomeException) -> action a >> MC.throwM e))
1176-- | Like 'finally' with following differences:
1178-- * action @a -> m c@ won't run if the stream is garbage collected
1179-- after partial evaluation.
1180-- * does not require a 'MonadAsync' constraint.
1182-- /Inhibits stream fusion/
1184-- /Pre-release/
1185{-# INLINE_NORMAL finally_ #-}
1186finally_ :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
1187finally_ action (Unfold step1 inject1) = Unfold step inject
1189 where
1191 inject x = do
1192 s <- inject1 x
1193 return (s, x)
1195 {-# INLINE_LATE step #-}
1196 step (st, v) = do
1197 res <- step1 st `MC.onException` action v
1198 case res of
1199 Yield x s -> return $ Yield x (s, v)
1200 Skip s -> return $ Skip (s, v)
1201 Stop -> action v >> return Stop
1203-- | Unfold the input @a@ using @Unfold m a b@, run an action on @a@ whenever
1204-- the unfold stops normally, aborts due to an exception or if it is garbage
1205-- collected after a partial lazy evaluation.
1207-- The semantics of the action @a -> m c@ are similar to the cleanup action
1208-- semantics in 'bracket'.
1210-- @
1211-- finally release = bracket return release
1212-- @
1214-- /See also 'finally_'/
1216-- /Inhibits stream fusion/
1218-- /Pre-release/
1219{-# INLINE_NORMAL finally #-}
1220finally :: (MonadAsync m, MonadCatch m)
1221 => (a -> m c) -> Unfold m a b -> Unfold m a b
1222finally action (Unfold step1 inject1) = Unfold step inject
1224 where
1226 inject x = do
1227 s <- inject1 x
1228 ref <- newIOFinalizer (action x)
1229 return (s, ref)
1231 {-# INLINE_LATE step #-}
1232 step (st, ref) = do
1233 res <- step1 st `MC.onException` runIOFinalizer ref
1234 case res of
1235 Yield x s -> return $ Yield x (s, ref)
1236 Skip s -> return $ Skip (s, ref)
1237 Stop -> do
1238 runIOFinalizer ref
1239 return Stop
1241{-# INLINE_NORMAL _bracket #-}
1242_bracket :: MonadCatch m
1243 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1244_bracket bef aft =
1245 gbracket_ bef MC.try aft (nilM (\(a, e :: MC.SomeException) -> aft a >>
1246 MC.throwM e))
1248-- | Like 'bracket' but with following differences:
1250-- * alloc action @a -> m c@ runs with async exceptions enabled
1251-- * cleanup action @c -> m d@ won't run if the stream is garbage collected
1252-- after partial evaluation.
1253-- * does not require a 'MonadAsync' constraint.
1255-- /Inhibits stream fusion/
1257-- /Pre-release/
1258{-# INLINE_NORMAL bracket_ #-}
1259bracket_ :: MonadCatch m
1260 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1261bracket_ bef aft (Unfold step1 inject1) = Unfold step inject
1263 where
1265 inject x = do
1266 r <- bef x
1267 s <- inject1 r
1268 return (s, r)
1270 {-# INLINE_LATE step #-}
1271 step (st, v) = do
1272 res <- step1 st `MC.onException` aft v
1273 case res of
1274 Yield x s -> return $ Yield x (s, v)
1275 Skip s -> return $ Skip (s, v)
1276 Stop -> aft v >> return Stop
1278-- | Run the alloc action @a -> m c@ with async exceptions disabled but keeping
1279-- blocking operations interruptible (see 'Control.Exception.mask'). Use the
1280-- output @c@ as input to @Unfold m c b@ to generate an output stream.
1282-- @c@ is usually a resource under the state of monad @m@, e.g. a file
1283-- handle, that requires a cleanup after use. The cleanup action @c -> m d@,
1284-- runs whenever the stream ends normally, due to a sync or async exception or
1285-- if it gets garbage collected after a partial lazy evaluation.
1287-- 'bracket' only guarantees that the cleanup action runs, and it runs with
1288-- async exceptions enabled. The action must ensure that it can successfully
1289-- cleanup the resource in the face of sync or async exceptions.
1291-- When the stream ends normally or on a sync exception, cleanup action runs
1292-- immediately in the current thread context, whereas in other cases it runs in
1293-- the GC context, therefore, cleanup may be delayed until the GC gets to run.
1295-- /See also: 'bracket_', 'gbracket'/
1297-- /Inhibits stream fusion/
1299-- /Pre-release/
1300{-# INLINE_NORMAL bracket #-}
1301bracket :: (MonadAsync m, MonadCatch m)
1302 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1303bracket bef aft (Unfold step1 inject1) = Unfold step inject
1305 where
1307 inject x = do
1308 -- Mask asynchronous exceptions to make the execution of 'bef' and
1309 -- the registration of 'aft' atomic. See comment in 'D.gbracketIO'.
1310 (r, ref) <- liftBaseOp_ mask_ $ do
1311 r <- bef x
1312 ref <- newIOFinalizer (aft r)
1313 return (r, ref)
1314 s <- inject1 r
1315 return (s, ref)
1317 {-# INLINE_LATE step #-}
1318 step (st, ref) = do
1319 res <- step1 st `MC.onException` runIOFinalizer ref
1320 case res of
1321 Yield x s -> return $ Yield x (s, ref)
1322 Skip s -> return $ Skip (s, ref)
1323 Stop -> do
1324 runIOFinalizer ref
1325 return Stop
1327-- | When unfolding @Unfold m a b@ if an exception @e@ occurs, unfold @e@ using
1328-- @Unfold m e b@.
1330-- /Inhibits stream fusion/
1332-- /Pre-release/
1333{-# INLINE_NORMAL handle #-}
1334handle :: (MonadCatch m, Exception e)
1335 => Unfold m e b -> Unfold m a b -> Unfold m a b
1336handle exc =
1337 gbracket_ return MC.try (\_ -> return ()) (discardFirst exc)