on 796bb22 help sign in
1335 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 , yieldM
133 , yield
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-- this is horribly slow
428{-# INLINE_NORMAL fromStreamD #-}
429fromStreamD :: Monad m => Unfold m (Stream m a) a
430fromStreamD = Unfold step return
431 where
433 {-# INLINE_LATE step #-}
434 step (UnStream step1 state1) = do
435 r <- step1 defState state1
436 return $ case r of
437 Yield x s -> Yield x (Stream step1 s)
438 Skip s -> Skip (Stream step1 s)
439 Stop -> Stop
441{-# INLINE_NORMAL fromStreamK #-}
442fromStreamK :: Monad m => Unfold m (K.Stream m a) a
443fromStreamK = Unfold step return
445 where
447 {-# INLINE_LATE step #-}
448 step stream = do
449 r <- K.uncons stream
450 return $ case r of
451 Just (x, xs) -> Yield x xs
452 Nothing -> Stop
454-- | Convert a stream into an 'Unfold'. Note that a stream converted to an
455-- 'Unfold' may not be as efficient as an 'Unfold' in some situations.
456--
457-- /Since: 0.8.0/
458--
459{-# INLINE_NORMAL fromStream #-}
460fromStream :: (K.IsStream t, Monad m) => Unfold m (t m a) a
461fromStream = lmap K.toStream fromStreamK
463-------------------------------------------------------------------------------
464-- Unfolds
465-------------------------------------------------------------------------------
467-- | Lift a monadic function into an unfold generating a nil stream with a side
468-- effect.
469--
470{-# INLINE nilM #-}
471nilM :: Monad m => (a -> m c) -> Unfold m a b
472nilM f = Unfold step return
473 where
474 {-# INLINE_LATE step #-}
475 step x = f x >> return Stop
477-- | Prepend a monadic single element generator function to an 'Unfold'. The
478-- same seed is used in the action as well as the unfold.
479--
480-- /Pre-release/
481{-# INLINE_NORMAL consM #-}
482consM :: Monad m => (a -> m b) -> Unfold m a b -> Unfold m a b
483consM action unf = Unfold step inject
485 where
487 inject = return . Left
489 {-# INLINE_LATE step #-}
490 step (Left a) =
491 action a >>= \r -> return $ Yield r (Right (D.unfold unf a))
492 step (Right (UnStream step1 st)) = do
493 res <- step1 defState st
494 case res of
495 Yield x s -> return $ Yield x (Right (Stream step1 s))
496 Skip s -> return $ Skip (Right (Stream step1 s))
497 Stop -> return Stop
499-- | Convert a list of pure values to a 'Stream'
500--
501-- /Since: 0.8.0/
502--
503{-# INLINE_LATE fromList #-}
504fromList :: Monad m => Unfold m [a] a
505fromList = Unfold step inject
506 where
507 inject = return
508 {-# INLINE_LATE step #-}
509 step (x:xs) = return $ Yield x xs
510 step [] = return Stop
512-- | Convert a list of monadic values to a 'Stream'
513--
514-- /Since: 0.8.0/
515--
516{-# INLINE_LATE fromListM #-}
517fromListM :: Monad m => Unfold m [m a] a
518fromListM = Unfold step inject
519 where
520 inject = return
521 {-# INLINE_LATE step #-}
522 step (x:xs) = x >>= \r -> return $ Yield r xs
523 step [] = return Stop
525------------------------------------------------------------------------------
526-- Specialized Generation
527------------------------------------------------------------------------------
529-- | Generates a stream replicating the seed @n@ times.
530--
531-- /Since: 0.8.0/
532--
533{-# INLINE replicateM #-}
534replicateM :: Monad m => Int -> Unfold m (m a) a
535replicateM n = Unfold step inject
537 where
539 inject x = return (x, n)
541 {-# INLINE_LATE step #-}
542 step (x, i) =
543 if i <= 0
544 then return Stop
545 else do
546 x1 <- x
547 return $ Yield x1 (x, i - 1)
549-- | Generates an infinite stream repeating the seed.
550--
551-- /Since: 0.8.0/
552--
553{-# INLINE repeatM #-}
554repeatM :: Monad m => Unfold m (m a) a
555repeatM = Unfold step return
556 where
557 {-# INLINE_LATE step #-}
558 step x = x >>= \x1 -> return $ Yield x1 x
560-- | Generates an infinite stream starting with the given seed and applying the
561-- given function repeatedly.
562--
563-- /Since: 0.8.0/
564--
565{-# INLINE iterateM #-}
566iterateM :: Monad m => (a -> m a) -> Unfold m (m a) a
567iterateM f = Unfold step id
568 where
569 {-# INLINE_LATE step #-}
570 step x = do
571 fx <- f x
572 return $ Yield x fx
574-- | @fromIndicesM gen@ generates an infinite stream of values using @gen@
575-- starting from the seed.
576--
577-- @
578-- fromIndicesM f = Unfold.mapM f $ Unfold.enumerateFrom 0
579-- @
580--
581-- /Pre-release/
582--
583{-# INLINE_NORMAL fromIndicesM #-}
584fromIndicesM :: Monad m => (Int -> m a) -> Unfold m Int a
585fromIndicesM gen = Unfold step return
586 where
587 {-# INLINE_LATE step #-}
588 step i = do
589 x <- gen i
590 return $ Yield x (i + 1)
592-------------------------------------------------------------------------------
593-- Filtering
594-------------------------------------------------------------------------------
596-- |
597-- >>> u = Unfold.take 2 Unfold.fromList
598-- >>> Unfold.fold Fold.toList u [1..100]
599-- [1,2]
600--
601-- /Since: 0.8.0/
602--
603{-# INLINE_NORMAL take #-}
604take :: Monad m => Int -> Unfold m a b -> Unfold m a b
605take n (Unfold step1 inject1) = Unfold step inject
606 where
607 inject x = do
608 s <- inject1 x
609 return (s, 0)
610 {-# INLINE_LATE step #-}
611 step (st, i) | i < n = do
612 r <- step1 st
613 return $ case r of
614 Yield x s -> Yield x (s, i + 1)
615 Skip s -> Skip (s, i)
616 Stop -> Stop
617 step (_, _) = return Stop
619-- | Same as 'takeWhile' but with a monadic predicate.
620--
621-- /Since: 0.8.0/
622--
623{-# INLINE_NORMAL takeWhileM #-}
624takeWhileM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
625takeWhileM f (Unfold step1 inject1) = Unfold step inject1
626 where
627 {-# INLINE_LATE step #-}
628 step st = do
629 r <- step1 st
630 case r of
631 Yield x s -> do
632 b <- f x
633 return $ if b then Yield x s else Stop
634 Skip s -> return $ Skip s
635 Stop -> return Stop
637-- | End the stream generated by the 'Unfold' as soon as the predicate fails
638-- on an element.
639--
640-- /Since: 0.8.0/
641--
642{-# INLINE takeWhile #-}
643takeWhile :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
644takeWhile f = takeWhileM (return . f)
646-- | Same as 'filter' but with a monadic predicate.
647--
648-- /Since: 0.8.0/
649--
650{-# INLINE_NORMAL filterM #-}
651filterM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
652filterM f (Unfold step1 inject1) = Unfold step inject1
653 where
654 {-# INLINE_LATE step #-}
655 step st = do
656 r <- step1 st
657 case r of
658 Yield x s -> do
659 b <- f x
660 return $ if b then Yield x s else Skip s
661 Skip s -> return $ Skip s
662 Stop -> return Stop
664-- | Include only those elements that pass a predicate.
665--
666-- /Since: 0.8.0/
667--
668{-# INLINE filter #-}
669filter :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
670filter f = filterM (return . f)
672-- | @drop n unf@ drops @n@ elements from the stream generated by @unf@.
673--
674-- /Since: 0.8.0/
675--
676{-# INLINE_NORMAL drop #-}
677drop :: Monad m => Int -> Unfold m a b -> Unfold m a b
678drop n (Unfold step inject) = Unfold step' inject'
680 where
682 inject' a = do
683 b <- inject a
684 return (b, n)
686 {-# INLINE_LATE step' #-}
687 step' (st, i)
688 | i > 0 = do
689 r <- step st
690 return
691 $ case r of
692 Yield _ s -> Skip (s, i - 1)
693 Skip s -> Skip (s, i)
694 Stop -> Stop
695 | otherwise = do
696 r <- step st
697 return
698 $ case r of
699 Yield x s -> Yield x (s, 0)
700 Skip s -> Skip (s, 0)
701 Stop -> Stop
703-- | @dropWhileM f unf@ drops elements from the stream generated by @unf@ while
704-- the condition holds true. The condition function @f@ is /monadic/ in nature.
705--
706-- /Since: 0.8.0/
707--
708{-# INLINE_NORMAL dropWhileM #-}
709dropWhileM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
710dropWhileM f (Unfold step inject) = Unfold step' inject'
712 where
714 inject' a = do
715 b <- inject a
716 return $ Left b
718 {-# INLINE_LATE step' #-}
719 step' (Left st) = do
720 r <- step st
721 case r of
722 Yield x s -> do
723 b <- f x
724 return
725 $ if b
726 then Skip (Left s)
727 else Yield x (Right s)
728 Skip s -> return $ Skip (Left s)
729 Stop -> return Stop
730 step' (Right st) = do
731 r <- step st
732 return
733 $ case r of
734 Yield x s -> Yield x (Right s)
735 Skip s -> Skip (Right s)
736 Stop -> Stop
738-- | Similar to 'dropWhileM' but with a pure condition function.
739--
740-- /Since: 0.8.0/
741--
742{-# INLINE dropWhile #-}
743dropWhile :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
744dropWhile f = dropWhileM (return . f)
746------------------------------------------------------------------------------
747-- Enumeration of Num
748------------------------------------------------------------------------------
750-- | Generate an infinite stream starting from a starting value with increments
751-- of the given stride. The implementation is numerically stable for floating
752-- point values.
753--
754-- Note 'enumerateFromStepIntegral' is faster for integrals.
755--
756-- /Pre-release/
757--
758{-# INLINE enumerateFromStepNum #-}
759enumerateFromStepNum :: (Monad m, Num a) => a -> Unfold m a a
760enumerateFromStepNum stride = Unfold step inject
762 where
764 inject !from = return (from, 0)
766 {-# INLINE_LATE step #-}
767 step (from, !i) = return $ (Yield $! (from + i * stride)) $! (from, i + 1)
769-- | @numFrom = enumerateFromStepNum 1@
770--
771-- /Pre-release/
772--
773{-# INLINE_NORMAL numFrom #-}
774numFrom :: (Monad m, Num a) => Unfold m a a
775numFrom = enumerateFromStepNum 1
777------------------------------------------------------------------------------
778-- Enumeration of Integrals
779------------------------------------------------------------------------------
781-- | Can be used to enumerate unbounded integrals. This does not check for
782-- overflow or underflow for bounded integrals.
783{-# INLINE_NORMAL enumerateFromStepIntegral #-}
784enumerateFromStepIntegral :: (Integral a, Monad m) => Unfold m (a, a) a
785enumerateFromStepIntegral = Unfold step inject
786 where
787 inject (from, stride) = from `seq` stride `seq` return (from, stride)
788 {-# INLINE_LATE step #-}
789 step (x, stride) = return $ Yield x $! (x + stride, stride)
791-- We are assuming that "to" is constrained by the type to be within
792-- max/min bounds.
793{-# INLINE enumerateFromToIntegral #-}
794enumerateFromToIntegral :: (Monad m, Integral a) => a -> Unfold m a a
795enumerateFromToIntegral to =
796 takeWhile (<= to) $ supplySecond 1 enumerateFromStepIntegral
798{-# INLINE enumerateFromIntegral #-}
799enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => Unfold m a a
800enumerateFromIntegral = enumerateFromToIntegral maxBound
802------------------------------------------------------------------------------
803-- Enumeration of Fractionals
804------------------------------------------------------------------------------
806-- | /Internal/
807--
808-- > enumerateFromToFractional to = takeWhile (<= to + 1 / 2) $ enumerateFromStepNum 1
809--
810{-# INLINE_NORMAL enumerateFromToFractional #-}
811enumerateFromToFractional :: (Monad m, Fractional a, Ord a) => a -> Unfold m a a
812enumerateFromToFractional to =
813 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum 1
815-------------------------------------------------------------------------------
816-- Generation from SVar
817-------------------------------------------------------------------------------
819data FromSVarState t m a =
820 FromSVarInit (SVar t m a)
821 | FromSVarRead (SVar t m a)
822 | FromSVarLoop (SVar t m a) [ChildEvent a]
823 | FromSVarDone (SVar t m a)
825-- | /Internal/
826--
827{-# INLINE_NORMAL fromSVar #-}
828fromSVar :: MonadAsync m => Unfold m (SVar t m a) a
829fromSVar = Unfold step (return . FromSVarInit)
830 where
832 {-# INLINE_LATE step #-}
833 step (FromSVarInit svar) = do
834 ref <- liftIO $ newIORef ()
835 _ <- liftIO $ mkWeakIORef ref hook
836 -- when this copy of svar gets garbage collected "ref" will get
837 -- garbage collected and our GC hook will be called.
838 let sv = svar{svarRef = Just ref}
839 return $ Skip (FromSVarRead sv)
841 where
843 {-# NOINLINE hook #-}
844 hook = do
845 when (svarInspectMode svar) $ do
846 r <- liftIO $ readIORef (svarStopTime (svarStats svar))
847 when (isNothing r) $
848 printSVar svar "SVar Garbage Collected"
849 cleanupSVar svar
850 -- If there are any SVars referenced by this SVar a GC will prompt
851 -- them to be cleaned up quickly.
852 when (svarInspectMode svar) performMajorGC
854 step (FromSVarRead sv) = do
855 list <- readOutputQ sv
856 -- Reversing the output is important to guarantee that we process the
857 -- outputs in the same order as they were generated by the constituent
858 -- streams.
859 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
861 step (FromSVarLoop sv []) = do
862 done <- postProcess sv
863 return $ Skip $ if done
864 then FromSVarDone sv
865 else FromSVarRead sv
867 step (FromSVarLoop sv (ev : es)) = do
868 case ev of
869 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
870 ChildStop tid e -> do
871 accountThread sv tid
872 case e of
873 Nothing -> do
874 stop <- shouldStop tid
875 if stop
876 then do
877 liftIO (cleanupSVar sv)
878 return $ Skip (FromSVarDone sv)
879 else return $ Skip (FromSVarLoop sv es)
880 Just ex ->
881 case fromException ex of
882 Just ThreadAbort ->
883 return $ Skip (FromSVarLoop sv es)
884 Nothing -> liftIO (cleanupSVar sv) >> MC.throwM ex
885 where
887 shouldStop tid =
888 case svarStopStyle sv of
889 StopNone -> return False
890 StopAny -> return True
891 StopBy -> do
892 sid <- liftIO $ readIORef (svarStopBy sv)
893 return $ tid == sid
895 step (FromSVarDone sv) = do
896 when (svarInspectMode sv) $ do
897 t <- liftIO $ getTime Monotonic
898 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
899 liftIO $ printSVar sv "SVar Done"
900 return Stop
902-------------------------------------------------------------------------------
903-- Process events received by a fold consumer from a stream producer
904-------------------------------------------------------------------------------
906-- XXX Have an error for "FromSVarInit" instead of undefined. The error can
907-- redirect the user to report the failure to the developers.
908-- | /Internal/
909--
910{-# INLINE_NORMAL fromProducer #-}
911fromProducer :: MonadAsync m => Unfold m (SVar t m a) a
912fromProducer = Unfold step (return . FromSVarRead)
913 where
915 {-# INLINE_LATE step #-}
916 step (FromSVarRead sv) = do
917 list <- readOutputQ sv
918 -- Reversing the output is important to guarantee that we process the
919 -- outputs in the same order as they were generated by the constituent
920 -- streams.
921 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
923 step (FromSVarLoop sv []) = return $ Skip $ FromSVarRead sv
924 step (FromSVarLoop sv (ev : es)) = do
925 case ev of
926 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
927 ChildStop tid e -> do
928 accountThread sv tid
929 case e of
930 Nothing -> do
931 sendStopToProducer sv
932 return $ Skip (FromSVarDone sv)
933 Just _ -> error "Bug: fromProducer: received exception"
935 step (FromSVarDone sv) = do
936 when (svarInspectMode sv) $ do
937 t <- liftIO $ getTime Monotonic
938 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
939 liftIO $ printSVar sv "SVar Done"
940 return Stop
942 step (FromSVarInit _) = undefined
944------------------------------------------------------------------------------
945-- Exceptions
946------------------------------------------------------------------------------
948-- | Like 'gbracket' but with following differences:
949--
950-- * alloc action @a -> m c@ runs with async exceptions enabled
951-- * cleanup action @c -> m d@ won't run if the stream is garbage collected
952-- after partial evaluation.
953-- * does not require a 'MonadAsync' constraint.
954--
955-- /Inhibits stream fusion/
956--
957-- /Pre-release/
958--
959{-# INLINE_NORMAL gbracket_ #-}
960gbracket_
961 :: Monad m
962 => (a -> m c) -- ^ before
963 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
964 -> (c -> m d) -- ^ after, on normal stop
965 -> Unfold m (c, e) b -- ^ on exception
966 -> Unfold m c b -- ^ unfold to run
967 -> Unfold m a b
968gbracket_ bef exc aft (Unfold estep einject) (Unfold step1 inject1) =
969 Unfold step inject
971 where
973 inject x = do
974 r <- bef x
975 s <- inject1 r
976 return $ Right (s, r)
978 {-# INLINE_LATE step #-}
979 step (Right (st, v)) = do
980 res <- exc $ step1 st
981 case res of
982 Right r -> case r of
983 Yield x s -> return $ Yield x (Right (s, v))
984 Skip s -> return $ Skip (Right (s, v))
985 Stop -> aft v >> return Stop
986 -- XXX Do not handle async exceptions, just rethrow them.
987 Left e -> do
988 r <- einject (v, e)
989 return $ Skip (Left r)
990 step (Left st) = do
991 res <- estep st
992 case res of
993 Yield x s -> return $ Yield x (Left s)
994 Skip s -> return $ Skip (Left s)
995 Stop -> return Stop
997-- | Run the alloc action @a -> m c@ with async exceptions disabled but keeping
998-- blocking operations interruptible (see 'Control.Exception.mask'). Use the
999-- output @c@ as input to @Unfold m c b@ to generate an output stream. When
1000-- unfolding use the supplied @try@ operation @forall s. m s -> m (Either e s)@
1001-- to catch synchronous exceptions. If an exception occurs run the exception
1002-- handling unfold @Unfold m (c, e) b@.
1004-- The cleanup action @c -> m d@, runs whenever the stream ends normally, due
1005-- to a sync or async exception or if it gets garbage collected after a partial
1006-- lazy evaluation. See 'bracket' for the semantics of the cleanup action.
1008-- 'gbracket' can express all other exception handling combinators.
1010-- /Inhibits stream fusion/
1012-- /Pre-release/
1013{-# INLINE_NORMAL gbracket #-}
1014gbracket
1015 :: (MonadIO m, MonadBaseControl IO m)
1016 => (a -> m c) -- ^ before
1017 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
1018 -> (c -> m d) -- ^ after, on normal stop, or GC
1019 -> Unfold m (c, e) b -- ^ on exception
1020 -> Unfold m c b -- ^ unfold to run
1021 -> Unfold m a b
1022gbracket bef exc aft (Unfold estep einject) (Unfold step1 inject1) =
1023 Unfold step inject
1025 where
1027 inject x = do
1028 -- Mask asynchronous exceptions to make the execution of 'bef' and
1029 -- the registration of 'aft' atomic. See comment in 'D.gbracketIO'.
1030 (r, ref) <- liftBaseOp_ mask_ $ do
1031 r <- bef x
1032 ref <- newIOFinalizer (aft r)
1033 return (r, ref)
1034 s <- inject1 r
1035 return $ Right (s, r, ref)
1037 {-# INLINE_LATE step #-}
1038 step (Right (st, v, ref)) = do
1039 res <- exc $ step1 st
1040 case res of
1041 Right r -> case r of
1042 Yield x s -> return $ Yield x (Right (s, v, ref))
1043 Skip s -> return $ Skip (Right (s, v, ref))
1044 Stop -> do
1045 runIOFinalizer ref
1046 return Stop
1047 -- XXX Do not handle async exceptions, just rethrow them.
1048 Left e -> do
1049 -- Clearing of finalizer and running of exception handler must
1050 -- be atomic wrt async exceptions. Otherwise if we have cleared
1051 -- the finalizer and have not run the exception handler then we
1052 -- may leak the resource.
1053 r <- clearingIOFinalizer ref (einject (v, e))
1054 return $ Skip (Left r)
1055 step (Left st) = do
1056 res <- estep st
1057 case res of
1058 Yield x s -> return $ Yield x (Left s)
1059 Skip s -> return $ Skip (Left s)
1060 Stop -> return Stop
1062-- | Run a side effect @a -> m c@ on the input @a@ before unfolding it using
1063-- @Unfold m a b@.
1065-- > before f = lmapM (\a -> f a >> return a)
1067-- /Pre-release/
1068{-# INLINE_NORMAL before #-}
1069before :: (a -> m c) -> Unfold m a b -> Unfold m a b
1070before action (Unfold step inject) = Unfold step (action >> inject)
1072-- The custom implementation of "after_" is slightly faster (5-7%) than
1073-- "_after". This is just to document and make sure that we can always use
1074-- gbracket to implement after_ The same applies to other combinators as well.
1076{-# INLINE_NORMAL _after #-}
1077_after :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
1078_after aft = gbracket_ return (fmap Right) aft undefined
1080-- | Like 'after' with following differences:
1082-- * action @a -> m c@ won't run if the stream is garbage collected
1083-- after partial evaluation.
1084-- * Monad @m@ does not require any other constraints.
1086-- /Pre-release/
1087{-# INLINE_NORMAL after_ #-}
1088after_ :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
1089after_ action (Unfold step1 inject1) = Unfold step inject
1091 where
1093 inject x = do
1094 s <- inject1 x
1095 return (s, x)
1097 {-# INLINE_LATE step #-}
1098 step (st, v) = do
1099 res <- step1 st
1100 case res of
1101 Yield x s -> return $ Yield x (s, v)
1102 Skip s -> return $ Skip (s, v)
1103 Stop -> action v >> return Stop
1105-- | Unfold the input @a@ using @Unfold m a b@, run an action on @a@ whenever
1106-- the unfold stops normally, or if it is garbage collected after a partial
1107-- lazy evaluation.
1109-- The semantics of the action @a -> m c@ are similar to the cleanup action
1110-- semantics in 'bracket'.
1112-- /See also 'after_'/
1114-- /Pre-release/
1115{-# INLINE_NORMAL after #-}
1116after :: (MonadIO m, MonadBaseControl IO m)
1117 => (a -> m c) -> Unfold m a b -> Unfold m a b
1118after action (Unfold step1 inject1) = Unfold step inject
1120 where
1122 inject x = do
1123 s <- inject1 x
1124 ref <- newIOFinalizer (action x)
1125 return (s, ref)
1127 {-# INLINE_LATE step #-}
1128 step (st, ref) = do
1129 res <- step1 st
1130 case res of
1131 Yield x s -> return $ Yield x (s, ref)
1132 Skip s -> return $ Skip (s, ref)
1133 Stop -> do
1134 runIOFinalizer ref
1135 return Stop
1137{-# INLINE_NORMAL _onException #-}
1138_onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
1139_onException action =
1140 gbracket_ return MC.try
1141 (\_ -> return ())
1142 (nilM (\(a, e :: MC.SomeException) -> action a >> MC.throwM e))
1144-- | Unfold the input @a@ using @Unfold m a b@, run the action @a -> m c@ on
1145-- @a@ if the unfold aborts due to an exception.
1147-- /Inhibits stream fusion/
1149-- /Pre-release/
1150{-# INLINE_NORMAL onException #-}
1151onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
1152onException action (Unfold step1 inject1) = Unfold step inject
1154 where
1156 inject x = do
1157 s <- inject1 x
1158 return (s, x)
1160 {-# INLINE_LATE step #-}
1161 step (st, v) = do
1162 res <- step1 st `MC.onException` action v
1163 case res of
1164 Yield x s -> return $ Yield x (s, v)
1165 Skip s -> return $ Skip (s, v)
1166 Stop -> return Stop
1168{-# INLINE_NORMAL _finally #-}
1169_finally :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
1170_finally action =
1171 gbracket_ return MC.try action
1172 (nilM (\(a, e :: MC.SomeException) -> action a >> MC.throwM e))
1174-- | Like 'finally' with following differences:
1176-- * action @a -> m c@ won't run if the stream is garbage collected
1177-- after partial evaluation.
1178-- * does not require a 'MonadAsync' constraint.
1180-- /Inhibits stream fusion/
1182-- /Pre-release/
1183{-# INLINE_NORMAL finally_ #-}
1184finally_ :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
1185finally_ action (Unfold step1 inject1) = Unfold step inject
1187 where
1189 inject x = do
1190 s <- inject1 x
1191 return (s, x)
1193 {-# INLINE_LATE step #-}
1194 step (st, v) = do
1195 res <- step1 st `MC.onException` action v
1196 case res of
1197 Yield x s -> return $ Yield x (s, v)
1198 Skip s -> return $ Skip (s, v)
1199 Stop -> action v >> return Stop
1201-- | Unfold the input @a@ using @Unfold m a b@, run an action on @a@ whenever
1202-- the unfold stops normally, aborts due to an exception or if it is garbage
1203-- collected after a partial lazy evaluation.
1205-- The semantics of the action @a -> m c@ are similar to the cleanup action
1206-- semantics in 'bracket'.
1208-- @
1209-- finally release = bracket return release
1210-- @
1212-- /See also 'finally_'/
1214-- /Inhibits stream fusion/
1216-- /Pre-release/
1217{-# INLINE_NORMAL finally #-}
1218finally :: (MonadAsync m, MonadCatch m)
1219 => (a -> m c) -> Unfold m a b -> Unfold m a b
1220finally action (Unfold step1 inject1) = Unfold step inject
1222 where
1224 inject x = do
1225 s <- inject1 x
1226 ref <- newIOFinalizer (action x)
1227 return (s, ref)
1229 {-# INLINE_LATE step #-}
1230 step (st, ref) = do
1231 res <- step1 st `MC.onException` runIOFinalizer ref
1232 case res of
1233 Yield x s -> return $ Yield x (s, ref)
1234 Skip s -> return $ Skip (s, ref)
1235 Stop -> do
1236 runIOFinalizer ref
1237 return Stop
1239{-# INLINE_NORMAL _bracket #-}
1240_bracket :: MonadCatch m
1241 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1242_bracket bef aft =
1243 gbracket_ bef MC.try aft (nilM (\(a, e :: MC.SomeException) -> aft a >>
1244 MC.throwM e))
1246-- | Like 'bracket' but with following differences:
1248-- * alloc action @a -> m c@ runs with async exceptions enabled
1249-- * cleanup action @c -> m d@ won't run if the stream is garbage collected
1250-- after partial evaluation.
1251-- * does not require a 'MonadAsync' constraint.
1253-- /Inhibits stream fusion/
1255-- /Pre-release/
1256{-# INLINE_NORMAL bracket_ #-}
1257bracket_ :: MonadCatch m
1258 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1259bracket_ bef aft (Unfold step1 inject1) = Unfold step inject
1261 where
1263 inject x = do
1264 r <- bef x
1265 s <- inject1 r
1266 return (s, r)
1268 {-# INLINE_LATE step #-}
1269 step (st, v) = do
1270 res <- step1 st `MC.onException` aft v
1271 case res of
1272 Yield x s -> return $ Yield x (s, v)
1273 Skip s -> return $ Skip (s, v)
1274 Stop -> aft v >> return Stop
1276-- | Run the alloc action @a -> m c@ with async exceptions disabled but keeping
1277-- blocking operations interruptible (see 'Control.Exception.mask'). Use the
1278-- output @c@ as input to @Unfold m c b@ to generate an output stream.
1280-- @c@ is usually a resource under the state of monad @m@, e.g. a file
1281-- handle, that requires a cleanup after use. The cleanup action @c -> m d@,
1282-- runs whenever the stream ends normally, due to a sync or async exception or
1283-- if it gets garbage collected after a partial lazy evaluation.
1285-- 'bracket' only guarantees that the cleanup action runs, and it runs with
1286-- async exceptions enabled. The action must ensure that it can successfully
1287-- cleanup the resource in the face of sync or async exceptions.
1289-- When the stream ends normally or on a sync exception, cleanup action runs
1290-- immediately in the current thread context, whereas in other cases it runs in
1291-- the GC context, therefore, cleanup may be delayed until the GC gets to run.
1293-- /See also: 'bracket_', 'gbracket'/
1295-- /Inhibits stream fusion/
1297-- /Pre-release/
1298{-# INLINE_NORMAL bracket #-}
1299bracket :: (MonadAsync m, MonadCatch m)
1300 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1301bracket bef aft (Unfold step1 inject1) = Unfold step inject
1303 where
1305 inject x = do
1306 -- Mask asynchronous exceptions to make the execution of 'bef' and
1307 -- the registration of 'aft' atomic. See comment in 'D.gbracketIO'.
1308 (r, ref) <- liftBaseOp_ mask_ $ do
1309 r <- bef x
1310 ref <- newIOFinalizer (aft r)
1311 return (r, ref)
1312 s <- inject1 r
1313 return (s, ref)
1315 {-# INLINE_LATE step #-}
1316 step (st, ref) = do
1317 res <- step1 st `MC.onException` runIOFinalizer ref
1318 case res of
1319 Yield x s -> return $ Yield x (s, ref)
1320 Skip s -> return $ Skip (s, ref)
1321 Stop -> do
1322 runIOFinalizer ref
1323 return Stop
1325-- | When unfolding @Unfold m a b@ if an exception @e@ occurs, unfold @e@ using
1326-- @Unfold m e b@.
1328-- /Inhibits stream fusion/
1330-- /Pre-release/
1331{-# INLINE_NORMAL handle #-}
1332handle :: (MonadCatch m, Exception e)
1333 => Unfold m e b -> Unfold m a b -> Unfold m a b
1334handle exc =
1335 gbracket_ return MC.try (\_ -> return ()) (discardFirst exc)