Type Safe State Machines in TypeScript

State machines are everywhere but I don't think they get enough air time. To remedy the situation I'm going to present an implementation in TypeScript and use it to model a very simple elevator.

I could point you to the definition of what a state machine is on Wikipedia but the name gives it away. State machines consist of states and a description of how to go from one state to the next. So let's start spelling things out.

The first thing we need is a description of the states. I'm going to use basic types like strings, numbers, and symbols because it greatly simplifies the rest of the implementation and we don't lose any generality

type Primitives = string | number | symbol;

Next we need a description of the transition function/map

type EventTypeMapping<I extends Primitives> = { [K in I]: unknown };

type TransitionMap<I extends Primitives, S, E extends EventTypeMapping<I>> = {
  [K in I]: (event: E[I], currentState: I, extraState: S) => I
};

That definition looks a bit elaborate but in plain English it says for every state we have a function that takes an event (E[I]), the current state ( I), some extra state ( S), and gives us back the next state. The events are indexed by the possible states the machine can be in because not every event is valid in every state. There is no way to actually enforce this constraint at compile time so our transition functions must deal with all event types which is what E[I] is about.

Now we can construct the machine. We'll use a class to model the machine with the above ingredients

export class Machine<I extends Primitives, S, E extends EventTypeMapping<I>> {

  constructor(
    readonly state: S,
    readonly initialState: I,
    readonly transitions: TransitionMap<I, S, E>,
    public currentState: I = initialState
  ) { }
}

The above machine doesn't do much because we are not leveraging the transition functions so let's remedy that.

// ...
  step(event: E[I]): [I, I] {
    const currentState = this.currentState;
    const newState = this.transitions[currentState](
      event, currentState, this.state);
    this.currentState = newState;
    return [currentState, newState];
  }
// ...

That's it. We just implemented an abstract state machine. Abstract state machines aren't that useful so let's implement a concrete one that models an elevator in a 3 story building.

The elevator will start on the 1st floor and then go all the way to the 3rd floor. After reaching the 3rd floor it will go all the way down to the 1st floor and then repeat this cycle. Like I said, it's very simple.

The first thing we need is the description of the states. The most relevant part of the state is the floor the elevator is on so that's what we'll use for the state description.

type Floor = 1 | 2 | 3;

This elevator doesn't take any inputs so the event mapping for each state is also very simple.

type AllowedEvents = { 1: void; 2: void; 3: void };

We need one more bit of information. The elevator needs to know which direction it's going.

type ExtraState = { direction: 'up' | 'down' };

The transitions functions are a runtime concern so to fill those in we're going to instantiate our abstract machine and fill them in.

const elevator = new Machine<Floor, ExtraState, AllowedEvents>(
  { direction: 'up' },
  1,
  {
    1: (e, f, s) => { return (s.direction = 'up', 2); }, // Can only go up
    2: (e, f, s) => { return s.direction === 'up' ? 3 : 1 }, // Can go up or down
    3: (e, f, s) => { return (s.direction = 'down', 2); } // Can only go down
  }
);

Let's also step through the transitions a few times to verify that it's working as we expect.

console.log(`*starting*`);
for (let i = 0; i < 12; i++) {
  const [previousFloor, nextFloor] = elevator.step(void(0));
  console.log(`Elevator going ${elevator.state.direction}: ${previousFloor} -> ${nextFloor}`);
  if (elevator.currentState === 1) {
    console.log(`---Cycle complete---`);
  }
}
console.log(`*done*`);

*starting*
Elevator going up: 1 -> 2
Elevator going up: 2 -> 3
Elevator going down: 3 -> 2
Elevator going down: 2 -> 1
---Cycle complete---
Elevator going up: 1 -> 2
Elevator going up: 2 -> 3
Elevator going down: 3 -> 2
Elevator going down: 2 -> 1
---Cycle complete---
Elevator going up: 1 -> 2
Elevator going up: 2 -> 3
Elevator going down: 3 -> 2
Elevator going down: 2 -> 1
---Cycle complete---
*done*

Looks correct to me. Making the elevator less simple is left as an exercise for the reader. See if you can model people requesting to be taken to a specific floor as a starting point.

All the code along with the above example lives at GitHub: davidk01/state-machine.