on bc4dba9 help sign in
1341 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-- $setup
260-- >>> import qualified Streamly.Data.Fold as Fold
261-- >>> import qualified Streamly.Internal.Data.Unfold as Unfold
263-------------------------------------------------------------------------------
264-- Input operations
265-------------------------------------------------------------------------------
267-- | Map an action on the input argument of the 'Unfold'.
268--
269-- @
270-- lmapM f = Unfold.many (Unfold.functionM f)
271-- @
272--
273-- /Since: 0.8.0/
274{-# INLINE_NORMAL lmapM #-}
275lmapM :: Monad m => (a -> m c) -> Unfold m c b -> Unfold m a b
276lmapM f (Unfold ustep uinject) = Unfold ustep (f >=> uinject)
278-- | Supply the seed to an unfold closing the input end of the unfold.
279--
280-- @
281-- supply a = Unfold.lmap (Prelude.const a)
282-- @
283--
284-- /Pre-release/
285--
286{-# INLINE_NORMAL supply #-}
287supply :: a -> Unfold m a b -> Unfold m Void b
288supply a = lmap (Prelude.const a)
290-- | Supply the first component of the tuple to an unfold that accepts a tuple
291-- as a seed resulting in a fold that accepts the second component of the tuple
292-- as a seed.
293--
294-- @
295-- supplyFirst a = Unfold.lmap (a, )
296-- @
297--
298-- /Pre-release/
299--
300{-# INLINE_NORMAL supplyFirst #-}
301supplyFirst :: a -> Unfold m (a, b) c -> Unfold m b c
302supplyFirst a = lmap (a, )
304-- | Supply the second component of the tuple to an unfold that accepts a tuple
305-- as a seed resulting in a fold that accepts the first component of the tuple
306-- as a seed.
307--
308-- @
309-- supplySecond b = Unfold.lmap (, b)
310-- @
311--
312-- /Pre-release/
313--
314{-# INLINE_NORMAL supplySecond #-}
315supplySecond :: b -> Unfold m (a, b) c -> Unfold m a c
316supplySecond b = lmap (, b)
318-- | Convert an 'Unfold' into an unfold accepting a tuple as an argument,
319-- using the argument of the original fold as the second element of tuple and
320-- discarding the first element of the tuple.
321--
322-- @
323-- discardFirst = Unfold.lmap snd
324-- @
325--
326-- /Pre-release/
327--
328{-# INLINE_NORMAL discardFirst #-}
329discardFirst :: Unfold m a b -> Unfold m (c, a) b
330discardFirst = lmap snd
332-- | Convert an 'Unfold' into an unfold accepting a tuple as an argument,
333-- using the argument of the original fold as the first element of tuple and
334-- discarding the second element of the tuple.
335--
336-- @
337-- discardSecond = Unfold.lmap fst
338-- @
339--
340-- /Pre-release/
341--
342{-# INLINE_NORMAL discardSecond #-}
343discardSecond :: Unfold m a b -> Unfold m (a, c) b
344discardSecond = lmap fst
346-- | Convert an 'Unfold' that accepts a tuple as an argument into an unfold
347-- that accepts a tuple with elements swapped.
348--
349-- @
350-- swap = Unfold.lmap Tuple.swap
351-- @
352--
353-- /Pre-release/
354--
355{-# INLINE_NORMAL swap #-}
356swap :: Unfold m (a, c) b -> Unfold m (c, a) b
357swap = lmap Tuple.swap
359-------------------------------------------------------------------------------
360-- Output operations
361-------------------------------------------------------------------------------
363-- | Compose an 'Unfold' and a 'Fold'. Given an @Unfold m a b@ and a
364-- @Fold m b c@, returns a monadic action @a -> m c@ representing the
365-- application of the fold on the unfolded stream.
366--
367-- >>> Unfold.fold Fold.sum Unfold.fromList [1..100]
368-- 5050
369--
370-- /Pre-release/
371--
372{-# INLINE_NORMAL fold #-}
373fold :: Monad m => Fold m b c -> Unfold m a b -> a -> m c
374fold (Fold fstep initial extract) (Unfold ustep inject) a = do
375 res <- initial
376 case res of
377 FL.Partial x -> inject a >>= go SPEC x
378 FL.Done b -> return b
380 where
382 {-# INLINE_LATE go #-}
383 go !_ !fs st = do
384 r <- ustep st
385 case r of
386 Yield x s -> do
387 res <- fstep fs x
388 case res of
389 FL.Partial fs1 -> go SPEC fs1 s
390 FL.Done c -> return c
391 Skip s -> go SPEC fs s
392 Stop -> extract fs
394-- | Apply a monadic function to each element of the stream and replace it
395-- with the output of the resulting action.
396--
397-- /Since: 0.8.0/
398--
399{-# INLINE_NORMAL mapM #-}
400mapM :: Monad m => (b -> m c) -> Unfold m a b -> Unfold m a c
401mapM f (Unfold ustep uinject) = Unfold step uinject
402 where
403 {-# INLINE_LATE step #-}
404 step st = do
405 r <- ustep st
406 case r of
407 Yield x s -> f x >>= \a -> return $ Yield a s
408 Skip s -> return $ Skip s
409 Stop -> return Stop
411{-# INLINE_NORMAL mapMWithInput #-}
412mapMWithInput :: Monad m => (a -> b -> m c) -> Unfold m a b -> Unfold m a c
413mapMWithInput f (Unfold ustep uinject) = Unfold step inject
414 where
415 inject a = do
416 r <- uinject a
417 return (a, r)
419 {-# INLINE_LATE step #-}
420 step (inp, st) = do
421 r <- ustep st
422 case r of
423 Yield x s -> f inp x >>= \a -> return $ Yield a (inp, s)
424 Skip s -> return $ Skip (inp, s)
425 Stop -> return Stop
427-------------------------------------------------------------------------------
428-- Convert streams into unfolds
429-------------------------------------------------------------------------------
431{-# INLINE_NORMAL fromStreamD #-}
432fromStreamD :: Monad m => Unfold m (Stream m a) a
433fromStreamD = Unfold step return
434 where
436 {-# INLINE_LATE step #-}
437 step (UnStream step1 state1) = do
438 r <- step1 defState state1
439 return $ case r of
440 Yield x s -> Yield x (Stream step1 s)
441 Skip s -> Skip (Stream step1 s)
442 Stop -> Stop
444{-# INLINE_NORMAL fromStreamK #-}
445fromStreamK :: Monad m => Unfold m (K.Stream m a) a
446fromStreamK = Unfold step return
448 where
450 {-# INLINE_LATE step #-}
451 step stream = do
452 r <- K.uncons stream
453 return $ case r of
454 Just (x, xs) -> Yield x xs
455 Nothing -> Stop
457-- XXX Using Unfold.fromStreamD seems to be faster (using cross product test
458-- case) than using fromStream even if it is implemented using fromStreamD.
459--
460-- | Convert a stream into an 'Unfold'. Note that a stream converted to an
461-- 'Unfold' may not be as efficient as an 'Unfold' in some situations.
462--
463-- /Since: 0.8.0/
464--
465{-# INLINE_NORMAL fromStream #-}
466fromStream :: (K.IsStream t, Monad m) => Unfold m (t m a) a
467fromStream = lmap K.toStream fromStreamK
469-------------------------------------------------------------------------------
470-- Unfolds
471-------------------------------------------------------------------------------
473-- | Lift a monadic function into an unfold generating a nil stream with a side
474-- effect.
475--
476{-# INLINE nilM #-}
477nilM :: Monad m => (a -> m c) -> Unfold m a b
478nilM f = Unfold step return
479 where
480 {-# INLINE_LATE step #-}
481 step x = f x >> return Stop
483-- | Prepend a monadic single element generator function to an 'Unfold'. The
484-- same seed is used in the action as well as the unfold.
485--
486-- /Pre-release/
487{-# INLINE_NORMAL consM #-}
488consM :: Monad m => (a -> m b) -> Unfold m a b -> Unfold m a b
489consM action unf = Unfold step inject
491 where
493 inject = return . Left
495 {-# INLINE_LATE step #-}
496 step (Left a) =
497 action a >>= \r -> return $ Yield r (Right (D.unfold unf a))
498 step (Right (UnStream step1 st)) = do
499 res <- step1 defState st
500 case res of
501 Yield x s -> return $ Yield x (Right (Stream step1 s))
502 Skip s -> return $ Skip (Right (Stream step1 s))
503 Stop -> return Stop
505-- | Convert a list of pure values to a 'Stream'
506--
507-- /Since: 0.8.0/
508--
509{-# INLINE_LATE fromList #-}
510fromList :: Monad m => Unfold m [a] a
511fromList = Unfold step inject
512 where
513 inject = return
514 {-# INLINE_LATE step #-}
515 step (x:xs) = return $ Yield x xs
516 step [] = return Stop
518-- | Convert a list of monadic values to a 'Stream'
519--
520-- /Since: 0.8.0/
521--
522{-# INLINE_LATE fromListM #-}
523fromListM :: Monad m => Unfold m [m a] a
524fromListM = Unfold step inject
525 where
526 inject = return
527 {-# INLINE_LATE step #-}
528 step (x:xs) = x >>= \r -> return $ Yield r xs
529 step [] = return Stop
531------------------------------------------------------------------------------
532-- Specialized Generation
533------------------------------------------------------------------------------
535-- | Generates a stream replicating the seed @n@ times.
536--
537-- /Since: 0.8.0/
538--
539{-# INLINE replicateM #-}
540replicateM :: Monad m => Int -> Unfold m (m a) a
541replicateM n = Unfold step inject
543 where
545 inject x = return (x, n)
547 {-# INLINE_LATE step #-}
548 step (x, i) =
549 if i <= 0
550 then return Stop
551 else do
552 x1 <- x
553 return $ Yield x1 (x, i - 1)
555-- | Generates an infinite stream repeating the seed.
556--
557-- /Since: 0.8.0/
558--
559{-# INLINE repeatM #-}
560repeatM :: Monad m => Unfold m (m a) a
561repeatM = Unfold step return
562 where
563 {-# INLINE_LATE step #-}
564 step x = x >>= \x1 -> return $ Yield x1 x
566-- | Generates an infinite stream starting with the given seed and applying the
567-- given function repeatedly.
568--
569-- /Since: 0.8.0/
570--
571{-# INLINE iterateM #-}
572iterateM :: Monad m => (a -> m a) -> Unfold m (m a) a
573iterateM f = Unfold step id
574 where
575 {-# INLINE_LATE step #-}
576 step x = do
577 fx <- f x
578 return $ Yield x fx
580-- | @fromIndicesM gen@ generates an infinite stream of values using @gen@
581-- starting from the seed.
582--
583-- @
584-- fromIndicesM f = Unfold.mapM f $ Unfold.enumerateFrom 0
585-- @
586--
587-- /Pre-release/
588--
589{-# INLINE_NORMAL fromIndicesM #-}
590fromIndicesM :: Monad m => (Int -> m a) -> Unfold m Int a
591fromIndicesM gen = Unfold step return
592 where
593 {-# INLINE_LATE step #-}
594 step i = do
595 x <- gen i
596 return $ Yield x (i + 1)
598-------------------------------------------------------------------------------
599-- Filtering
600-------------------------------------------------------------------------------
602-- |
603-- >>> u = Unfold.take 2 Unfold.fromList
604-- >>> Unfold.fold Fold.toList u [1..100]
605-- [1,2]
606--
607-- /Since: 0.8.0/
608--
609{-# INLINE_NORMAL take #-}
610take :: Monad m => Int -> Unfold m a b -> Unfold m a b
611take n (Unfold step1 inject1) = Unfold step inject
612 where
613 inject x = do
614 s <- inject1 x
615 return (s, 0)
616 {-# INLINE_LATE step #-}
617 step (st, i) | i < n = do
618 r <- step1 st
619 return $ case r of
620 Yield x s -> Yield x (s, i + 1)
621 Skip s -> Skip (s, i)
622 Stop -> Stop
623 step (_, _) = return Stop
625-- | Same as 'takeWhile' but with a monadic predicate.
626--
627-- /Since: 0.8.0/
628--
629{-# INLINE_NORMAL takeWhileM #-}
630takeWhileM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
631takeWhileM f (Unfold step1 inject1) = Unfold step inject1
632 where
633 {-# INLINE_LATE step #-}
634 step st = do
635 r <- step1 st
636 case r of
637 Yield x s -> do
638 b <- f x
639 return $ if b then Yield x s else Stop
640 Skip s -> return $ Skip s
641 Stop -> return Stop
643-- | End the stream generated by the 'Unfold' as soon as the predicate fails
644-- on an element.
645--
646-- /Since: 0.8.0/
647--
648{-# INLINE takeWhile #-}
649takeWhile :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
650takeWhile f = takeWhileM (return . f)
652-- | Same as 'filter' but with a monadic predicate.
653--
654-- /Since: 0.8.0/
655--
656{-# INLINE_NORMAL filterM #-}
657filterM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
658filterM f (Unfold step1 inject1) = Unfold step inject1
659 where
660 {-# INLINE_LATE step #-}
661 step st = do
662 r <- step1 st
663 case r of
664 Yield x s -> do
665 b <- f x
666 return $ if b then Yield x s else Skip s
667 Skip s -> return $ Skip s
668 Stop -> return Stop
670-- | Include only those elements that pass a predicate.
671--
672-- /Since: 0.8.0/
673--
674{-# INLINE filter #-}
675filter :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
676filter f = filterM (return . f)
678-- | @drop n unf@ drops @n@ elements from the stream generated by @unf@.
679--
680-- /Since: 0.8.0/
681--
682{-# INLINE_NORMAL drop #-}
683drop :: Monad m => Int -> Unfold m a b -> Unfold m a b
684drop n (Unfold step inject) = Unfold step' inject'
686 where
688 inject' a = do
689 b <- inject a
690 return (b, n)
692 {-# INLINE_LATE step' #-}
693 step' (st, i)
694 | i > 0 = do
695 r <- step st
696 return
697 $ case r of
698 Yield _ s -> Skip (s, i - 1)
699 Skip s -> Skip (s, i)
700 Stop -> Stop
701 | otherwise = do
702 r <- step st
703 return
704 $ case r of
705 Yield x s -> Yield x (s, 0)
706 Skip s -> Skip (s, 0)
707 Stop -> Stop
709-- | @dropWhileM f unf@ drops elements from the stream generated by @unf@ while
710-- the condition holds true. The condition function @f@ is /monadic/ in nature.
711--
712-- /Since: 0.8.0/
713--
714{-# INLINE_NORMAL dropWhileM #-}
715dropWhileM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
716dropWhileM f (Unfold step inject) = Unfold step' inject'
718 where
720 inject' a = do
721 b <- inject a
722 return $ Left b
724 {-# INLINE_LATE step' #-}
725 step' (Left st) = do
726 r <- step st
727 case r of
728 Yield x s -> do
729 b <- f x
730 return
731 $ if b
732 then Skip (Left s)
733 else Yield x (Right s)
734 Skip s -> return $ Skip (Left s)
735 Stop -> return Stop
736 step' (Right st) = do
737 r <- step st
738 return
739 $ case r of
740 Yield x s -> Yield x (Right s)
741 Skip s -> Skip (Right s)
742 Stop -> Stop
744-- | Similar to 'dropWhileM' but with a pure condition function.
745--
746-- /Since: 0.8.0/
747--
748{-# INLINE dropWhile #-}
749dropWhile :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
750dropWhile f = dropWhileM (return . f)
752------------------------------------------------------------------------------
753-- Enumeration of Num
754------------------------------------------------------------------------------
756-- | Generate an infinite stream starting from a starting value with increments
757-- of the given stride. The implementation is numerically stable for floating
758-- point values.
759--
760-- Note 'enumerateFromStepIntegral' is faster for integrals.
761--
762-- /Pre-release/
763--
764{-# INLINE enumerateFromStepNum #-}
765enumerateFromStepNum :: (Monad m, Num a) => a -> Unfold m a a
766enumerateFromStepNum stride = Unfold step inject
768 where
770 inject !from = return (from, 0)
772 {-# INLINE_LATE step #-}
773 step (from, !i) = return $ (Yield $! (from + i * stride)) $! (from, i + 1)
775-- | @numFrom = enumerateFromStepNum 1@
776--
777-- /Pre-release/
778--
779{-# INLINE_NORMAL numFrom #-}
780numFrom :: (Monad m, Num a) => Unfold m a a
781numFrom = enumerateFromStepNum 1
783------------------------------------------------------------------------------
784-- Enumeration of Integrals
785------------------------------------------------------------------------------
787-- | Can be used to enumerate unbounded integrals. This does not check for
788-- overflow or underflow for bounded integrals.
789{-# INLINE_NORMAL enumerateFromStepIntegral #-}
790enumerateFromStepIntegral :: (Integral a, Monad m) => Unfold m (a, a) a
791enumerateFromStepIntegral = Unfold step inject
792 where
793 inject (from, stride) = from `seq` stride `seq` return (from, stride)
794 {-# INLINE_LATE step #-}
795 step (x, stride) = return $ Yield x $! (x + stride, stride)
797-- We are assuming that "to" is constrained by the type to be within
798-- max/min bounds.
799{-# INLINE enumerateFromToIntegral #-}
800enumerateFromToIntegral :: (Monad m, Integral a) => a -> Unfold m a a
801enumerateFromToIntegral to =
802 takeWhile (<= to) $ supplySecond 1 enumerateFromStepIntegral
804{-# INLINE enumerateFromIntegral #-}
805enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => Unfold m a a
806enumerateFromIntegral = enumerateFromToIntegral maxBound
808------------------------------------------------------------------------------
809-- Enumeration of Fractionals
810------------------------------------------------------------------------------
812-- | /Internal/
813--
814-- > enumerateFromToFractional to = takeWhile (<= to + 1 / 2) $ enumerateFromStepNum 1
815--
816{-# INLINE_NORMAL enumerateFromToFractional #-}
817enumerateFromToFractional :: (Monad m, Fractional a, Ord a) => a -> Unfold m a a
818enumerateFromToFractional to =
819 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum 1
821-------------------------------------------------------------------------------
822-- Generation from SVar
823-------------------------------------------------------------------------------
825data FromSVarState t m a =
826 FromSVarInit (SVar t m a)
827 | FromSVarRead (SVar t m a)
828 | FromSVarLoop (SVar t m a) [ChildEvent a]
829 | FromSVarDone (SVar t m a)
831-- | /Internal/
832--
833{-# INLINE_NORMAL fromSVar #-}
834fromSVar :: MonadAsync m => Unfold m (SVar t m a) a
835fromSVar = Unfold step (return . FromSVarInit)
836 where
838 {-# INLINE_LATE step #-}
839 step (FromSVarInit svar) = do
840 ref <- liftIO $ newIORef ()
841 _ <- liftIO $ mkWeakIORef ref hook
842 -- when this copy of svar gets garbage collected "ref" will get
843 -- garbage collected and our GC hook will be called.
844 let sv = svar{svarRef = Just ref}
845 return $ Skip (FromSVarRead sv)
847 where
849 {-# NOINLINE hook #-}
850 hook = do
851 when (svarInspectMode svar) $ do
852 r <- liftIO $ readIORef (svarStopTime (svarStats svar))
853 when (isNothing r) $
854 printSVar svar "SVar Garbage Collected"
855 cleanupSVar svar
856 -- If there are any SVars referenced by this SVar a GC will prompt
857 -- them to be cleaned up quickly.
858 when (svarInspectMode svar) performMajorGC
860 step (FromSVarRead sv) = do
861 list <- readOutputQ sv
862 -- Reversing the output is important to guarantee that we process the
863 -- outputs in the same order as they were generated by the constituent
864 -- streams.
865 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
867 step (FromSVarLoop sv []) = do
868 done <- postProcess sv
869 return $ Skip $ if done
870 then FromSVarDone sv
871 else FromSVarRead sv
873 step (FromSVarLoop sv (ev : es)) = do
874 case ev of
875 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
876 ChildStop tid e -> do
877 accountThread sv tid
878 case e of
879 Nothing -> do
880 stop <- shouldStop tid
881 if stop
882 then do
883 liftIO (cleanupSVar sv)
884 return $ Skip (FromSVarDone sv)
885 else return $ Skip (FromSVarLoop sv es)
886 Just ex ->
887 case fromException ex of
888 Just ThreadAbort ->
889 return $ Skip (FromSVarLoop sv es)
890 Nothing -> liftIO (cleanupSVar sv) >> MC.throwM ex
891 where
893 shouldStop tid =
894 case svarStopStyle sv of
895 StopNone -> return False
896 StopAny -> return True
897 StopBy -> do
898 sid <- liftIO $ readIORef (svarStopBy sv)
899 return $ tid == sid
901 step (FromSVarDone sv) = do
902 when (svarInspectMode sv) $ do
903 t <- liftIO $ getTime Monotonic
904 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
905 liftIO $ printSVar sv "SVar Done"
906 return Stop
908-------------------------------------------------------------------------------
909-- Process events received by a fold consumer from a stream producer
910-------------------------------------------------------------------------------
912-- XXX Have an error for "FromSVarInit" instead of undefined. The error can
913-- redirect the user to report the failure to the developers.
914-- | /Internal/
915--
916{-# INLINE_NORMAL fromProducer #-}
917fromProducer :: MonadAsync m => Unfold m (SVar t m a) a
918fromProducer = Unfold step (return . FromSVarRead)
919 where
921 {-# INLINE_LATE step #-}
922 step (FromSVarRead sv) = do
923 list <- readOutputQ sv
924 -- Reversing the output is important to guarantee that we process the
925 -- outputs in the same order as they were generated by the constituent
926 -- streams.
927 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
929 step (FromSVarLoop sv []) = return $ Skip $ FromSVarRead sv
930 step (FromSVarLoop sv (ev : es)) = do
931 case ev of
932 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
933 ChildStop tid e -> do
934 accountThread sv tid
935 case e of
936 Nothing -> do
937 sendStopToProducer sv
938 return $ Skip (FromSVarDone sv)
939 Just _ -> error "Bug: fromProducer: received exception"
941 step (FromSVarDone sv) = do
942 when (svarInspectMode sv) $ do
943 t <- liftIO $ getTime Monotonic
944 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
945 liftIO $ printSVar sv "SVar Done"
946 return Stop
948 step (FromSVarInit _) = undefined
950------------------------------------------------------------------------------
951-- Exceptions
952------------------------------------------------------------------------------
954-- | Like 'gbracket' but with following differences:
955--
956-- * alloc action @a -> m c@ runs with async exceptions enabled
957-- * cleanup action @c -> m d@ won't run if the stream is garbage collected
958-- after partial evaluation.
959-- * does not require a 'MonadAsync' constraint.
960--
961-- /Inhibits stream fusion/
962--
963-- /Pre-release/
964--
965{-# INLINE_NORMAL gbracket_ #-}
966gbracket_
967 :: Monad m
968 => (a -> m c) -- ^ before
969 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
970 -> (c -> m d) -- ^ after, on normal stop
971 -> Unfold m (c, e) b -- ^ on exception
972 -> Unfold m c b -- ^ unfold to run
973 -> Unfold m a b
974gbracket_ bef exc aft (Unfold estep einject) (Unfold step1 inject1) =
975 Unfold step inject
977 where
979 inject x = do
980 r <- bef x
981 s <- inject1 r
982 return $ Right (s, r)
984 {-# INLINE_LATE step #-}
985 step (Right (st, v)) = do
986 res <- exc $ step1 st
987 case res of
988 Right r -> case r of
989 Yield x s -> return $ Yield x (Right (s, v))
990 Skip s -> return $ Skip (Right (s, v))
991 Stop -> aft v >> return Stop
992 -- XXX Do not handle async exceptions, just rethrow them.
993 Left e -> do
994 r <- einject (v, e)
995 return $ Skip (Left r)
996 step (Left st) = do
997 res <- estep st
998 case res of
999 Yield x s -> return $ Yield x (Left s)
1000 Skip s -> return $ Skip (Left s)
1001 Stop -> return Stop
1003-- | Run the alloc action @a -> m c@ with async exceptions disabled but keeping
1004-- blocking operations interruptible (see 'Control.Exception.mask'). Use the
1005-- output @c@ as input to @Unfold m c b@ to generate an output stream. When
1006-- unfolding use the supplied @try@ operation @forall s. m s -> m (Either e s)@
1007-- to catch synchronous exceptions. If an exception occurs run the exception
1008-- handling unfold @Unfold m (c, e) b@.
1010-- The cleanup action @c -> m d@, runs whenever the stream ends normally, due
1011-- to a sync or async exception or if it gets garbage collected after a partial
1012-- lazy evaluation. See 'bracket' for the semantics of the cleanup action.
1014-- 'gbracket' can express all other exception handling combinators.
1016-- /Inhibits stream fusion/
1018-- /Pre-release/
1019{-# INLINE_NORMAL gbracket #-}
1020gbracket
1021 :: (MonadIO m, MonadBaseControl IO m)
1022 => (a -> m c) -- ^ before
1023 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
1024 -> (c -> m d) -- ^ after, on normal stop, or GC
1025 -> Unfold m (c, e) b -- ^ on exception
1026 -> Unfold m c b -- ^ unfold to run
1027 -> Unfold m a b
1028gbracket bef exc aft (Unfold estep einject) (Unfold step1 inject1) =
1029 Unfold step inject
1031 where
1033 inject x = do
1034 -- Mask asynchronous exceptions to make the execution of 'bef' and
1035 -- the registration of 'aft' atomic. See comment in 'D.gbracketIO'.
1036 (r, ref) <- liftBaseOp_ mask_ $ do
1037 r <- bef x
1038 ref <- newIOFinalizer (aft r)
1039 return (r, ref)
1040 s <- inject1 r
1041 return $ Right (s, r, ref)
1043 {-# INLINE_LATE step #-}
1044 step (Right (st, v, ref)) = do
1045 res <- exc $ step1 st
1046 case res of
1047 Right r -> case r of
1048 Yield x s -> return $ Yield x (Right (s, v, ref))
1049 Skip s -> return $ Skip (Right (s, v, ref))
1050 Stop -> do
1051 runIOFinalizer ref
1052 return Stop
1053 -- XXX Do not handle async exceptions, just rethrow them.
1054 Left e -> do
1055 -- Clearing of finalizer and running of exception handler must
1056 -- be atomic wrt async exceptions. Otherwise if we have cleared
1057 -- the finalizer and have not run the exception handler then we
1058 -- may leak the resource.
1059 r <- clearingIOFinalizer ref (einject (v, e))
1060 return $ Skip (Left r)
1061 step (Left st) = do
1062 res <- estep st
1063 case res of
1064 Yield x s -> return $ Yield x (Left s)
1065 Skip s -> return $ Skip (Left s)
1066 Stop -> return Stop
1068-- | Run a side effect @a -> m c@ on the input @a@ before unfolding it using
1069-- @Unfold m a b@.
1071-- > before f = lmapM (\a -> f a >> return a)
1073-- /Pre-release/
1074{-# INLINE_NORMAL before #-}
1075before :: (a -> m c) -> Unfold m a b -> Unfold m a b
1076before action (Unfold step inject) = Unfold step (action >> inject)
1078-- The custom implementation of "after_" is slightly faster (5-7%) than
1079-- "_after". This is just to document and make sure that we can always use
1080-- gbracket to implement after_ The same applies to other combinators as well.
1082{-# INLINE_NORMAL _after #-}
1083_after :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
1084_after aft = gbracket_ return (fmap Right) aft undefined
1086-- | Like 'after' with following differences:
1088-- * action @a -> m c@ won't run if the stream is garbage collected
1089-- after partial evaluation.
1090-- * Monad @m@ does not require any other constraints.
1092-- /Pre-release/
1093{-# INLINE_NORMAL after_ #-}
1094after_ :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
1095after_ action (Unfold step1 inject1) = Unfold step inject
1097 where
1099 inject x = do
1100 s <- inject1 x
1101 return (s, x)
1103 {-# INLINE_LATE step #-}
1104 step (st, v) = do
1105 res <- step1 st
1106 case res of
1107 Yield x s -> return $ Yield x (s, v)
1108 Skip s -> return $ Skip (s, v)
1109 Stop -> action v >> return Stop
1111-- | Unfold the input @a@ using @Unfold m a b@, run an action on @a@ whenever
1112-- the unfold stops normally, or if it is garbage collected after a partial
1113-- lazy evaluation.
1115-- The semantics of the action @a -> m c@ are similar to the cleanup action
1116-- semantics in 'bracket'.
1118-- /See also 'after_'/
1120-- /Pre-release/
1121{-# INLINE_NORMAL after #-}
1122after :: (MonadIO m, MonadBaseControl IO m)
1123 => (a -> m c) -> Unfold m a b -> Unfold m a b
1124after action (Unfold step1 inject1) = Unfold step inject
1126 where
1128 inject x = do
1129 s <- inject1 x
1130 ref <- newIOFinalizer (action x)
1131 return (s, ref)
1133 {-# INLINE_LATE step #-}
1134 step (st, ref) = do
1135 res <- step1 st
1136 case res of
1137 Yield x s -> return $ Yield x (s, ref)
1138 Skip s -> return $ Skip (s, ref)
1139 Stop -> do
1140 runIOFinalizer ref
1141 return Stop
1143{-# INLINE_NORMAL _onException #-}
1144_onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
1145_onException action =
1146 gbracket_ return MC.try
1147 (\_ -> return ())
1148 (nilM (\(a, e :: MC.SomeException) -> action a >> MC.throwM e))
1150-- | Unfold the input @a@ using @Unfold m a b@, run the action @a -> m c@ on
1151-- @a@ if the unfold aborts due to an exception.
1153-- /Inhibits stream fusion/
1155-- /Pre-release/
1156{-# INLINE_NORMAL onException #-}
1157onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
1158onException action (Unfold step1 inject1) = Unfold step inject
1160 where
1162 inject x = do
1163 s <- inject1 x
1164 return (s, x)
1166 {-# INLINE_LATE step #-}
1167 step (st, v) = do
1168 res <- step1 st `MC.onException` action v
1169 case res of
1170 Yield x s -> return $ Yield x (s, v)
1171 Skip s -> return $ Skip (s, v)
1172 Stop -> return Stop
1174{-# INLINE_NORMAL _finally #-}
1175_finally :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
1176_finally action =
1177 gbracket_ return MC.try action
1178 (nilM (\(a, e :: MC.SomeException) -> action a >> MC.throwM e))
1180-- | Like 'finally' with following differences:
1182-- * action @a -> m c@ won't run if the stream is garbage collected
1183-- after partial evaluation.
1184-- * does not require a 'MonadAsync' constraint.
1186-- /Inhibits stream fusion/
1188-- /Pre-release/
1189{-# INLINE_NORMAL finally_ #-}
1190finally_ :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
1191finally_ action (Unfold step1 inject1) = Unfold step inject
1193 where
1195 inject x = do
1196 s <- inject1 x
1197 return (s, x)
1199 {-# INLINE_LATE step #-}
1200 step (st, v) = do
1201 res <- step1 st `MC.onException` action v
1202 case res of
1203 Yield x s -> return $ Yield x (s, v)
1204 Skip s -> return $ Skip (s, v)
1205 Stop -> action v >> return Stop
1207-- | Unfold the input @a@ using @Unfold m a b@, run an action on @a@ whenever
1208-- the unfold stops normally, aborts due to an exception or if it is garbage
1209-- collected after a partial lazy evaluation.
1211-- The semantics of the action @a -> m c@ are similar to the cleanup action
1212-- semantics in 'bracket'.
1214-- @
1215-- finally release = bracket return release
1216-- @
1218-- /See also 'finally_'/
1220-- /Inhibits stream fusion/
1222-- /Pre-release/
1223{-# INLINE_NORMAL finally #-}
1224finally :: (MonadAsync m, MonadCatch m)
1225 => (a -> m c) -> Unfold m a b -> Unfold m a b
1226finally action (Unfold step1 inject1) = Unfold step inject
1228 where
1230 inject x = do
1231 s <- inject1 x
1232 ref <- newIOFinalizer (action x)
1233 return (s, ref)
1235 {-# INLINE_LATE step #-}
1236 step (st, ref) = do
1237 res <- step1 st `MC.onException` runIOFinalizer ref
1238 case res of
1239 Yield x s -> return $ Yield x (s, ref)
1240 Skip s -> return $ Skip (s, ref)
1241 Stop -> do
1242 runIOFinalizer ref
1243 return Stop
1245{-# INLINE_NORMAL _bracket #-}
1246_bracket :: MonadCatch m
1247 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1248_bracket bef aft =
1249 gbracket_ bef MC.try aft (nilM (\(a, e :: MC.SomeException) -> aft a >>
1250 MC.throwM e))
1252-- | Like 'bracket' but with following differences:
1254-- * alloc action @a -> m c@ runs with async exceptions enabled
1255-- * cleanup action @c -> m d@ won't run if the stream is garbage collected
1256-- after partial evaluation.
1257-- * does not require a 'MonadAsync' constraint.
1259-- /Inhibits stream fusion/
1261-- /Pre-release/
1262{-# INLINE_NORMAL bracket_ #-}
1263bracket_ :: MonadCatch m
1264 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1265bracket_ bef aft (Unfold step1 inject1) = Unfold step inject
1267 where
1269 inject x = do
1270 r <- bef x
1271 s <- inject1 r
1272 return (s, r)
1274 {-# INLINE_LATE step #-}
1275 step (st, v) = do
1276 res <- step1 st `MC.onException` aft v
1277 case res of
1278 Yield x s -> return $ Yield x (s, v)
1279 Skip s -> return $ Skip (s, v)
1280 Stop -> aft v >> return Stop
1282-- | Run the alloc action @a -> m c@ with async exceptions disabled but keeping
1283-- blocking operations interruptible (see 'Control.Exception.mask'). Use the
1284-- output @c@ as input to @Unfold m c b@ to generate an output stream.
1286-- @c@ is usually a resource under the state of monad @m@, e.g. a file
1287-- handle, that requires a cleanup after use. The cleanup action @c -> m d@,
1288-- runs whenever the stream ends normally, due to a sync or async exception or
1289-- if it gets garbage collected after a partial lazy evaluation.
1291-- 'bracket' only guarantees that the cleanup action runs, and it runs with
1292-- async exceptions enabled. The action must ensure that it can successfully
1293-- cleanup the resource in the face of sync or async exceptions.
1295-- When the stream ends normally or on a sync exception, cleanup action runs
1296-- immediately in the current thread context, whereas in other cases it runs in
1297-- the GC context, therefore, cleanup may be delayed until the GC gets to run.
1299-- /See also: 'bracket_', 'gbracket'/
1301-- /Inhibits stream fusion/
1303-- /Pre-release/
1304{-# INLINE_NORMAL bracket #-}
1305bracket :: (MonadAsync m, MonadCatch m)
1306 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1307bracket bef aft (Unfold step1 inject1) = Unfold step inject
1309 where
1311 inject x = do
1312 -- Mask asynchronous exceptions to make the execution of 'bef' and
1313 -- the registration of 'aft' atomic. See comment in 'D.gbracketIO'.
1314 (r, ref) <- liftBaseOp_ mask_ $ do
1315 r <- bef x
1316 ref <- newIOFinalizer (aft r)
1317 return (r, ref)
1318 s <- inject1 r
1319 return (s, ref)
1321 {-# INLINE_LATE step #-}
1322 step (st, ref) = do
1323 res <- step1 st `MC.onException` runIOFinalizer ref
1324 case res of
1325 Yield x s -> return $ Yield x (s, ref)
1326 Skip s -> return $ Skip (s, ref)
1327 Stop -> do
1328 runIOFinalizer ref
1329 return Stop
1331-- | When unfolding @Unfold m a b@ if an exception @e@ occurs, unfold @e@ using
1332-- @Unfold m e b@.
1334-- /Inhibits stream fusion/
1336-- /Pre-release/
1337{-# INLINE_NORMAL handle #-}
1338handle :: (MonadCatch m, Exception e)
1339 => Unfold m e b -> Unfold m a b -> Unfold m a b
1340handle exc =
1341 gbracket_ return MC.try (\_ -> return ()) (discardFirst exc)