State Monad in Coq Tutorial

In this post we will review usage of StateMonad in Coq using ExtLib on an adapted example from  a Haskell State Monad tutorial.

We will implement a simple game function which takes a string from dictionary {a,b,c} and produces a score number. The game could have two modes: on and off. It consumes an input string character by character and updates the mode and the score at each step. Initially the mode is off. The character 'c' toggles the on/off mode flag. In off mode, consuming characters does not affect the score. In on mode, each occurrence of the character 'a' increases the score by 1, and each occurrence of the character 'b' decreases it by 1. So for example, the input sequence "ab" will result in a score of 0, "ca" a score of 1, "cabca" also a 0, and "abcaaacbbcabbab" will result in a score of 2.

Let us look the implementation of this game in Coq using State Monad:

1 Require Import Coq.ZArith.ZArith_base Coq.Strings.String.
2 Require Import ExtLib.Data.Monads.StateMonad ExtLib.Structures.Monads.
3 
4 Section StateGame.
5   Import MonadNotation.
6   Local Open Scope Z_scope.
7   Local Open Scope char_scope.
8   Local Open Scope monad_scope.
9 
10   Definition GameValue : Type := Z.
11   Definition GameState : Type := (prod bool Z).
12 
13   Variable m : Type → Type.
14   Context {Monad_m: Monad m}.
15   Context {State_m: MonadState GameState m}.
16   
17   Fixpoint playGame (s: string): m GameValue :=
18     match s with
19     |  EmptyString =>
20        v <- get ;;
21          let '(on, score) := v in ret score
22     |  String x xs =>
23        v <- get ;;
24          let '(on, score) := v in
25          match x, on with
26          | "a", true =>  put (on, score + 1)
27          | "b", true => put (on, score - 1)
28          | "c", _ =>   put (negb on, score)
29          |  _, _  =>    put (on, score)
30          end ;; playGame xs
31     end.
32 
33   Definition startState: GameState := (false, 0).
34 End StateGame.
35 
36 Definition main : GameValue :=
37   (evalState  
38            (playGame (state GameState) "abcaaacbbcabbab") startState).

I will omit description of the Coq nuts and bolts of Require, Import, and Open Scope instructions in lines 1-8.

In lines 10 and 11, we declare type aliases for our game score value as a signed integer (Z), and game state as a product of the boolean game mode and and an integer game score.

Variable 'm' in line 13 identifies our Monad by it's type constructor. Line 14 states that 'm' must be an instance of type class Monad. If we look at ExtLib's definition of this type class we will see that it just postulates the presence of the standard Monad bind and return (called ret here) functions:

Class Monad (m : Type → Type) : Type := { 
    ret : ∀ {t}, t → m t ;
    bind : ∀ {t u}, m t → (t → m u) → m u
}.

The next class, mentioned in line 15 brings additional get and put definitions, specific for State Monad:

Class MonadState (T : Type) (m : Type → Type) : Type := {
    get : m T ;
    put : T → m unit
}.

The actual game function definition begins on line 17. It is a recursive function consuming an input string and returning a monadic value as a result. The string is an inductive type and we match its two constructors on lines 19 and 22, respectively. If the string is empty, we use syntax x <- c1 ;; c2 which is a Coq's rendering of well a known do-notation. Informally: we extract a score element from the current state tuple and return it as a monadic value using ret. While in Haskell, we can write a much terser expression (_, score) <- get, in Coq we have to use an intermediate value 'v' followed by destructive let, which unfortunately does not allow us to use a wildcard for the first element (mode), which we do not need anyway.

In the case of the second constructor on line 22 we again use do-notation and destructing let to break down the current state into on and score components and then analyze their combinations via match. Again, we use slightly different syntax from Haskell's version and, instead of matching on x with guards, we match on both x and on. Depending on their combination, we update the score and on/off flag according to our game rules using the state monad's put function. Then, we proceed with the next recursive iteration chained by c1 ;; c2 notation which is sometimes called then operator (>>) in Haskell, and used to to chain monadic computations where c2 does not depend on the outcome of c1. This could be expanded into our familiar do-notation of:  _ <- c1 ;; c2.

On line 33, we define our initial game state (game off, score=0) as startState.

On line 36, we define our 'main' function, which runs a game for a specific string and returns an integer score. It is a bit difficult to parse, so for clarity let us split it into several definitions, specifying return types and making all arguments explicit via '@':

Definition m : Type → Type := state GameState.
Definition pg :  state GameState GameValue := @playGame m _  _ "abcaaacbbcabbab".
Definition main : GameValue := @evalState GameState GameValue pg startState.

The first definition 'm' uses the state record type to create our State Monad's constructor. It is defined as:

state: Type → Type → Type
Record state (S t : Type) : Type := mkState { runState : S → t * S }

We partially specialize state with GameState, which gives us a new type constructor (Type → Type), which instantiates a type which holds a function with a type GameState -> t * GameState.

Th next definition 'pg' holds a monadic value returned as a result of running our game. You may notice that its type is state GameState GameValue which is the same as m GameValue, the type of  playGame. Please note that playGame was defined inside a section. Consequently, when seen outside one, it is additionally universally quantified with all section variables. In our case playGame now takes 4 arguments. The first is a type constructor for the monad, the second is the type class instance Monad m that carries the monad functions, the third is the type class instance MonadState GameState m that carries the monad state functions, and the fourth is the string input to the function. We explicitly provide 'm' but use wildcards for type class parameters, letting the type class resolution mechanism find instances for parameters marked as _ . Coq will look for a type class instance of Monad (state GameState) and MonadState (state GameState). Both of these are defined as Global Instances in the ExtLib.Data.Monads.StateMonad library so type class resolution finds the definitions. Thus, after class resolution, the 'pg' definition will expand to:

(@playGame (state GameState) (Monad_state GameState)
     (MonadState_state GameState) "abcaaacbbcabbab"

Finally in the 'main' definition, we use the evalState function to get our result. It is defined as:

Definition evalState {t} (c : state t) (s : S) : t :=
    fst (runState c s).

The semantics are pretty straightforward; we apply State Monad to the initial state (run state transition function), and project out the first element of the returned tuple (the score).

To test our implementation we can use Compute main instruction to get the expected result (2).

P.S. I would like to thank Gregory Malecha for the explanations, comments, and suggestions he provided.