{-# LANGUAGE AllowAmbiguousTypes  #-}
{-# LANGUAGE BlockArguments       #-}
{-# LANGUAGE DataKinds            #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE GADTs                #-}
{-# LANGUAGE KindSignatures       #-}
{-# LANGUAGE LambdaCase           #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE TemplateHaskell      #-}
{-# LANGUAGE TypeApplications     #-}
{-# LANGUAGE TypeOperators        #-}
{-# LANGUAGE UndecidableInstances #-}

{-# OPTIONS_GHC -Wno-orphans #-}

module Polysemy.Managed
  ( Managed(..)

  , runManagedResource
  , runManaged
  , runManagedFinal

  , managedAsk
  , managedLocal
  ) where

import Control.Exception qualified as E
import Control.Monad.Trans.Resource (createInternalState, MonadResource, InternalState)
import Control.Monad.Trans.Resource qualified as MTR
import Control.Monad.Trans.Resource.Internal qualified as RI
import Polysemy (Embed, Final, Sem, Member)
import Polysemy qualified as P
import Polysemy.Final qualified as PF
import Polysemy.Reader (Reader, runReader)

data Managed m a where
  ManagedAsk :: Managed m InternalState
  ManagedLocal :: m a -> Managed m a

P.makeSem ''Managed

runManaged :: forall a r. ()
  => Member (Embed IO) r
  => Sem (Managed ': r) a
  -> Sem r a
runManaged :: Sem (Managed : r) a -> Sem r a
runManaged Sem (Managed : r) a
f = do
  InternalState
state <- Sem r InternalState
forall (m :: * -> *). MonadIO m => m InternalState
createInternalState
  InternalState -> Sem r a -> Sem r a
forall a (r :: EffectRow).
Member (Embed IO) r =>
InternalState -> Sem r a -> Sem r a
managedBracketImpl InternalState
state (Sem r a -> Sem r a) -> Sem r a -> Sem r a
forall a b. (a -> b) -> a -> b
$ InternalState -> Sem (Managed : r) a -> Sem r a
forall a (r :: EffectRow).
Member (Embed IO) r =>
InternalState -> Sem (Managed : r) a -> Sem r a
runManagedImpl InternalState
state Sem (Managed : r) a
f
{-# INLINE runManaged #-}

runManagedFinal :: forall a r. ()
  => Member (Final IO) r
  => Sem (Managed ': r) a
  -> Sem r a
runManagedFinal :: Sem (Managed : r) a -> Sem r a
runManagedFinal Sem (Managed : r) a
f = do
  InternalState
state <- IO InternalState -> Sem r InternalState
forall (m :: * -> *) (r :: EffectRow) a.
(Member (Final m) r, Functor m) =>
m a -> Sem r a
P.embedFinal @IO IO InternalState
forall (m :: * -> *). MonadIO m => m InternalState
createInternalState
  InternalState -> Sem r a -> Sem r a
forall a (r :: EffectRow).
Member (Final IO) r =>
InternalState -> Sem r a -> Sem r a
managedBracketFinalImpl InternalState
state (Sem r a -> Sem r a) -> Sem r a -> Sem r a
forall a b. (a -> b) -> a -> b
$ InternalState -> Sem (Managed : r) a -> Sem r a
forall a (r :: EffectRow).
Member (Final IO) r =>
InternalState -> Sem (Managed : r) a -> Sem r a
runManagedFinalImpl InternalState
state Sem (Managed : r) a
f
{-# INLINE runManagedFinal #-}

runManagedResource :: ()
  => Member (Embed IO) r
  => Sem (Reader MTR.InternalState ': r) a
  -> Sem r a
runManagedResource :: Sem (Reader InternalState : r) a -> Sem r a
runManagedResource Sem (Reader InternalState : r) a
f = do
  InternalState
istate <- Sem r InternalState
forall (m :: * -> *). MonadIO m => m InternalState
createInternalState
  InternalState -> Sem (Reader InternalState : r) a -> Sem r a
forall i (r :: EffectRow) a. i -> Sem (Reader i : r) a -> Sem r a
runReader InternalState
istate Sem (Reader InternalState : r) a
f

--------------------------------------------------------------------------------
-- Internal

runManagedImpl :: forall a r. ()
  => Member (Embed IO) r
  => InternalState
  -> Sem (Managed ': r) a
  -> Sem r a
runManagedImpl :: InternalState -> Sem (Managed : r) a -> Sem r a
runManagedImpl InternalState
state = (forall (rInitial :: EffectRow) x.
 Managed (Sem rInitial) x -> Tactical Managed (Sem rInitial) r x)
-> Sem (Managed : r) a -> Sem r a
forall (e :: (* -> *) -> * -> *) (r :: EffectRow) a.
(forall (rInitial :: EffectRow) x.
 e (Sem rInitial) x -> Tactical e (Sem rInitial) r x)
-> Sem (e : r) a -> Sem r a
P.interpretH ((forall (rInitial :: EffectRow) x.
  Managed (Sem rInitial) x -> Tactical Managed (Sem rInitial) r x)
 -> Sem (Managed : r) a -> Sem r a)
-> (forall (rInitial :: EffectRow) x.
    Managed (Sem rInitial) x -> Tactical Managed (Sem rInitial) r x)
-> Sem (Managed : r) a
-> Sem r a
forall a b. (a -> b) -> a -> b
$ \case
  Managed (Sem rInitial) x
ManagedAsk -> InternalState
-> Sem (WithTactics Managed f (Sem rInitial) r) (f InternalState)
forall (f :: * -> *) a (e :: (* -> *) -> * -> *) (m :: * -> *)
       (r :: EffectRow).
Functor f =>
a -> Sem (WithTactics e f m r) (f a)
P.pureT InternalState
state
  ManagedLocal m -> do
    Sem (Managed : r) (f x)
mm <- Sem rInitial x
-> Sem
     (WithTactics Managed f (Sem rInitial) r) (Sem (Managed : r) (f x))
forall (m :: * -> *) a (e :: (* -> *) -> * -> *) (f :: * -> *)
       (r :: EffectRow).
m a -> Sem (WithTactics e f m r) (Sem (e : r) (f a))
P.runT Sem rInitial x
m
    InternalState
newState <- Sem (WithTactics Managed f (Sem rInitial) r) InternalState
forall (m :: * -> *). MonadIO m => m InternalState
createInternalState
    InternalState
-> Sem (WithTactics Managed f (Sem rInitial) r) (f x)
-> Sem (WithTactics Managed f (Sem rInitial) r) (f x)
forall a (r :: EffectRow).
Member (Embed IO) r =>
InternalState -> Sem r a -> Sem r a
managedBracketImpl InternalState
newState (Sem (WithTactics Managed f (Sem rInitial) r) (f x)
 -> Sem (WithTactics Managed f (Sem rInitial) r) (f x))
-> Sem (WithTactics Managed f (Sem rInitial) r) (f x)
-> Sem (WithTactics Managed f (Sem rInitial) r) (f x)
forall a b. (a -> b) -> a -> b
$ Sem r (f x) -> Sem (WithTactics Managed f (Sem rInitial) r) (f x)
forall (e :: (* -> *) -> * -> *) (r :: EffectRow) a.
Sem r a -> Sem (e : r) a
P.raise (Sem r (f x) -> Sem (WithTactics Managed f (Sem rInitial) r) (f x))
-> Sem r (f x)
-> Sem (WithTactics Managed f (Sem rInitial) r) (f x)
forall a b. (a -> b) -> a -> b
$ InternalState -> Sem (Managed : r) (f x) -> Sem r (f x)
forall a (r :: EffectRow).
Member (Embed IO) r =>
InternalState -> Sem (Managed : r) a -> Sem r a
runManagedImpl InternalState
newState Sem (Managed : r) (f x)
mm
{-# INLINE runManagedImpl #-}

runManagedFinalImpl :: forall a r. ()
  => Member (Final IO) r
  => InternalState
  -> Sem (Managed ': r) a
  -> Sem r a
runManagedFinalImpl :: InternalState -> Sem (Managed : r) a -> Sem r a
runManagedFinalImpl InternalState
state = (forall (rInitial :: EffectRow) x.
 Managed (Sem rInitial) x -> Tactical Managed (Sem rInitial) r x)
-> Sem (Managed : r) a -> Sem r a
forall (e :: (* -> *) -> * -> *) (r :: EffectRow) a.
(forall (rInitial :: EffectRow) x.
 e (Sem rInitial) x -> Tactical e (Sem rInitial) r x)
-> Sem (e : r) a -> Sem r a
P.interpretH ((forall (rInitial :: EffectRow) x.
  Managed (Sem rInitial) x -> Tactical Managed (Sem rInitial) r x)
 -> Sem (Managed : r) a -> Sem r a)
-> (forall (rInitial :: EffectRow) x.
    Managed (Sem rInitial) x -> Tactical Managed (Sem rInitial) r x)
-> Sem (Managed : r) a
-> Sem r a
forall a b. (a -> b) -> a -> b
$ \case
  Managed (Sem rInitial) x
ManagedAsk -> InternalState
-> Sem (WithTactics Managed f (Sem rInitial) r) (f InternalState)
forall (f :: * -> *) a (e :: (* -> *) -> * -> *) (m :: * -> *)
       (r :: EffectRow).
Functor f =>
a -> Sem (WithTactics e f m r) (f a)
P.pureT InternalState
state
  ManagedLocal m -> do
    Sem (Managed : r) (f x)
mm <- Sem rInitial x
-> Sem
     (WithTactics Managed f (Sem rInitial) r) (Sem (Managed : r) (f x))
forall (m :: * -> *) a (e :: (* -> *) -> * -> *) (f :: * -> *)
       (r :: EffectRow).
m a -> Sem (WithTactics e f m r) (Sem (e : r) (f a))
P.runT Sem rInitial x
m
    InternalState
newState <- IO InternalState
-> Sem (WithTactics Managed f (Sem rInitial) r) InternalState
forall (m :: * -> *) (r :: EffectRow) a.
(Member (Final m) r, Functor m) =>
m a -> Sem r a
P.embedFinal @IO IO InternalState
forall (m :: * -> *). MonadIO m => m InternalState
createInternalState
    InternalState
-> Sem (WithTactics Managed f (Sem rInitial) r) (f x)
-> Sem (WithTactics Managed f (Sem rInitial) r) (f x)
forall a (r :: EffectRow).
Member (Final IO) r =>
InternalState -> Sem r a -> Sem r a
managedBracketFinalImpl InternalState
newState (Sem (WithTactics Managed f (Sem rInitial) r) (f x)
 -> Sem (WithTactics Managed f (Sem rInitial) r) (f x))
-> Sem (WithTactics Managed f (Sem rInitial) r) (f x)
-> Sem (WithTactics Managed f (Sem rInitial) r) (f x)
forall a b. (a -> b) -> a -> b
$ Sem r (f x) -> Sem (WithTactics Managed f (Sem rInitial) r) (f x)
forall (e :: (* -> *) -> * -> *) (r :: EffectRow) a.
Sem r a -> Sem (e : r) a
P.raise (Sem r (f x) -> Sem (WithTactics Managed f (Sem rInitial) r) (f x))
-> Sem r (f x)
-> Sem (WithTactics Managed f (Sem rInitial) r) (f x)
forall a b. (a -> b) -> a -> b
$ InternalState -> Sem (Managed : r) (f x) -> Sem r (f x)
forall a (r :: EffectRow).
Member (Final IO) r =>
InternalState -> Sem (Managed : r) a -> Sem r a
runManagedFinalImpl InternalState
newState Sem (Managed : r) (f x)
mm
{-# INLINE runManagedFinalImpl #-}

managedBracketImpl :: forall a r. ()
  => Member (Embed IO) r
  => MTR.InternalState
  -> Sem r a
  -> Sem r a
managedBracketImpl :: InternalState -> Sem r a -> Sem r a
managedBracketImpl InternalState
istate Sem r a
f = do
  ((forall x. Sem r x -> IO x) -> IO () -> IO a) -> Sem r a
forall (r :: EffectRow) a.
Member (Embed IO) r =>
((forall x. Sem r x -> IO x) -> IO () -> IO a) -> Sem r a
P.withLowerToIO (((forall x. Sem r x -> IO x) -> IO () -> IO a) -> Sem r a)
-> ((forall x. Sem r x -> IO x) -> IO () -> IO a) -> Sem r a
forall a b. (a -> b) -> a -> b
$ \forall x. Sem r x -> IO x
lower IO ()
finish -> do
    ((forall a. IO a -> IO a) -> IO a) -> IO a
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
E.mask (((forall a. IO a -> IO a) -> IO a) -> IO a)
-> ((forall a. IO a -> IO a) -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
restore -> do
      a
res <- IO a -> IO a
forall a. IO a -> IO a
restore (Sem r a -> IO a
forall x. Sem r x -> IO x
lower Sem r a
f) IO a -> (SomeException -> IO a) -> IO a
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`E.catch` \SomeException
e -> do
        Maybe SomeException -> InternalState -> IO ()
RI.stateCleanupChecked (SomeException -> Maybe SomeException
forall a. a -> Maybe a
Just SomeException
e) InternalState
istate
        SomeException -> IO a
forall e a. Exception e => e -> IO a
E.throwIO SomeException
e
      Maybe SomeException -> InternalState -> IO ()
RI.stateCleanupChecked Maybe SomeException
forall a. Maybe a
Nothing InternalState
istate
      IO ()
finish
      a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
res

managedBracketFinalImpl :: forall a r. ()
  => Member (Final IO) r
  => MTR.InternalState
  -> Sem r a
  -> Sem r a
managedBracketFinalImpl :: InternalState -> Sem r a -> Sem r a
managedBracketFinalImpl InternalState
istate Sem r a
f = forall (r :: EffectRow) a.
Member (Final IO) r =>
Strategic IO (Sem r) a -> Sem r a
forall (m :: * -> *) (r :: EffectRow) a.
Member (Final m) r =>
Strategic m (Sem r) a -> Sem r a
PF.withStrategicToFinal @IO (Strategic IO (Sem r) a -> Sem r a)
-> Strategic IO (Sem r) a -> Sem r a
forall a b. (a -> b) -> a -> b
$ do
  IO (f a)
f' <- Sem r a -> Sem (WithStrategy IO f (Sem r)) (IO (f a))
forall (n :: * -> *) a (m :: * -> *) (f :: * -> *).
n a -> Sem (WithStrategy m f n) (m (f a))
PF.runS Sem r a
f
  IO (f a) -> Sem (WithStrategy IO f (Sem r)) (IO (f a))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IO (f a) -> Sem (WithStrategy IO f (Sem r)) (IO (f a)))
-> IO (f a) -> Sem (WithStrategy IO f (Sem r)) (IO (f a))
forall a b. (a -> b) -> a -> b
$ ((forall a. IO a -> IO a) -> IO (f a)) -> IO (f a)
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
E.mask (((forall a. IO a -> IO a) -> IO (f a)) -> IO (f a))
-> ((forall a. IO a -> IO a) -> IO (f a)) -> IO (f a)
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
restore -> do
    f a
res <- IO (f a) -> IO (f a)
forall a. IO a -> IO a
restore IO (f a)
f' IO (f a) -> (SomeException -> IO (f a)) -> IO (f a)
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`E.catch` \SomeException
e -> do
      Maybe SomeException -> InternalState -> IO ()
RI.stateCleanupChecked  (SomeException -> Maybe SomeException
forall a. a -> Maybe a
Just SomeException
e) InternalState
istate
      SomeException -> IO (f a)
forall e a. Exception e => e -> IO a
E.throwIO SomeException
e
    Maybe SomeException -> InternalState -> IO ()
RI.stateCleanupChecked Maybe SomeException
forall a. Maybe a
Nothing InternalState
istate
    f a -> IO (f a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure f a
res

instance
  ( Member (Embed IO) r
  , Member Managed r
  ) => MonadResource (Sem r) where
  liftResourceT :: ResourceT IO a -> Sem r a
liftResourceT (RI.ResourceT InternalState -> IO a
r) = Sem r InternalState
forall (r :: EffectRow). Member Managed r => Sem r InternalState
managedAsk Sem r InternalState -> (InternalState -> Sem r a) -> Sem r a
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= IO a -> Sem r a
forall (m :: * -> *) (r :: EffectRow) a.
Member (Embed m) r =>
m a -> Sem r a
P.embed (IO a -> Sem r a)
-> (InternalState -> IO a) -> InternalState -> Sem r a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. InternalState -> IO a
r