Understanding monad transformers (3): stateful is beautiful

In the first and second post of this series I describe a test case (a function evaluator) and a method to tackle its implementation by writing a complex monad in plain Haskell, built as a pile of transformers. All this from the ground up, without importing any library (except basic stuff like Data.Map and alike) and following the lines described in a presentation by Henrik Nilsson.

All code is available in GitHub, fully covered by tests.

The implementation we’ve come to, handles errors returning Nothing in all cases of impossible evaluation.

This time we are to provide the capability to weave a mutable state through the computation and modify it to record useful happenings. To start with, state will be an type Int to be increased by one at each recursive call of the evaluation function.

Teaching a monad to handle mutable state

In order to give any monad m the capability to pass a state of type s through each step of its own computation we have to follow the example of the state monad and let m become part of a function of s.

The type that expresses this transformation is called StateT in mtl, whereas Nilsson calls it ST:

newtype ST s m a = ST (s -> m (a, s))
deriving instance Show (s -> m (a, s)) => Show (ST s m a)
deriving instance Eq (s -> m (a, s)) => Eq (ST s m a)

The monad instance that results by adding state management to m is the following:

instance (Monad m) => Monad (ST s m) where
  return x = ST (\s -> return (x, s))
  stsma >>= fastsmb = ST $ \s -> do (aaa, sss) <- unST stsma s
                                    unST (fastsmb aaa) sss

On Github you may find also the corresponding functor and applicative instance definitions, along with the unwrapper and the runner, as Nilsson does not want us to use the pesky Haskell record syntax.

Standing on the shoulders of the giants that came here before us, we accept their warm suggestion to fit the ST transformer right around the inner I. That’s to say that ET (seen in the previous post) will wrap ST. The type of the resulting evaluation will therefore be (please follow carefully the substitutions):

ET (ST Int I) Value 
  = ET (ST Int I) (Maybe Value) 
  = ET (ST (Int -> I (Maybe Value, Int)))

Preparing the monad and running it afterwards

As it is typical with the state monad, each stateful computation is prepared by building the monad and running it afterwards, after providing the initial state. The evaluation function must therefore have type:

Env -> Exp -> ET (ST Int I) Value


and it will have to be run by a function that closes on an Int (the state monad value at the beginning of the run) and returns a couple (Maybe Value, Int). The couple will contain either a Just Value or a Nothing (the latter in case of impossibility to evaluate the given expression) and the final state value. We decide to call that function runEval3:

runEval3 :: Int -> ET (ST Int I) Value -> (Maybe Value, Int)
runEval3 s esm = unI $ unST (unET esm) s


In order to fully appreciate this last step you gotta remember that for Nilsson unST is actually a runner (corresponding to runStateT in mtl) and keep in mind the sequence of the onion peeling operations.

A state monad to record each invocation of a given function

In the body of any monadic function we may always use the do-notation abstraction to mimic an imperative step; in the case of a state monad the imitation of giving orders nears perfection. We want to carry a counter through the evaluation of an expression, and we want it to increase each time the evaluation function is invoked, both at the very start and afterwards through recursion. Our guide, Martin Grabmüller calls that state monad tick.

The way a simpler version of tick works is explained thoroughly in another post. In this case also I is involved and the function is:

tick :: ST Int I Int
tick = ST $ \n -> I (n, n+1)


Each time tick is invoked inside a do-notation or used as body of a function passed to bind (if you like no sugar in your coffee), the state become the value for an instant, it gets increased and passed further down the computation.

Our expression evaluator begins to get interesting

Following the steps of the Grabmüller paper and using Nilsson’s Lego blocks we come to the following evaluation function:

eval3 :: Env -> Exp -> ET (ST Int I) Value
eval3 env (Lit i)      = do (lift tick)
                            return $ IntVal i
eval3 env (Var name)   = do (lift tick)
                            ET $ return $ Map.lookup name env  
eval3 env (Plus e1 e2) = do (lift tick)
                            v1 <- eval3 env e1
                            v2 <- eval3 env e2
                            case (v1, v2) of
                              (IntVal i1, IntVal i2) -> return $ IntVal $ i1 + i2
                              _ -> eFail
eval3 env (Lambda argname body) = do (lift tick)
                                     return $ FunVal argname body env
eval3 env (App lambda expr) = do (lift tick)
                                 FunVal argname body env' <- eval3 env lambda
                                   val <- eval3 env expr
                                   eval3 (Map.insert argname val env') body

There are two important observation to make:

• the function runs at the ET level, therefore tick must be lifted (details follow), whereas eFail must not
• there is actually no calculation involved so far; as we said before, this is a Int -> I (Maybe Value, Int) function at heart

We are kind of winding a spring mouse that will run as soon as it’s fed an Int. This is not only true in the notoriously lazy Haskell: state monads keep quiet even when written in the eager JavaScript.

Here you may admire the function at work in several tests.

Giving ST all the traits it deserves

The list of typeclasses that our state transformer ST may want to implement becomes rather long:

• MonadTransformer: gotta implement lift
• S (MonadState for mtl): this requires implementing a getState/sGet and a setState/sSet methods
• become itself an error handler, by implementing eFail and eHandle

On its turn, ET may become a state handler. This way we wouldn’t have to lift tick inside the eval3 function and tick would work also at the outer ET layer.

But we started this whole endeavor saying how confusing it is when every transformer is able to do everything. All these implementations are to be found in the Github code, but from now on we will never mention mix-ups of monad effects spanning multiple levels.

Time to write and read, now

The next posts will illustrate the usage of more monad transformers that will change our types into:

• a writer monad that logs the various processing steps
• a reader monad that carries implicitly the evaluation environment and serves it to any computation step that may need to consult it
• a list monad that evaluates n expressions without a single trace of a loop inside the code