{-# LANGUAGE CPP #-}
-- |
-- Module      : Streamly.Internal.Data.Unfold
-- Copyright   : (c) 2019 Composewell Technologies
-- License     : BSD3
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
module Streamly.Internal.Data.Unfold
    (
    -- * Setup
    -- | To execute the code examples provided in this module in ghci, please
    -- run the following commands first.
    --
    -- $setup

    -- * Unfold Type
      module Streamly.Internal.Data.Unfold.Type

    -- * Unfolds
    -- One to one correspondence with
    -- "Streamly.Internal.Data.Stream.Generate"
    -- ** Basic Constructors
    , nilM
    , nil
    , consM

    -- ** Generators
    -- | Generate a monadic stream from a seed.
    , repeatM
    , replicateM
    , fromIndicesM
    , iterateM

    -- ** Enumerations
    , module Streamly.Internal.Data.Unfold.Enumeration

    -- ** From Containers
    , fromListM

    -- ** From Memory
    , fromPtr

    -- ** From Stream
    , fromStreamK
    , fromStreamD
    , fromStream

    -- * Combinators
    -- ** Mapping on Input
    , discardFirst
    , discardSecond
    , swap
    -- coapply
    -- comonad

    -- * Folding
    , fold

    -- XXX Add "WithInput" versions of all these, map2, postscan2, takeWhile2,
    -- filter2 etc.  Alternatively, we can use the default operations with
    -- input, but that might make the common case more inconvenient.

    -- ** Mapping on Output
    , postscanlM'
    , postscan
    , scan
    , scanMany
    , foldMany
    -- pipe

    -- ** Either Wrapped Input
    , either

    -- ** Filtering
    , take
    , filter
    , filterM
    , drop
    , dropWhile
    , dropWhileM

    -- ** Cross product
    , joinInnerGeneric

    -- ** Resource Management
    -- | 'bracket' is the most general resource management operation, all other
    -- operations can be expressed using it. These functions have IO suffix
    -- because the allocation and cleanup functions are IO actions. For
    -- generalized allocation and cleanup functions see the functions without
    -- the IO suffix in the "streamly" package.
    , gbracket_
    , gbracketIO
    , before
    , afterIO
    , after_
    , finallyIO
    , finally_
    , bracketIO
    , bracket_

    -- ** Exceptions
    -- | Most of these combinators inhibit stream fusion, therefore, when
    -- possible, they should be called in an outer loop to mitigate the cost.
    -- For example, instead of calling them on a stream of chars call them on a
    -- stream of arrays before flattening it to a stream of chars.
    , onException
    , handle
    )
where

#include "inline.hs"
#include "ArrayMacros.h"

import Control.Exception (Exception, mask_)
import Control.Monad.Catch (MonadCatch)
import Data.Functor (($>))
import GHC.Types (SPEC(..))
import Streamly.Internal.Data.Fold.Type (Fold(..))
import Streamly.Internal.Data.IOFinalizer
    (newIOFinalizer, runIOFinalizer, clearingIOFinalizer)
import Streamly.Internal.Data.Stream.Type (Stream(..))
import Streamly.Internal.Data.SVar.Type (defState)

import qualified Control.Monad.Catch as MC
import qualified Data.Tuple as Tuple
import qualified Streamly.Internal.Data.Fold.Type as FL
import qualified Streamly.Internal.Data.Stream.Type as D
import qualified Streamly.Internal.Data.StreamK.Type as K
import qualified Prelude

import Streamly.Internal.Data.Unfold.Enumeration
import Streamly.Internal.Data.Unfold.Type
import Prelude
       hiding (map, mapM, takeWhile, take, filter, const, zipWith
              , drop, dropWhile, either)
import Control.Monad.IO.Class (MonadIO (liftIO))
import Foreign (Storable, peek, sizeOf)
import Foreign.Ptr

#include "DocTestDataUnfold.hs"

-------------------------------------------------------------------------------
-- Input operations
-------------------------------------------------------------------------------

-- | Convert an 'Unfold' into an unfold accepting a tuple as an argument,
-- using the argument of the original fold as the second element of tuple and
-- discarding the first element of the tuple.
--
-- @
-- discardFirst = Unfold.lmap snd
-- @
--
-- /Pre-release/
--
{-# INLINE_NORMAL discardFirst #-}
discardFirst :: Unfold m a b -> Unfold m (c, a) b
discardFirst :: forall (m :: * -> *) a b c. Unfold m a b -> Unfold m (c, a) b
discardFirst = ((c, a) -> a) -> Unfold m a b -> Unfold m (c, a) b
forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
lmap (c, a) -> a
forall a b. (a, b) -> b
snd

-- | Convert an 'Unfold' into an unfold accepting a tuple as an argument,
-- using the argument of the original fold as the first element of tuple and
-- discarding the second element of the tuple.
--
-- @
-- discardSecond = Unfold.lmap fst
-- @
--
-- /Pre-release/
--
{-# INLINE_NORMAL discardSecond #-}
discardSecond :: Unfold m a b -> Unfold m (a, c) b
discardSecond :: forall (m :: * -> *) a b c. Unfold m a b -> Unfold m (a, c) b
discardSecond = ((a, c) -> a) -> Unfold m a b -> Unfold m (a, c) b
forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
lmap (a, c) -> a
forall a b. (a, b) -> a
fst

-- | Convert an 'Unfold' that accepts a tuple as an argument into an unfold
-- that accepts a tuple with elements swapped.
--
-- @
-- swap = Unfold.lmap Tuple.swap
-- @
--
-- /Pre-release/
--
{-# INLINE_NORMAL swap #-}
swap :: Unfold m (a, c) b -> Unfold m (c, a) b
swap :: forall (m :: * -> *) a c b. Unfold m (a, c) b -> Unfold m (c, a) b
swap = ((c, a) -> (a, c)) -> Unfold m (a, c) b -> Unfold m (c, a) b
forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
lmap (c, a) -> (a, c)
forall a b. (a, b) -> (b, a)
Tuple.swap

-------------------------------------------------------------------------------
-- Output operations
-------------------------------------------------------------------------------

-- XXX Do we need this combinator or the stream based idiom is enough?

-- | Compose an 'Unfold' and a 'Fold'. Given an @Unfold m a b@ and a
-- @Fold m b c@, returns a monadic action @a -> m c@ representing the
-- application of the fold on the unfolded stream.
--
-- >>> Unfold.fold Fold.sum Unfold.fromList [1..100]
-- 5050
--
-- >>> fold f u = Stream.fold f . Stream.unfold u
--
-- /Pre-release/
--
{-# INLINE_NORMAL fold #-}
fold :: Monad m => Fold m b c -> Unfold m a b -> a -> m c
fold :: forall (m :: * -> *) b c a.
Monad m =>
Fold m b c -> Unfold m a b -> a -> m c
fold (Fold s -> b -> m (Step s c)
fstep m (Step s c)
initial s -> m c
_ s -> m c
final) (Unfold s -> m (Step s b)
ustep a -> m s
inject) a
a = do
    Step s c
res <- m (Step s c)
initial
    case Step s c
res of
        FL.Partial s
x -> a -> m s
inject a
a m s -> (s -> m c) -> m c
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= SPEC -> s -> s -> m c
go SPEC
SPEC s
x
        FL.Done c
b -> c -> m c
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return c
b

    where

    {-# INLINE_LATE go #-}
    go :: SPEC -> s -> s -> m c
go !SPEC
_ !s
fs s
st = do
        Step s b
r <- s -> m (Step s b)
ustep s
st
        case Step s b
r of
            Yield b
x s
s -> do
                Step s c
res <- s -> b -> m (Step s c)
fstep s
fs b
x
                case Step s c
res of
                    FL.Partial s
fs1 -> SPEC -> s -> s -> m c
go SPEC
SPEC s
fs1 s
s
                    FL.Done c
c -> c -> m c
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return c
c
            Skip s
s -> SPEC -> s -> s -> m c
go SPEC
SPEC s
fs s
s
            Step s b
Stop -> s -> m c
final s
fs

-- {-# ANN type FoldMany Fuse #-}
data FoldMany s fs b a
    = FoldManyStart s
    | FoldManyFirst fs s
    | FoldManyLoop s fs
    | FoldManyYield b (FoldMany s fs b a)
    | FoldManyDone

-- | Apply a fold multiple times on the output of an unfold.
--
-- /Pre-release/
{-# INLINE_NORMAL foldMany #-}
foldMany :: Monad m => Fold m b c -> Unfold m a b -> Unfold m a c
foldMany :: forall (m :: * -> *) b c a.
Monad m =>
Fold m b c -> Unfold m a b -> Unfold m a c
foldMany (Fold s -> b -> m (Step s c)
fstep m (Step s c)
initial s -> m c
_ s -> m c
final) (Unfold s -> m (Step s b)
ustep a -> m s
inject1) =
    (FoldMany s s c Any -> m (Step (FoldMany s s c Any) c))
-> (a -> m (FoldMany s s c Any)) -> Unfold m a c
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold FoldMany s s c Any -> m (Step (FoldMany s s c Any) c)
forall {a}. FoldMany s s c a -> m (Step (FoldMany s s c a) c)
step a -> m (FoldMany s s c Any)
forall {fs} {b} {a}. a -> m (FoldMany s fs b a)
inject

    where

    inject :: a -> m (FoldMany s fs b a)
inject a
x = do
        s
r <- a -> m s
inject1 a
x
        FoldMany s fs b a -> m (FoldMany s fs b a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (s -> FoldMany s fs b a
forall s fs b a. s -> FoldMany s fs b a
FoldManyStart s
r)

    {-# INLINE consume #-}
    consume :: b -> s -> s -> m (Step (FoldMany s s c a) a)
consume b
x s
s s
fs = do
        Step s c
res <- s -> b -> m (Step s c)
fstep s
fs b
x
        Step (FoldMany s s c a) a -> m (Step (FoldMany s s c a) a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return
            (Step (FoldMany s s c a) a -> m (Step (FoldMany s s c a) a))
-> Step (FoldMany s s c a) a -> m (Step (FoldMany s s c a) a)
forall a b. (a -> b) -> a -> b
$ FoldMany s s c a -> Step (FoldMany s s c a) a
forall s a. s -> Step s a
Skip
            (FoldMany s s c a -> Step (FoldMany s s c a) a)
-> FoldMany s s c a -> Step (FoldMany s s c a) a
forall a b. (a -> b) -> a -> b
$ case Step s c
res of
                  FL.Done c
b -> c -> FoldMany s s c a -> FoldMany s s c a
forall s fs b a. b -> FoldMany s fs b a -> FoldMany s fs b a
FoldManyYield c
b (s -> FoldMany s s c a
forall s fs b a. s -> FoldMany s fs b a
FoldManyStart s
s)
                  FL.Partial s
ps -> s -> s -> FoldMany s s c a
forall s fs b a. s -> fs -> FoldMany s fs b a
FoldManyLoop s
s s
ps

    {-# INLINE_LATE step #-}
    step :: FoldMany s s c a -> m (Step (FoldMany s s c a) c)
step (FoldManyStart s
st) = do
        Step s c
r <- m (Step s c)
initial
        Step (FoldMany s s c a) c -> m (Step (FoldMany s s c a) c)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return
            (Step (FoldMany s s c a) c -> m (Step (FoldMany s s c a) c))
-> Step (FoldMany s s c a) c -> m (Step (FoldMany s s c a) c)
forall a b. (a -> b) -> a -> b
$ FoldMany s s c a -> Step (FoldMany s s c a) c
forall s a. s -> Step s a
Skip
            (FoldMany s s c a -> Step (FoldMany s s c a) c)
-> FoldMany s s c a -> Step (FoldMany s s c a) c
forall a b. (a -> b) -> a -> b
$ case Step s c
r of
                  FL.Done c
b -> c -> FoldMany s s c a -> FoldMany s s c a
forall s fs b a. b -> FoldMany s fs b a -> FoldMany s fs b a
FoldManyYield c
b (s -> FoldMany s s c a
forall s fs b a. s -> FoldMany s fs b a
FoldManyStart s
st)
                  FL.Partial s
fs -> s -> s -> FoldMany s s c a
forall s fs b a. fs -> s -> FoldMany s fs b a
FoldManyFirst s
fs s
st
    step (FoldManyFirst s
fs s
st) = do
        Step s b
r <- s -> m (Step s b)
ustep s
st
        case Step s b
r of
            Yield b
x s
s -> b -> s -> s -> m (Step (FoldMany s s c a) c)
forall {s} {a} {a}. b -> s -> s -> m (Step (FoldMany s s c a) a)
consume b
x s
s s
fs
            Skip s
s -> Step (FoldMany s s c a) c -> m (Step (FoldMany s s c a) c)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (FoldMany s s c a) c -> m (Step (FoldMany s s c a) c))
-> Step (FoldMany s s c a) c -> m (Step (FoldMany s s c a) c)
forall a b. (a -> b) -> a -> b
$ FoldMany s s c a -> Step (FoldMany s s c a) c
forall s a. s -> Step s a
Skip (s -> s -> FoldMany s s c a
forall s fs b a. fs -> s -> FoldMany s fs b a
FoldManyFirst s
fs s
s)
            Step s b
Stop -> s -> m c
final s
fs m c
-> m (Step (FoldMany s s c a) c) -> m (Step (FoldMany s s c a) c)
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Step (FoldMany s s c a) c -> m (Step (FoldMany s s c a) c)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Step (FoldMany s s c a) c
forall s a. Step s a
Stop
    step (FoldManyLoop s
st s
fs) = do
        Step s b
r <- s -> m (Step s b)
ustep s
st
        case Step s b
r of
            Yield b
x s
s -> b -> s -> s -> m (Step (FoldMany s s c a) c)
forall {s} {a} {a}. b -> s -> s -> m (Step (FoldMany s s c a) a)
consume b
x s
s s
fs
            Skip s
s -> Step (FoldMany s s c a) c -> m (Step (FoldMany s s c a) c)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (FoldMany s s c a) c -> m (Step (FoldMany s s c a) c))
-> Step (FoldMany s s c a) c -> m (Step (FoldMany s s c a) c)
forall a b. (a -> b) -> a -> b
$ FoldMany s s c a -> Step (FoldMany s s c a) c
forall s a. s -> Step s a
Skip (s -> s -> FoldMany s s c a
forall s fs b a. s -> fs -> FoldMany s fs b a
FoldManyLoop s
s s
fs)
            Step s b
Stop -> do
                c
b <- s -> m c
final s
fs
                Step (FoldMany s s c a) c -> m (Step (FoldMany s s c a) c)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (FoldMany s s c a) c -> m (Step (FoldMany s s c a) c))
-> Step (FoldMany s s c a) c -> m (Step (FoldMany s s c a) c)
forall a b. (a -> b) -> a -> b
$ FoldMany s s c a -> Step (FoldMany s s c a) c
forall s a. s -> Step s a
Skip (c -> FoldMany s s c a -> FoldMany s s c a
forall s fs b a. b -> FoldMany s fs b a -> FoldMany s fs b a
FoldManyYield c
b FoldMany s s c a
forall s fs b a. FoldMany s fs b a
FoldManyDone)
    step (FoldManyYield c
b FoldMany s s c a
next) = Step (FoldMany s s c a) c -> m (Step (FoldMany s s c a) c)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (FoldMany s s c a) c -> m (Step (FoldMany s s c a) c))
-> Step (FoldMany s s c a) c -> m (Step (FoldMany s s c a) c)
forall a b. (a -> b) -> a -> b
$ c -> FoldMany s s c a -> Step (FoldMany s s c a) c
forall s a. a -> s -> Step s a
Yield c
b FoldMany s s c a
next
    step FoldMany s s c a
FoldManyDone = Step (FoldMany s s c a) c -> m (Step (FoldMany s s c a) c)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Step (FoldMany s s c a) c
forall s a. Step s a
Stop

-------------------------------------------------------------------------------
-- Either
-------------------------------------------------------------------------------

-- | Choose left or right unfold based on an either input.
--
-- /Pre-release/
{-# INLINE_NORMAL either #-}
either :: Applicative m =>
    Unfold m a c -> Unfold m b c -> Unfold m (Either a b) c
either :: forall (m :: * -> *) a c b.
Applicative m =>
Unfold m a c -> Unfold m b c -> Unfold m (Either a b) c
either (Unfold s -> m (Step s c)
stepL a -> m s
injectL) (Unfold s -> m (Step s c)
stepR b -> m s
injectR) = (Either s s -> m (Step (Either s s) c))
-> (Either a b -> m (Either s s)) -> Unfold m (Either a b) c
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold Either s s -> m (Step (Either s s) c)
step Either a b -> m (Either s s)
inject

    where

    inject :: Either a b -> m (Either s s)
inject (Left a
x) = s -> Either s s
forall a b. a -> Either a b
Left (s -> Either s s) -> m s -> m (Either s s)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> m s
injectL a
x
    inject (Right b
x) = s -> Either s s
forall a b. b -> Either a b
Right (s -> Either s s) -> m s -> m (Either s s)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> b -> m s
injectR b
x

    {-# INLINE_LATE step #-}
    step :: Either s s -> m (Step (Either s s) c)
step (Left s
st) = do
        (\case
            Yield c
x s
s -> c -> Either s s -> Step (Either s s) c
forall s a. a -> s -> Step s a
Yield c
x (s -> Either s s
forall a b. a -> Either a b
Left s
s)
            Skip s
s -> Either s s -> Step (Either s s) c
forall s a. s -> Step s a
Skip (s -> Either s s
forall a b. a -> Either a b
Left s
s)
            Step s c
Stop -> Step (Either s s) c
forall s a. Step s a
Stop) (Step s c -> Step (Either s s) c)
-> m (Step s c) -> m (Step (Either s s) c)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> s -> m (Step s c)
stepL s
st
    step (Right s
st) = do
        (\case
            Yield c
x s
s -> c -> Either s s -> Step (Either s s) c
forall s a. a -> s -> Step s a
Yield c
x (s -> Either s s
forall a b. b -> Either a b
Right s
s)
            Skip s
s -> Either s s -> Step (Either s s) c
forall s a. s -> Step s a
Skip (s -> Either s s
forall a b. b -> Either a b
Right s
s)
            Step s c
Stop -> Step (Either s s) c
forall s a. Step s a
Stop) (Step s c -> Step (Either s s) c)
-> m (Step s c) -> m (Step (Either s s) c)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> s -> m (Step s c)
stepR s
st

-- postscan2 :: Monad m => Refold m a b c -> Unfold m a b -> Unfold m a c

-- | Scan the output of an 'Unfold' to change it in a stateful manner.
--
-- /Pre-release/
{-# INLINE_NORMAL postscan #-}
postscan :: Monad m => Fold m b c -> Unfold m a b -> Unfold m a c
postscan :: forall (m :: * -> *) b c a.
Monad m =>
Fold m b c -> Unfold m a b -> Unfold m a c
postscan (Fold s -> b -> m (Step s c)
stepF m (Step s c)
initial s -> m c
extract s -> m c
final) (Unfold s -> m (Step s b)
stepU a -> m s
injectU) =
    (Maybe (s, s) -> m (Step (Maybe (s, s)) c))
-> (a -> m (Maybe (s, s))) -> Unfold m a c
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold Maybe (s, s) -> m (Step (Maybe (s, s)) c)
step a -> m (Maybe (s, s))
inject

    where

    inject :: a -> m (Maybe (s, s))
inject a
a =  do
        Step s c
r <- m (Step s c)
initial
        case Step s c
r of
            FL.Partial s
fs -> (s, s) -> Maybe (s, s)
forall a. a -> Maybe a
Just ((s, s) -> Maybe (s, s)) -> (s -> (s, s)) -> s -> Maybe (s, s)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (s
fs,) (s -> Maybe (s, s)) -> m s -> m (Maybe (s, s))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> m s
injectU a
a
            FL.Done c
_ -> Maybe (s, s) -> m (Maybe (s, s))
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (s, s)
forall a. Maybe a
Nothing

    {-# INLINE_LATE step #-}
    step :: Maybe (s, s) -> m (Step (Maybe (s, s)) c)
step (Just (s
fs, s
us)) = do
        Step s b
ru <- s -> m (Step s b)
stepU s
us
        case Step s b
ru of
            Yield b
x s
s -> do
                Step s c
rf <- s -> b -> m (Step s c)
stepF s
fs b
x
                case Step s c
rf of
                    FL.Done c
v -> Step (Maybe (s, s)) c -> m (Step (Maybe (s, s)) c)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Maybe (s, s)) c -> m (Step (Maybe (s, s)) c))
-> Step (Maybe (s, s)) c -> m (Step (Maybe (s, s)) c)
forall a b. (a -> b) -> a -> b
$ c -> Maybe (s, s) -> Step (Maybe (s, s)) c
forall s a. a -> s -> Step s a
Yield c
v Maybe (s, s)
forall a. Maybe a
Nothing
                    FL.Partial s
fs1 -> do
                        c
v <- s -> m c
extract s
fs1
                        Step (Maybe (s, s)) c -> m (Step (Maybe (s, s)) c)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Maybe (s, s)) c -> m (Step (Maybe (s, s)) c))
-> Step (Maybe (s, s)) c -> m (Step (Maybe (s, s)) c)
forall a b. (a -> b) -> a -> b
$ c -> Maybe (s, s) -> Step (Maybe (s, s)) c
forall s a. a -> s -> Step s a
Yield c
v ((s, s) -> Maybe (s, s)
forall a. a -> Maybe a
Just (s
fs1, s
s))
            Skip s
s -> Step (Maybe (s, s)) c -> m (Step (Maybe (s, s)) c)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Maybe (s, s)) c -> m (Step (Maybe (s, s)) c))
-> Step (Maybe (s, s)) c -> m (Step (Maybe (s, s)) c)
forall a b. (a -> b) -> a -> b
$ Maybe (s, s) -> Step (Maybe (s, s)) c
forall s a. s -> Step s a
Skip ((s, s) -> Maybe (s, s)
forall a. a -> Maybe a
Just (s
fs, s
s))
            Step s b
Stop -> s -> m c
final s
fs m c -> m (Step (Maybe (s, s)) c) -> m (Step (Maybe (s, s)) c)
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Step (Maybe (s, s)) c -> m (Step (Maybe (s, s)) c)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Step (Maybe (s, s)) c
forall s a. Step s a
Stop

    step Maybe (s, s)
Nothing = Step (Maybe (s, s)) c -> m (Step (Maybe (s, s)) c)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Step (Maybe (s, s)) c
forall s a. Step s a
Stop

data ScanState s f = ScanInit s | ScanDo s !f | ScanDone

{-# INLINE_NORMAL scanWith #-}
scanWith :: Monad m => Bool -> Fold m b c -> Unfold m a b -> Unfold m a c
scanWith :: forall (m :: * -> *) b c a.
Monad m =>
Bool -> Fold m b c -> Unfold m a b -> Unfold m a c
scanWith Bool
restart (Fold s -> b -> m (Step s c)
fstep m (Step s c)
initial s -> m c
extract s -> m c
final) (Unfold s -> m (Step s b)
stepU a -> m s
injectU) =
    (ScanState s s -> m (Step (ScanState s s) c))
-> (a -> m (ScanState s s)) -> Unfold m a c
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold ScanState s s -> m (Step (ScanState s s) c)
step a -> m (ScanState s s)
forall {f}. a -> m (ScanState s f)
inject

    where

    inject :: a -> m (ScanState s f)
inject a
a = s -> ScanState s f
forall s f. s -> ScanState s f
ScanInit (s -> ScanState s f) -> m s -> m (ScanState s f)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> m s
injectU a
a

    {-# INLINE runStep #-}
    runStep :: s -> m (Step s c) -> m (Step (ScanState s s) c)
runStep s
us m (Step s c)
action = do
        Step s c
r <- m (Step s c)
action
        case Step s c
r of
            FL.Partial s
fs -> do
                !c
b <- s -> m c
extract s
fs
                Step (ScanState s s) c -> m (Step (ScanState s s) c)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (ScanState s s) c -> m (Step (ScanState s s) c))
-> Step (ScanState s s) c -> m (Step (ScanState s s) c)
forall a b. (a -> b) -> a -> b
$ c -> ScanState s s -> Step (ScanState s s) c
forall s a. a -> s -> Step s a
Yield c
b (s -> s -> ScanState s s
forall s f. s -> f -> ScanState s f
ScanDo s
us s
fs)
            FL.Done c
b ->
                let next :: ScanState s f
next = if Bool
restart then s -> ScanState s f
forall s f. s -> ScanState s f
ScanInit s
us else ScanState s f
forall s f. ScanState s f
ScanDone
                 in Step (ScanState s s) c -> m (Step (ScanState s s) c)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (ScanState s s) c -> m (Step (ScanState s s) c))
-> Step (ScanState s s) c -> m (Step (ScanState s s) c)
forall a b. (a -> b) -> a -> b
$ c -> ScanState s s -> Step (ScanState s s) c
forall s a. a -> s -> Step s a
Yield c
b ScanState s s
forall {f}. ScanState s f
next

    {-# INLINE_LATE step #-}
    step :: ScanState s s -> m (Step (ScanState s s) c)
step (ScanInit s
us) = s -> m (Step s c) -> m (Step (ScanState s s) c)
forall {s}. s -> m (Step s c) -> m (Step (ScanState s s) c)
runStep s
us m (Step s c)
initial
    step (ScanDo s
us s
fs) = do
        Step s b
res <- s -> m (Step s b)
stepU s
us
        case Step s b
res of
            Yield b
x s
s -> s -> m (Step s c) -> m (Step (ScanState s s) c)
forall {s}. s -> m (Step s c) -> m (Step (ScanState s s) c)
runStep s
s (s -> b -> m (Step s c)
fstep s
fs b
x)
            Skip s
s -> Step (ScanState s s) c -> m (Step (ScanState s s) c)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (ScanState s s) c -> m (Step (ScanState s s) c))
-> Step (ScanState s s) c -> m (Step (ScanState s s) c)
forall a b. (a -> b) -> a -> b
$ ScanState s s -> Step (ScanState s s) c
forall s a. s -> Step s a
Skip (ScanState s s -> Step (ScanState s s) c)
-> ScanState s s -> Step (ScanState s s) c
forall a b. (a -> b) -> a -> b
$ s -> s -> ScanState s s
forall s f. s -> f -> ScanState s f
ScanDo s
s s
fs
            Step s b
Stop -> s -> m c
final s
fs m c -> m (Step (ScanState s s) c) -> m (Step (ScanState s s) c)
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Step (ScanState s s) c -> m (Step (ScanState s s) c)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Step (ScanState s s) c
forall s a. Step s a
Stop
    step ScanState s s
ScanDone = Step (ScanState s s) c -> m (Step (ScanState s s) c)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Step (ScanState s s) c
forall s a. Step s a
Stop

-- | Scan the output of an 'Unfold' to change it in a stateful manner.
-- Once fold is done it will restart from its initial state.
--
-- >>> u = Unfold.scanMany (Fold.take 2 Fold.sum) Unfold.fromList
-- >>> Unfold.fold Fold.toList u [1,2,3,4,5]
-- [0,1,3,0,3,7,0,5]
--
-- /Pre-release/
{-# INLINE_NORMAL scanMany #-}
scanMany :: Monad m => Fold m b c -> Unfold m a b -> Unfold m a c
scanMany :: forall (m :: * -> *) b c a.
Monad m =>
Fold m b c -> Unfold m a b -> Unfold m a c
scanMany = Bool -> Fold m b c -> Unfold m a b -> Unfold m a c
forall (m :: * -> *) b c a.
Monad m =>
Bool -> Fold m b c -> Unfold m a b -> Unfold m a c
scanWith Bool
True

-- scan2 :: Monad m => Refold m a b c -> Unfold m a b -> Unfold m a c

-- | Scan the output of an 'Unfold' to change it in a stateful manner.
-- Once fold is done it will stop.
--
-- >>> u = Unfold.scan (Fold.take 2 Fold.sum) Unfold.fromList
-- >>> Unfold.fold Fold.toList u [1,2,3,4,5]
-- [0,1,3]
--
-- /Pre-release/
{-# INLINE_NORMAL scan #-}
scan :: Monad m => Fold m b c -> Unfold m a b -> Unfold m a c
scan :: forall (m :: * -> *) b c a.
Monad m =>
Fold m b c -> Unfold m a b -> Unfold m a c
scan = Bool -> Fold m b c -> Unfold m a b -> Unfold m a c
forall (m :: * -> *) b c a.
Monad m =>
Bool -> Fold m b c -> Unfold m a b -> Unfold m a c
scanWith Bool
False

-- | Scan the output of an 'Unfold' to change it in a stateful manner.
--
-- /Pre-release/
{-# INLINE_NORMAL postscanlM' #-}
postscanlM' :: Monad m => (b -> a -> m b) -> m b -> Unfold m c a -> Unfold m c b
postscanlM' :: forall (m :: * -> *) b a c.
Monad m =>
(b -> a -> m b) -> m b -> Unfold m c a -> Unfold m c b
postscanlM' b -> a -> m b
f m b
z = Fold m a b -> Unfold m c a -> Unfold m c b
forall (m :: * -> *) b c a.
Monad m =>
Fold m b c -> Unfold m a b -> Unfold m a c
postscan ((b -> a -> m b) -> m b -> Fold m a b
forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Fold m a b
FL.foldlM' b -> a -> m b
f m b
z)

-------------------------------------------------------------------------------
-- Convert streams into unfolds
-------------------------------------------------------------------------------

{-# INLINE_NORMAL fromStreamD #-}
fromStreamD :: Applicative m => Unfold m (Stream m a) a
fromStreamD :: forall (m :: * -> *) a. Applicative m => Unfold m (Stream m a) a
fromStreamD = (Stream m a -> m (Step (Stream m a) a))
-> (Stream m a -> m (Stream m a)) -> Unfold m (Stream m a) a
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold Stream m a -> m (Step (Stream m a) a)
forall {m :: * -> *} {a}.
Functor m =>
Stream m a -> m (Step (Stream m a) a)
step Stream m a -> m (Stream m a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

    where

    {-# INLINE_LATE step #-}
    step :: Stream m a -> m (Step (Stream m a) a)
step (UnStream State StreamK m a -> s -> m (Step s a)
step1 s
state1) =
        (\case
            Yield a
x s
s -> a -> Stream m a -> Step (Stream m a) a
forall s a. a -> s -> Step s a
Yield a
x ((State StreamK m a -> s -> m (Step s a)) -> s -> Stream m a
forall (m :: * -> *) a s.
(State StreamK m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State StreamK m a -> s -> m (Step s a)
step1 s
s)
            Skip s
s    -> Stream m a -> Step (Stream m a) a
forall s a. s -> Step s a
Skip ((State StreamK m a -> s -> m (Step s a)) -> s -> Stream m a
forall (m :: * -> *) a s.
(State StreamK m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State StreamK m a -> s -> m (Step s a)
step1 s
s)
            Step s a
Stop      -> Step (Stream m a) a
forall s a. Step s a
Stop) (Step s a -> Step (Stream m a) a)
-> m (Step s a) -> m (Step (Stream m a) a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> State StreamK m a -> s -> m (Step s a)
step1 State StreamK m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
state1

{-# INLINE_NORMAL fromStreamK #-}
fromStreamK :: Applicative m => Unfold m (K.StreamK m a) a
fromStreamK :: forall (m :: * -> *) a. Applicative m => Unfold m (StreamK m a) a
fromStreamK = (StreamK m a -> m (Step (StreamK m a) a))
-> (StreamK m a -> m (StreamK m a)) -> Unfold m (StreamK m a) a
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold StreamK m a -> m (Step (StreamK m a) a)
forall {f :: * -> *} {a}.
Applicative f =>
StreamK f a -> f (Step (StreamK f a) a)
step StreamK m a -> m (StreamK m a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

    where

    {-# INLINE_LATE step #-}
    step :: StreamK f a -> f (Step (StreamK f a) a)
step StreamK f a
stream = do
        (\case
            Just (a
x, StreamK f a
xs) -> a -> StreamK f a -> Step (StreamK f a) a
forall s a. a -> s -> Step s a
Yield a
x StreamK f a
xs
            Maybe (a, StreamK f a)
Nothing -> Step (StreamK f a) a
forall s a. Step s a
Stop) (Maybe (a, StreamK f a) -> Step (StreamK f a) a)
-> f (Maybe (a, StreamK f a)) -> f (Step (StreamK f a) a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StreamK f a -> f (Maybe (a, StreamK f a))
forall (m :: * -> *) a.
Applicative m =>
StreamK m a -> m (Maybe (a, StreamK m a))
K.uncons StreamK f a
stream

{-# INLINE fromStream #-}
fromStream :: Applicative m => Unfold m (Stream m a) a
fromStream :: forall (m :: * -> *) a. Applicative m => Unfold m (Stream m a) a
fromStream = Unfold m (Stream m a) a
forall (m :: * -> *) a. Applicative m => Unfold m (Stream m a) a
fromStreamD

-------------------------------------------------------------------------------
-- Unfolds
-------------------------------------------------------------------------------

-- | Lift a monadic function into an unfold generating a nil stream with a side
-- effect.
--
{-# INLINE nilM #-}
nilM :: Applicative m => (a -> m c) -> Unfold m a b
nilM :: forall (m :: * -> *) a c b.
Applicative m =>
(a -> m c) -> Unfold m a b
nilM a -> m c
f = (a -> m (Step a b)) -> (a -> m a) -> Unfold m a b
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold a -> m (Step a b)
forall {s} {a}. a -> m (Step s a)
step a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

    where

    {-# INLINE_LATE step #-}
    step :: a -> m (Step s a)
step a
x = a -> m c
f a
x m c -> Step s a -> m (Step s a)
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Step s a
forall s a. Step s a
Stop

-- | An empty stream.
{-# INLINE nil #-}
nil :: Applicative m => Unfold m a b
nil :: forall (m :: * -> *) a b. Applicative m => Unfold m a b
nil = (a -> m (Step a b)) -> (a -> m a) -> Unfold m a b
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold (m (Step a b) -> a -> m (Step a b)
forall a b. a -> b -> a
Prelude.const (Step a b -> m (Step a b)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Step a b
forall s a. Step s a
Stop)) a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

-- | Prepend a monadic single element generator function to an 'Unfold'. The
-- same seed is used in the action as well as the unfold.
--
-- /Pre-release/
{-# INLINE_NORMAL consM #-}
consM :: Applicative m => (a -> m b) -> Unfold m a b -> Unfold m a b
consM :: forall (m :: * -> *) a b.
Applicative m =>
(a -> m b) -> Unfold m a b -> Unfold m a b
consM a -> m b
action Unfold m a b
unf = (Either a (Stream m b) -> m (Step (Either a (Stream m b)) b))
-> (a -> m (Either a (Stream m b))) -> Unfold m a b
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold Either a (Stream m b) -> m (Step (Either a (Stream m b)) b)
forall {a}.
Either a (Stream m b) -> m (Step (Either a (Stream m b)) b)
step a -> m (Either a (Stream m b))
forall {a} {b}. a -> m (Either a b)
inject

    where

    inject :: a -> m (Either a b)
inject = Either a b -> m (Either a b)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either a b -> m (Either a b))
-> (a -> Either a b) -> a -> m (Either a b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Either a b
forall a b. a -> Either a b
Left

    {-# INLINE_LATE step #-}
    step :: Either a (Stream m b) -> m (Step (Either a (Stream m b)) b)
step (Left a
a) = (b -> Either a (Stream m b) -> Step (Either a (Stream m b)) b
forall s a. a -> s -> Step s a
`Yield` Stream m b -> Either a (Stream m b)
forall a b. b -> Either a b
Right (Unfold m a b -> a -> Stream m b
forall (m :: * -> *) a b.
Applicative m =>
Unfold m a b -> a -> Stream m b
D.unfold Unfold m a b
unf a
a)) (b -> Step (Either a (Stream m b)) b)
-> m b -> m (Step (Either a (Stream m b)) b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> m b
action a
a
    step (Right (UnStream State StreamK m b -> s -> m (Step s b)
step1 s
st)) = do
        (\case
            Yield b
x s
s -> b -> Either a (Stream m b) -> Step (Either a (Stream m b)) b
forall s a. a -> s -> Step s a
Yield b
x (Stream m b -> Either a (Stream m b)
forall a b. b -> Either a b
Right ((State StreamK m b -> s -> m (Step s b)) -> s -> Stream m b
forall (m :: * -> *) a s.
(State StreamK m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State StreamK m b -> s -> m (Step s b)
step1 s
s))
            Skip s
s -> Either a (Stream m b) -> Step (Either a (Stream m b)) b
forall s a. s -> Step s a
Skip (Stream m b -> Either a (Stream m b)
forall a b. b -> Either a b
Right ((State StreamK m b -> s -> m (Step s b)) -> s -> Stream m b
forall (m :: * -> *) a s.
(State StreamK m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State StreamK m b -> s -> m (Step s b)
step1 s
s))
            Step s b
Stop -> Step (Either a (Stream m b)) b
forall s a. Step s a
Stop) (Step s b -> Step (Either a (Stream m b)) b)
-> m (Step s b) -> m (Step (Either a (Stream m b)) b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> State StreamK m b -> s -> m (Step s b)
step1 State StreamK m b
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st

-- | Convert a list of monadic values to a 'Stream'
--
{-# INLINE_LATE fromListM #-}
fromListM :: Applicative m => Unfold m [m a] a
fromListM :: forall (m :: * -> *) a. Applicative m => Unfold m [m a] a
fromListM = ([m a] -> m (Step [m a] a))
-> ([m a] -> m [m a]) -> Unfold m [m a] a
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold [m a] -> m (Step [m a] a)
forall {f :: * -> *} {a}.
Applicative f =>
[f a] -> f (Step [f a] a)
step [m a] -> m [m a]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

    where

    {-# INLINE_LATE step #-}
    step :: [f a] -> f (Step [f a] a)
step (f a
x:[f a]
xs) = (a -> [f a] -> Step [f a] a
forall s a. a -> s -> Step s a
`Yield` [f a]
xs) (a -> Step [f a] a) -> f a -> f (Step [f a] a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f a
x
    step [] = Step [f a] a -> f (Step [f a] a)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Step [f a] a
forall s a. Step s a
Stop

{-# INLINE fromPtr #-}
fromPtr :: forall m a. (MonadIO m, Storable a) => Unfold m (Ptr a) a
fromPtr :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Unfold m (Ptr a) a
fromPtr = (Ptr a -> m (Step (Ptr a) a))
-> (Ptr a -> m (Ptr a)) -> Unfold m (Ptr a) a
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold Ptr a -> m (Step (Ptr a) a)
forall {m :: * -> *} {a} {b}.
(MonadIO m, Storable a) =>
Ptr a -> m (Step (Ptr b) a)
step Ptr a -> m (Ptr a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return

    where

    {-# INLINE_LATE step #-}
    step :: Ptr a -> m (Step (Ptr b) a)
step Ptr a
p = do
        a
x <- IO a -> m a
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO a -> m a) -> IO a -> m a
forall a b. (a -> b) -> a -> b
$ Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek Ptr a
p
        Step (Ptr b) a -> m (Step (Ptr b) a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Ptr b) a -> m (Step (Ptr b) a))
-> Step (Ptr b) a -> m (Step (Ptr b) a)
forall a b. (a -> b) -> a -> b
$ a -> Ptr b -> Step (Ptr b) a
forall s a. a -> s -> Step s a
Yield a
x (PTR_NEXT(p, a))

------------------------------------------------------------------------------
-- Specialized Generation
------------------------------------------------------------------------------

-- | Given a seed @(n, action)@, generates a stream replicating the @action@ @n@
-- times.
--
{-# INLINE replicateM #-}
replicateM :: Applicative m => Unfold m (Int, m a) a
replicateM :: forall (m :: * -> *) a. Applicative m => Unfold m (Int, m a) a
replicateM = ((Int, m a) -> m (Step (Int, m a) a))
-> ((Int, m a) -> m (Int, m a)) -> Unfold m (Int, m a) a
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold (Int, m a) -> m (Step (Int, m a) a)
forall {a} {f :: * -> *} {a}.
(Ord a, Num a, Applicative f) =>
(a, f a) -> f (Step (a, f a) a)
step (Int, m a) -> m (Int, m a)
forall {f :: * -> *} {a}. Applicative f => a -> f a
inject

    where

    inject :: a -> f a
inject a
seed = a -> f a
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
seed

    {-# INLINE_LATE step #-}
    step :: (a, f a) -> f (Step (a, f a) a)
step (a
i, f a
action) =
        if a
i a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
0
        then Step (a, f a) a -> f (Step (a, f a) a)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Step (a, f a) a
forall s a. Step s a
Stop
        else (\a
x -> a -> (a, f a) -> Step (a, f a) a
forall s a. a -> s -> Step s a
Yield a
x (a
i a -> a -> a
forall a. Num a => a -> a -> a
- a
1, f a
action)) (a -> Step (a, f a) a) -> f a -> f (Step (a, f a) a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f a
action

-- | Generates an infinite stream repeating the seed.
--
{-# INLINE repeatM #-}
repeatM :: Applicative m => Unfold m (m a) a
repeatM :: forall (m :: * -> *) a. Applicative m => Unfold m (m a) a
repeatM = (m a -> m (Step (m a) a)) -> (m a -> m (m a)) -> Unfold m (m a) a
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold m a -> m (Step (m a) a)
forall {f :: * -> *} {a}. Functor f => f a -> f (Step (f a) a)
step m a -> m (m a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

    where

    {-# INLINE_LATE step #-}
    step :: f a -> f (Step (f a) a)
step f a
action = (a -> f a -> Step (f a) a
forall s a. a -> s -> Step s a
`Yield` f a
action) (a -> Step (f a) a) -> f a -> f (Step (f a) a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f a
action

-- | Generates an infinite stream starting with the given seed and applying the
-- given function repeatedly.
--
{-# INLINE iterateM #-}
iterateM :: Applicative m => (a -> m a) -> Unfold m (m a) a
iterateM :: forall (m :: * -> *) a.
Applicative m =>
(a -> m a) -> Unfold m (m a) a
iterateM a -> m a
f = (a -> m (Step a a)) -> (m a -> m a) -> Unfold m (m a) a
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold a -> m (Step a a)
step m a -> m a
forall a. a -> a
id

    where

    {-# INLINE_LATE step #-}
    step :: a -> m (Step a a)
step a
x = a -> a -> Step a a
forall s a. a -> s -> Step s a
Yield a
x (a -> Step a a) -> m a -> m (Step a a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> m a
f a
x

-- | @fromIndicesM gen@ generates an infinite stream of values using @gen@
-- starting from the seed.
--
-- @
-- fromIndicesM f = Unfold.mapM f $ Unfold.enumerateFrom 0
-- @
--
-- /Pre-release/
--
{-# INLINE_NORMAL fromIndicesM #-}
fromIndicesM :: Applicative m => (Int -> m a) -> Unfold m Int a
fromIndicesM :: forall (m :: * -> *) a.
Applicative m =>
(Int -> m a) -> Unfold m Int a
fromIndicesM Int -> m a
gen = (Int -> m (Step Int a)) -> (Int -> m Int) -> Unfold m Int a
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold Int -> m (Step Int a)
step Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

    where

    {-# INLINE_LATE step #-}
    step :: Int -> m (Step Int a)
step Int
i = (a -> Int -> Step Int a
forall s a. a -> s -> Step s a
`Yield` (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) (a -> Step Int a) -> m a -> m (Step Int a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> m a
gen Int
i

-------------------------------------------------------------------------------
-- Filtering
-------------------------------------------------------------------------------

-- |
-- >>> u = Unfold.take 2 Unfold.fromList
-- >>> Unfold.fold Fold.toList u [1..100]
-- [1,2]
--
{-# INLINE_NORMAL take #-}
take :: Applicative m => Int -> Unfold m a b -> Unfold m a b
take :: forall (m :: * -> *) a b.
Applicative m =>
Int -> Unfold m a b -> Unfold m a b
take Int
n (Unfold s -> m (Step s b)
step1 a -> m s
inject1) = ((s, Int) -> m (Step (s, Int) b))
-> (a -> m (s, Int)) -> Unfold m a b
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold (s, Int) -> m (Step (s, Int) b)
step a -> m (s, Int)
forall {t}. Num t => a -> m (s, t)
inject

    where

    inject :: a -> m (s, t)
inject a
x = (, t
0) (s -> (s, t)) -> m s -> m (s, t)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> m s
inject1 a
x

    {-# INLINE_LATE step #-}
    step :: (s, Int) -> m (Step (s, Int) b)
step (s
st, Int
i) | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n = do
        (\case
            Yield b
x s
s -> b -> (s, Int) -> Step (s, Int) b
forall s a. a -> s -> Step s a
Yield b
x (s
s, Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
            Skip s
s -> (s, Int) -> Step (s, Int) b
forall s a. s -> Step s a
Skip (s
s, Int
i)
            Step s b
Stop   -> Step (s, Int) b
forall s a. Step s a
Stop) (Step s b -> Step (s, Int) b)
-> m (Step s b) -> m (Step (s, Int) b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> s -> m (Step s b)
step1 s
st
    step (s
_, Int
_) = Step (s, Int) b -> m (Step (s, Int) b)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Step (s, Int) b
forall s a. Step s a
Stop

-- | Same as 'filter' but with a monadic predicate.
--
{-# INLINE_NORMAL filterM #-}
filterM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
filterM :: forall (m :: * -> *) b a.
Monad m =>
(b -> m Bool) -> Unfold m a b -> Unfold m a b
filterM b -> m Bool
f (Unfold s -> m (Step s b)
step1 a -> m s
inject1) = (s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold s -> m (Step s b)
step a -> m s
inject1
  where
    {-# INLINE_LATE step #-}
    step :: s -> m (Step s b)
step s
st = do
        Step s b
r <- s -> m (Step s b)
step1 s
st
        case Step s b
r of
            Yield b
x s
s -> do
                Bool
b <- b -> m Bool
f b
x
                Step s b -> m (Step s b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step s b -> m (Step s b)) -> Step s b -> m (Step s b)
forall a b. (a -> b) -> a -> b
$ if Bool
b then b -> s -> Step s b
forall s a. a -> s -> Step s a
Yield b
x s
s else s -> Step s b
forall s a. s -> Step s a
Skip s
s
            Skip s
s -> Step s b -> m (Step s b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step s b -> m (Step s b)) -> Step s b -> m (Step s b)
forall a b. (a -> b) -> a -> b
$ s -> Step s b
forall s a. s -> Step s a
Skip s
s
            Step s b
Stop   -> Step s b -> m (Step s b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Step s b
forall s a. Step s a
Stop

-- | Include only those elements that pass a predicate.
--
{-# INLINE filter #-}
filter :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
filter :: forall (m :: * -> *) b a.
Monad m =>
(b -> Bool) -> Unfold m a b -> Unfold m a b
filter b -> Bool
f = (b -> m Bool) -> Unfold m a b -> Unfold m a b
forall (m :: * -> *) b a.
Monad m =>
(b -> m Bool) -> Unfold m a b -> Unfold m a b
filterM (Bool -> m Bool
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> m Bool) -> (b -> Bool) -> b -> m Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> Bool
f)

-- | @drop n unf@ drops @n@ elements from the stream generated by @unf@.
--
{-# INLINE_NORMAL drop #-}
drop :: Applicative m => Int -> Unfold m a b -> Unfold m a b
drop :: forall (m :: * -> *) a b.
Applicative m =>
Int -> Unfold m a b -> Unfold m a b
drop Int
n (Unfold s -> m (Step s b)
step a -> m s
inject) = ((s, Int) -> m (Step (s, Int) b))
-> (a -> m (s, Int)) -> Unfold m a b
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold (s, Int) -> m (Step (s, Int) b)
forall {b}. (Ord b, Num b) => (s, b) -> m (Step (s, b) b)
step' a -> m (s, Int)
inject'

    where

    inject' :: a -> m (s, Int)
inject' a
a = (, Int
n) (s -> (s, Int)) -> m s -> m (s, Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> m s
inject a
a

    {-# INLINE_LATE step' #-}
    step' :: (s, b) -> m (Step (s, b) b)
step' (s
st, b
i)
        | b
i b -> b -> Bool
forall a. Ord a => a -> a -> Bool
> b
0 = do
            (\case
                  Yield b
_ s
s -> (s, b) -> Step (s, b) b
forall s a. s -> Step s a
Skip (s
s, b
i b -> b -> b
forall a. Num a => a -> a -> a
- b
1)
                  Skip s
s -> (s, b) -> Step (s, b) b
forall s a. s -> Step s a
Skip (s
s, b
i)
                  Step s b
Stop -> Step (s, b) b
forall s a. Step s a
Stop) (Step s b -> Step (s, b) b) -> m (Step s b) -> m (Step (s, b) b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> s -> m (Step s b)
step s
st
        | Bool
otherwise = do
            (\case
                  Yield b
x s
s -> b -> (s, b) -> Step (s, b) b
forall s a. a -> s -> Step s a
Yield b
x (s
s, b
0)
                  Skip s
s -> (s, b) -> Step (s, b) b
forall s a. s -> Step s a
Skip (s
s, b
0)
                  Step s b
Stop -> Step (s, b) b
forall s a. Step s a
Stop) (Step s b -> Step (s, b) b) -> m (Step s b) -> m (Step (s, b) b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> s -> m (Step s b)
step s
st

-- | @dropWhileM f unf@ drops elements from the stream generated by @unf@ while
-- the condition holds true. The condition function @f@ is /monadic/ in nature.
--
{-# INLINE_NORMAL dropWhileM #-}
dropWhileM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
dropWhileM :: forall (m :: * -> *) b a.
Monad m =>
(b -> m Bool) -> Unfold m a b -> Unfold m a b
dropWhileM b -> m Bool
f (Unfold s -> m (Step s b)
step a -> m s
inject) = (Either s s -> m (Step (Either s s) b))
-> (a -> m (Either s s)) -> Unfold m a b
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold Either s s -> m (Step (Either s s) b)
step' a -> m (Either s s)
forall {b}. a -> m (Either s b)
inject'

    where

    inject' :: a -> m (Either s b)
inject' a
a = do
        s
b <- a -> m s
inject a
a
        Either s b -> m (Either s b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either s b -> m (Either s b)) -> Either s b -> m (Either s b)
forall a b. (a -> b) -> a -> b
$ s -> Either s b
forall a b. a -> Either a b
Left s
b

    {-# INLINE_LATE step' #-}
    step' :: Either s s -> m (Step (Either s s) b)
step' (Left s
st) = do
        Step s b
r <- s -> m (Step s b)
step s
st
        case Step s b
r of
            Yield b
x s
s -> do
                Bool
b <- b -> m Bool
f b
x
                Step (Either s s) b -> m (Step (Either s s) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return
                    (Step (Either s s) b -> m (Step (Either s s) b))
-> Step (Either s s) b -> m (Step (Either s s) b)
forall a b. (a -> b) -> a -> b
$ if Bool
b
                      then Either s s -> Step (Either s s) b
forall s a. s -> Step s a
Skip (s -> Either s s
forall a b. a -> Either a b
Left s
s)
                      else b -> Either s s -> Step (Either s s) b
forall s a. a -> s -> Step s a
Yield b
x (s -> Either s s
forall a b. b -> Either a b
Right s
s)
            Skip s
s -> Step (Either s s) b -> m (Step (Either s s) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Either s s) b -> m (Step (Either s s) b))
-> Step (Either s s) b -> m (Step (Either s s) b)
forall a b. (a -> b) -> a -> b
$ Either s s -> Step (Either s s) b
forall s a. s -> Step s a
Skip (s -> Either s s
forall a b. a -> Either a b
Left s
s)
            Step s b
Stop -> Step (Either s s) b -> m (Step (Either s s) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Step (Either s s) b
forall s a. Step s a
Stop
    step' (Right s
st) = do
        Step s b
r <- s -> m (Step s b)
step s
st
        Step (Either s s) b -> m (Step (Either s s) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return
            (Step (Either s s) b -> m (Step (Either s s) b))
-> Step (Either s s) b -> m (Step (Either s s) b)
forall a b. (a -> b) -> a -> b
$ case Step s b
r of
                  Yield b
x s
s -> b -> Either s s -> Step (Either s s) b
forall s a. a -> s -> Step s a
Yield b
x (s -> Either s s
forall a b. b -> Either a b
Right s
s)
                  Skip s
s -> Either s s -> Step (Either s s) b
forall s a. s -> Step s a
Skip (s -> Either s s
forall a b. b -> Either a b
Right s
s)
                  Step s b
Stop -> Step (Either s s) b
forall s a. Step s a
Stop

-- | Similar to 'dropWhileM' but with a pure condition function.
--
{-# INLINE dropWhile #-}
dropWhile :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
dropWhile :: forall (m :: * -> *) b a.
Monad m =>
(b -> Bool) -> Unfold m a b -> Unfold m a b
dropWhile b -> Bool
f = (b -> m Bool) -> Unfold m a b -> Unfold m a b
forall (m :: * -> *) b a.
Monad m =>
(b -> m Bool) -> Unfold m a b -> Unfold m a b
dropWhileM (Bool -> m Bool
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> m Bool) -> (b -> Bool) -> b -> m Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> Bool
f)

{-# INLINE_NORMAL joinInnerGeneric #-}
joinInnerGeneric :: Monad m =>
    (b -> c -> Bool) -> Unfold m a b -> Unfold m a c -> Unfold m a (b, c)
joinInnerGeneric :: forall (m :: * -> *) b c a.
Monad m =>
(b -> c -> Bool)
-> Unfold m a b -> Unfold m a c -> Unfold m a (b, c)
joinInnerGeneric b -> c -> Bool
eq Unfold m a b
s1 Unfold m a c
s2 = ((b, c) -> Bool) -> Unfold m a (b, c) -> Unfold m a (b, c)
forall (m :: * -> *) b a.
Monad m =>
(b -> Bool) -> Unfold m a b -> Unfold m a b
filter (\(b
a, c
b) -> b
a b -> c -> Bool
`eq` c
b) (Unfold m a (b, c) -> Unfold m a (b, c))
-> Unfold m a (b, c) -> Unfold m a (b, c)
forall a b. (a -> b) -> a -> b
$ Unfold m a b -> Unfold m a c -> Unfold m a (b, c)
forall (m :: * -> *) a b c.
Monad m =>
Unfold m a b -> Unfold m a c -> Unfold m a (b, c)
cross Unfold m a b
s1 Unfold m a c
s2

------------------------------------------------------------------------------
-- Exceptions
------------------------------------------------------------------------------

-- | Like 'gbracketIO' but with following differences:
--
-- * alloc action @a -> m c@ runs with async exceptions enabled
-- * cleanup action @c -> m d@ won't run if the stream is garbage collected
--   after partial evaluation.
--
-- /Inhibits stream fusion/
--
-- /Pre-release/
--
{-# INLINE_NORMAL gbracket_ #-}
gbracket_
    :: Monad m
    => (a -> m c)                           -- ^ before
    -> (forall s. m s -> m (Either e s))    -- ^ try (exception handling)
    -> (c -> m d)                           -- ^ after, on normal stop
    -> Unfold m (c, e) b                    -- ^ on exception
    -> Unfold m c b                         -- ^ unfold to run
    -> Unfold m a b
gbracket_ :: forall (m :: * -> *) a c e d b.
Monad m =>
(a -> m c)
-> (forall s. m s -> m (Either e s))
-> (c -> m d)
-> Unfold m (c, e) b
-> Unfold m c b
-> Unfold m a b
gbracket_ a -> m c
bef forall s. m s -> m (Either e s)
exc c -> m d
aft (Unfold s -> m (Step s b)
estep (c, e) -> m s
einject) (Unfold s -> m (Step s b)
step1 c -> m s
inject1) =
    (Either s (s, c) -> m (Step (Either s (s, c)) b))
-> (a -> m (Either s (s, c))) -> Unfold m a b
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold Either s (s, c) -> m (Step (Either s (s, c)) b)
step a -> m (Either s (s, c))
forall {a}. a -> m (Either a (s, c))
inject

    where

    inject :: a -> m (Either a (s, c))
inject a
x = do
        c
r <- a -> m c
bef a
x
        s
s <- c -> m s
inject1 c
r
        Either a (s, c) -> m (Either a (s, c))
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either a (s, c) -> m (Either a (s, c)))
-> Either a (s, c) -> m (Either a (s, c))
forall a b. (a -> b) -> a -> b
$ (s, c) -> Either a (s, c)
forall a b. b -> Either a b
Right (s
s, c
r)

    {-# INLINE_LATE step #-}
    step :: Either s (s, c) -> m (Step (Either s (s, c)) b)
step (Right (s
st, c
v)) = do
        Either e (Step s b)
res <- m (Step s b) -> m (Either e (Step s b))
forall s. m s -> m (Either e s)
exc (m (Step s b) -> m (Either e (Step s b)))
-> m (Step s b) -> m (Either e (Step s b))
forall a b. (a -> b) -> a -> b
$ s -> m (Step s b)
step1 s
st
        case Either e (Step s b)
res of
            Right Step s b
r -> case Step s b
r of
                Yield b
x s
s -> Step (Either s (s, c)) b -> m (Step (Either s (s, c)) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Either s (s, c)) b -> m (Step (Either s (s, c)) b))
-> Step (Either s (s, c)) b -> m (Step (Either s (s, c)) b)
forall a b. (a -> b) -> a -> b
$ b -> Either s (s, c) -> Step (Either s (s, c)) b
forall s a. a -> s -> Step s a
Yield b
x ((s, c) -> Either s (s, c)
forall a b. b -> Either a b
Right (s
s, c
v))
                Skip s
s    -> Step (Either s (s, c)) b -> m (Step (Either s (s, c)) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Either s (s, c)) b -> m (Step (Either s (s, c)) b))
-> Step (Either s (s, c)) b -> m (Step (Either s (s, c)) b)
forall a b. (a -> b) -> a -> b
$ Either s (s, c) -> Step (Either s (s, c)) b
forall s a. s -> Step s a
Skip ((s, c) -> Either s (s, c)
forall a b. b -> Either a b
Right (s
s, c
v))
                Step s b
Stop      -> c -> m d
aft c
v m d -> m (Step (Either s (s, c)) b) -> m (Step (Either s (s, c)) b)
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Step (Either s (s, c)) b -> m (Step (Either s (s, c)) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Step (Either s (s, c)) b
forall s a. Step s a
Stop
            -- XXX Do not handle async exceptions, just rethrow them.
            Left e
e -> do
                s
r <- (c, e) -> m s
einject (c
v, e
e)
                Step (Either s (s, c)) b -> m (Step (Either s (s, c)) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Either s (s, c)) b -> m (Step (Either s (s, c)) b))
-> Step (Either s (s, c)) b -> m (Step (Either s (s, c)) b)
forall a b. (a -> b) -> a -> b
$ Either s (s, c) -> Step (Either s (s, c)) b
forall s a. s -> Step s a
Skip (s -> Either s (s, c)
forall a b. a -> Either a b
Left s
r)
    step (Left s
st) = do
        Step s b
res <- s -> m (Step s b)
estep s
st
        Step (Either s (s, c)) b -> m (Step (Either s (s, c)) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Either s (s, c)) b -> m (Step (Either s (s, c)) b))
-> Step (Either s (s, c)) b -> m (Step (Either s (s, c)) b)
forall a b. (a -> b) -> a -> b
$ case Step s b
res of
            Yield b
x s
s -> b -> Either s (s, c) -> Step (Either s (s, c)) b
forall s a. a -> s -> Step s a
Yield b
x (s -> Either s (s, c)
forall a b. a -> Either a b
Left s
s)
            Skip s
s    -> Either s (s, c) -> Step (Either s (s, c)) b
forall s a. s -> Step s a
Skip (s -> Either s (s, c)
forall a b. a -> Either a b
Left s
s)
            Step s b
Stop      -> Step (Either s (s, c)) b
forall s a. Step s a
Stop

-- | Run the alloc action @a -> m c@ with async exceptions disabled but keeping
-- blocking operations interruptible (see 'Control.Exception.mask').  Use the
-- output @c@ as input to @Unfold m c b@ to generate an output stream. When
-- unfolding use the supplied @try@ operation @forall s. m s -> m (Either e s)@
-- to catch synchronous exceptions. If an exception occurs run the exception
-- handling unfold @Unfold m (c, e) b@.
--
-- The cleanup action @c -> m d@, runs whenever the stream ends normally, due
-- to a sync or async exception or if it gets garbage collected after a partial
-- lazy evaluation.  See 'bracket' for the semantics of the cleanup action.
--
-- 'gbracket' can express all other exception handling combinators.
--
-- /Inhibits stream fusion/
--
-- /Pre-release/
{-# INLINE_NORMAL gbracketIO #-}
gbracketIO
    :: MonadIO m
    => (a -> IO c)                           -- ^ before
    -> (c -> IO d)                           -- ^ after, on normal stop, or GC
    -> (c -> IO ())                          -- ^ action on exception
    -> Unfold m e b                          -- ^ stream on exception
    -> (forall s. m s -> IO (Either e s))    -- ^ try (exception handling)
    -> Unfold m c b                         -- ^ unfold to run
    -> Unfold m a b
gbracketIO :: forall (m :: * -> *) a c d e b.
MonadIO m =>
(a -> IO c)
-> (c -> IO d)
-> (c -> IO ())
-> Unfold m e b
-> (forall s. m s -> IO (Either e s))
-> Unfold m c b
-> Unfold m a b
gbracketIO a -> IO c
bef c -> IO d
aft c -> IO ()
onExc (Unfold s -> m (Step s b)
estep e -> m s
einject) forall s. m s -> IO (Either e s)
ftry (Unfold s -> m (Step s b)
step1 c -> m s
inject1) =
    (Either s (s, c, IOFinalizer)
 -> m (Step (Either s (s, c, IOFinalizer)) b))
-> (a -> m (Either s (s, c, IOFinalizer))) -> Unfold m a b
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold Either s (s, c, IOFinalizer)
-> m (Step (Either s (s, c, IOFinalizer)) b)
step a -> m (Either s (s, c, IOFinalizer))
forall {a}. a -> m (Either a (s, c, IOFinalizer))
inject

    where

    inject :: a -> m (Either a (s, c, IOFinalizer))
inject a
x = do
        -- Mask asynchronous exceptions to make the execution of 'bef' and
        -- the registration of 'aft' atomic. See comment in 'D.gbracketIO'.
        (c
r, IOFinalizer
ref) <- IO (c, IOFinalizer) -> m (c, IOFinalizer)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (c, IOFinalizer) -> m (c, IOFinalizer))
-> IO (c, IOFinalizer) -> m (c, IOFinalizer)
forall a b. (a -> b) -> a -> b
$ IO (c, IOFinalizer) -> IO (c, IOFinalizer)
forall a. IO a -> IO a
mask_ (IO (c, IOFinalizer) -> IO (c, IOFinalizer))
-> IO (c, IOFinalizer) -> IO (c, IOFinalizer)
forall a b. (a -> b) -> a -> b
$ do
            c
r <- a -> IO c
bef a
x
            IOFinalizer
ref <- IO d -> IO IOFinalizer
forall (m :: * -> *) a. MonadIO m => IO a -> m IOFinalizer
newIOFinalizer (c -> IO d
aft c
r)
            (c, IOFinalizer) -> IO (c, IOFinalizer)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (c
r, IOFinalizer
ref)
        s
s <- c -> m s
inject1 c
r
        Either a (s, c, IOFinalizer) -> m (Either a (s, c, IOFinalizer))
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either a (s, c, IOFinalizer) -> m (Either a (s, c, IOFinalizer)))
-> Either a (s, c, IOFinalizer) -> m (Either a (s, c, IOFinalizer))
forall a b. (a -> b) -> a -> b
$ (s, c, IOFinalizer) -> Either a (s, c, IOFinalizer)
forall a b. b -> Either a b
Right (s
s, c
r, IOFinalizer
ref)

    {-# INLINE_LATE step #-}
    step :: Either s (s, c, IOFinalizer)
-> m (Step (Either s (s, c, IOFinalizer)) b)
step (Right (s
st, c
v, IOFinalizer
ref)) = do
        Either e (Step s b)
res <- IO (Either e (Step s b)) -> m (Either e (Step s b))
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Either e (Step s b)) -> m (Either e (Step s b)))
-> IO (Either e (Step s b)) -> m (Either e (Step s b))
forall a b. (a -> b) -> a -> b
$ m (Step s b) -> IO (Either e (Step s b))
forall s. m s -> IO (Either e s)
ftry (m (Step s b) -> IO (Either e (Step s b)))
-> m (Step s b) -> IO (Either e (Step s b))
forall a b. (a -> b) -> a -> b
$ s -> m (Step s b)
step1 s
st
        case Either e (Step s b)
res of
            Right Step s b
r -> case Step s b
r of
                Yield b
x s
s -> Step (Either s (s, c, IOFinalizer)) b
-> m (Step (Either s (s, c, IOFinalizer)) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Either s (s, c, IOFinalizer)) b
 -> m (Step (Either s (s, c, IOFinalizer)) b))
-> Step (Either s (s, c, IOFinalizer)) b
-> m (Step (Either s (s, c, IOFinalizer)) b)
forall a b. (a -> b) -> a -> b
$ b
-> Either s (s, c, IOFinalizer)
-> Step (Either s (s, c, IOFinalizer)) b
forall s a. a -> s -> Step s a
Yield b
x ((s, c, IOFinalizer) -> Either s (s, c, IOFinalizer)
forall a b. b -> Either a b
Right (s
s, c
v, IOFinalizer
ref))
                Skip s
s    -> Step (Either s (s, c, IOFinalizer)) b
-> m (Step (Either s (s, c, IOFinalizer)) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Either s (s, c, IOFinalizer)) b
 -> m (Step (Either s (s, c, IOFinalizer)) b))
-> Step (Either s (s, c, IOFinalizer)) b
-> m (Step (Either s (s, c, IOFinalizer)) b)
forall a b. (a -> b) -> a -> b
$ Either s (s, c, IOFinalizer)
-> Step (Either s (s, c, IOFinalizer)) b
forall s a. s -> Step s a
Skip ((s, c, IOFinalizer) -> Either s (s, c, IOFinalizer)
forall a b. b -> Either a b
Right (s
s, c
v, IOFinalizer
ref))
                Step s b
Stop      -> do
                    IOFinalizer -> m ()
forall (m :: * -> *). MonadIO m => IOFinalizer -> m ()
runIOFinalizer IOFinalizer
ref
                    Step (Either s (s, c, IOFinalizer)) b
-> m (Step (Either s (s, c, IOFinalizer)) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Step (Either s (s, c, IOFinalizer)) b
forall s a. Step s a
Stop
            -- XXX Do not handle async exceptions, just rethrow them.
            Left e
e -> do
                -- Clearing of finalizer and running of exception handler must
                -- be atomic wrt async exceptions. Otherwise if we have cleared
                -- the finalizer and have not run the exception handler then we
                -- may leak the resource.
                IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IOFinalizer -> IO () -> IO ()
forall (m :: * -> *) a. MonadIO m => IOFinalizer -> IO a -> m a
clearingIOFinalizer IOFinalizer
ref (c -> IO ()
onExc c
v)
                s
r <- e -> m s
einject e
e
                Step (Either s (s, c, IOFinalizer)) b
-> m (Step (Either s (s, c, IOFinalizer)) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Either s (s, c, IOFinalizer)) b
 -> m (Step (Either s (s, c, IOFinalizer)) b))
-> Step (Either s (s, c, IOFinalizer)) b
-> m (Step (Either s (s, c, IOFinalizer)) b)
forall a b. (a -> b) -> a -> b
$ Either s (s, c, IOFinalizer)
-> Step (Either s (s, c, IOFinalizer)) b
forall s a. s -> Step s a
Skip (s -> Either s (s, c, IOFinalizer)
forall a b. a -> Either a b
Left s
r)
    step (Left s
st) = do
        Step s b
res <- s -> m (Step s b)
estep s
st
        Step (Either s (s, c, IOFinalizer)) b
-> m (Step (Either s (s, c, IOFinalizer)) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Either s (s, c, IOFinalizer)) b
 -> m (Step (Either s (s, c, IOFinalizer)) b))
-> Step (Either s (s, c, IOFinalizer)) b
-> m (Step (Either s (s, c, IOFinalizer)) b)
forall a b. (a -> b) -> a -> b
$ case Step s b
res of
            Yield b
x s
s -> b
-> Either s (s, c, IOFinalizer)
-> Step (Either s (s, c, IOFinalizer)) b
forall s a. a -> s -> Step s a
Yield b
x (s -> Either s (s, c, IOFinalizer)
forall a b. a -> Either a b
Left s
s)
            Skip s
s    -> Either s (s, c, IOFinalizer)
-> Step (Either s (s, c, IOFinalizer)) b
forall s a. s -> Step s a
Skip (s -> Either s (s, c, IOFinalizer)
forall a b. a -> Either a b
Left s
s)
            Step s b
Stop      -> Step (Either s (s, c, IOFinalizer)) b
forall s a. Step s a
Stop

-- | Run a side effect @a -> m c@ on the input @a@ before unfolding it using
-- @Unfold m a b@.
--
-- > before f = lmapM (\a -> f a >> return a)
--
-- /Pre-release/
{-# INLINE_NORMAL before #-}
before :: (a -> m c) -> Unfold m a b -> Unfold m a b
before :: forall a (m :: * -> *) c b.
(a -> m c) -> Unfold m a b -> Unfold m a b
before a -> m c
action (Unfold s -> m (Step s b)
step a -> m s
inject) = (s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold s -> m (Step s b)
step (a -> m c
action (a -> m c) -> (a -> m s) -> a -> m s
forall a b. (a -> a) -> (a -> b) -> a -> b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> a -> m s
inject)

-- The custom implementation of "after_" is slightly faster (5-7%) than
-- "_after".  This is just to document and make sure that we can always use
-- gbracket to implement after_ The same applies to other combinators as well.
--
{-# INLINE_NORMAL _after #-}
_after :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
_after :: forall (m :: * -> *) a c b.
Monad m =>
(a -> m c) -> Unfold m a b -> Unfold m a b
_after a -> m c
aft = (a -> m a)
-> (forall s. m s -> m (Either Any s))
-> (a -> m c)
-> Unfold m (a, Any) b
-> Unfold m a b
-> Unfold m a b
forall (m :: * -> *) a c e d b.
Monad m =>
(a -> m c)
-> (forall s. m s -> m (Either e s))
-> (c -> m d)
-> Unfold m (c, e) b
-> Unfold m c b
-> Unfold m a b
gbracket_ a -> m a
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ((s -> Either Any s) -> m s -> m (Either Any s)
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap s -> Either Any s
forall a b. b -> Either a b
Right) a -> m c
aft Unfold m (a, Any) b
forall a. HasCallStack => a
undefined

-- | Like 'after' with following differences:
--
-- * action @a -> m c@ won't run if the stream is garbage collected
--   after partial evaluation.
-- * Monad @m@ does not require any other constraints.
--
-- /Pre-release/
{-# INLINE_NORMAL after_ #-}
after_ :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b
after_ :: forall (m :: * -> *) a c b.
Monad m =>
(a -> m c) -> Unfold m a b -> Unfold m a b
after_ a -> m c
action (Unfold s -> m (Step s b)
step1 a -> m s
inject1) = ((s, a) -> m (Step (s, a) b)) -> (a -> m (s, a)) -> Unfold m a b
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold (s, a) -> m (Step (s, a) b)
step a -> m (s, a)
inject

    where

    inject :: a -> m (s, a)
inject a
x = do
        s
s <- a -> m s
inject1 a
x
        (s, a) -> m (s, a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (s
s, a
x)

    {-# INLINE_LATE step #-}
    step :: (s, a) -> m (Step (s, a) b)
step (s
st, a
v) = do
        Step s b
res <- s -> m (Step s b)
step1 s
st
        case Step s b
res of
            Yield b
x s
s -> Step (s, a) b -> m (Step (s, a) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (s, a) b -> m (Step (s, a) b))
-> Step (s, a) b -> m (Step (s, a) b)
forall a b. (a -> b) -> a -> b
$ b -> (s, a) -> Step (s, a) b
forall s a. a -> s -> Step s a
Yield b
x (s
s, a
v)
            Skip s
s    -> Step (s, a) b -> m (Step (s, a) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (s, a) b -> m (Step (s, a) b))
-> Step (s, a) b -> m (Step (s, a) b)
forall a b. (a -> b) -> a -> b
$ (s, a) -> Step (s, a) b
forall s a. s -> Step s a
Skip (s
s, a
v)
            Step s b
Stop      -> a -> m c
action a
v m c -> m (Step (s, a) b) -> m (Step (s, a) b)
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Step (s, a) b -> m (Step (s, a) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Step (s, a) b
forall s a. Step s a
Stop

-- | Unfold the input @a@ using @Unfold m a b@, run an action on @a@ whenever
-- the unfold stops normally, or if it is garbage collected after a partial
-- lazy evaluation.
--
-- The semantics of the action @a -> m c@ are similar to the cleanup action
-- semantics in 'bracket'.
--
-- /See also 'after_'/
--
-- /Pre-release/
{-# INLINE_NORMAL afterIO #-}
afterIO :: MonadIO m
    => (a -> IO c) -> Unfold m a b -> Unfold m a b
afterIO :: forall (m :: * -> *) a c b.
MonadIO m =>
(a -> IO c) -> Unfold m a b -> Unfold m a b
afterIO a -> IO c
action (Unfold s -> m (Step s b)
step1 a -> m s
inject1) = ((s, IOFinalizer) -> m (Step (s, IOFinalizer) b))
-> (a -> m (s, IOFinalizer)) -> Unfold m a b
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold (s, IOFinalizer) -> m (Step (s, IOFinalizer) b)
step a -> m (s, IOFinalizer)
inject

    where

    inject :: a -> m (s, IOFinalizer)
inject a
x = do
        s
s <- a -> m s
inject1 a
x
        IOFinalizer
ref <- IO IOFinalizer -> m IOFinalizer
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO IOFinalizer -> m IOFinalizer)
-> IO IOFinalizer -> m IOFinalizer
forall a b. (a -> b) -> a -> b
$ IO c -> IO IOFinalizer
forall (m :: * -> *) a. MonadIO m => IO a -> m IOFinalizer
newIOFinalizer (a -> IO c
action a
x)
        (s, IOFinalizer) -> m (s, IOFinalizer)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (s
s, IOFinalizer
ref)

    {-# INLINE_LATE step #-}
    step :: (s, IOFinalizer) -> m (Step (s, IOFinalizer) b)
step (s
st, IOFinalizer
ref) = do
        Step s b
res <- s -> m (Step s b)
step1 s
st
        case Step s b
res of
            Yield b
x s
s -> Step (s, IOFinalizer) b -> m (Step (s, IOFinalizer) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (s, IOFinalizer) b -> m (Step (s, IOFinalizer) b))
-> Step (s, IOFinalizer) b -> m (Step (s, IOFinalizer) b)
forall a b. (a -> b) -> a -> b
$ b -> (s, IOFinalizer) -> Step (s, IOFinalizer) b
forall s a. a -> s -> Step s a
Yield b
x (s
s, IOFinalizer
ref)
            Skip s
s    -> Step (s, IOFinalizer) b -> m (Step (s, IOFinalizer) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (s, IOFinalizer) b -> m (Step (s, IOFinalizer) b))
-> Step (s, IOFinalizer) b -> m (Step (s, IOFinalizer) b)
forall a b. (a -> b) -> a -> b
$ (s, IOFinalizer) -> Step (s, IOFinalizer) b
forall s a. s -> Step s a
Skip (s
s, IOFinalizer
ref)
            Step s b
Stop      -> do
                IOFinalizer -> m ()
forall (m :: * -> *). MonadIO m => IOFinalizer -> m ()
runIOFinalizer IOFinalizer
ref
                Step (s, IOFinalizer) b -> m (Step (s, IOFinalizer) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Step (s, IOFinalizer) b
forall s a. Step s a
Stop

{-# INLINE_NORMAL _onException #-}
_onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
_onException :: forall (m :: * -> *) a c b.
MonadCatch m =>
(a -> m c) -> Unfold m a b -> Unfold m a b
_onException a -> m c
action =
    (a -> m a)
-> (forall s. m s -> m (Either SomeException s))
-> (a -> m ())
-> Unfold m (a, SomeException) b
-> Unfold m a b
-> Unfold m a b
forall (m :: * -> *) a c e d b.
Monad m =>
(a -> m c)
-> (forall s. m s -> m (Either e s))
-> (c -> m d)
-> Unfold m (c, e) b
-> Unfold m c b
-> Unfold m a b
gbracket_ a -> m a
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return m s -> m (Either SomeException s)
forall s. m s -> m (Either SomeException s)
forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
m a -> m (Either e a)
MC.try
        (\a
_ -> () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ())
        (((a, SomeException) -> m Any) -> Unfold m (a, SomeException) b
forall (m :: * -> *) a c b.
Applicative m =>
(a -> m c) -> Unfold m a b
nilM (\(a
a, SomeException
e :: MC.SomeException) -> a -> m c
action a
a m c -> m Any -> m Any
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> SomeException -> m Any
forall e a. Exception e => e -> m a
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
MC.throwM SomeException
e))

-- | Unfold the input @a@ using @Unfold m a b@, run the action @a -> m c@ on
-- @a@ if the unfold aborts due to an exception.
--
-- /Inhibits stream fusion/
--
-- /Pre-release/
{-# INLINE_NORMAL onException #-}
onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
onException :: forall (m :: * -> *) a c b.
MonadCatch m =>
(a -> m c) -> Unfold m a b -> Unfold m a b
onException a -> m c
action (Unfold s -> m (Step s b)
step1 a -> m s
inject1) = ((s, a) -> m (Step (s, a) b)) -> (a -> m (s, a)) -> Unfold m a b
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold (s, a) -> m (Step (s, a) b)
step a -> m (s, a)
inject

    where

    inject :: a -> m (s, a)
inject a
x = do
        s
s <- a -> m s
inject1 a
x
        (s, a) -> m (s, a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (s
s, a
x)

    {-# INLINE_LATE step #-}
    step :: (s, a) -> m (Step (s, a) b)
step (s
st, a
v) = do
        Step s b
res <- s -> m (Step s b)
step1 s
st m (Step s b) -> m c -> m (Step s b)
forall (m :: * -> *) a b. MonadCatch m => m a -> m b -> m a
`MC.onException` a -> m c
action a
v
        Step (s, a) b -> m (Step (s, a) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (s, a) b -> m (Step (s, a) b))
-> Step (s, a) b -> m (Step (s, a) b)
forall a b. (a -> b) -> a -> b
$ case Step s b
res of
            Yield b
x s
s -> b -> (s, a) -> Step (s, a) b
forall s a. a -> s -> Step s a
Yield b
x (s
s, a
v)
            Skip s
s    -> (s, a) -> Step (s, a) b
forall s a. s -> Step s a
Skip (s
s, a
v)
            Step s b
Stop      -> Step (s, a) b
forall s a. Step s a
Stop

{-# INLINE_NORMAL _finally #-}
_finally :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
_finally :: forall (m :: * -> *) a c b.
MonadCatch m =>
(a -> m c) -> Unfold m a b -> Unfold m a b
_finally a -> m c
action =
    (a -> m a)
-> (forall s. m s -> m (Either SomeException s))
-> (a -> m c)
-> Unfold m (a, SomeException) b
-> Unfold m a b
-> Unfold m a b
forall (m :: * -> *) a c e d b.
Monad m =>
(a -> m c)
-> (forall s. m s -> m (Either e s))
-> (c -> m d)
-> Unfold m (c, e) b
-> Unfold m c b
-> Unfold m a b
gbracket_ a -> m a
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return m s -> m (Either SomeException s)
forall s. m s -> m (Either SomeException s)
forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
m a -> m (Either e a)
MC.try a -> m c
action
        (((a, SomeException) -> m Any) -> Unfold m (a, SomeException) b
forall (m :: * -> *) a c b.
Applicative m =>
(a -> m c) -> Unfold m a b
nilM (\(a
a, SomeException
e :: MC.SomeException) -> a -> m c
action a
a m c -> m Any -> m Any
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> SomeException -> m Any
forall e a. Exception e => e -> m a
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
MC.throwM SomeException
e))

-- | Like 'finallyIO' with following differences:
--
-- * action @a -> m c@ won't run if the stream is garbage collected
--   after partial evaluation.
--
-- /Inhibits stream fusion/
--
-- /Pre-release/
{-# INLINE_NORMAL finally_ #-}
finally_ :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b
finally_ :: forall (m :: * -> *) a c b.
MonadCatch m =>
(a -> m c) -> Unfold m a b -> Unfold m a b
finally_ a -> m c
action (Unfold s -> m (Step s b)
step1 a -> m s
inject1) = ((s, a) -> m (Step (s, a) b)) -> (a -> m (s, a)) -> Unfold m a b
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold (s, a) -> m (Step (s, a) b)
step a -> m (s, a)
inject

    where

    inject :: a -> m (s, a)
inject a
x = do
        s
s <- a -> m s
inject1 a
x
        (s, a) -> m (s, a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (s
s, a
x)

    {-# INLINE_LATE step #-}
    step :: (s, a) -> m (Step (s, a) b)
step (s
st, a
v) = do
        Step s b
res <- s -> m (Step s b)
step1 s
st m (Step s b) -> m c -> m (Step s b)
forall (m :: * -> *) a b. MonadCatch m => m a -> m b -> m a
`MC.onException` a -> m c
action a
v
        case Step s b
res of
            Yield b
x s
s -> Step (s, a) b -> m (Step (s, a) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (s, a) b -> m (Step (s, a) b))
-> Step (s, a) b -> m (Step (s, a) b)
forall a b. (a -> b) -> a -> b
$ b -> (s, a) -> Step (s, a) b
forall s a. a -> s -> Step s a
Yield b
x (s
s, a
v)
            Skip s
s    -> Step (s, a) b -> m (Step (s, a) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (s, a) b -> m (Step (s, a) b))
-> Step (s, a) b -> m (Step (s, a) b)
forall a b. (a -> b) -> a -> b
$ (s, a) -> Step (s, a) b
forall s a. s -> Step s a
Skip (s
s, a
v)
            Step s b
Stop      -> a -> m c
action a
v m c -> m (Step (s, a) b) -> m (Step (s, a) b)
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Step (s, a) b -> m (Step (s, a) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Step (s, a) b
forall s a. Step s a
Stop

-- | Unfold the input @a@ using @Unfold m a b@, run an action on @a@ whenever
-- the unfold stops normally, aborts due to an exception or if it is garbage
-- collected after a partial lazy evaluation.
--
-- The semantics of the action @a -> m c@ are similar to the cleanup action
-- semantics in 'bracket'.
--
-- @
-- finally release = bracket return release
-- @
--
-- /See also 'finally_'/
--
-- /Inhibits stream fusion/
--
-- /Pre-release/
{-# INLINE_NORMAL finallyIO #-}
finallyIO :: (MonadIO m, MonadCatch m)
    => (a -> IO c) -> Unfold m a b -> Unfold m a b
finallyIO :: forall (m :: * -> *) a c b.
(MonadIO m, MonadCatch m) =>
(a -> IO c) -> Unfold m a b -> Unfold m a b
finallyIO a -> IO c
action (Unfold s -> m (Step s b)
step1 a -> m s
inject1) = ((s, IOFinalizer) -> m (Step (s, IOFinalizer) b))
-> (a -> m (s, IOFinalizer)) -> Unfold m a b
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold (s, IOFinalizer) -> m (Step (s, IOFinalizer) b)
step a -> m (s, IOFinalizer)
inject

    where

    inject :: a -> m (s, IOFinalizer)
inject a
x = do
        s
s <- a -> m s
inject1 a
x
        IOFinalizer
ref <- IO IOFinalizer -> m IOFinalizer
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO IOFinalizer -> m IOFinalizer)
-> IO IOFinalizer -> m IOFinalizer
forall a b. (a -> b) -> a -> b
$ IO c -> IO IOFinalizer
forall (m :: * -> *) a. MonadIO m => IO a -> m IOFinalizer
newIOFinalizer (a -> IO c
action a
x)
        (s, IOFinalizer) -> m (s, IOFinalizer)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (s
s, IOFinalizer
ref)

    {-# INLINE_LATE step #-}
    step :: (s, IOFinalizer) -> m (Step (s, IOFinalizer) b)
step (s
st, IOFinalizer
ref) = do
        Step s b
res <- s -> m (Step s b)
step1 s
st m (Step s b) -> m () -> m (Step s b)
forall (m :: * -> *) a b. MonadCatch m => m a -> m b -> m a
`MC.onException` IOFinalizer -> m ()
forall (m :: * -> *). MonadIO m => IOFinalizer -> m ()
runIOFinalizer IOFinalizer
ref
        case Step s b
res of
            Yield b
x s
s -> Step (s, IOFinalizer) b -> m (Step (s, IOFinalizer) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (s, IOFinalizer) b -> m (Step (s, IOFinalizer) b))
-> Step (s, IOFinalizer) b -> m (Step (s, IOFinalizer) b)
forall a b. (a -> b) -> a -> b
$ b -> (s, IOFinalizer) -> Step (s, IOFinalizer) b
forall s a. a -> s -> Step s a
Yield b
x (s
s, IOFinalizer
ref)
            Skip s
s    -> Step (s, IOFinalizer) b -> m (Step (s, IOFinalizer) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (s, IOFinalizer) b -> m (Step (s, IOFinalizer) b))
-> Step (s, IOFinalizer) b -> m (Step (s, IOFinalizer) b)
forall a b. (a -> b) -> a -> b
$ (s, IOFinalizer) -> Step (s, IOFinalizer) b
forall s a. s -> Step s a
Skip (s
s, IOFinalizer
ref)
            Step s b
Stop      -> do
                IOFinalizer -> m ()
forall (m :: * -> *). MonadIO m => IOFinalizer -> m ()
runIOFinalizer IOFinalizer
ref
                Step (s, IOFinalizer) b -> m (Step (s, IOFinalizer) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Step (s, IOFinalizer) b
forall s a. Step s a
Stop

{-# INLINE_NORMAL _bracket #-}
_bracket :: MonadCatch m
    => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
_bracket :: forall (m :: * -> *) a c d b.
MonadCatch m =>
(a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
_bracket a -> m c
bef c -> m d
aft =
    (a -> m c)
-> (forall s. m s -> m (Either SomeException s))
-> (c -> m d)
-> Unfold m (c, SomeException) b
-> Unfold m c b
-> Unfold m a b
forall (m :: * -> *) a c e d b.
Monad m =>
(a -> m c)
-> (forall s. m s -> m (Either e s))
-> (c -> m d)
-> Unfold m (c, e) b
-> Unfold m c b
-> Unfold m a b
gbracket_ a -> m c
bef m s -> m (Either SomeException s)
forall s. m s -> m (Either SomeException s)
forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
m a -> m (Either e a)
MC.try c -> m d
aft (((c, SomeException) -> m Any) -> Unfold m (c, SomeException) b
forall (m :: * -> *) a c b.
Applicative m =>
(a -> m c) -> Unfold m a b
nilM (\(c
a, SomeException
e :: MC.SomeException) -> c -> m d
aft c
a m d -> m Any -> m Any
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
    SomeException -> m Any
forall e a. Exception e => e -> m a
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
MC.throwM SomeException
e))

-- | Like 'bracketIO' but with following differences:
--
-- * alloc action @a -> m c@ runs with async exceptions enabled
-- * cleanup action @c -> m d@ won't run if the stream is garbage collected
--   after partial evaluation.
--
-- /Inhibits stream fusion/
--
-- /Pre-release/
{-# INLINE_NORMAL bracket_ #-}
bracket_ :: MonadCatch m
    => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
bracket_ :: forall (m :: * -> *) a c d b.
MonadCatch m =>
(a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
bracket_ a -> m c
bef c -> m d
aft (Unfold s -> m (Step s b)
step1 c -> m s
inject1) = ((s, c) -> m (Step (s, c) b)) -> (a -> m (s, c)) -> Unfold m a b
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold (s, c) -> m (Step (s, c) b)
step a -> m (s, c)
inject

    where

    inject :: a -> m (s, c)
inject a
x = do
        c
r <- a -> m c
bef a
x
        s
s <- c -> m s
inject1 c
r
        (s, c) -> m (s, c)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (s
s, c
r)

    {-# INLINE_LATE step #-}
    step :: (s, c) -> m (Step (s, c) b)
step (s
st, c
v) = do
        Step s b
res <- s -> m (Step s b)
step1 s
st m (Step s b) -> m d -> m (Step s b)
forall (m :: * -> *) a b. MonadCatch m => m a -> m b -> m a
`MC.onException` c -> m d
aft c
v
        case Step s b
res of
            Yield b
x s
s -> Step (s, c) b -> m (Step (s, c) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (s, c) b -> m (Step (s, c) b))
-> Step (s, c) b -> m (Step (s, c) b)
forall a b. (a -> b) -> a -> b
$ b -> (s, c) -> Step (s, c) b
forall s a. a -> s -> Step s a
Yield b
x (s
s, c
v)
            Skip s
s    -> Step (s, c) b -> m (Step (s, c) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (s, c) b -> m (Step (s, c) b))
-> Step (s, c) b -> m (Step (s, c) b)
forall a b. (a -> b) -> a -> b
$ (s, c) -> Step (s, c) b
forall s a. s -> Step s a
Skip (s
s, c
v)
            Step s b
Stop      -> c -> m d
aft c
v m d -> m (Step (s, c) b) -> m (Step (s, c) b)
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Step (s, c) b -> m (Step (s, c) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Step (s, c) b
forall s a. Step s a
Stop

-- | Run the alloc action @a -> m c@ with async exceptions disabled but keeping
-- blocking operations interruptible (see 'Control.Exception.mask').  Use the
-- output @c@ as input to @Unfold m c b@ to generate an output stream.
--
-- @c@ is usually a resource under the state of monad @m@, e.g. a file
-- handle, that requires a cleanup after use. The cleanup action @c -> m d@,
-- runs whenever the stream ends normally, due to a sync or async exception or
-- if it gets garbage collected after a partial lazy evaluation.
--
-- 'bracket' only guarantees that the cleanup action runs, and it runs with
-- async exceptions enabled. The action must ensure that it can successfully
-- cleanup the resource in the face of sync or async exceptions.
--
-- When the stream ends normally or on a sync exception, cleanup action runs
-- immediately in the current thread context, whereas in other cases it runs in
-- the GC context, therefore, cleanup may be delayed until the GC gets to run.
--
-- /See also: 'bracket_', 'gbracket'/
--
-- /Inhibits stream fusion/
--
-- /Pre-release/
{-# INLINE_NORMAL bracketIO #-}
bracketIO :: (MonadIO m, MonadCatch m)
    => (a -> IO c) -> (c -> IO d) -> Unfold m c b -> Unfold m a b
bracketIO :: forall (m :: * -> *) a c d b.
(MonadIO m, MonadCatch m) =>
(a -> IO c) -> (c -> IO d) -> Unfold m c b -> Unfold m a b
bracketIO a -> IO c
bef c -> IO d
aft (Unfold s -> m (Step s b)
step1 c -> m s
inject1) = ((s, IOFinalizer) -> m (Step (s, IOFinalizer) b))
-> (a -> m (s, IOFinalizer)) -> Unfold m a b
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold (s, IOFinalizer) -> m (Step (s, IOFinalizer) b)
step a -> m (s, IOFinalizer)
inject

    where

    inject :: a -> m (s, IOFinalizer)
inject a
x = do
        -- Mask asynchronous exceptions to make the execution of 'bef' and
        -- the registration of 'aft' atomic. See comment in 'D.gbracketIO'.
        (c
r, IOFinalizer
ref) <- IO (c, IOFinalizer) -> m (c, IOFinalizer)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (c, IOFinalizer) -> m (c, IOFinalizer))
-> IO (c, IOFinalizer) -> m (c, IOFinalizer)
forall a b. (a -> b) -> a -> b
$ IO (c, IOFinalizer) -> IO (c, IOFinalizer)
forall a. IO a -> IO a
mask_ (IO (c, IOFinalizer) -> IO (c, IOFinalizer))
-> IO (c, IOFinalizer) -> IO (c, IOFinalizer)
forall a b. (a -> b) -> a -> b
$ do
            c
r <- a -> IO c
bef a
x
            IOFinalizer
ref <- IO d -> IO IOFinalizer
forall (m :: * -> *) a. MonadIO m => IO a -> m IOFinalizer
newIOFinalizer (c -> IO d
aft c
r)
            (c, IOFinalizer) -> IO (c, IOFinalizer)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (c
r, IOFinalizer
ref)
        s
s <- c -> m s
inject1 c
r
        (s, IOFinalizer) -> m (s, IOFinalizer)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (s
s, IOFinalizer
ref)

    {-# INLINE_LATE step #-}
    step :: (s, IOFinalizer) -> m (Step (s, IOFinalizer) b)
step (s
st, IOFinalizer
ref) = do
        Step s b
res <- s -> m (Step s b)
step1 s
st m (Step s b) -> m () -> m (Step s b)
forall (m :: * -> *) a b. MonadCatch m => m a -> m b -> m a
`MC.onException` IOFinalizer -> m ()
forall (m :: * -> *). MonadIO m => IOFinalizer -> m ()
runIOFinalizer IOFinalizer
ref
        case Step s b
res of
            Yield b
x s
s -> Step (s, IOFinalizer) b -> m (Step (s, IOFinalizer) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (s, IOFinalizer) b -> m (Step (s, IOFinalizer) b))
-> Step (s, IOFinalizer) b -> m (Step (s, IOFinalizer) b)
forall a b. (a -> b) -> a -> b
$ b -> (s, IOFinalizer) -> Step (s, IOFinalizer) b
forall s a. a -> s -> Step s a
Yield b
x (s
s, IOFinalizer
ref)
            Skip s
s    -> Step (s, IOFinalizer) b -> m (Step (s, IOFinalizer) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (s, IOFinalizer) b -> m (Step (s, IOFinalizer) b))
-> Step (s, IOFinalizer) b -> m (Step (s, IOFinalizer) b)
forall a b. (a -> b) -> a -> b
$ (s, IOFinalizer) -> Step (s, IOFinalizer) b
forall s a. s -> Step s a
Skip (s
s, IOFinalizer
ref)
            Step s b
Stop      -> do
                IOFinalizer -> m ()
forall (m :: * -> *). MonadIO m => IOFinalizer -> m ()
runIOFinalizer IOFinalizer
ref
                Step (s, IOFinalizer) b -> m (Step (s, IOFinalizer) b)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Step (s, IOFinalizer) b
forall s a. Step s a
Stop

-- | When unfolding @Unfold m a b@ if an exception @e@ occurs, unfold @e@ using
-- @Unfold m e b@.
--
-- /Inhibits stream fusion/
--
-- /Pre-release/
{-# INLINE_NORMAL handle #-}
handle :: (MonadCatch m, Exception e)
    => Unfold m e b -> Unfold m a b -> Unfold m a b
handle :: forall (m :: * -> *) e b a.
(MonadCatch m, Exception e) =>
Unfold m e b -> Unfold m a b -> Unfold m a b
handle Unfold m e b
exc =
    (a -> m a)
-> (forall s. m s -> m (Either e s))
-> (a -> m ())
-> Unfold m (a, e) b
-> Unfold m a b
-> Unfold m a b
forall (m :: * -> *) a c e d b.
Monad m =>
(a -> m c)
-> (forall s. m s -> m (Either e s))
-> (c -> m d)
-> Unfold m (c, e) b
-> Unfold m c b
-> Unfold m a b
gbracket_ a -> m a
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return m s -> m (Either e s)
forall s. m s -> m (Either e s)
forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
m a -> m (Either e a)
MC.try (\a
_ -> () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()) (Unfold m e b -> Unfold m (a, e) b
forall (m :: * -> *) a b c. Unfold m a b -> Unfold m (c, a) b
discardFirst Unfold m e b
exc)