Continuation monads for state machines (04 Jul 2007)

CPS (continuation-passing-style) is a code style which is often the result of the first step in compiling many Scheme like languages. Since I learned this stuff in Scheme, that's what I'm going to use in the beginning, switching to Haskell soon after.

So here's a top level Scheme program

(print (fact 10))

Rather than return, each function gets a function in its argument list which is its continuation. It's the function for the rest of the program which takes the result of the current function. It's always a tail-call.

(fact 10 (lambda (v) print v #exit#))

So here, fact runs and calls its continuation with the result. This continuation is a function which prints the value and calls its continuation, which terminates the program.

Easy, right?

So here's a continuation monad in Haskell; but a quick motivation first. What this will give us is something like a Python generator, but which we can pass values in to. So it's a state machine, but without the inverted flow of control and without threads.

newtype M o a = M ((a->o)->o)
nstance Monad (M o) where
  return x = M (\c -> c x)
  (M g)>>=f = M (\c -> g (\a -> let M h = f a in h c))

Here, o is the output type (the type of the values which are yielded) and a is the input type. The monad itself is a wrapper around a function of type ((a->o)->o) - a function which takes a continuation and returns the output type of that contination. The bind method is pretty scary, but I can't explain it any better here than the code already does - I'll just end up using more letters to say the same thing. (I have to work it through every time I read it anyway.)

Now we need a couple of helper functions, but first the example: we're going to build a lightswitch which takes three commands: set (with an Bool value), toggle and query:

data Result a = NewState (a->Result a) | Value Bool (a->Result a) | Final
data Input = Set Bool | Toggle | Query

So all our commands have to yield a value of type Result, querying will return a Value and the other two will return a NewState (which isn't really a result, it just gives the new continuation of the system). Final is there to be the value marking the end of the stream (it doesn't contain a next continuation).

yield x = M (\c -> Value x c)
wait = M (\c -> NewState c)

These functions are how we yeild values. Both are values in M which take a continuation and return a Result which passes that continuation back to the caller.

runM cm = \x -> let (M f) = cm x in f (\c -> Final)

This is a function which takes a first input, applies it to something which results in a continuation monad, unwraps that monad and gives it the final continuation, one which eats its given contination and gives Final

lightswitch state v = do
  case v of
    Set state -> wait >>= lightswitch state
    Toggle -> wait >>= lightswitch (not state)
    Query -> yield state >>= lightswitch state

This is our lightswitch, it takes an initial state and an Input and updates its state accordingly and returns some Result using either yield or wait. It recurses forever.

step cont = do
  line <- getLine
    case line of
	"toggle" -> case cont Toggle of NewState cont' -> step cont'
    	"query" -> case cont Query of Value x cont' -> putStrLn (show x) >> step cont'
	otherwise -> putStrLn "?" >> step cont

Here's the code that uses the state machine. It's in the IO monad and runs a little command line where you can type in commands and see the results:


It takes a continuation (which is the state of the state machine) and applies an Input to it to get another continuation (state of the machine). You can, of course, keep around any of these states and "undo" something by using an older state.

And to tie it all together:

main = step $ runM $ lightswitch False

We pass False as the initial state of the system and use runM to stick the final continuation on the end; then we have a continuation for the state machine and off we go.

Hopefully that made some kind of sense. To give credit where it's due: I stole the bind method (and motivation) from this paper.