Coding the state machines

2023-07-20

Intro

I coded my first state machine consciously while studying at the university. During the course of the operating systems, we were introduced to the topics of parsers implementation and were tasked to code a regular expression implementation using an approach of the deterministic finite automaton (DFA).

It was a really fun experience. Before that, in my mind, DFAs were mainly associated with the hardware. So it made me realize how approaches or methods may cross boundaries of disciplines they originate from. For the details on the DFA approach for the regular expression you refer to an excellent series of posts from Russ Cox that also explains the motivation.

When it comes to coding the state machines, there are several general approaches you might take. I will try to review three of them (switch statements, transition tables, and state functions) in this post. I'm illustrating these approaches using the code in Go, but they can be implemented in any language.

We'll consider an implementation of a turnstile state machine which is described in a Wikipedia article about the finite state machines. The image is taken from there.

Switch statements

To illustrate the approaches, we'll implement programs that simulate the turnstile by printing its current state to the terminal, then asking the user for an action. If the input from the user is either "push" or "coin", it will be transmitted to the state machine. Any other input will terminate our simulation.

To help with this task, let's introduce a function nextEvent() which will print the current state and get the user input.

type event string

// nextEvent returns event according to the user input.
// It can be "push", "coin", or an empty string to signal the termination.
func nextEvent(current state) event {
    fmt.Println("Current state:", current)

    s := bufio.NewScanner(os.Stdin)
    s.Split(bufio.ScanLines)
    if s.Scan() {
        switch t := strings.ToLower(s.Text()); t {
        case "push", "coin":
            return event(t)
        }
    }
    // Any other input means termination.
    return ""
}

The state here is either "locked" or "unlocked":

type state string

const (
    stateLocked   state = "locked"
    stateUnlocked state = "unlocked"
)

Our state machine code can get the following form using the switch statements:

func main() {
    s := stateLocked
    for e := nextEvent(s); e != ""; e = nextEvent(s) {

        switch s {
        case stateLocked:
            switch e {
            case "coin":
                s = stateUnlocked
            }

        case stateUnlocked:
            switch e {
            case "push":
                s = stateLocked
            }

        default:
            panic("unexpected state " + s)
        }
    }
    fmt.Println("Bye!")
}

You can play with a variant of this program on Go Playground. There is no stdin there, so inputs are provided via a hardcoded array.

In this program, we maintain the machine state in our program data (using the variable s) and represent the structure of the state machine transitions using the structure of our code. Nesting the switch and case statements follows the diagram of the turnstile state machine above.

I remember that Android Open Source Project had a similar implementation related to managing the Wi-Fi connection state a few years ago. I cannot find this code today (very likely it got rewritten), but see the next section for another example from the AOSP project.

Transitions table

I like the code above because its switch statements give me a very quick overview of the state machine I'm dealing with reviewing the code. However, for a bigger number of states/transitions such code structure may turn out to be a problem. Most likely, we would start splitting this code into multiple functions, each responsible for its own part of the states diagram.

Another approach would be moving the representation of transitions from the program code to the program data forming a table (a data structure) that describes our state machine.

We'll start from introducing the concept of a transition function: a function that checks the input (event) and decides what is the next state.

type transitionFunc func(current state, e event) state

func unlockOnCoin(current state, e event) state {
    if e == "coin" {
        return stateUnlocked
    }
    return current
}

func lockOnPush(current state, e event) state {
    if e == "push" {
        return stateLocked
    }
    return current
}

With these function, the main code may look like

func main() {
    transitions := map[state]transitionFunc{
        stateLocked:   unlockOnCoin,
        stateUnlocked: lockOnPush,
    }
    s := stateLocked
    for e := nextEvent(s); e != ""; e = nextEvent(s) {
        s = transitions[s](s, e)
    }
    fmt.Println("Bye!")
}

Check it on Go Playground.

Note how transition functions explicitly describe the arcs on the state diagram in the beginning of this article. Remaining in the same state is also explicit now as opposed to the variant with the switch statements where we just didn't change the value of the state variable.

This is the approach we used to implement the regular expression automaton during the OS course at the university.

In different code bases, engineers may choose slightly different variants of how this table is defined. For example, the table could directly have some input pattern matching to describe multiple transitions from the same state (as opposed to a simple variant above where we have one row per state).

More (advanced) examples.

Android Open Source Project code base has a class called com.android.internal.util.StateMachine which can be used to represent even more complicated hierarchical structures. This implementation leverages object oriented programming to represent the state machines: each state is an object, handling inputs/events is done by calling a method on this object. You can easily find a few places in AOSP where this approach is used.

Another example (now with Scheme language) is declaring a DFA using the automata/dfa package in Rocket.

State functions

Using the transition tables, we fully represented the state diagram in the program data using 2 data types:

type state string // Representing a node on the diagram.
type transitionFunc(state, event) state // Representing an arc on the diagram. 

Let's take a different turn an replace them with a single type: state function.

// stateFn represents a state in the state machine.
// Invoking this function leads to processing machine inputs relevant to the state.
// The returned value is the next state - a function that should be invoked next.
// Returning nil means reaching the terminal state (finishing the execution of the state machine).
type stateFn func() stateFn

For our turnstile state machine, the state functions could look like this.

func stateLocked() stateFn {
    switch e := nextEvent("locked"); e {
    case "coin":
        return stateUnlocked
    case "":
        return nil
    default:
        return stateLocked
    }
}

func stateUnlocked() stateFn {
    switch e := nextEvent("unlocked"); e {
    case "push":
        return stateLocked
    case "":
        return nil
    default:
        return stateUnlocked
    }
}

Each of these functions is responsible for getting the inputs that should be processed in the corresponding state, acting on them, and returning the next state function to call. Returning nil means we want to terminate the machine.

Our main function now gets the following form: a loop that begins from the initial state and invokes the current state function until nil is returned.

func main() {
    for s := stateLocked; s != nil; {
        s = s()
    }
    fmt.Println("Bye!")
}

Check it on Go Playground.

With this approach, we moved the state machine structure representation from the program data back to the code.

Rob Pike has a great presentation on this idea with a video on Youtube. At some point, it was leveraged to refactor the code of Go's templates parser. It fixed some problems with the code structure there, but also introduced another kind of issue :).

Experimenting a lot with this approach in production systems led me to working on a generalization in Go that can be used to implement state machines in such a way. I will try to cover the details in a follow-up post.