module Control.Monad.Trans.Managed
  ( ManagedT (..),
    allocate,
    allocate_,
    lowerManagedT,
    hoistManagedTReaderT,
  )
where

import Control.Concurrent qualified as C
import Control.Exception.Lifted (bracket, bracket_)
import Control.Monad.Codensity (Codensity (..))
import Control.Monad.Fix (MonadFix (..))
import Control.Monad.IO.Class (MonadIO, liftIO)
import Control.Monad.Reader.Class (MonadReader)
import Control.Monad.State.Class (MonadState)
import Control.Monad.Trans (MonadTrans (..))
import Control.Monad.Trans.Control (MonadBaseControl)
import Control.Monad.Trans.Reader (ReaderT (..))
import GHC.IO.Unsafe (unsafeDupableInterleaveIO)
import Prelude

-- | This type is like a transformer version of the @Managed@ monad from the
-- @managed@ library. It can be used to manage resources by pairing together
-- their allocation with their finalizers.
--
-- The documentation for the @managed@ library is an excellent introduction to
-- the idea here.
--
-- We could use 'Codensity' directly, but we'd have to define an orphan instance
-- for 'MonadFix'. This also gives us the opportunity to give it a slightly more
-- friendly name.
--
-- We could also have used @ResourceT@, but that would have involved writing
-- instances for @MonadUnliftIO@. That could still be a good option to consider
-- later, however.
newtype ManagedT m a = ManagedT {forall (m :: * -> *) a. ManagedT m a -> forall r. (a -> m r) -> m r
runManagedT :: forall r. (a -> m r) -> m r}
  deriving
    ( (forall a b. (a -> b) -> ManagedT m a -> ManagedT m b)
-> (forall a b. a -> ManagedT m b -> ManagedT m a)
-> Functor (ManagedT m)
forall a b. a -> ManagedT m b -> ManagedT m a
forall a b. (a -> b) -> ManagedT m a -> ManagedT m b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
forall (m :: * -> *) a b. a -> ManagedT m b -> ManagedT m a
forall (m :: * -> *) a b. (a -> b) -> ManagedT m a -> ManagedT m b
$cfmap :: forall (m :: * -> *) a b. (a -> b) -> ManagedT m a -> ManagedT m b
fmap :: forall a b. (a -> b) -> ManagedT m a -> ManagedT m b
$c<$ :: forall (m :: * -> *) a b. a -> ManagedT m b -> ManagedT m a
<$ :: forall a b. a -> ManagedT m b -> ManagedT m a
Functor,
      Functor (ManagedT m)
Functor (ManagedT m)
-> (forall a. a -> ManagedT m a)
-> (forall a b.
    ManagedT m (a -> b) -> ManagedT m a -> ManagedT m b)
-> (forall a b c.
    (a -> b -> c) -> ManagedT m a -> ManagedT m b -> ManagedT m c)
-> (forall a b. ManagedT m a -> ManagedT m b -> ManagedT m b)
-> (forall a b. ManagedT m a -> ManagedT m b -> ManagedT m a)
-> Applicative (ManagedT m)
forall a. a -> ManagedT m a
forall a b. ManagedT m a -> ManagedT m b -> ManagedT m a
forall a b. ManagedT m a -> ManagedT m b -> ManagedT m b
forall a b. ManagedT m (a -> b) -> ManagedT m a -> ManagedT m b
forall a b c.
(a -> b -> c) -> ManagedT m a -> ManagedT m b -> ManagedT m c
forall (m :: * -> *). Functor (ManagedT m)
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
forall (m :: * -> *) a. a -> ManagedT m a
forall (m :: * -> *) a b.
ManagedT m a -> ManagedT m b -> ManagedT m a
forall (m :: * -> *) a b.
ManagedT m a -> ManagedT m b -> ManagedT m b
forall (m :: * -> *) a b.
ManagedT m (a -> b) -> ManagedT m a -> ManagedT m b
forall (m :: * -> *) a b c.
(a -> b -> c) -> ManagedT m a -> ManagedT m b -> ManagedT m c
$cpure :: forall (m :: * -> *) a. a -> ManagedT m a
pure :: forall a. a -> ManagedT m a
$c<*> :: forall (m :: * -> *) a b.
ManagedT m (a -> b) -> ManagedT m a -> ManagedT m b
<*> :: forall a b. ManagedT m (a -> b) -> ManagedT m a -> ManagedT m b
$cliftA2 :: forall (m :: * -> *) a b c.
(a -> b -> c) -> ManagedT m a -> ManagedT m b -> ManagedT m c
liftA2 :: forall a b c.
(a -> b -> c) -> ManagedT m a -> ManagedT m b -> ManagedT m c
$c*> :: forall (m :: * -> *) a b.
ManagedT m a -> ManagedT m b -> ManagedT m b
*> :: forall a b. ManagedT m a -> ManagedT m b -> ManagedT m b
$c<* :: forall (m :: * -> *) a b.
ManagedT m a -> ManagedT m b -> ManagedT m a
<* :: forall a b. ManagedT m a -> ManagedT m b -> ManagedT m a
Applicative,
      Applicative (ManagedT m)
Applicative (ManagedT m)
-> (forall a b.
    ManagedT m a -> (a -> ManagedT m b) -> ManagedT m b)
-> (forall a b. ManagedT m a -> ManagedT m b -> ManagedT m b)
-> (forall a. a -> ManagedT m a)
-> Monad (ManagedT m)
forall a. a -> ManagedT m a
forall a b. ManagedT m a -> ManagedT m b -> ManagedT m b
forall a b. ManagedT m a -> (a -> ManagedT m b) -> ManagedT m b
forall (m :: * -> *). Applicative (ManagedT m)
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
forall (m :: * -> *) a. a -> ManagedT m a
forall (m :: * -> *) a b.
ManagedT m a -> ManagedT m b -> ManagedT m b
forall (m :: * -> *) a b.
ManagedT m a -> (a -> ManagedT m b) -> ManagedT m b
$c>>= :: forall (m :: * -> *) a b.
ManagedT m a -> (a -> ManagedT m b) -> ManagedT m b
>>= :: forall a b. ManagedT m a -> (a -> ManagedT m b) -> ManagedT m b
$c>> :: forall (m :: * -> *) a b.
ManagedT m a -> ManagedT m b -> ManagedT m b
>> :: forall a b. ManagedT m a -> ManagedT m b -> ManagedT m b
$creturn :: forall (m :: * -> *) a. a -> ManagedT m a
return :: forall a. a -> ManagedT m a
Monad,
      Monad (ManagedT m)
Monad (ManagedT m)
-> (forall a. IO a -> ManagedT m a) -> MonadIO (ManagedT m)
forall a. IO a -> ManagedT m a
forall (m :: * -> *).
Monad m -> (forall a. IO a -> m a) -> MonadIO m
forall {m :: * -> *}. MonadIO m => Monad (ManagedT m)
forall (m :: * -> *) a. MonadIO m => IO a -> ManagedT m a
$cliftIO :: forall (m :: * -> *) a. MonadIO m => IO a -> ManagedT m a
liftIO :: forall a. IO a -> ManagedT m a
MonadIO,
      MonadReader r,
      MonadState s
    )
    via (Codensity m)
  deriving ((forall (m :: * -> *) a. Monad m => m a -> ManagedT m a)
-> MonadTrans ManagedT
forall (m :: * -> *) a. Monad m => m a -> ManagedT m a
forall (t :: (* -> *) -> * -> *).
(forall (m :: * -> *) a. Monad m => m a -> t m a) -> MonadTrans t
$clift :: forall (m :: * -> *) a. Monad m => m a -> ManagedT m a
lift :: forall (m :: * -> *) a. Monad m => m a -> ManagedT m a
MonadTrans) via Codensity

-- | Allocate a resource by providing setup and finalizer actions.
allocate :: (MonadBaseControl IO m) => m a -> (a -> m b) -> ManagedT m a
allocate :: forall (m :: * -> *) a b.
MonadBaseControl IO m =>
m a -> (a -> m b) -> ManagedT m a
allocate m a
setup a -> m b
finalize = (forall r. (a -> m r) -> m r) -> ManagedT m a
forall (m :: * -> *) a.
(forall r. (a -> m r) -> m r) -> ManagedT m a
ManagedT (m a -> (a -> m b) -> (a -> m r) -> m r
forall (m :: * -> *) a b c.
MonadBaseControl IO m =>
m a -> (a -> m b) -> (a -> m c) -> m c
bracket m a
setup a -> m b
finalize)

-- | Allocate a resource but do not return a reference to it.
allocate_ :: (MonadBaseControl IO m) => m a -> m b -> ManagedT m ()
allocate_ :: forall (m :: * -> *) a b.
MonadBaseControl IO m =>
m a -> m b -> ManagedT m ()
allocate_ m a
setup m b
finalize = (forall r. (() -> m r) -> m r) -> ManagedT m ()
forall (m :: * -> *) a.
(forall r. (a -> m r) -> m r) -> ManagedT m a
ManagedT (\() -> m r
k -> m a -> m b -> m r -> m r
forall (m :: * -> *) a b c.
MonadBaseControl IO m =>
m a -> m b -> m c -> m c
bracket_ m a
setup m b
finalize (() -> m r
k ()))

-- | Run the provided computation by returning its result, and run any finalizers.
-- Watch out: this function might leak finalized resources.
lowerManagedT :: (Monad m) => ManagedT m a -> m a
lowerManagedT :: forall (m :: * -> *) a. Monad m => ManagedT m a -> m a
lowerManagedT ManagedT m a
m = ManagedT m a -> forall r. (a -> m r) -> m r
forall (m :: * -> *) a. ManagedT m a -> forall r. (a -> m r) -> m r
runManagedT ManagedT m a
m a -> m a
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return

hoistManagedTReaderT :: (Monad m) => r -> ManagedT (ReaderT r m) a -> ManagedT m a
hoistManagedTReaderT :: forall (m :: * -> *) r a.
Monad m =>
r -> ManagedT (ReaderT r m) a -> ManagedT m a
hoistManagedTReaderT r
r ManagedT (ReaderT r m) a
cod = (forall r. (a -> m r) -> m r) -> ManagedT m a
forall (m :: * -> *) a.
(forall r. (a -> m r) -> m r) -> ManagedT m a
ManagedT ((forall r. (a -> m r) -> m r) -> ManagedT m a)
-> (forall r. (a -> m r) -> m r) -> ManagedT m a
forall a b. (a -> b) -> a -> b
$ \a -> m r
k ->
  ReaderT r m r -> r -> m r
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT (ManagedT (ReaderT r m) a
-> forall r. (a -> ReaderT r m r) -> ReaderT r m r
forall (m :: * -> *) a. ManagedT m a -> forall r. (a -> m r) -> m r
runManagedT ManagedT (ReaderT r m) a
cod (m r -> ReaderT r m r
forall (m :: * -> *) a. Monad m => m a -> ReaderT r m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m r -> ReaderT r m r) -> (a -> m r) -> a -> ReaderT r m r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> m r
k)) r
r

-- | We need this instance to tie the knot when initializing resources.
-- It'd be nice if we could do this with a 'MonadFix' constraint on the underlying
-- monad, but here we just use 'MonadIO' to tie the knot using a lazily-evaluated
-- 'MVar'-based promise for the eventual result.
--
-- We need to be careful not to leak allocated resources via the use of
-- recursively-defined monadic actions when making use of this instance.
instance (MonadIO m) => MonadFix (ManagedT m) where
  mfix :: forall a. (a -> ManagedT m a) -> ManagedT m a
mfix a -> ManagedT m a
f = (forall r. (a -> m r) -> m r) -> ManagedT m a
forall (m :: * -> *) a.
(forall r. (a -> m r) -> m r) -> ManagedT m a
ManagedT \a -> m r
k -> do
    MVar a
m <- IO (MVar a) -> m (MVar a)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO (MVar a)
forall a. IO (MVar a)
C.newEmptyMVar
    a
ans <- 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
$ IO a -> IO a
forall a. IO a -> IO a
unsafeDupableInterleaveIO (MVar a -> IO a
forall a. MVar a -> IO a
C.readMVar MVar a
m)
    ManagedT m a -> forall r. (a -> m r) -> m r
forall (m :: * -> *) a. ManagedT m a -> forall r. (a -> m r) -> m r
runManagedT (a -> ManagedT m a
f a
ans) \a
a -> do
      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
$ MVar a -> a -> IO ()
forall a. MVar a -> a -> IO ()
C.putMVar MVar a
m a
a
      a -> m r
k a
a