{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE RankNTypes #-}

module Control.Monad.Catch.Scope (
  Scope,
  FinalizationException (..),
  withScope,
  withScopeLiftedTo,
  atEndOf,
  scopedTo,
  scoped
) where

import Control.Monad (void, (>=>))
import Control.Monad.Catch (Exception, MonadCatch, MonadMask, SomeException, bracket, mask_, catchAll, throwM)
import Control.Monad.IO.Class (MonadIO, liftIO)
import Data.IORef (IORef, newIORef, readIORef, atomicModifyIORef)
import Data.Functor (($>))
import Data.Maybe (catMaybes)
import Type.Reflection (Typeable)

newtype Scope m = Scope { addFinalizer :: m () -> IO () }

newtype FinalizationException = FinalizationException [SomeException] deriving (Show, Typeable)
instance Exception FinalizationException

withScope :: (MonadMask m, MonadIO m) => (Scope m -> m a) -> m a
withScope useScope = bracket (liftIO $ newIORef []) finalize (useScope . newScope)
  where
    newScope :: MonadIO m => IORef [m ()] -> Scope m
    newScope ref = Scope $ \finalizer -> liftIO $ atomicModifyIORef ref $ (, ()) . (finalizer :)

    finalize :: (MonadCatch m, MonadIO m) => IORef [m ()] -> m ()
    finalize =
      liftIO . readIORef >=>
      fmap catMaybes . mapM maybeCaughtException >=>
      \case
        [] -> return ()
        exceptions -> throwM $ FinalizationException exceptions

    maybeCaughtException :: MonadCatch m => m () -> m (Maybe SomeException)
    maybeCaughtException m = catchAll (m $> Nothing) (return . Just)

withScopeLiftedTo :: (MonadIO super, MonadIO sub) => Scope super -> (forall b. sub b -> super b) -> (Scope sub -> sub a) -> super a
withScopeLiftedTo superScope lift useSubScope = lift $ useSubScope $ Scope $ addFinalizer superScope . lift

atEndOf :: MonadIO n => Scope m -> m () -> n ()
atEndOf scope finalizer = liftIO $ addFinalizer scope finalizer

scoped :: (MonadMask m, MonadIO m) => m a -> (a -> m b) -> Scope m -> m a
scoped acquire release scope = mask_ $ do
  a <- acquire
  atEndOf scope (void $ release a)
  return a

scopedTo :: (MonadMask m, MonadIO m) => Scope m -> m a -> (a -> m b) -> m a
scopedTo scope acquire release = scoped acquire release scope