git.iko.soy/streamly
on af62054 help sign in
streamly/src/Streamly/Internal/Data/Unfold.hs
1411 lines 45 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-- Streams forcing a closed control flow loop can be categorized under
12-- two types, unfolds and folds, both of these are duals of each other.
13--
14-- Unfold streams are really generators of a sequence of elements, we can also
15-- call them pull style streams. These are lazy producers of streams. On each
16-- evaluation the producer generates the next element. A consumer can
17-- therefore pull elements from the stream whenever it wants to. A stream
18-- consumer can multiplex pull streams by pulling elements from the chosen
19-- streams, therefore, pull streams allow merging or multiplexing. On the
20-- other hand, with this representation we cannot split or demultiplex a
21-- stream. So really these are stream sources that can be generated from a
22-- seed and can be merged or zipped into a single stream.
23--
24-- The dual of Unfolds are Folds. Folds can also be called as push style
25-- streams or reducers. These are strict consumers of streams. We keep pushing
26-- elements to a fold and we can extract the result at any point. A driver can
27-- choose which fold to push to and can also push the same element to multiple
28-- folds. Therefore, folds allow splitting or demultiplexing a stream. On the
29-- other hand, we cannot merge streams using this representation. So really
30-- these are stream consumers that reduce the stream to a single value, these
31-- consumers can be composed such that a stream can be split over multiple
32-- consumers.
33--
34-- Performance:
35--
36-- Composing a tree or graph of computations with unfolds can be much more
37-- efficient compared to composing with the Monad instance. The reason is that
38-- unfolds allow the compiler to statically know the state and optimize it
39-- using stream fusion whereas it is not possible with the monad bind because
40-- the state is determined dynamically.
42-- Open control flow style streams can also have two representations. StreamK
43-- is a producer style representation. We can also have a consumer style
44-- representation. We can use that for composable folds in StreamK
45-- representation.
46--
47module Streamly.Internal.Data.Unfold
48 (
49 -- * Unfold Type
50 Unfold
52 -- * Operations on Input
53 , lmap
54 , lmapM
55 , supply
56 , supplyFirst
57 , supplySecond
58 , discardFirst
59 , discardSecond
60 , swap
61 -- coapply
62 -- comonad
64 -- * Operations on Output
65 , fold
66 -- pipe
68 -- * Unfolds
69 , fromStream
70 , fromStreamK
71 , fromStreamD
72 , nilM
73 , consM
74 , effect
75 , singletonM
76 , singleton
77 , identity
78 , const
79 , unfoldrM
81 , fromList
82 , fromListM
84 , fromSVar
85 , fromProducer
87 -- ** Specialized Generation
88 -- | Generate a monadic stream from a seed.
89 , replicateM
90 , repeatM
91 , iterateM
92 , fromIndicesM
94 -- ** Enumerations
95 , enumerateFromStepIntegral
96 , enumerateFromToIntegral
97 , enumerateFromIntegral
99 , enumerateFromStepNum
100 , numFrom
101 , enumerateFromToFractional
103 -- * Transformations
104 , map
105 , mapM
106 , mapMWithInput
108 -- * Filtering
109 , takeWhileM
110 , takeWhile
111 , take
112 , filter
113 , filterM
114 , drop
115 , dropWhile
116 , dropWhileM
118 -- * Zipping
119 , zipWithM
120 , zipWith
121 , teeZipWith
123 -- * Nesting
124 , ConcatState (..)
125 , concat
126 , concatMapM
127 , outerProduct
128 , ap
129 , apDiscardFst
130 , apDiscardSnd
132 -- * Exceptions
133 , gbracket_
134 , gbracket
135 , before
136 , after_
137 , after
138 , onException
139 , finally_
140 , finally
141 , bracket_
142 , bracket
143 , handle
144 )
145where
147import Control.Exception (Exception, fromException, mask_)
148import Control.Monad ((>=>), when)
149import Control.Monad.Catch (MonadCatch)
150import Control.Monad.IO.Class (MonadIO(..))
151import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_)
152import Data.Maybe (isNothing)
153import Data.Void (Void)
154import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef)
155import GHC.Types (SPEC(..))
156import Fusion.Plugin.Types (Fuse(..))
157import Streamly.Internal.Data.Stream.StreamD.Type (Stream(..), Step(..))
158import Streamly.Internal.Data.Unfold.Types (Unfold(..))
159import Streamly.Internal.Data.Fold.Types (Fold(..))
160import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
161import System.Mem (performMajorGC)
163import qualified Prelude
164import qualified Control.Monad.Catch as MC
165import qualified Data.Tuple as Tuple
166import qualified Streamly.Internal.Data.Stream.StreamK as K
167import qualified Streamly.Internal.Data.Stream.StreamD as D
169import Streamly.Internal.Data.SVar
170import Prelude
171 hiding (concat, map, mapM, takeWhile, take, filter, const, zipWith
172 , drop, dropWhile)
174-------------------------------------------------------------------------------
175-- Input operations
176-------------------------------------------------------------------------------
178-- | Map a function on the input argument of the 'Unfold'.
179--
180-- @
181-- lmap f = concat (singleton f)
182-- @
183--
184-- /Internal/
185{-# INLINE_NORMAL lmap #-}
186lmap :: (a -> c) -> Unfold m c b -> Unfold m a b
187lmap f (Unfold ustep uinject) = Unfold ustep (uinject . f)
189-- | Map an action on the input argument of the 'Unfold'.
190--
191-- @
192-- lmapM f = concat (singletonM f)
193-- @
194--
195-- /Internal/
196{-# INLINE_NORMAL lmapM #-}
197lmapM :: Monad m => (a -> m c) -> Unfold m c b -> Unfold m a b
198lmapM f (Unfold ustep uinject) = Unfold ustep (f >=> uinject)
200-- XXX change the signature to the following?
201-- supply :: a -> Unfold m a b -> Unfold m Void b
202--
203-- | Supply the seed to an unfold closing the input end of the unfold.
204--
205-- /Internal/
206--
207{-# INLINE_NORMAL supply #-}
208supply :: Unfold m a b -> a -> Unfold m Void b
209supply unf a = lmap (Prelude.const a) unf
211-- XXX change the signature to the following?
212-- supplyFirst :: a -> Unfold m (a, b) c -> Unfold m b c
213--
214-- | Supply the first component of the tuple to an unfold that accepts a tuple
215-- as a seed resulting in a fold that accepts the second component of the tuple
216-- as a seed.
217--
218-- /Internal/
219--
220{-# INLINE_NORMAL supplyFirst #-}
221supplyFirst :: Unfold m (a, b) c -> a -> Unfold m b c
222supplyFirst unf a = lmap (a, ) unf
224-- XXX change the signature to the following?
225-- supplySecond :: b -> Unfold m (a, b) c -> Unfold m a c
226--
227-- | Supply the second component of the tuple to an unfold that accepts a tuple
228-- as a seed resulting in a fold that accepts the first component of the tuple
229-- as a seed.
230--
231-- /Internal/
232--
233{-# INLINE_NORMAL supplySecond #-}
234supplySecond :: Unfold m (a, b) c -> b -> Unfold m a c
235supplySecond unf b = lmap (, b) unf
237-- | Convert an 'Unfold' into an unfold accepting a tuple as an argument,
238-- using the argument of the original fold as the second element of tuple and
239-- discarding the first element of the tuple.
240--
241-- /Internal/
242--
243{-# INLINE_NORMAL discardFirst #-}
244discardFirst :: Unfold m a b -> Unfold m (c, a) b
245discardFirst = lmap snd
247-- | Convert an 'Unfold' into an unfold accepting a tuple as an argument,
248-- using the argument of the original fold as the first element of tuple and
249-- discarding the second element of the tuple.
250--
251-- /Internal/
252--
253{-# INLINE_NORMAL discardSecond #-}
254discardSecond :: Unfold m a b -> Unfold m (a, c) b
255discardSecond = lmap fst
257-- | Convert an 'Unfold' that accepts a tuple as an argument into an unfold
258-- that accepts a tuple with elements swapped.
259--
260-- /Internal/
261--
262{-# INLINE_NORMAL swap #-}
263swap :: Unfold m (a, c) b -> Unfold m (c, a) b
264swap = lmap Tuple.swap
266-------------------------------------------------------------------------------
267-- Output operations
268-------------------------------------------------------------------------------
270-- | Compose an 'Unfold' and a 'Fold'. Given an @Unfold m a b@ and a
271-- @Fold m b c@, returns a monadic action @a -> m c@ representing the
272-- application of the fold on the unfolded stream.
273--
274-- /Internal/
275--
276{-# INLINE_NORMAL fold #-}
277fold :: Monad m => Unfold m a b -> Fold m b c -> a -> m c
278fold (Unfold ustep inject) (Fold fstep initial extract) a =
279 initial >>= \x -> inject a >>= go SPEC x
280 where
281 -- XXX !acc?
282 {-# INLINE_LATE go #-}
283 go !_ acc st = acc `seq` do
284 r <- ustep st
285 case r of
286 Yield x s -> do
287 acc' <- fstep acc x
288 go SPEC acc' s
289 Skip s -> go SPEC acc s
290 Stop -> extract acc
292{-# INLINE_NORMAL map #-}
293map :: Monad m => (b -> c) -> Unfold m a b -> Unfold m a c
294map f (Unfold ustep uinject) = Unfold step uinject
295 where
296 {-# INLINE_LATE step #-}
297 step st = do
298 r <- ustep st
299 return $ case r of
300 Yield x s -> Yield (f x) s
301 Skip s -> Skip s
302 Stop -> Stop
304{-# INLINE_NORMAL mapM #-}
305mapM :: Monad m => (b -> m c) -> Unfold m a b -> Unfold m a c
306mapM f (Unfold ustep uinject) = Unfold step uinject
307 where
308 {-# INLINE_LATE step #-}
309 step st = do
310 r <- ustep st
311 case r of
312 Yield x s -> f x >>= \a -> return $ Yield a s
313 Skip s -> return $ Skip s
314 Stop -> return Stop
316{-# INLINE_NORMAL mapMWithInput #-}
317mapMWithInput :: Monad m => (a -> b -> m c) -> Unfold m a b -> Unfold m a c
318mapMWithInput f (Unfold ustep uinject) = Unfold step inject
319 where
320 inject a = do
321 r <- uinject a
322 return (a, r)
324 {-# INLINE_LATE step #-}
325 step (inp, st) = do
326 r <- ustep st
327 case r of
328 Yield x s -> f inp x >>= \a -> return $ Yield a (inp, s)
329 Skip s -> return $ Skip (inp, s)
330 Stop -> return Stop
332-------------------------------------------------------------------------------
333-- Convert streams into unfolds
334-------------------------------------------------------------------------------
336-- this is horribly slow
337{-# INLINE_NORMAL fromStreamD #-}
338fromStreamD :: Monad m => Unfold m (Stream m a) a
339fromStreamD = Unfold step return
340 where
342 {-# INLINE_LATE step #-}
343 step (Stream step1 state1) = do
344 r <- step1 defState state1
345 return $ case r of
346 Yield x s -> Yield x (Stream step1 s)
347 Skip s -> Skip (Stream step1 s)
348 Stop -> Stop
350{-# INLINE_NORMAL fromStreamK #-}
351fromStreamK :: Monad m => Unfold m (K.Stream m a) a
352fromStreamK = Unfold step return
354 where
356 {-# INLINE_LATE step #-}
357 step stream = do
358 r <- K.uncons stream
359 return $ case r of
360 Just (x, xs) -> Yield x xs
361 Nothing -> Stop
363-- | Convert a stream into an 'Unfold'. Note that a stream converted to an
364-- 'Unfold' may not be as efficient as an 'Unfold' in some situations.
365--
366-- /Internal/
367{-# INLINE_NORMAL fromStream #-}
368fromStream :: (K.IsStream t, Monad m) => Unfold m (t m a) a
369fromStream = lmap K.toStream fromStreamK
371-------------------------------------------------------------------------------
372-- Unfolds
373-------------------------------------------------------------------------------
375-- | Lift a monadic function into an unfold generating a nil stream with a side
376-- effect.
377--
378{-# INLINE nilM #-}
379nilM :: Monad m => (a -> m c) -> Unfold m a b
380nilM f = Unfold step return
381 where
382 {-# INLINE_LATE step #-}
383 step x = f x >> return Stop
385-- | Prepend a monadic single element generator function to an 'Unfold'.
386--
387-- /Internal/
388{-# INLINE_NORMAL consM #-}
389consM :: Monad m => (a -> m b) -> Unfold m a b -> Unfold m a b
390consM action unf = Unfold step inject
392 where
394 inject = return . Left
396 {-# INLINE_LATE step #-}
397 step (Left a) =
398 action a >>= \r -> return $ Yield r (Right (D.unfold unf a))
399 step (Right (UnStream step1 st)) = do
400 res <- step1 defState st
401 case res of
402 Yield x s -> return $ Yield x (Right (Stream step1 s))
403 Skip s -> return $ Skip (Right (Stream step1 s))
404 Stop -> return Stop
406-- | Lift a monadic effect into an unfold generating a singleton stream.
407--
408{-# INLINE effect #-}
409effect :: Monad m => m b -> Unfold m Void b
410effect eff = Unfold step inject
411 where
412 inject _ = return True
413 {-# INLINE_LATE step #-}
414 step True = eff >>= \r -> return $ Yield r False
415 step False = return Stop
417-- XXX change it to yieldM or change yieldM in Prelude to singletonM
418--
419-- | Lift a monadic function into an unfold generating a singleton stream.
420--
421{-# INLINE singletonM #-}
422singletonM :: Monad m => (a -> m b) -> Unfold m a b
423singletonM f = Unfold step inject
424 where
425 inject x = return $ Just x
426 {-# INLINE_LATE step #-}
427 step (Just x) = f x >>= \r -> return $ Yield r Nothing
428 step Nothing = return Stop
430-- | Lift a pure function into an unfold generating a singleton stream.
431--
432{-# INLINE singleton #-}
433singleton :: Monad m => (a -> b) -> Unfold m a b
434singleton f = singletonM $ return . f
436-- | Identity unfold. Generates a singleton stream with the seed as the only
437-- element in the stream.
438--
439-- > identity = singletonM return
440--
441{-# INLINE identity #-}
442identity :: Monad m => Unfold m a a
443identity = singletonM return
445const :: Monad m => m b -> Unfold m a b
446const m = Unfold step inject
447 where
448 inject _ = return ()
449 step () = m >>= \r -> return $ Yield r ()
452-- | Convert a list of pure values to a 'Stream'
453{-# INLINE_LATE fromList #-}
454fromList :: Monad m => Unfold m [a] a
455fromList = Unfold step inject
456 where
457 inject x = return x
458 {-# INLINE_LATE step #-}
459 step (x:xs) = return $ Yield x xs
460 step [] = return Stop
462-- | Convert a list of monadic values to a 'Stream'
463{-# INLINE_LATE fromListM #-}
464fromListM :: Monad m => Unfold m [m a] a
465fromListM = Unfold step inject
466 where
467 inject x = return x
468 {-# INLINE_LATE step #-}
469 step (x:xs) = x >>= \r -> return $ Yield r xs
470 step [] = return Stop
472-- | Build a stream by unfolding a /monadic/ step function starting from a seed.
473-- The step function returns the next element in the stream and the next seed
474-- value. When it is done it returns 'Nothing' and the stream ends.
475--
476-- /Internal/
477--
478{-# INLINE unfoldrM #-}
479unfoldrM :: Monad m => (a -> m (Maybe (b, a))) -> Unfold m a b
480unfoldrM next = Unfold step return
481 where
482 {-# INLINE_LATE step #-}
483 step st = do
484 r <- next st
485 return $ case r of
486 Just (x, s) -> Yield x s
487 Nothing -> Stop
489------------------------------------------------------------------------------
490-- Specialized Generation
491------------------------------------------------------------------------------
493-- | Generates a stream replicating the seed @n@ times.
494--
495-- /Internal/
496--
497{-# INLINE replicateM #-}
498replicateM :: Monad m => Int -> Unfold m (m a) a
499replicateM n = Unfold step inject
501 where
503 inject x = return (x, n)
505 {-# INLINE_LATE step #-}
506 step (x, i) = do
507 x1 <- x
508 return
509 $ if i <= 0
510 then Stop
511 else Yield x1 (x, i - 1)
513-- | Generates an infinite stream repeating the seed.
514--
515-- /Internal/
516--
517{-# INLINE repeatM #-}
518repeatM :: Monad m => Unfold m (m a) a
519repeatM = Unfold step return
520 where
521 {-# INLINE_LATE step #-}
522 step x = x >>= \x1 -> return $ Yield x1 x
524-- | Generates an infinite stream starting with the given seed and applying the
525-- given function repeteadly.
526--
527-- /Internal/
528--
529{-# INLINE iterateM #-}
530iterateM :: Monad m => (a -> m a) -> Unfold m a a
531iterateM f = Unfold step f
532 where
533 {-# INLINE_LATE step #-}
534 step x = do
535 fx <- f x
536 return $ Yield fx fx
538-- | @fromIndicesM gen@ generates an infinite stream of values using @gen@
539-- starting from the seed.
540--
541-- /Internal/
542--
543{-# INLINE_NORMAL fromIndicesM #-}
544fromIndicesM :: Monad m => (Int -> m a) -> Unfold m Int a
545fromIndicesM gen = Unfold step return
546 where
547 {-# INLINE_LATE step #-}
548 step i = do
549 x <- gen i
550 return $ Yield x (i + 1)
552-------------------------------------------------------------------------------
553-- Filtering
554-------------------------------------------------------------------------------
556{-# INLINE_NORMAL take #-}
557take :: Monad m => Int -> Unfold m a b -> Unfold m a b
558take n (Unfold step1 inject1) = Unfold step inject
559 where
560 inject x = do
561 s <- inject1 x
562 return (s, 0)
563 {-# INLINE_LATE step #-}
564 step (st, i) | i < n = do
565 r <- step1 st
566 return $ case r of
567 Yield x s -> Yield x (s, i + 1)
568 Skip s -> Skip (s, i)
569 Stop -> Stop
570 step (_, _) = return Stop
572{-# INLINE_NORMAL takeWhileM #-}
573takeWhileM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
574takeWhileM f (Unfold step1 inject1) = Unfold step inject1
575 where
576 {-# INLINE_LATE step #-}
577 step st = do
578 r <- step1 st
579 case r of
580 Yield x s -> do
581 b <- f x
582 return $ if b then Yield x s else Stop
583 Skip s -> return $ Skip s
584 Stop -> return Stop
586{-# INLINE takeWhile #-}
587takeWhile :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
588takeWhile f = takeWhileM (return . f)
590{-# INLINE_NORMAL filterM #-}
591filterM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
592filterM f (Unfold step1 inject1) = Unfold step inject1
593 where
594 {-# INLINE_LATE step #-}
595 step st = do
596 r <- step1 st
597 case r of
598 Yield x s -> do
599 b <- f x
600 return $ if b then Yield x s else Skip s
601 Skip s -> return $ Skip s
602 Stop -> return Stop
604{-# INLINE filter #-}
605filter :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
606filter f = filterM (return . f)
608-- | @drop n unf@ drops @n@ elements from the stream generated by @unf@.
609--
610-- /Internal/
611--
612{-# INLINE_NORMAL drop #-}
613drop :: Monad m => Int -> Unfold m a b -> Unfold m a b
614drop n (Unfold step inject) = Unfold step' inject'
616 where
618 inject' a = do
619 b <- inject a
620 return (b, n)
622 {-# INLINE_LATE step' #-}
623 step' (st, i)
624 | i > 0 = do
625 r <- step st
626 return
627 $ case r of
628 Yield _ s -> Skip (s, i - 1)
629 Skip s -> Skip (s, i)
630 Stop -> Stop
631 | otherwise = do
632 r <- step st
633 return
634 $ case r of
635 Yield x s -> Yield x (s, 0)
636 Skip s -> Skip (s, 0)
637 Stop -> Stop
639-- | @dropWhileM f unf@ drops elements from the stream generated by @unf@ while
640-- the condition holds true. The condition function @f@ is /monadic/ in nature.
641--
642-- /Internal/
643--
644{-# INLINE_NORMAL dropWhileM #-}
645dropWhileM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
646dropWhileM f (Unfold step inject) = Unfold step' inject'
648 where
650 inject' a = do
651 b <- inject a
652 return $ Left b
654 {-# INLINE_LATE step' #-}
655 step' (Left st) = do
656 r <- step st
657 case r of
658 Yield x s -> do
659 b <- f x
660 return
661 $ if b
662 then Skip (Left s)
663 else Yield x (Right s)
664 Skip s -> return $ Skip (Left s)
665 Stop -> return Stop
666 step' (Right st) = do
667 r <- step st
668 return
669 $ case r of
670 Yield x s -> Yield x (Right s)
671 Skip s -> Skip (Right s)
672 Stop -> Stop
674-- | Similar to 'dropWhileM' but with a pure condition function.
675--
676-- /Internal/
677--
678{-# INLINE dropWhile #-}
679dropWhile :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
680dropWhile f = dropWhileM (return . f)
682-------------------------------------------------------------------------------
683-- Enumeration
684-------------------------------------------------------------------------------
686-- | Can be used to enumerate unbounded integrals. This does not check for
687-- overflow or underflow for bounded integrals.
688{-# INLINE_NORMAL enumerateFromStepIntegral #-}
689enumerateFromStepIntegral :: (Integral a, Monad m) => Unfold m (a, a) a
690enumerateFromStepIntegral = Unfold step inject
691 where
692 inject (from, stride) = from `seq` stride `seq` return (from, stride)
693 {-# INLINE_LATE step #-}
694 step (x, stride) = return $ Yield x $! (x + stride, stride)
696-- We are assuming that "to" is constrained by the type to be within
697-- max/min bounds.
698{-# INLINE enumerateFromToIntegral #-}
699enumerateFromToIntegral :: (Monad m, Integral a) => a -> Unfold m a a
700enumerateFromToIntegral to =
701 takeWhile (<= to) $ supplySecond enumerateFromStepIntegral 1
703{-# INLINE enumerateFromIntegral #-}
704enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => Unfold m a a
705enumerateFromIntegral = enumerateFromToIntegral maxBound
707-- | Generate an infinite stream starting from the seed with increments of the
708-- given stride.
709--
710-- /Internal/
711--
712{-# INLINE enumerateFromStepNum #-}
713enumerateFromStepNum :: (Monad m, Num a) => a -> Unfold m a a
714enumerateFromStepNum stride = Unfold step return
715 where
716 {-# INLINE_LATE step #-}
717 step !s = return $ (Yield $! s) $! (s + stride)
719-- | @numFrom = enumerateFromStepNum 1@
720--
721-- /Internal/
722--
723{-# INLINE_NORMAL numFrom #-}
724numFrom :: (Monad m, Num a) => Unfold m a a
725numFrom = enumerateFromStepNum 1
727-- | /Internal/
728--
729{-# INLINE_NORMAL enumerateFromToFractional #-}
730enumerateFromToFractional :: (Monad m, Fractional a, Ord a) => a -> Unfold m a a
731enumerateFromToFractional to =
732 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum 1
734-------------------------------------------------------------------------------
735-- Generation from SVar
736-------------------------------------------------------------------------------
738data FromSVarState t m a =
739 FromSVarInit (SVar t m a)
740 | FromSVarRead (SVar t m a)
741 | FromSVarLoop (SVar t m a) [ChildEvent a]
742 | FromSVarDone (SVar t m a)
744-- | /Internal/
745--
746{-# INLINE_NORMAL fromSVar #-}
747fromSVar :: MonadAsync m => Unfold m (SVar t m a) a
748fromSVar = Unfold step (return . FromSVarInit)
749 where
751 {-# INLINE_LATE step #-}
752 step (FromSVarInit svar) = do
753 ref <- liftIO $ newIORef ()
754 _ <- liftIO $ mkWeakIORef ref hook
755 -- when this copy of svar gets garbage collected "ref" will get
756 -- garbage collected and our GC hook will be called.
757 let sv = svar{svarRef = Just ref}
758 return $ Skip (FromSVarRead sv)
760 where
762 {-# NOINLINE hook #-}
763 hook = do
764 when (svarInspectMode svar) $ do
765 r <- liftIO $ readIORef (svarStopTime (svarStats svar))
766 when (isNothing r) $
767 printSVar svar "SVar Garbage Collected"
768 cleanupSVar svar
769 -- If there are any SVars referenced by this SVar a GC will prompt
770 -- them to be cleaned up quickly.
771 when (svarInspectMode svar) performMajorGC
773 step (FromSVarRead sv) = do
774 list <- readOutputQ sv
775 -- Reversing the output is important to guarantee that we process the
776 -- outputs in the same order as they were generated by the constituent
777 -- streams.
778 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
780 step (FromSVarLoop sv []) = do
781 done <- postProcess sv
782 return $ Skip $ if done
783 then FromSVarDone sv
784 else FromSVarRead sv
786 step (FromSVarLoop sv (ev : es)) = do
787 case ev of
788 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
789 ChildStop tid e -> do
790 accountThread sv tid
791 case e of
792 Nothing -> do
793 stop <- shouldStop tid
794 if stop
795 then do
796 liftIO (cleanupSVar sv)
797 return $ Skip (FromSVarDone sv)
798 else return $ Skip (FromSVarLoop sv es)
799 Just ex ->
800 case fromException ex of
801 Just ThreadAbort ->
802 return $ Skip (FromSVarLoop sv es)
803 Nothing -> liftIO (cleanupSVar sv) >> MC.throwM ex
804 where
806 shouldStop tid =
807 case svarStopStyle sv of
808 StopNone -> return False
809 StopAny -> return True
810 StopBy -> do
811 sid <- liftIO $ readIORef (svarStopBy sv)
812 return $ tid == sid
814 step (FromSVarDone sv) = do
815 when (svarInspectMode sv) $ do
816 t <- liftIO $ getTime Monotonic
817 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
818 liftIO $ printSVar sv "SVar Done"
819 return Stop
821-------------------------------------------------------------------------------
822-- Process events received by a fold consumer from a stream producer
823-------------------------------------------------------------------------------
825-- XXX Have an error for "FromSVarInit" instead of undefined. The error can
826-- redirect the user to report the failure to the developers.
827-- | /Internal/
828--
829{-# INLINE_NORMAL fromProducer #-}
830fromProducer :: MonadAsync m => Unfold m (SVar t m a) a
831fromProducer = Unfold step (return . FromSVarRead)
832 where
834 {-# INLINE_LATE step #-}
835 step (FromSVarRead sv) = do
836 list <- readOutputQ sv
837 -- Reversing the output is important to guarantee that we process the
838 -- outputs in the same order as they were generated by the constituent
839 -- streams.
840 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
842 step (FromSVarLoop sv []) = return $ Skip $ FromSVarRead sv
843 step (FromSVarLoop sv (ev : es)) = do
844 case ev of
845 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
846 ChildStop tid e -> do
847 accountThread sv tid
848 case e of
849 Nothing -> do
850 sendStopToProducer sv
851 return $ Skip (FromSVarDone sv)
852 Just _ -> error "Bug: fromProducer: received exception"
854 step (FromSVarDone sv) = do
855 when (svarInspectMode sv) $ do
856 t <- liftIO $ getTime Monotonic
857 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
858 liftIO $ printSVar sv "SVar Done"
859 return Stop
861 step (FromSVarInit _) = undefined
863-------------------------------------------------------------------------------
864-- Zipping
865-------------------------------------------------------------------------------
867{-# INLINE_NORMAL zipWithM #-}
868zipWithM :: Monad m
869 => (a -> b -> m c) -> Unfold m x a -> Unfold m y b -> Unfold m (x, y) c
870zipWithM f (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject
872 where
874 inject (x, y) = do
875 s1 <- inject1 x
876 s2 <- inject2 y
877 return (s1, s2, Nothing)
879 {-# INLINE_LATE step #-}
880 step (s1, s2, Nothing) = do
881 r <- step1 s1
882 return $
883 case r of
884 Yield x s -> Skip (s, s2, Just x)
885 Skip s -> Skip (s, s2, Nothing)
886 Stop -> Stop
888 step (s1, s2, Just x) = do
889 r <- step2 s2
890 case r of
891 Yield y s -> do
892 z <- f x y
893 return $ Yield z (s1, s, Nothing)
894 Skip s -> return $ Skip (s1, s, Just x)
895 Stop -> return Stop
897-- | Divide the input into two unfolds and then zip the outputs to a single
898-- stream.
899--
900-- @
901-- S.mapM_ print
902-- $ S.concatUnfold (UF.zipWith (,) UF.identity (UF.singleton sqrt))
903-- $ S.map (\x -> (x,x))
904-- $ S.fromList [1..10]
905-- @
906--
907-- /Internal/
908--
909{-# INLINE zipWith #-}
910zipWith :: Monad m
911 => (a -> b -> c) -> Unfold m x a -> Unfold m y b -> Unfold m (x, y) c
912zipWith f = zipWithM (\a b -> return (f a b))
914-- | Distribute the input to two unfolds and then zip the outputs to a single
915-- stream.
916--
917-- @
918-- S.mapM_ print $ S.concatUnfold (UF.teeZipWith (,) UF.identity (UF.singleton sqrt)) $ S.fromList [1..10]
919-- @
920--
921-- /Internal/
922--
923{-# INLINE_NORMAL teeZipWith #-}
924teeZipWith :: Monad m
925 => (a -> b -> c) -> Unfold m x a -> Unfold m x b -> Unfold m x c
926teeZipWith f unf1 unf2 = lmap (\x -> (x,x)) $ zipWith f unf1 unf2
928-------------------------------------------------------------------------------
929-- Nested
930-------------------------------------------------------------------------------
932{-# ANN type ConcatState Fuse #-}
933data ConcatState s1 s2 = ConcatOuter s1 | ConcatInner s1 s2
935-- | Apply the second unfold to each output element of the first unfold and
936-- flatten the output in a single stream.
937--
938-- /Internal/
939--
940{-# INLINE_NORMAL concat #-}
941concat :: Monad m => Unfold m a b -> Unfold m b c -> Unfold m a c
942concat (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject
943 where
944 inject x = do
945 s <- inject1 x
946 return $ ConcatOuter s
948 {-# INLINE_LATE step #-}
949 step (ConcatOuter st) = do
950 r <- step1 st
951 case r of
952 Yield x s -> do
953 innerSt <- inject2 x
954 return $ Skip (ConcatInner s innerSt)
955 Skip s -> return $ Skip (ConcatOuter s)
956 Stop -> return Stop
958 step (ConcatInner ost ist) = do
959 r <- step2 ist
960 return $ case r of
961 Yield x s -> Yield x (ConcatInner ost s)
962 Skip s -> Skip (ConcatInner ost s)
963 Stop -> Skip (ConcatOuter ost)
965data OuterProductState s1 s2 sy x y =
966 OuterProductOuter s1 y | OuterProductInner s1 sy s2 x
968-- | Create an outer product (vector product or cartesian product) of the
969-- output streams of two unfolds.
970--
971{-# INLINE_NORMAL outerProduct #-}
972outerProduct :: Monad m
973 => Unfold m a b -> Unfold m c d -> Unfold m (a, c) (b, d)
974outerProduct (Unfold step1 inject1) (Unfold step2 inject2) = Unfold step inject
975 where
976 inject (x, y) = do
977 s1 <- inject1 x
978 return $ OuterProductOuter s1 y
980 {-# INLINE_LATE step #-}
981 step (OuterProductOuter st1 sy) = do
982 r <- step1 st1
983 case r of
984 Yield x s -> do
985 s2 <- inject2 sy
986 return $ Skip (OuterProductInner s sy s2 x)
987 Skip s -> return $ Skip (OuterProductOuter s sy)
988 Stop -> return Stop
990 step (OuterProductInner ost sy ist x) = do
991 r <- step2 ist
992 return $ case r of
993 Yield y s -> Yield (x, y) (OuterProductInner ost sy s x)
994 Skip s -> Skip (OuterProductInner ost sy s x)
995 Stop -> Skip (OuterProductOuter ost sy)
997-- Special cases of outer product
998-- | Outer product with a function application.
999--
1000-- /Unimplemented/
1002{-# INLINE_NORMAL ap #-}
1003ap :: -- Monad m =>
1004 Unfold m a (b -> c) -> Unfold m d b -> Unfold m (a, d) c
1005ap (Unfold _step1 _inject1) (Unfold _step2 _inject2) = undefined
1007-- | Outer product discarding the first element.
1009-- /Unimplemented/
1011{-# INLINE_NORMAL apDiscardFst #-}
1012apDiscardFst :: -- Monad m =>
1013 Unfold m a b -> Unfold m c d -> Unfold m (a, c) d
1014apDiscardFst (Unfold _step1 _inject1) (Unfold _step2 _inject2) = undefined
1016-- | Outer product discarding the second element.
1018-- /Unimplemented/
1020{-# INLINE_NORMAL apDiscardSnd #-}
1021apDiscardSnd :: -- Monad m =>
1022 Unfold m a b -> Unfold m c d -> Unfold m (a, c) b
1023apDiscardSnd (Unfold _step1 _inject1) (Unfold _step2 _inject2) = undefined
1025-- XXX This can be used to implement a Monad instance for "Unfold m ()".
1027data ConcatMapState s1 s2 = ConcatMapOuter s1 | ConcatMapInner s1 s2
1029-- | Map an unfold generating action to each element of an unfold and
1030-- flattern the results into a single stream.
1032{-# INLINE_NORMAL concatMapM #-}
1033concatMapM :: Monad m
1034 => (b -> m (Unfold m () c)) -> Unfold m a b -> Unfold m a c
1035concatMapM f (Unfold step1 inject1) = Unfold step inject
1036 where
1037 inject x = do
1038 s <- inject1 x
1039 return $ ConcatMapOuter s
1041 {-# INLINE_LATE step #-}
1042 step (ConcatMapOuter st) = do
1043 r <- step1 st
1044 case r of
1045 Yield x s -> do
1046 Unfold step2 inject2 <- f x
1047 innerSt <- inject2 ()
1048 return $ Skip (ConcatMapInner s (Stream (\_ ss -> step2 ss)
1049 innerSt))
1050 Skip s -> return $ Skip (ConcatMapOuter s)
1051 Stop -> return Stop
1053 step (ConcatMapInner ost (UnStream istep ist)) = do
1054 r <- istep defState ist
1055 return $ case r of
1056 Yield x s -> Yield x (ConcatMapInner ost (Stream istep s))
1057 Skip s -> Skip (ConcatMapInner ost (Stream istep s))
1058 Stop -> Skip (ConcatMapOuter ost)
1060------------------------------------------------------------------------------
1061-- Exceptions
1062------------------------------------------------------------------------------
1064-- | The most general bracketing and exception combinator. All other
1065-- combinators can be expressed in terms of this combinator. This can also be
1066-- used for cases which are not covered by the standard combinators.
1068-- /Internal/
1070{-# INLINE_NORMAL gbracket_ #-}
1071gbracket_
1072 :: Monad m
1073 => (a -> m c) -- ^ before
1074 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
1075 -> (c -> m d) -- ^ after, on normal stop
1076 -> Unfold m (c, e) b -- ^ on exception
1077 -> Unfold m c b -- ^ unfold to run
1078 -> Unfold m a b
1079gbracket_ bef exc aft (Unfold estep einject) (Unfold step1 inject1) =
1080 Unfold step inject
1082 where
1084 inject x = do
1085 r <- bef x
1086 s <- inject1 r
1087 return $ Right (s, r)
1089 {-# INLINE_LATE step #-}
1090 step (Right (st, v)) = do
1091 res <- exc $ step1 st
1092 case res of
1093 Right r -> case r of
1094 Yield x s -> return $ Yield x (Right (s, v))
1095 Skip s -> return $ Skip (Right (s, v))
1096 Stop -> aft v >> return Stop
1097 Left e -> do
1098 r <- einject (v, e)
1099 return $ Skip (Left r)
1100 step (Left st) = do
1101 res <- estep st
1102 case res of
1103 Yield x s -> return $ Yield x (Left s)
1104 Skip s -> return $ Skip (Left s)
1105 Stop -> return Stop
1107-- | The most general bracketing and exception combinator. All other
1108-- combinators can be expressed in terms of this combinator. This can also be
1109-- used for cases which are not covered by the standard combinators.
1111-- /Internal/
1113{-# INLINE_NORMAL gbracket #-}
1114gbracket
1115 :: (MonadIO m, MonadBaseControl IO m)
1116 => (a -> m c) -- ^ before
1117 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
1118 -> (c -> m d) -- ^ after, on normal stop, or GC
1119 -> Unfold m (c, e) b -- ^ on exception
1120 -> Unfold m c b -- ^ unfold to run
1121 -> Unfold m a b
1122gbracket bef exc aft (Unfold estep einject) (Unfold step1 inject1) =
1123 Unfold step inject
1125 where
1127 inject x = do
1128 -- Mask asynchronous exceptions to make the execution of 'bef' and
1129 -- the registration of 'aft' atomic. See comment in 'D.gbracketIO'.
1130 (r, ref) <- liftBaseOp_ mask_ $ do
1131 r <- bef x
1132 ref <- D.newFinalizedIORef (aft r)
1133 return (r, ref)
1134 s <- inject1 r
1135 return $ Right (s, r, ref)
1137 {-# INLINE_LATE step #-}
1138 step (Right (st, v, ref)) = do
1139 res <- exc $ step1 st
1140 case res of
1141 Right r -> case r of
1142 Yield x s -> return $ Yield x (Right (s, v, ref))
1143 Skip s -> return $ Skip (Right (s, v, ref))
1144 Stop -> do
1145 D.runIORefFinalizer ref
1146 return Stop
1147 Left e -> do
1148 -- Clearing of finalizer and running of exception handler must
1149 -- be atomic wrt async exceptions. Otherwise if we have cleared
1150 -- the finalizer and have not run the exception handler then we
1151 -- may leak the resource.
1152 r <- D.withIORefFinalizer ref (einject (v, e))
1153 return $ Skip (Left r)
1154 step (Left st) = do
1155 res <- estep st
1156 case res of
1157 Yield x s -> return $ Yield x (Left s)
1158 Skip s -> return $ Skip (Left s)
1159 Stop -> return Stop
1161-- The custom implementation of "before" is slightly faster (5-7%) than
1162-- "_before". This is just to document and make sure that we can always use
1163-- gbracket to implement before. The same applies to other combinators as well.
1165{-# INLINE_NORMAL _before #-}
1166_before :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
1167_before action = gbracket_ (\x -> action x >> return x) (fmap Right)
1168 (\_ -> return ()) undefined
1170-- | Run a side effect before the unfold yields its first element.
1172-- /Internal/
1173{-# INLINE_NORMAL before #-}
1174before :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
1175before action (Unfold step1 inject1) = Unfold step inject
1177 where
1179 inject x = do
1180 _ <- action x
1181 inject1 x
1183 {-# INLINE_LATE step #-}
1184 step st = do
1185 res <- step1 st
1186 case res of
1187 Yield x s -> return $ Yield x s
1188 Skip s -> return $ Skip s
1189 Stop -> return Stop
1191{-# INLINE_NORMAL _after #-}
1192_after :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
1193_after aft = gbracket_ return (fmap Right) aft undefined
1195-- | Run a side effect whenever the unfold stops normally.
1197-- Prefer 'after' over this as the @after@ action in this combinator is not
1198-- executed if the unfold is partially evaluated lazily and then garbage
1199-- collected.
1201-- /Internal/
1202{-# INLINE_NORMAL after_ #-}
1203after_ :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
1204after_ action (Unfold step1 inject1) = Unfold step inject
1206 where
1208 inject x = do
1209 s <- inject1 x
1210 return (s, x)
1212 {-# INLINE_LATE step #-}
1213 step (st, v) = do
1214 res <- step1 st
1215 case res of
1216 Yield x s -> return $ Yield x (s, v)
1217 Skip s -> return $ Skip (s, v)
1218 Stop -> action v >> return Stop
1220-- | Run a side effect whenever the unfold stops normally
1221-- or is garbage collected after a partial lazy evaluation.
1223-- /Internal/
1224{-# INLINE_NORMAL after #-}
1225after :: (MonadIO m, MonadBaseControl IO m)
1226 => (a -> m c) -> Unfold m a b -> Unfold m a b
1227after action (Unfold step1 inject1) = Unfold step inject
1229 where
1231 inject x = do
1232 s <- inject1 x
1233 ref <- D.newFinalizedIORef (action x)
1234 return (s, ref)
1236 {-# INLINE_LATE step #-}
1237 step (st, ref) = do
1238 res <- step1 st
1239 case res of
1240 Yield x s -> return $ Yield x (s, ref)
1241 Skip s -> return $ Skip (s, ref)
1242 Stop -> do
1243 D.runIORefFinalizer ref
1244 return Stop
1246{-# INLINE_NORMAL _onException #-}
1247_onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
1248_onException action =
1249 gbracket_ return MC.try
1250 (\_ -> return ())
1251 (nilM (\(a, e :: MC.SomeException) -> action a >> MC.throwM e))
1253-- | Run a side effect whenever the unfold aborts due to an exception.
1255-- /Internal/
1256{-# INLINE_NORMAL onException #-}
1257onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
1258onException action (Unfold step1 inject1) = Unfold step inject
1260 where
1262 inject x = do
1263 s <- inject1 x
1264 return (s, x)
1266 {-# INLINE_LATE step #-}
1267 step (st, v) = do
1268 res <- step1 st `MC.onException` action v
1269 case res of
1270 Yield x s -> return $ Yield x (s, v)
1271 Skip s -> return $ Skip (s, v)
1272 Stop -> return Stop
1274{-# INLINE_NORMAL _finally #-}
1275_finally :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
1276_finally action =
1277 gbracket_ return MC.try action
1278 (nilM (\(a, e :: MC.SomeException) -> action a >> MC.throwM e))
1280-- | Run a side effect whenever the unfold stops normally or aborts due to an
1281-- exception.
1283-- Prefer 'finally' over this as the @after@ action in this combinator is not
1284-- executed if the unfold is partially evaluated lazily and then garbage
1285-- collected.
1287-- /Internal/
1288{-# INLINE_NORMAL finally_ #-}
1289finally_ :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
1290finally_ action (Unfold step1 inject1) = Unfold step inject
1292 where
1294 inject x = do
1295 s <- inject1 x
1296 return (s, x)
1298 {-# INLINE_LATE step #-}
1299 step (st, v) = do
1300 res <- step1 st `MC.onException` action v
1301 case res of
1302 Yield x s -> return $ Yield x (s, v)
1303 Skip s -> return $ Skip (s, v)
1304 Stop -> action v >> return Stop
1306-- | Run a side effect whenever the unfold stops normally, aborts due to an
1307-- exception or if it is garbage collected after a partial lazy evaluation.
1309-- /Internal/
1310{-# INLINE_NORMAL finally #-}
1311finally :: (MonadAsync m, MonadCatch m)
1312 => (a -> m c) -> Unfold m a b -> Unfold m a b
1313finally action (Unfold step1 inject1) = Unfold step inject
1315 where
1317 inject x = do
1318 s <- inject1 x
1319 ref <- D.newFinalizedIORef (action x)
1320 return (s, ref)
1322 {-# INLINE_LATE step #-}
1323 step (st, ref) = do
1324 res <- step1 st `MC.onException` D.runIORefFinalizer ref
1325 case res of
1326 Yield x s -> return $ Yield x (s, ref)
1327 Skip s -> return $ Skip (s, ref)
1328 Stop -> do
1329 D.runIORefFinalizer ref
1330 return Stop
1332{-# INLINE_NORMAL _bracket #-}
1333_bracket :: MonadCatch m
1334 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1335_bracket bef aft =
1336 gbracket_ bef MC.try aft (nilM (\(a, e :: MC.SomeException) -> aft a >>
1337 MC.throwM e))
1339-- | @bracket_ before after between@ runs the @before@ action and then unfolds
1340-- its output using the @between@ unfold. When the @between@ unfold is done or
1341-- if an exception occurs then the @after@ action is run with the output of
1342-- @before@ as argument.
1344-- Prefer 'bracket' over this as the @after@ action in this combinator is not
1345-- executed if the unfold is partially evaluated lazily and then garbage
1346-- collected.
1348-- /Internal/
1349{-# INLINE_NORMAL bracket_ #-}
1350bracket_ :: MonadCatch m
1351 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1352bracket_ bef aft (Unfold step1 inject1) = Unfold step inject
1354 where
1356 inject x = do
1357 r <- bef x
1358 s <- inject1 r
1359 return (s, r)
1361 {-# INLINE_LATE step #-}
1362 step (st, v) = do
1363 res <- step1 st `MC.onException` aft v
1364 case res of
1365 Yield x s -> return $ Yield x (s, v)
1366 Skip s -> return $ Skip (s, v)
1367 Stop -> aft v >> return Stop
1369-- | @bracket before after between@ runs the @before@ action and then unfolds
1370-- its output using the @between@ unfold. When the @between@ unfold is done or
1371-- if an exception occurs then the @after@ action is run with the output of
1372-- @before@ as argument. The after action is also executed if the unfold is
1373-- paritally evaluated and then garbage collected.
1375-- /Internal/
1376{-# INLINE_NORMAL bracket #-}
1377bracket :: (MonadAsync m, MonadCatch m)
1378 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
1379bracket bef aft (Unfold step1 inject1) = Unfold step inject
1381 where
1383 inject x = do
1384 -- Mask asynchronous exceptions to make the execution of 'bef' and
1385 -- the registration of 'aft' atomic. See comment in 'D.gbracketIO'.
1386 (r, ref) <- liftBaseOp_ mask_ $ do
1387 r <- bef x
1388 ref <- D.newFinalizedIORef (aft r)
1389 return (r, ref)
1390 s <- inject1 r
1391 return (s, ref)
1393 {-# INLINE_LATE step #-}
1394 step (st, ref) = do
1395 res <- step1 st `MC.onException` D.runIORefFinalizer ref
1396 case res of
1397 Yield x s -> return $ Yield x (s, ref)
1398 Skip s -> return $ Skip (s, ref)
1399 Stop -> do
1400 D.runIORefFinalizer ref
1401 return Stop
1403-- | When unfolding if an exception occurs, unfold the exception using the
1404-- exception unfold supplied as the first argument to 'handle'.
1406-- /Internal/
1407{-# INLINE_NORMAL handle #-}
1408handle :: (MonadCatch m, Exception e)
1409 => Unfold m e b -> Unfold m a b -> Unfold m a b
1410handle exc =
1411 gbracket_ return MC.try (\_ -> return ()) (discardFirst exc)