More Monads: Stateful Programming

Time for more monads. In this article, I'm going to show you how to implement a very simple
state monad - it's a trivial monad which allows you to use a mutable state
consisting of a single integer. Then we'll expand it to allow a more interesting
notion of state.

Let's get the trivial stuff out of the way. To start off, we'll use a simple fixed
state type which is just a wrapper for a single integer

>data State = State Int deriving (Show)

A monad over that would be written:

>data SimpleStateMonad a = SSMon (State -> (a,State))

To understand that, just remember that the monad is really a transition function for
executing a sequence of computations that use a state. In fact, we'll describe the sequence of steps
chained together that way as a single monad operation as a monadic sequence. In a monadic sequence, each step of the computation is going to take the current state, do something with it, and return both a value and an updated state.

Given the monad type - the basic type of the step function for the sequences we want to build - the first thing that we need to do is implement the fundamental monadic operators. There are two of them: the sequencing operator (>>=), and the return function that brings a non-monadic value into a monad.

We'll implement the (>>=) operator, which we'll implement as a function named chain:

>chain :: (SimpleStateMonad a) -> (a -> SimpleStateMonad b) 
>          -> SimpleStateMonad b
>chain (SSMon state) stepfun =
>      SSMon (\s -> let (carried, stateBefore) = state s
>                       SSMon stateMonadAfter = stepfun carried
>                   in stateMonadAfter stateBefore)

The way that chain works is crucial, and so we'll take some time and tease it apart. Chain is a function that takes the result of the previous computation step of the sequence, and
chains it into the next step of the sequence. The way to understand the code is to look at the basic
pieces, and how they're put together:

• The first binding of the let: (carried,stateBefore) = state s. This looks
  mysterious - so we'll take the time to look at it carefully. state is the
  function inside of the monad that carries hidden monadic state and the result of the previous
  step. So all we're really doing here is unwrapping the results of the previous step - extracting the state and the return values, and binding them to separate variables. carried is the result of the previous step, and stateBefore is the
  state returned by the previous step - that is, the state before we chain in
  the next step of the monadic sequence.
• The second binding of the let: SSMon stateMonadAfter = stepfun carried. Again,
  it looks mysterious, but it's really pretty simple. The monad is a step function - we're
  creating a step function from the result of the previous step, so that we can
  use that to run the next step. stepfun carried is just a function which
  passes the state and return value of the first step into the second step of our
  chain.
• The main expression of the let: stateMonadAfter stateBefore. This just
  runs the step function we generated on the result of the previous step.
• The λ wrapping the whole shebang: monads are a bit of a trap - you don't
  get to get out of them without using a type-specific operation that isn't part of the
  monad class. The monadic chain operator needs to result in a monad value - the wrapping
  lambda expression just encapsulates the stepping sequence we've built into a monadic
  function wrapper.

Return is trivial, so we'll just write it inline in the declaration of our monad type as an instance of the monad typeclass. The typeclass declaration looks like:

>instance Monad SimpleStateMonad where
>    statemonad >>= statefun = chain statemonad statefun
>    return nonMonadValue = SSMon (\state -> (nonMonadValue, state))

Nothing complicated there. The next thing we need to do is to write some functions usable
inside of a monadic computation to access and update the state hidden inside the
monad. Let's start with retrieving the state value:

>getState :: SimpleStateMonad Int
>getState = SSMon (\s@(State i) -> (i,s))

That's it! It's just a monadic step function that takes the value hidden in the state,
and puts it into the "return value" slot of the monad result. So the monadic state is
unmodified, and the result passed to the next state is the value that was wrapped up in the monad.

Updating is almost as easy:

>putState ::  (Int -> Int) -> SimpleStateMonad ()
>putState updateFun = 
>     let stateUpdateFun = (\(State i)  -> State (updateFun i)) 
>     in SSMon (\s -> ((), stateUpdateFun s))

putState takes a function that performs a modification of the state value. Since
we want people to treat the state as an integer, we do a little bit of wrapping, to translate an Int -> Int function to a State -> State function. The real meat is just
executing the state update function to produce the new state value, and wrapping it up in a monad
that returns the unit value and the updated state.

Finally, we need a function that takes an initial state value, and a monadic sequence of operations, and executes the monadic sequence starting with the state value.

>runWithState :: Int -> SimpleStateMonad a ->  (a,State)
>runWithState f (SSMon c) = c (State f)

And voila! We have a state monad. Let's try a simple test of it. Here's
a monadic sequence that plays with the state a bit.

>testState :: SimpleStateMonad Int 
>testState = do
>           putState (3*)
>           putState (1+)
>           i <- getState
>           return (i+2)

Running that in GHCI gives us:

 *Main> runWithState 3 testState 
(12,State 10) 
*Main> runWithState 10 testState (33,State 31)
*Main> 

I ran it twice, with different values for the initial state. The code takes whatever the initial
state is, multiplies it by three, and adds one. That's the final value of the state. And then the do
block finishes with a value of the current state plus two. What's actually returned at the end of the
do block is a pair of the return value of the last statement in the block, and the final state. As you can see, it does exactly the right thing. It looks almost like a piece of imperative code,
except that the "assignment" operator, putState takes a function on the current
value of the state instead of just a new value.

So, now we've seen the basic method that we can use t make a monad that can encapsulate a kind of
state. But so far, it's a pretty trivial kind of state: how much good is one integer? Suppose we
wanted something more than that. In fact, suppose we wanted an arbitrary set of named variables
encapsulated in a state. We do almost the same thing as the simple state monad; we just use a more
complicated value inside the monad, and a different set of functions for working with the value in
the state. We'll start by just building a trivial non-monadic implementation of a name/value
list.

>type NameValueMap = [(String,Int)]
>
>getValueFromList :: NameValueMap -> String -> Int
>getValueFromList ((n,v):rest) name 
>             | n == name = v
>              otherwise = getValueFromList rest name
>getValueFromList [] name = -1
>
>setValueInList :: NameValueMap -> String 
                    -> Int -> [(String,Int)]
>setValueInList ((n,v):rest) name newval 
>          | n == name = ((n,newval):rest)
>          | otherwise = ((n,v):(setValueInList rest name newval))
>setValueInList [] name newval = [(name,newval)]

Now, we're going to wrap that as a monad. The declaration of the monad type
and the chain and return functions are going to be essentially identical to the integer
state monad.

>data MapStateMonad a = MapMon (NameValueMap -> (a,NameValueMap))
>
>mchain :: (MapStateMonad a) -> 
>          (a -> MapStateMonad b) -> 
>          MapStateMonad b
>mchain (MapMon state) stepfun = 
>        MapMon (\s -> let (carried, stateBefore) = state s
>                          MapMon stateMonadAfter = stepfun carried
>                      in stateMonadAfter stateBefore)
>
>
>instance Monad MapStateMonad where
>    statemonad >>= statefun = mchain statemonad statefun
>    return nonMonadValue = 
>            MapMon (\state -> (nonMonadValue, state))

Here's where we do something a little bit different. Instead of having "getState" and "putState" functions for working on the state in the monad, we need operations that fit the new name-map state.

>getFromState :: String -> MapStateMonad Int
>getFromState name = 
>      MapMon (\statemap -> (getValueFromList statemap name, 
>                            statemap))
>
>setInState :: String -> Int -> MapStateMonad ()
>setInState  name newval = 
>       MapMon (\statemap -> 
>                   ((), 
                    setValueInList statemap name newval))

And, just as in the case of the simple integer state monad, we need a way of running a monadic sequence with a particular initial state. And while we're at it, we'll write a test sequence.

>runWithMapState :: NameValueMap -> 
>                   MapStateMonad a ->  
>                   (a,NameValueMap)
>runWithMapState f (MapMon c) = c f
>
>testMapMonad = do
>                  setInState "x" 3
>                  setInState "y" 5
>                  z <- getFromState "a"
>                  x <- getFromState "x"
>                  return (x*z)

Now, aside from a bit of syntax, that looks like an imperative program, doesn't it? Except that we've implemented it in a pure functional way - using a hidden state parameter. That's the basic idea of what we do with monads - they allow us to create mini-languages with different
semantics based on the idea of chaining things together into monadic sequences. We write most of the program in the general functional language; but when we need to, we have multiple custom languages at our disposal in the form of the monad system. We can easily use monads for IO, for imperative programming, for Prolog-style backtracking, for non-determinism... All there, in Monads.

Next up, we'll look at the theory of monads: what they are in category theory, and how that relates to what we've seen in programming, and what rules are required by the monadic semantics to make a new monad work correctly.