Learning Monads & Monad Transformers

{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE ExplicitForAll #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}

{-
This is an exercise in learning Haskell. I wrote these mostly in a clean-room fashion
as a way to figure out how to do monads, monad transformers, a Reader monad (RT) and
a state monad (ST).
-}

module Main where

import Control.Applicative
import Control.Monad.Identity
import Control.Monad.IO.Class
import Control.Monad.Reader
import Control.Monad.State
import Control.Monad.Trans
import Data.Monoid

newtype RT r m a = RT { runRT :: r -> m a }

instance Functor m => Functor (RT r m) where
    fmap f (RT a) = RT $ fmap f . a

instance Applicative m => Applicative (RT r m) where
    pure = RT . const . pure
    (RT f) <*> (RT a) = RT $ \r -> (f r) <*> (a r)

instance Monad m => Monad (RT r m) where
    return = pure
    (RT a) >>= f = RT $ \r -> do
        a' <- a r
        runRT (f a') r

instance MonadTrans (RT r) where
    lift = RT . const

instance MonadIO m => MonadIO (RT r m) where
    liftIO = lift . liftIO

instance Monad m => MonadReader r (RT r m) where
    ask = RT return
    local t (RT a) = RT $ \r -> a (t r)

type R r = RT r Identity

runR :: R r a -> r -> a
runR a = runIdentity . runRT a

newtype ST s m a = ST { runST :: s -> m (a, s) }

evalST a = fst . runST a
execST a = snd . runST a

instance Functor m => Functor (ST s m) where
    fmap f (ST a) = ST $ fmap (\(a', s') -> (f a', s')) . a

instance Monad m => Applicative (ST s m) where
    pure a = ST $ pure . (a,)
    (ST f) <*> (ST a) = ST $ \s -> do
        ~(f', s') <- f s
        ~(a', s'') <- a s'
        return $ (f' a', s'')

instance Monad m => Monad (ST s m) where
    return = pure
    (ST a) >>= f = ST $ \s -> do
        ~(a', s') <- (a s)
        runST (f a') s'

instance Monad m => MonadState s (ST s m) where
    get = ST $ \s -> return (s, s)
    put s = ST $ \_ -> return ((), s)

instance MonadTrans (ST s) where
    lift a = ST $ \s -> a >>= return . (,s)

instance MonadIO m => MonadIO (ST s m) where
    liftIO = lift . liftIO

type S s = ST s Identity

runS :: S s a -> s -> (a, s)
runS a = runIdentity . runST a

evalS a = fst . runS a
execS a = snd . runS a

type RS a = RT a (S a) a

runRS :: Int -> Int -> RS Int -> (Int, Int)
runRS r s a = runS (runRT a r) s

evalRS r s a = fst $ runRS r s a
execRS r s a = snd $ runRS r s a

x :: Int -> RT Int (S Int) Int
x n = do
    r <- ask
    s <- lift $ do
        modify (+n)
        get
    return $ (n * n) + (r * s)

main = print $ runS (runRT (liftA2 (+) (x 2) (x 1)) 2) 1