-- Copyright (c) 2008-2010
--         The President and Fellows of Harvard College.
--
-- Redistribution and use in source and binary forms, with or without
-- modification, are permitted provided that the following conditions
-- are met:
-- 1. Redistributions of source code must retain the above copyright
--    notice, this list of conditions and the following disclaimer.
-- 2. Redistributions in binary form must reproduce the above copyright
--    notice, this list of conditions and the following disclaimer in the
--    documentation and/or other materials provided with the distribution.
-- 3. Neither the name of the University nor the names of its contributors
--    may be used to endorse or promote products derived from this software
--    without specific prior written permission.
--
-- THIS SOFTWARE IS PROVIDED BY THE UNIVERSITY AND CONTRIBUTORS ``AS IS'' AND
-- ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-- IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
-- ARE DISCLAIMED.  IN NO EVENT SHALL THE UNIVERSITY OR CONTRIBUTORS BE LIABLE
-- FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-- DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
-- OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
-- HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
-- LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
-- OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
-- SUCH DAMAGE.

--------------------------------------------------------------------------------
-- |
-- Module      :  Control.Monad.Exception
-- Copyright   :  (c) Harvard University 2008-2010
-- License     :  BSD-style
-- Maintainer  :  mainland@eecs.harvard.edu
--
--------------------------------------------------------------------------------

{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE UndecidableInstances #-}

module Control.Monad.Exception (
    E.Exception(..),
    E.SomeException,

    MonadException(..),
    MonadAsyncException(..),

    ExceptionT(..),
    mapExceptionT,
    liftException
  ) where

import Prelude hiding (catch)

import qualified Control.Exception as E (Exception(..),
                                         SomeException,
                                         block,
                                         catch,
                                         throw,
                                         unblock)
import Control.Monad (MonadPlus(..))
import Control.Monad.Cont (MonadCont(..))
import Control.Monad.Fix (MonadFix(..))
import Control.Monad.RWS.Class (MonadRWS)
import Control.Monad.RWS.Lazy as Lazy (RWST(..),
                                       runRWST)
import Control.Monad.RWS.Strict as Strict (RWST(..),
                                           runRWST)
import Control.Monad.Reader (ReaderT(..))
import Control.Monad.Reader.Class (MonadReader(..))
import Control.Monad.State.Class (MonadState(..))
import Control.Monad.State.Lazy as Lazy (StateT(..),
                                         runStateT)
import Control.Monad.State.Strict as Strict (StateT(..),
                                             runStateT)
import Control.Monad.Trans (MonadIO(..),
                            MonadTrans(..))
import Control.Monad.Writer.Class (MonadWriter(..))
import Control.Monad.Writer.Lazy as Lazy (WriterT(..),
                                          runWriterT)
import Control.Monad.Writer.Strict as Strict (WriterT(..),
                                              runWriterT)
import Data.Monoid (Monoid)
import GHC.Base (RealWorld,
                 State#,
                 catchSTM#,
                 raiseIO#)
import GHC.Conc (STM(..))

class (Monad m) => MonadException m where
    -- |Throw an exception.
    throw :: E.Exception e => e -> m a
    -- |Catch an exception.
    catch :: E.Exception e => m a -> (e -> m a) -> m a

class (MonadIO m, MonadException m) => MonadAsyncException m where
    -- |Applying 'block' to a computation will execute that computation with
    -- asynchronous exceptions /blocked/.
    block :: m a -> m a
    -- |To re-enable asynchronous exceptions inside the scope of 'block',
    -- 'unblock' can be used.
    unblock :: m a -> m a

--
-- The ExceptionT monad transformer.
--

newtype ExceptionT m a =
    ExceptionT { runExceptionT :: m (Either E.SomeException a) }

mapExceptionT :: (m (Either E.SomeException a) -> n (Either E.SomeException b))
              -> ExceptionT m a
              -> ExceptionT n b
mapExceptionT f m = ExceptionT $ f (runExceptionT m)

-- |Lift the result of running a computation in a monad transformed by
-- 'ExceptionT' into another monad that supports exceptions.
liftException :: MonadException m => Either E.SomeException a -> m a
liftException (Left e)  = throw e
liftException (Right a) = return a

instance MonadTrans ExceptionT where
    lift m = ExceptionT $ do
        a <- m
        return (Right a)

instance (Monad m) => Functor (ExceptionT m) where
    fmap f m = ExceptionT $ do
        a <- runExceptionT m
        case a of
            Left  l -> return (Left  l)
            Right r -> return (Right (f r))

instance (Monad m) => Monad (ExceptionT m) where
    return a = ExceptionT $ return (Right a)
    m >>= k  = ExceptionT $ do
        a <- runExceptionT m
        case a of
          Left l  -> return (Left l)
          Right r -> runExceptionT (k r)
    fail msg = ExceptionT $ return (Left (E.toException (userError msg)))

instance (Monad m) => MonadPlus (ExceptionT m) where
    mzero       = ExceptionT $ return (Left (E.toException (userError "")))
    m `mplus` n = ExceptionT $ do
        a <- runExceptionT m
        case a of
          Left _  -> runExceptionT n
          Right r -> return (Right r)

instance (MonadFix m) => MonadFix (ExceptionT m) where
    mfix f = ExceptionT $ mfix $ \a -> runExceptionT $ f $ case a of
        Right r -> r
        _       -> error "empty mfix argument"

instance (Monad m) => MonadException (ExceptionT m) where
    throw e     = ExceptionT $ return (Left (E.toException e))
    m `catch` h = ExceptionT $ do
        a <- runExceptionT m
        case a of
          Left l  ->  case E.fromException l of
                        Just e  -> runExceptionT (h e)
                        Nothing -> return (Left l)
          Right r -> return (Right r)

instance (MonadIO m) => MonadIO (ExceptionT m) where
    liftIO m = ExceptionT $ liftIO $
        (m >>= return . Right)
        `E.catch` \(e :: E.SomeException) -> return (Left e)

instance (MonadAsyncException m) => MonadAsyncException (ExceptionT m) where
    block   = ExceptionT . block . runExceptionT
    unblock = ExceptionT . unblock . runExceptionT

--
-- Instances for the IO monad.
--

instance MonadException IO where
    catch = E.catch
    throw = E.throw

instance MonadAsyncException IO where
    block   = E.block
    unblock = E.unblock

--
-- Instances for the STM monad.
--

instance MonadException STM where
    catch = catchSTM
    throw = throwSTM

unSTM :: STM a -> (State# RealWorld -> (# State# RealWorld, a #))
unSTM (STM a) = a

catchSTM :: E.Exception e => STM a -> (e -> STM a) -> STM a
catchSTM (STM m) handler = STM $ catchSTM# m handler'
  where
    handler' e = case E.fromException e of
                   Just e' -> unSTM (handler e')
                   Nothing -> raiseIO# e

throwSTM :: E.Exception e => e -> STM a
throwSTM e = STM $ raiseIO# (E.toException e)

--
-- MonadException instances for mtl monad transformers.
--

instance (Monoid w, MonadException m) =>
    MonadException (Lazy.RWST r w s m) where
    throw       = lift . throw
    m `catch` h = Lazy.RWST $ \r s ->
                  Lazy.runRWST m r s `catch` \e -> Lazy.runRWST (h e) r s

instance (Monoid w, MonadException m) =>
    MonadException (Strict.RWST r w s m) where
    throw       = lift . throw
    m `catch` h = Strict.RWST $ \r s ->
                  Strict.runRWST m r s `catch` \e -> Strict.runRWST (h e) r s

instance (MonadException m) =>
    MonadException (ReaderT r m) where
    throw       = lift . throw
    m `catch` h = ReaderT $ \r ->
                  runReaderT m r `catch` \e -> runReaderT (h e) r

instance (MonadException m) =>
    MonadException (Lazy.StateT s m) where
    throw       = lift . throw
    m `catch` h = Lazy.StateT $ \s ->
                  Lazy.runStateT m s `catch` \e -> Lazy.runStateT (h e) s

instance (MonadException m) =>
    MonadException (Strict.StateT s m) where
    throw       = lift . throw
    m `catch` h = Strict.StateT $ \s ->
                  Strict.runStateT m s `catch` \e -> Strict.runStateT (h e) s

instance (Monoid w, MonadException m) =>
    MonadException (Lazy.WriterT w m) where
    throw       = lift . throw
    m `catch` h = Lazy.WriterT $
                  Lazy.runWriterT m `catch` \e -> Lazy.runWriterT (h e)

instance (Monoid w, MonadException m) =>
    MonadException (Strict.WriterT w m) where
    throw       = lift . throw
    m `catch` h = Strict.WriterT $
                  Strict.runWriterT m `catch` \e -> Strict.runWriterT (h e)

--
-- MonadAsyncException instances for mtl monad transformers.
--

instance (Monoid w, MonadAsyncException m) =>
    MonadAsyncException (Lazy.RWST r w s m) where
    block m    = Lazy.RWST $ \r s -> block (Lazy.runRWST m r s)
    unblock m  = Lazy.RWST $ \r s -> unblock (Lazy.runRWST m r s)

instance (Monoid w, MonadAsyncException m) =>
    MonadAsyncException (Strict.RWST r w s m) where
    block m    = Strict.RWST $ \r s -> block (Strict.runRWST m r s)
    unblock m  = Strict.RWST $ \r s -> unblock (Strict.runRWST m r s)

instance (MonadAsyncException m) =>
    MonadAsyncException (ReaderT r m) where
    block m    = ReaderT $ \r -> block (runReaderT m r)
    unblock m  = ReaderT $ \r -> unblock (runReaderT m r)

instance (MonadAsyncException m) =>
    MonadAsyncException (Lazy.StateT s m) where
    block m    = Lazy.StateT $ \s -> block (Lazy.runStateT m s)
    unblock m  = Lazy.StateT $ \s -> unblock (Lazy.runStateT m s)

instance (MonadAsyncException m) =>
    MonadAsyncException (Strict.StateT s m) where
    block m    = Strict.StateT $ \s -> block (Strict.runStateT m s)
    unblock m  = Strict.StateT $ \s -> unblock (Strict.runStateT m s)

instance (Monoid w, MonadAsyncException m) =>
    MonadAsyncException (Lazy.WriterT w m) where
    block m    = Lazy.WriterT $ block (Lazy.runWriterT m)
    unblock m  = Lazy.WriterT $ unblock (Lazy.runWriterT m)

instance (Monoid w, MonadAsyncException m) =>
    MonadAsyncException (Strict.WriterT w m) where
    block m    = Strict.WriterT $ block (Strict.runWriterT m)
    unblock m  = Strict.WriterT $ unblock (Strict.runWriterT m)

--
-- mtl instances for transformed monads.
--

instance (MonadCont m) => MonadCont (ExceptionT m) where
    callCC f = ExceptionT $
        callCC $ \c ->
        runExceptionT (f (\a -> ExceptionT $ c (Right a)))

instance (MonadRWS r w s m) => MonadRWS r w s (ExceptionT m)

instance (MonadReader r m) => MonadReader r (ExceptionT m) where
    ask       = lift ask
    local f m = ExceptionT $ local f (runExceptionT m)

instance (MonadState s m) => MonadState s (ExceptionT m) where
    get = lift get
    put = lift . put

instance (MonadWriter w m) => MonadWriter w (ExceptionT m) where
    tell     = lift . tell
    listen m = ExceptionT $ do
        (a, w) <- listen (runExceptionT m)
        case a of
            Left  l -> return $ Left  l
            Right r -> return $ Right (r, w)
    pass   m = ExceptionT $ pass $ do
        a <- runExceptionT m
        case a of
            Left  l      -> return (Left  l, id)
            Right (r, f) -> return (Right r, f)