-- |
-- Module      : Streamly.Internal.Data.Unfold.Exception
-- Copyright   : (c) 2019 Composewell Technologies
-- License     : BSD3
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
-- Lifted resource management primitives.

module Streamly.Internal.Data.Unfold.Exception
    (
      gbracket
    , after
    , finally
    , bracket
    )
where

#include "inline.hs"

import Control.Exception (mask_)
import Control.Monad.Catch (MonadCatch)
import Streamly.Internal.Control.Concurrent
    (MonadRunInIO, MonadAsync, withRunInIO)
import Streamly.Internal.Data.IOFinalizer.Lifted
    (newIOFinalizer, runIOFinalizer, clearingIOFinalizer)
import qualified Control.Monad.Catch as MC

import Streamly.Internal.Data.Unfold

-- | 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 gbracket #-}
gbracket
    :: MonadRunInIO m
    => (a -> m c)                           -- ^ before
    -> (c -> m d)                           -- ^ after, on normal stop, or GC
    -> Unfold m (c, e) b                    -- ^ on exception
    -> (forall s. m s -> m (Either e s))    -- ^ try (exception handling)
    -> Unfold m c b                         -- ^ unfold to run
    -> Unfold m a b
gbracket :: forall (m :: * -> *) a c d e b.
MonadRunInIO m =>
(a -> m c)
-> (c -> m d)
-> Unfold m (c, e) b
-> (forall s. m s -> m (Either e s))
-> Unfold m c b
-> Unfold m a b
gbracket a -> m c
bef c -> m d
aft (Unfold s -> m (Step s b)
estep (c, e) -> m s
einject) forall s. m s -> m (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) <- ((forall a. m a -> IO (StM m a)) -> IO (StM m (c, IOFinalizer)))
-> m (c, IOFinalizer)
forall (m :: * -> *) b.
MonadRunInIO m =>
((forall a. m a -> IO (StM m a)) -> IO (StM m b)) -> m b
withRunInIO (((forall a. m a -> IO (StM m a)) -> IO (StM m (c, IOFinalizer)))
 -> m (c, IOFinalizer))
-> ((forall a. m a -> IO (StM m a)) -> IO (StM m (c, IOFinalizer)))
-> m (c, IOFinalizer)
forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO (StM m a)
run -> IO (StM m (c, IOFinalizer)) -> IO (StM m (c, IOFinalizer))
forall a. IO a -> IO a
mask_ (IO (StM m (c, IOFinalizer)) -> IO (StM m (c, IOFinalizer)))
-> IO (StM m (c, IOFinalizer)) -> IO (StM m (c, IOFinalizer))
forall a b. (a -> b) -> a -> b
$ m (c, IOFinalizer) -> IO (StM m (c, IOFinalizer))
forall a. m a -> IO (StM m a)
run (m (c, IOFinalizer) -> IO (StM m (c, IOFinalizer)))
-> m (c, IOFinalizer) -> IO (StM m (c, IOFinalizer))
forall a b. (a -> b) -> a -> b
$ do
            c
r <- a -> m c
bef a
x
            IOFinalizer
ref <- m d -> m IOFinalizer
forall (m :: * -> *) a. MonadRunInIO m => m a -> m IOFinalizer
newIOFinalizer (c -> m d
aft c
r)
            (c, IOFinalizer) -> m (c, IOFinalizer)
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 (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 <- m (Step s b) -> m (Either e (Step s b))
forall s. m s -> m (Either e s)
ftry (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, IOFinalizer)) b
-> m (Step (Either s (s, c, IOFinalizer)) b)
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 (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 (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.
                s
r <- IOFinalizer -> m s -> m s
forall (m :: * -> *) a. MonadRunInIO m => IOFinalizer -> m a -> m a
clearingIOFinalizer IOFinalizer
ref ((c, e) -> m s
einject (c
v, e
e))
                Step (Either s (s, c, IOFinalizer)) b
-> m (Step (Either s (s, c, IOFinalizer)) b)
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 (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

-- | 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 after #-}
after :: MonadRunInIO m
    => (a -> m c) -> Unfold m a b -> Unfold m a b
after :: forall (m :: * -> *) a c b.
MonadRunInIO 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, 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 <- m c -> m IOFinalizer
forall (m :: * -> *) a. MonadRunInIO m => m a -> m IOFinalizer
newIOFinalizer (a -> m c
action a
x)
        (s, IOFinalizer) -> m (s, IOFinalizer)
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 (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 (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 (m :: * -> *) a. Monad m => a -> m a
return Step (s, IOFinalizer) 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 finally #-}
finally :: (MonadAsync m, MonadCatch m)
    => (a -> m c) -> Unfold m a b -> Unfold m a b
finally :: forall (m :: * -> *) a c b.
(MonadAsync m, 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, 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 <- m c -> m IOFinalizer
forall (m :: * -> *) a. MonadRunInIO m => m a -> m IOFinalizer
newIOFinalizer (a -> m c
action a
x)
        (s, IOFinalizer) -> m (s, IOFinalizer)
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 (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 (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 (m :: * -> *) a. Monad m => a -> m a
return Step (s, IOFinalizer) 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 bracket #-}
bracket :: (MonadAsync m, MonadCatch m)
    => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b
bracket :: forall (m :: * -> *) a c d b.
(MonadAsync m, 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, 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) <- ((forall a. m a -> IO (StM m a)) -> IO (StM m (c, IOFinalizer)))
-> m (c, IOFinalizer)
forall (m :: * -> *) b.
MonadRunInIO m =>
((forall a. m a -> IO (StM m a)) -> IO (StM m b)) -> m b
withRunInIO (((forall a. m a -> IO (StM m a)) -> IO (StM m (c, IOFinalizer)))
 -> m (c, IOFinalizer))
-> ((forall a. m a -> IO (StM m a)) -> IO (StM m (c, IOFinalizer)))
-> m (c, IOFinalizer)
forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO (StM m a)
run -> IO (StM m (c, IOFinalizer)) -> IO (StM m (c, IOFinalizer))
forall a. IO a -> IO a
mask_ (IO (StM m (c, IOFinalizer)) -> IO (StM m (c, IOFinalizer)))
-> IO (StM m (c, IOFinalizer)) -> IO (StM m (c, IOFinalizer))
forall a b. (a -> b) -> a -> b
$ m (c, IOFinalizer) -> IO (StM m (c, IOFinalizer))
forall a. m a -> IO (StM m a)
run (m (c, IOFinalizer) -> IO (StM m (c, IOFinalizer)))
-> m (c, IOFinalizer) -> IO (StM m (c, IOFinalizer))
forall a b. (a -> b) -> a -> b
$ do
            c
r <- a -> m c
bef a
x
            IOFinalizer
ref <- m d -> m IOFinalizer
forall (m :: * -> *) a. MonadRunInIO m => m a -> m IOFinalizer
newIOFinalizer (c -> m d
aft c
r)
            (c, IOFinalizer) -> m (c, IOFinalizer)
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 (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 (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 (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 (m :: * -> *) a. Monad m => a -> m a
return Step (s, IOFinalizer) b
forall s a. Step s a
Stop