-- 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 MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

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

    MonadException(..),
    MonadIOException(..),

    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.IO.Class (MonadIO(..))
import Control.Monad.RWS.Class (MonadRWS)
import Control.Monad.Reader.Class (MonadReader(..))
import Control.Monad.State.Class (MonadState(..))
import Control.Monad.Writer.Class (MonadWriter(..))
import Control.Monad.Trans.Class (MonadTrans(..))
import Control.Monad.Trans.RWS.Lazy as Lazy (RWST(..),
                                             runRWST)
import Control.Monad.Trans.RWS.Strict as Strict (RWST(..),
                                                 runRWST)
import Control.Monad.Trans.Reader (ReaderT(..))
import Control.Monad.Trans.State.Lazy as Lazy (StateT(..),
                                               runStateT)
import Control.Monad.Trans.State.Strict as Strict (StateT(..),
                                                   runStateT)
import Control.Monad.Trans.Writer.Lazy as Lazy (WriterT(..),
                                                runWriterT)
import Control.Monad.Trans.Writer.Strict as Strict (WriterT(..),
                                                    runWriterT)
import Data.Monoid (Monoid)

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) => MonadIOException 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

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)

liftException :: MonadException m => Either E.SomeException a -> m a
liftException (Left e)  = throw e
liftException (Right a) = return 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 MonadTrans ExceptionT where
    lift m = ExceptionT $ do
        a <- m
        return (Right a)

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

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

instance (MonadRWS m) => MonadRWS (ExceptionT m)

instance (MonadReader m) => MonadReader (ExceptionT m) where
    type EnvType (ExceptionT m) = EnvType m

    ask       = lift ask
    local f m = ExceptionT $ local f (runExceptionT m)

instance (MonadState m) => MonadState (ExceptionT m) where
    type StateType (ExceptionT m) = StateType m

    get = lift get
    put = lift . put

instance (MonadWriter m) => MonadWriter (ExceptionT m) where
    type WriterType (ExceptionT m) = WriterType m

    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)

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

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)

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

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

instance (Monoid w, MonadIOException m) => MonadIOException (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, MonadIOException m) => MonadIOException (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 (MonadIOException m) => MonadIOException (ReaderT r m) where
    block m    = ReaderT $ \r -> block (runReaderT m r)
    unblock m  = ReaderT $ \r -> unblock (runReaderT m r)

instance (MonadIOException m) => MonadIOException (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 (MonadIOException m) => MonadIOException (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, MonadIOException m) => MonadIOException (Lazy.WriterT w m) where
    block m    = Lazy.WriterT $ block (Lazy.runWriterT m)
    unblock m  = Lazy.WriterT $ unblock (Lazy.runWriterT m)

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