RWS.hs

module Trans.RWS (
  module X,
  RWS,
  runRWS,
  evalRWS,
  execRWS,
  mapRWS,
  RWST,
  runRWST,
  evalRWST,
  execRWST,
  mapRWST,
) where

import Prelude
import Control.Monad.Identity
import Control.Monad.Trans
import qualified Control.Monad.Fail as Fail
import Control.Monad.Fix (MonadFix(..))
import Control.Applicative (Alternative(..))

import Control.Monad.RWS as X (MonadRWS)
import Control.Monad.Writer as X (MonadWriter(writer, tell, listen, pass))
import Control.Monad.Reader as X (MonadReader(ask, local, reader))
import Control.Monad.State as X (MonadState(get, put, state))

newtype RWST r w s m a = RWST { unRWST :: r -> s -> w -> m (a, s, w) }
type RWS r w s = RWST r w s Identity

instance Functor m => Functor (RWST r w s m) where
  fmap f m = RWST $ \r s w -> (\(a, s', w') -> (f a, s', w')) <$> unRWST m r s w
  {-# INLINE fmap #-}

instance Monad m => Applicative (RWST r w s m) where
  pure a = RWST $ \_ s w -> pure (a, s, w)
  {-# INLINE pure #-}

  RWST mf <*> RWST mx = RWST $ \r s w -> do
    (f, s', w')  <- mf r s w
    (x, s'', w'') <- mx r s' w'
    pure (f x, s'', w'')
  {-# INLINE (<*>) #-}

instance MonadPlus m => Alternative (RWST r w s m) where
  empty = RWST $ \_ _ _ -> mzero
  {-# INLINE empty #-}

  RWST m <|> RWST n = RWST $ \r s w -> m r s w `mplus` n r s w
  {-# INLINE (<|>) #-}

instance Monad m => Monad (RWST r w s m) where
  m >>= k = RWST $ \r s w -> do
    (a, s', w')  <- unRWST m r s w
    unRWST (k a) r s' w'
  {-# INLINE (>>=) #-}

instance Fail.MonadFail m => Fail.MonadFail (RWST r w s m) where
  fail msg = RWST $ \_ _ _ -> Fail.fail msg
  {-# INLINE fail #-}

instance MonadPlus m => MonadPlus (RWST r w s m) where
  mzero = empty
  {-# INLINE mzero #-}
  mplus = (<|>)
  {-# INLINE mplus #-}

instance MonadFix m => MonadFix (RWST r w s m) where
  mfix f = RWST $ \r s w -> mfix $ \ ~(a, _, _) -> unRWST (f a) r s w
  {-# INLINE mfix #-}

instance MonadTrans (RWST r w s) where
  lift m = RWST $ \_ s w -> (,s,w) <$> m
  {-# INLINE lift #-}

instance MonadIO m => MonadIO (RWST r w s m) where
  liftIO = lift . liftIO
  {-# INLINE liftIO #-}

instance (Monoid w, Monad m) => MonadWriter w (RWST r w s m) where
  writer (a, w') = RWST $ \_ s w -> let wt = w `mappend` w' in wt `seq` pure (a, s, wt)
  {-# INLINE writer #-}

  listen m = RWST $ \r s w -> do
    (a, s', w') <- runRWST m r s
    let wt = w `mappend` w'
    wt `seq` pure ((a, w'), s', wt)
  {-# INLINE listen #-}

  pass m = RWST $ \r s w -> do
    ((a, f), s', w') <- runRWST m r s
    let wt = w `mappend` f w'
    wt `seq` pure (a, s', wt)
  {-# INLINE pass #-}

instance Monad m => MonadReader r (RWST r w s m) where
  ask = RWST $ \r s w -> pure (r, s, w)
  {-# INLINE ask #-}

  local f m = RWST $ \r s w -> unRWST m (f r) s w
  {-# INLINE local #-}

  reader f = RWST $ \r s w -> pure (f r, s, w)
  {-# INLINE reader #-}

instance Monad m => MonadState s (RWST r w s m) where
  get = RWST $ \_ s w -> pure (s, s, w)
  {-# INLINE get #-}

  put s = RWST $ \_ _ w -> pure ((), s, w)
  {-# INLINE put #-}

  state f = RWST $ \_ s w -> let (a, s') = f s in pure (a, s', w)
  {-# INLINE state #-}

runRWST :: Monoid w => RWST r w s m a -> r -> s -> m (a, s, w)
runRWST m r s = unRWST m r s mempty
{-# INLINE runRWST #-}

evalRWST :: (Functor m, Monoid w) => RWST r w s m a -> r -> s -> m (a, w)
evalRWST m r s = (\(a, _, w) -> (a, w)) <$> runRWST m r s
{-# INLINE evalRWST #-}

execRWST :: (Functor m, Monoid w) => RWST r w s m a -> r -> s -> m (s, w)
execRWST m r s = (\(_, s', w) -> (s', w)) <$> runRWST m r s
{-# INLINE execRWST #-}

mapRWST :: (Functor n, Monoid w, Monoid w') =>
  (m (a, s, w) -> n (b, s, w')) -> RWST r w s m a -> RWST r w' s n b
mapRWST f m = RWST $ \r s w -> (\(b, s', w') -> (b, s', w `mappend` w')) <$> f (runRWST m r s)
{-# INLINE mapRWST #-}

runRWS :: Monoid w => RWS r w s a -> r -> s -> (a, s, w)
runRWS m r s = runIdentity $ runRWST m r s
{-# INLINE runRWS #-}

execRWS :: Monoid w => RWS r w s a -> r -> s -> (s, w)
execRWS m r s = runIdentity $ execRWST m r s
{-# INLINE execRWS #-}

evalRWS :: Monoid w => RWS r w s a -> r -> s -> (a, w)
evalRWS m r s = runIdentity $ evalRWST m r s
{-# INLINE evalRWS #-}

mapRWS :: (Monoid w, Monoid w') => ((a, s, w) -> (b, s, w')) -> RWS r w s a -> RWS r w' s b
mapRWS f = mapRWST (Identity . f . runIdentity)
{-# INLINE mapRWS #-}

Writer.hs

module Trans.Writer (
  module X,
  Writer,
  runWriter,
  execWriter,
  mapWriter,
  WriterT,
  runWriterT,
  execWriterT,
  mapWriterT,
) where

import Prelude
import Control.Monad.Identity
import Control.Monad.Trans
import Control.Arrow (first, second)
import qualified Control.Monad.Fail as Fail
import Control.Monad.Fix (MonadFix(..))
import Control.Applicative (Alternative(..))

import Control.Monad.Writer as X (MonadWriter(writer, tell, listen, pass))

newtype WriterT w m a = WriterT { unWriterT :: w -> m (a, w) }
type Writer w = WriterT w Identity

instance Functor m => Functor (WriterT w m) where
  fmap f m = WriterT $ \w -> first f <$> unWriterT m w
  {-# INLINE fmap #-}

instance Monad m => Applicative (WriterT w m) where
  pure a = WriterT $ pure . (a,)
  {-# INLINE pure #-}

  WriterT mf <*> WriterT mx = WriterT $ \w -> do
    (f, w') <- mf w
    (x, w'') <- mx w'
    pure (f x, w'')
  {-# INLINE (<*>) #-}

instance MonadPlus m => Alternative (WriterT w m) where
  empty = WriterT $ const mzero
  {-# INLINE empty #-}

  WriterT m <|> WriterT n = WriterT $ \w -> m w `mplus` n w
  {-# INLINE (<|>) #-}

instance Monad m => Monad (WriterT w m) where
  m >>= k = WriterT $ \w -> do
    (a, w') <- unWriterT m w
    unWriterT (k a) w'
  {-# INLINE (>>=) #-}

instance Fail.MonadFail m => Fail.MonadFail (WriterT w m) where
  fail msg = WriterT $ \_ -> Fail.fail msg
  {-# INLINE fail #-}

instance MonadPlus m => MonadPlus (WriterT w m) where
  mzero = empty
  {-# INLINE mzero #-}
  mplus = (<|>)
  {-# INLINE mplus #-}

instance MonadFix m => MonadFix (WriterT w m) where
  mfix f = WriterT $ \w -> mfix $ \ ~(a, _) -> unWriterT (f a) w
  {-# INLINE mfix #-}

instance MonadTrans (WriterT s) where
  lift m = WriterT $ \w -> (,w) <$> m
  {-# INLINE lift #-}

instance MonadIO m => MonadIO (WriterT w m) where
  liftIO = lift . liftIO
  {-# INLINE liftIO #-}

instance (Monoid w, Monad m) => MonadWriter w (WriterT w m) where
  writer (a, w') = WriterT $ \w -> let wt = w `mappend` w' in wt `seq` pure (a, wt)
  {-# INLINE writer #-}

  listen m = WriterT $ \w -> do
    (a, w') <- runWriterT m
    let wt = w `mappend` w'
    wt `seq` pure ((a, w'), wt)
  {-# INLINE listen #-}

  pass m = WriterT $ \w -> do
    ((a, f), w') <- runWriterT m
    let wt = w `mappend` f w'
    wt `seq` pure (a, wt)
  {-# INLINE pass #-}

runWriterT :: Monoid w => WriterT w m a -> m (a, w)
runWriterT m = unWriterT m mempty
{-# INLINE runWriterT #-}

execWriterT :: (Functor m, Monoid w) => WriterT w m a -> m w
execWriterT = fmap snd . runWriterT
{-# INLINE execWriterT #-}

mapWriterT :: (Functor n, Monoid w, Monoid w') =>
  (m (a, w) -> n (b, w')) -> WriterT w m a -> WriterT w' n b
mapWriterT f m = WriterT $ \w -> second (mappend w) <$> f (runWriterT m)
{-# INLINE mapWriterT #-}

runWriter :: Monoid w => Writer w a -> (a, w)
runWriter = runIdentity . runWriterT
{-# INLINE runWriter #-}

execWriter :: Monoid w => Writer w a -> w
execWriter = runIdentity . execWriterT
{-# INLINE execWriter #-}

mapWriter :: (Monoid w, Monoid w') => ((a, w) -> (b, w')) -> Writer w a -> Writer w' b
mapWriter f = mapWriterT (Identity . f . runIdentity)
{-# INLINE mapWriter #-}