StateT

Let's play this game one more time. Since State was defined as

newtype State s a = State { runState :: s -> (a, s) }

StateT s m a needs to wrap a function of type a -> m (a, s):

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

and the Monad instance of StateT once again needs to make sure that we pass the correct state from function to function:

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

Our put, get, gets, and modify functions for the State monad need to be rewritten to work in the undelying monad m instead of with pure functions:

get :: Monad m => StateT s m s
get = StateT $ \s -> return (s, s)

gets :: Monad m => (s -> a) -> StateT s m a
gets f = StateT $ \s -> return (f s, s)

put :: Monad m => s -> StateT s m ()
put s = StateT $ \_ -> return ((), s)

modify :: Monad m => (s -> s) -> StateT s m ()
modify f = StateT $ \s -> return ((), f s)

And the actual implementation of State once again uses StateT to equip the Identity monad with some state:

type State s = StateT s Identity

runState f = runIdentity . runStateT f