Implement State Monad transformer in Haskell from scratch

Question

When I was studying Monad Transformer, I decided to create StateT s m a from scratch with instances for Functor, Applicative and Monad.

This is what I have:

newtype StateT s m a = StateT { runStateT :: (s -> m (a, s)) }

instance Functor m => Functor (StateT s m) where
    -- fmap :: (a -> b) -> StateT s m a -> StateT s m b
    -- which is (a -> b) -> (s -> m (a, s)) -> (s -> m (b, s))
    f `fmap` (StateT x) = StateT $ \ s -> fmap run (x s)
        where run (a, s) = (f a, s)

instance Monad m => Applicative (StateT s m) where
    -- pure :: a -> StateT s m a
    pure a = StateT $ \ s -> pure (a, s)
    -- <*> :: f (a -> b) -> f a -> f b
    -- which is StateT s m (a -> b) -> StateT s m a -> State s m b
    k <*> x = StateT $ \ s -> do    
        (f, s1) <- runStateT k s -- :: m ((a -> b), s)
        (a, s2) <- runStateT x s1
        return (f a, s2)

instance (Monad m) => Monad (StateT s m) where
    return a = StateT $ \ s -> return (a, s)
    -- >>= :: StateT s m a -> (a -> StateT s m b) -> StateT s m b
    (StateT x) >>= f = StateT $ \ s -> do
        (v, s') <- x s
        runStateT (f v) s'

My original intention is to implement Functor (StateT s m) with Functor m restriction, Applicative (StateT s m) with Applicative m restriction, and Monad (StateT s m) withMonad m) restriction. However I couldn't do the Applicative case and had to use Monad m restriction instead. Is there a way to do it with Applicative m?

Thank you in advance.

Answer 1

Your implementation is correct. The additional comments are also welcome. I would write <*> without runState, but that's personal preference.

Keep in mind that you've re-implementend Control.Monad.Trans.State.Strict.StateT. For the lazy variant Control.Monad.Trans.State.Lazy.StateT, you would have to use lazy pattern matches in several places, so I assume you wanted to implement the strict variant.

To answer your question: no, we cannot implement Applicative (StateT s m) in terms of Applicative m. A Q&A on StackOverflow contains some hints, but let's reiterate them:

Suppose you have two values f, x and want y with the following types:

f :: StateT m s (a -> b)
x :: StateT m s a
y :: StateT m s b

We first remove the newtype:

f :: s -> m (s, (a -> b))
x :: s -> m (s, a)
y :: s -> m (s, b)

If our initial state is p, we have

f p :: m (s, (a -> b))
x   :: s -> m (s, a)

We want to use the s from f p in x to feed into x and the function to use the a. Now, let's suppose that m is an Applicative. This gives us only pure, fmap and <*> to work with. Let's try to write an expression that only uses those three functionspure, fmap and <*>:

z s = fmap (\(s', f') -> fmap (\(s'', x') -> (s'', f' x')) x s') (f s)

Let's check whether that implementation is sane. It's easier to deciper if we swap fmap's arguments:

x ~~> f = fmap f x
z s = f s                                   --  1
      ~~> (\(s', f') ->                     --  2
           x s'                             --  3
           ~~> (\(s'', x') -> (s'', f' x')  --  4
         )

1. We run f s to get our new state and our function.
2. Both are in m, so we use fmap to get in there.
3. We run x s' to get our new state again and our value.
4. Both are in m, so we use fmap yet again. However, we're using fmap inside fmap, which already tells us that we're not going to end up with a simple m ....

In (4), we end up with the correct new state s'' and the correct value f' x'. However, our type is s -> m (m (s, b)). Applicative does not provide any methods to reduce the number of m's, so we're stuck. We need to use join :: Monad m => m (m x) -> m x.

If we go back, we see that the problem arises due to x's type s -> m (s, a). If it was m (s -> (s, a)), we could simply use Applicative. Petr provides a detailed answer on SO.