{-# LANGUAGE UnicodeSyntax #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}

-- | Guarding monadic computations with cleanup actions.
module Control.Monad.Finally
  ( MonadFinally(..)
  , chainCleanups
  , finallyMany
  , onEscape
  , onEscapeMany
  , bracket_
  , bracketOnEscape
  ) where

#if !MIN_VERSION_base(4,8,0)
import Prelude hiding (mapM, forM)
#endif
import Data.Functor.Identity
import Data.Monoid (Monoid(..))
import Data.Traversable (mapM, forM)
import Control.Applicative (Applicative(..), (<$>), (<$))
import Control.Monad (join)
import Control.Monad.Trans.Maybe
import Control.Monad.Trans.Error
import Control.Monad.Trans.Except
import Control.Monad.Trans.Reader
import qualified Control.Monad.Trans.State.Lazy as L
import qualified Control.Monad.Trans.State.Strict as S
import qualified Control.Monad.Trans.Writer.Lazy as LW
import qualified Control.Monad.Trans.Writer.Strict as SW
import qualified Control.Monad.Trans.RWS.Lazy as L
import qualified Control.Monad.Trans.RWS.Strict as S
import Control.Monad.Trans.Abort
import Control.Monad.Trans.Finish
import qualified Control.Exception as E

-- | Class of monads that support guarding computations with cleanup actions.
class (Applicative μ, Monad μ) ⇒ MonadFinally μ where
#if __GLASGOW_HASKELL__ >= 707
  {-# MINIMAL finally' | bracket' #-}
#endif
  -- | @'finally'' m f@ runs computation @m@ and then
  --
  -- 1. runs @f ('Just' x)@ if @m@ produced a result @x@. The result of that is
  -- returned alongside with @x@.
  -- 2. runs @f 'Nothing'@ otherwise.
  finally' ∷ μ α → (Maybe α → μ β) → μ (α, β)
  finally' m f = bracket' (return ()) (const f) (const m)
  -- | A simplified version of 'finally'' in which the cleanup action
  -- does not care about the result of the main computation. The default
  -- implementation is
  --
  -- @
  --     'finally' m = 'fmap' 'fst' . 'finally'' m . 'const'
  -- @
  finally  ∷ μ α → μ β → μ α
  finally m = fmap fst . finally' m . const
  -- | Safely acquire a resource and use it in a computation, releasing it
  -- even when the computation does not produce a result.
  bracket' ∷ μ r                 -- ^ Acquire resource
           → (r → Maybe α → μ β) -- ^ Release resource
           → (r → μ α)           -- ^ Main computation
           → μ (α, β)
  bracket' acquire release m = do
    r ← acquire
    finally' (m r) (release r)
  -- | A simplified version of 'bracket'' in which the releasing action
  -- does not care about the result of the main computation. The default
  -- implementation is
  --
  -- @
  --     'bracket' acquire release =
  --       'fmap' 'fst' . 'bracket'' acquire ('const' . release)
  -- @
  bracket ∷ μ r → (r → μ β) → (r → μ α) → μ α
  bracket acquire release = fmap fst . bracket' acquire (const . release)

instance MonadFinally Identity where
  finally' m f = do
    mr ← m
    return (mr, runIdentity $ f $ Just mr)

instance MonadFinally IO where
  finally' m f = E.mask $ \restore → do
    mr ← restore m `E.onException` f Nothing
    fr ← f $ Just mr
    return (mr, fr)
  bracket' acquire release m = E.mask $ \restore → do
    ar ← acquire
    mr ← restore (m ar) `E.onException` release ar Nothing
    rr ← release ar $ Just mr
    return (mr, rr)

instance MonadFinally μ ⇒ MonadFinally (MaybeT μ) where
  finally' m f = MaybeT $ do
    (mr, fr) ← finally' (runMaybeT m) (runMaybeT . f . join)
    return $ (,) <$> mr <*> fr
  bracket' acquire release m = MaybeT $ do
    (mr, rr) ← bracket' (runMaybeT acquire)
                        (\case
                           Just ar → runMaybeT . release ar . join
                           Nothing → const (return Nothing))
                        (fmap join . mapM (runMaybeT . m))
    return $ (,) <$> mr <*> rr

justRight ∷ Maybe (Either e α) → Maybe α
justRight (Just (Right a)) = Just a
justRight _                = Nothing

instance MonadFinally μ ⇒ MonadFinally (AbortT e μ) where
  finally' m f = AbortT $ do
    (mr, fr) ← finally' (runAbortT m) (runAbortT . f . justRight)
    return $ (,) <$> mr <*> fr
  bracket' acquire release m = AbortT $ do
    (mr, rr) ← bracket' (runAbortT acquire)
                        (\case
                          Right ar → runAbortT . release ar . justRight
                          Left e → const (return (Left e)))
                        (fmap join . mapM (runAbortT . m))
    return $ (,) <$> mr <*> rr

instance MonadFinally μ ⇒ MonadFinally (FinishT β μ) where
  finally' m f = FinishT $ do
    (mr, fr) ← finally' (runFinishT m) (runFinishT . f . justRight)
    return $ (,) <$> mr <*> fr
  bracket' acquire release m = FinishT $ do
    (mr, rr) ← bracket' (runFinishT acquire)
                        (\case
                          Right ar → runFinishT . release ar . justRight
                          Left e → const (return (Left e)))
                        (fmap join . mapM (runFinishT . m))
    return $ (,) <$> mr <*> rr

instance (MonadFinally μ, Error e) ⇒ MonadFinally (ErrorT e μ) where
  finally' m f = ErrorT $ do
    (mr, fr) ← finally' (runErrorT m) (runErrorT . f . justRight)
    return $ (,) <$> mr <*> fr
  bracket' acquire release m = ErrorT $ do
    (mr, rr) ← bracket' (runErrorT acquire)
                        (\case
                          Right ar → runErrorT . release ar . justRight
                          Left e → const (return (Left e)))
                        (fmap join . mapM (runErrorT . m))
    return $ (,) <$> mr <*> rr

instance MonadFinally μ ⇒ MonadFinally (ExceptT e μ) where
  finally' m f = ExceptT $ do
    (mr, fr) ← finally' (runExceptT m) (runExceptT . f . justRight)
    return $ (,) <$> mr <*> fr
  bracket' acquire release m = ExceptT $ do
    (mr, rr) ← bracket' (runExceptT acquire)
                        (\case
                          Right ar → runExceptT . release ar . justRight
                          Left e → const (return (Left e)))
                        (fmap join . mapM (runExceptT . m))
    return $ (,) <$> mr <*> rr

instance MonadFinally μ ⇒ MonadFinally (ReaderT r μ) where
  finally' m f = ReaderT $ \r →
    finally' (runReaderT m r) ((`runReaderT` r) . f)
  bracket' acquire release m = ReaderT $ \r →
    bracket' (runReaderT acquire r)
             (fmap (`runReaderT` r) . release)
             ((`runReaderT` r) . m)

instance MonadFinally μ ⇒ MonadFinally (L.StateT s μ) where
  finally' m f = L.StateT $ \s → do
    ~(~(mr, _), ~(fr, s'')) ← finally' (L.runStateT m s) $ \case
      Just ~(a, s') → L.runStateT (f $ Just a) s'
      Nothing       → L.runStateT (f Nothing) s
    return ((mr, fr), s'')
  bracket' acquire release m = L.StateT $ \s → do
    ~(~(mr, _), ~(fr, s'')) ←
      bracket' (L.runStateT acquire s)
               (\ ~(ar, s') → \case
                  Just ~(mr, s'') → L.runStateT (release ar (Just mr)) s''
                  Nothing         → L.runStateT (release ar Nothing) s')
               (\ ~(ar, s') → L.runStateT (m ar) s')
    return ((mr, fr), s'')

instance MonadFinally μ ⇒ MonadFinally (S.StateT s μ) where
  finally' m f = S.StateT $ \s → do
    ((mr, _), (fr, s''')) ← finally' (S.runStateT m s) $ \case
      Just (a, s') → S.runStateT (f $ Just a) s'
      Nothing      → S.runStateT (f Nothing) s
    return ((mr, fr), s''')
  bracket' acquire release m = S.StateT $ \s → do
    ((mr, _), (fr, s''')) ←
      bracket' (S.runStateT acquire s)
               (\(ar, s') → \case
                  Just ~(mr, s'') → S.runStateT (release ar (Just mr)) s''
                  Nothing         → S.runStateT (release ar Nothing) s')
               (\(ar, s') → S.runStateT (m ar) s')
    return ((mr, fr), s''')

instance (MonadFinally μ, Monoid w) ⇒ MonadFinally (LW.WriterT w μ) where
  finally' m f = LW.WriterT $ do
    ~(~(mr, w), ~(fr, w')) ← finally' (LW.runWriterT m) $
      LW.runWriterT . f . fmap fst
    return ((mr, fr), w `mappend` w')
  bracket' acquire release m = LW.WriterT $ do
    ~(~(mr, w), ~(fr, w')) ←
      bracket' (LW.runWriterT acquire)
               (\ ~(ar, _) → \case
                  Just ~(mr, _) → LW.runWriterT (release ar (Just mr))
                  Nothing → LW.runWriterT (release ar Nothing))
               (\ ~(ar, aw) → LW.runWriterT (LW.tell aw >> m ar))
    return ((mr, fr), w `mappend` w')

instance (MonadFinally μ, Monoid w) ⇒ MonadFinally (SW.WriterT w μ) where
  finally' m f = SW.WriterT $ do
    ((mr, w), (fr, w')) ← finally' (SW.runWriterT m) $ \mbr → case mbr of
      Just (a, _) → SW.runWriterT $ f $ Just a
      Nothing     → SW.runWriterT $ f Nothing
    return ((mr, fr), w `mappend` w')
  bracket' acquire release m = SW.WriterT $ do
    ((mr, w), (fr, w')) ←
      bracket' (SW.runWriterT acquire)
               (\(ar, _) → \case
                  Just (mr, _) → SW.runWriterT (release ar (Just mr))
                  Nothing → SW.runWriterT (release ar Nothing))
               (\(ar, aw) → SW.runWriterT (SW.tell aw >> m ar))
    return ((mr, fr), w `mappend` w')

instance (MonadFinally μ, Monoid w) ⇒ MonadFinally (L.RWST r w s μ) where
  finally' m f = L.RWST $ \r s → do
    ~(~(mr, _, w), ~(fr, s'', w')) ← finally' (L.runRWST m r s) $ \case
      Just ~(mr, s', _) → L.runRWST (f $ Just mr) r s'
      Nothing → L.runRWST (f Nothing) r s
    return ((mr, fr), s'', w `mappend` w')
  bracket' acquire release m = L.RWST $ \r s → do
    ~(~(mr, _, w), ~(fr, s''', w')) ←
      bracket' (L.runRWST acquire r s)
               (\ ~(ar, s', _) → \case 
                  Just ~(mr, s'', _) → L.runRWST (release ar $ Just mr) r s''
                  Nothing → L.runRWST (release ar Nothing) r s')
               (\ ~(ar, s', aw) → L.runRWST (L.tell aw >> m ar) r s')
    return ((mr, fr), s''', w `mappend` w')

instance (MonadFinally μ, Monoid w) ⇒ MonadFinally (S.RWST r w s μ) where
  finally' m f = S.RWST $ \r s → do
    ((mr, _, w), (fr, s'', w')) ← finally' (S.runRWST m r s) $ \case
      Just (a, s', _) → S.runRWST (f $ Just a) r s'
      Nothing         → S.runRWST (f Nothing) r s
    return ((mr, fr), s'', w `mappend` w')
  bracket' acquire release m = S.RWST $ \r s → do
    ((mr, _, w), (fr, s''', w')) ←
      bracket' (S.runRWST acquire r s)
               (\(ar, s', _) → \case 
                  Just (mr, s'', _) → S.runRWST (release ar $ Just mr) r s''
                  Nothing → S.runRWST (release ar Nothing) r s')
               (\(ar, s', aw) → S.runRWST (S.tell aw >> m ar) r s')
    return ((mr, fr), s''', w `mappend` w')

-- | Run the provided list of cleanup actions sequentually, attempting to run
-- the next action even if the previous one did not produce a result.
chainCleanups ∷ MonadFinally μ ⇒ [μ α] → μ ()
chainCleanups []       = return ()
chainCleanups (m : ms) = finally (() <$ m) $ chainCleanups ms

-- | A variant of 'finally' that combines multiple cleanup actions with
-- 'chainCleanups'.
finallyMany ∷ MonadFinally μ ⇒ μ α → [μ β] → μ α
finallyMany m = finally m . chainCleanups
{-# INLINE finallyMany #-}

-- | @'onEscape' m c@ runs computation @m@ and then, if it did not produce
-- a result, runs computation @c@.
onEscape ∷ MonadFinally μ ⇒ μ α → μ β → μ α
onEscape m f = fmap fst $ finally' m $ maybe (() <$ f) (const $ return ())
{-# INLINE onEscape #-}

-- | A variant of `onEscape` that combines multiple cleanup actions with
-- 'chainCleanups'.
onEscapeMany ∷ MonadFinally μ ⇒ μ α → [μ β] → μ α
onEscapeMany m = onEscape m . chainCleanups

-- | A variant of 'bracket' where acquired value is not needed (e.g. using
-- a static resource).
bracket_ ∷ MonadFinally μ
         ⇒ μ r -- ^ Acquire resource
         → μ β -- ^ Release resource
         → μ α -- ^ Main computation
         → μ α
bracket_ acquire release = bracket acquire (const release) . const
{-# INLINE bracket_ #-}

-- | A variant of 'bracket' that releases the acquired resource only when
-- the main computation does not produce a value.
bracketOnEscape ∷ MonadFinally μ
                ⇒ μ r       -- ^ Acquire resource
                → (r → μ β) -- ^ Release resource
                → (r → μ α) -- ^ Main computation
                → μ α
bracketOnEscape acquire release = fmap fst . bracket' acquire release'
  where release' _  (Just _) = return Nothing
        release' ar Nothing  = Just <$> release ar