{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE UndecidableInstances #-}

{-
   Module Control.Monad.Wrap
   Copyright 2009 David Mazieres (http://www.scs.stanford.edu/~dm/)
   BSD Licence

   Non-portable.  Requires MultiParamTypeClasses, FlexibleInstances,
   FunctionalDependencies, and (for MonadWrapIO) UndecidableInstances.

   Process with haddock for documentation.
-}

{- |

   Class of monad transformers whose computations can be wrapped by
   functions such as 'catch' and 'finally' that operate on inner
   monadic types.  This works for all standard monad transformers
   except for 'ContT'.

   The main method in this module is 'wrap', which wraps a function on
   one monad around a transformed version of that monad.  This is
   useful in many situations, but it is important to keep in mind that
   this will affect the semantics of the monad.  For example, if you
   wrap 'forkIO' around a @'StateT' X 'IO'@ monad, it will fork the
   state of the monad into two independently updated copies of the
   state, one for each thread.  Similarly, if you wrap 'catch' around
   a @'WriterT' X 'IO'@ monad, you can catch an exception and return a
   value, but the writer state will be re-wound to the point at which
   you computed the result value (using 'result' or 'resultF').

-}


module Control.Monad.Wrap (
                 -- * The MonadWrap type
                  MonadWrap(..)
                 -- * Wrapping IO functions around other monads
                 , MonadWrapIO(..)
                 -- * Example
                 -- $examples
                 ) where

import Control.Monad.Trans (MonadTrans(..), MonadIO(..))
import Control.Monad.Cont (MonadCont(..))
import Control.Monad.Error (Error(..), ErrorT(..))
import Control.Monad.List (ListT(..))
import Control.Monad.Reader (ReaderT(..))
import Control.Monad.RWS (RWST(..))
import Control.Monad.State (StateT(..))
import Control.Monad.Writer (WriterT(..))
import Data.Monoid (Monoid(..))

class MonadWrap t a r | t a -> r where
    -- | This function is called to wrap a computation in a tansformed
    -- monad, @t m@, with a function that expects an argument and
    -- returns a result of the inner monadic type, @m@.  The canonical
    -- example of this is using a monad transformer such as 'ReaderT',
    -- having a computation @x :: 'ReaderT' MyConfig 'IO' a@, and
    -- wanting to wrap a function such as
    -- @'finally' :: 'IO' b -> 'IO' b@
    -- around @x@ to run a @cleanup@ function even when an
    -- exception is thrown.  If @x@ invokes a 'ReaderT' function such
    -- as 'ask', this cannot be accomplished with the ordinary monad
    -- transformer method 'lift'.  Instead, it can be achieved with:
    --
    -- >    wrap (\op -> op `finally` cleanup) x
    --
    -- If the inner monad is nested within multiple levels of
    -- transformer, you can invoke @wrap@ multiple times, e.g.:
    --
    -- >    wrap (wrap (\op -> op `finally` cleanup)) x
    --
    wrap :: (Monad m) => (m r -> m r) -> t m a -> t m a

    -- | Sometimes the wrapping function passed as the first argument
    -- of 'wrap' needs to produce its own return value rather than
    -- passing one straight through.  This is the case with 'catch'.
    -- To do this, use @result@ to tranlate a pure value into the
    -- inner-monad value required by the outer monad transformer.
    --
    -- (Note that @result@ is basically a pure transofmation, but it
    -- produces a value in the outer monad primarily just so that it
    -- can infer from context what type of result to produce.  The
    -- pure value returned must be lifted into the inner monad with
    -- @return@.)
    --
    -- @result@ is perhaps best illustrated by example.  With a
    -- single level of nesting, use something like:
    --
    -- >    err <- result Nothing
    -- >    wrap (handle (\(SomeException _) -> return err))
    -- >          $ liftM Just $ someIOComputation
    --
    -- For multiple levels of nesting, you will need to invoke
    -- @result@ multiple times, e.g.:
    --
    -- >    outerErr <- result Nothing
    -- >    innerErr <- lift $ result outerErr
    -- >    wrap (wrap (handle (\(SomeException _) -> return innerErr)))
    -- >          $ liftM Just $ someIOComputation
    --
    -- Note the use of 'lift' so as to compute the inner result
    -- transformation in the inner monad.  Further levels of nesting
    -- require increasing numbers of 'lift's.
    --
    -- It is important to keep in mind that where you invoke @result@
    -- (or 'resultF') may affect the result.  For instance, with the
    -- 'StateT' or 'WriterT' transformer, if you return a value
    -- computed with @result@, the state will be re-wound to the
    -- point at which you called @result@, discarding any changes that
    -- may have taken place in the mean time.
    result :: (Monad m, Monad (t m)) => a -> t m r
    result a = resultF >>= return . ($ a)

    -- | @resultF@ returns a function that can be used to perform the
    -- same transformation as 'result'.  This is necessary if you need
    -- to compute the result type dynamically within the wrapping
    -- function.  'result' can be implemented in terms of @resultF@
    -- as:
    --
    -- >  result a = do f <- resultF
    -- >                return (f a)
    --
    -- The example given above for 'result' could altenratively have
    -- been implemented as:
    --
    -- >    fout <- resultF
    -- >    fin <- lift resultF
    -- >    let f = fin . fout
    -- >    wrap (wrap (handle (\(SomeException _) -> return $ f Nothing)))
    -- >          $ liftM Just $ someIOComputation
    --
    -- Note that the order of composition is @f = fin . fout@, and
    -- /not/ @f = fout . fin@.  This is because types and result
    -- values nest in opposite directions with monad transformers.
    -- Even though 'IO' is the inner monad in a type such as
    -- @'WriterT' 'String' 'IO' ()@, running any computation of this
    -- type produces a value of type @'IO' ((), 'String')@.
    resultF :: (Monad m) => t m (a -> r)

instance (Error e) => MonadWrap (ErrorT e) a (Either e a) where
    wrap f = ErrorT . f . runErrorT
    resultF = return Right

instance MonadWrap ListT a [a] where
    wrap f = ListT . f . runListT
    resultF = return (: [])

instance MonadWrap (ReaderT r) a a where
    wrap f m = ReaderT $ f . runReaderT m
    resultF = return id

instance (Monoid w) => MonadWrap (RWST r w s) a (a, s, w) where
    wrap f m = RWST $ \r s -> f $ runRWST m r s
    resultF = RWST $ \r s -> return (\a -> (a, s, mempty), s, mempty)

instance MonadWrap (StateT s) a (a, s) where
    wrap f m = StateT $ f . runStateT m
    resultF = StateT $ \s -> return (\a -> (a, s), s)

instance (Monoid w) => MonadWrap (WriterT w) a (a, w) where
    wrap f = WriterT . f . runWriterT
    resultF = return $ \a -> (a, mempty)

-- | MonadWrapIO is analogous to 'MonadWrap', but where the wrapping
-- function is always of type @'IO' r -> 'IO' r@.  The point of
-- @MonadWrapIO@ is to go through as many nested monad transformers as
-- necessary to reach the 'IO' monad, so you don't have to keep track
-- of where you are in terms of monad nesting depth.
class (Monad m) => MonadWrapIO m a r | m a -> r where
    -- | @wrapIO@ is to 'wrap' as 'liftIO' is to 'lift'.
    wrapIO :: (IO r -> IO r) -> m a -> m a
    -- | @resultFIO@ is to 'resultF' as 'liftIO' is to 'lift'.
    resultFIO :: m (a -> r)
    -- | @resultIO@ is to 'result' as 'liftIO' is to 'lift'.
    resultIO :: a -> m r
    resultIO a = resultFIO >>= return . ($ a)

instance MonadWrapIO IO a a where
    wrapIO f = f
    resultFIO = return id
    resultIO = return

-- This implementation requires -XUndecidableInstances.  Note that it
-- would be possible to avoid -XUndecidableInstances by implementing
-- an instance of MonadWrapIO for each Monad Transformer, as MonadIO
-- does.
instance (Monad m, MonadTrans t, Monad (t m), MonadWrapIO m ar r,
          MonadWrap t a ar) => MonadWrapIO (t m) a r where
    wrapIO f = wrap (wrapIO f)
    resultFIO = do outer <- resultF
                   inner <- lift resultFIO
                   return $ inner . outer
    resultIO a = result a >>= lift . resultIO


{- $examples

   Here is a longer example showing @finally@ and @catch@ used within
   the @WriterT@ and @StateT@ monads.  (Note that it would be easier
   to use 'resultIO' and 'wrapIO' in @middle@, but here we show how to
   wrap through multiple monads manually.)

   @
    \{\-\# LANGUAGE DeriveDataTypeable \#\-\}
   @
   
   >    module Main where
   >
   >    import Prelude hiding (catch)
   >    import Control.Exception
   >    import Control.Monad.State
   >    import Control.Monad.Writer
   >    import Data.Typeable
   >
   >    import Control.Monad.Wrap
   >
   >    type OuterMonad = WriterT String IO
   >    
   >    type MyState = Int
   >    type InnerMonad = StateT MyState OuterMonad
   >    
   >    data Trap = Trap deriving (Typeable, Show)
   >    instance Exception Trap
   >    
   >    handler :: String -> IO a -> Trap -> IO a
   >    handler place a e = do
   >      putStrLn $ "caught " ++ show e ++ " in " ++ place
   >      a
   >    
   >    inner :: InnerMonad ()
   >    inner = do
   >      liftIO $ putStrLn "running inner"
   >      liftIO $ throwIO Trap
   >    
   >    middle :: InnerMonad ()
   >    middle = do
   >      put 1                         -- Can do StateT operations
   >      liftIO $ putStrLn "running middle"
   >      x <- result ()
   >      y <- lift $ result x
   >      wrap (wrap (handle $ handler "middle" $ return y)) inner
   >      wrap (wrap do_finally) inner
   >        where
   >          do_finally = flip finally $ putStrLn "middle finally!"
   >    
   >    outer :: OuterMonad ()
   >    outer = do
   >      tell "This is outer"          -- Can do WriteT operations
   >      liftIO $ putStrLn "About to run middle"
   >      x <- result ()
   >      wrap (handle $ handler "outer" $ return x) (evalStateT middle 0)
   >      liftIO $ putStrLn "Just ran middle"
   >      wrap do_finally (evalStateT middle 0)
   >      liftIO $ putStrLn "This line won't be reached"
   >        where
   >          do_finally = flip finally $ putStrLn "outer finally!"
   >    
   >    main :: IO ((), String)
   >    main = runWriterT outer

   The above code should produce the following output: 

   >    *Main> main
   >    About to run middle
   >    running middle
   >    running inner
   >    caught Trap in middle
   >    running inner
   >    middle finally!
   >    caught Trap in outer
   >    Just ran middle
   >    running middle
   >    running inner
   >    caught Trap in middle
   >    running inner
   >    middle finally!
   >    outer finally!
   >    *** Exception: Trap
   >    *Main> 

-}