Finite State Machines in Swift

Finite state machines are a powerful abstraction that can help you wrangle difficult flows of logic into a single, testable unit.

As you learn about finite state machines – or FSMs – you’ll see them pop up everywhere in real life. Vending machines, turnstiles and elevators can all be modeled by using this technique.

Let’s get familiar with state machines by looking at a common example: the traffic light.

Mapping out a finite state machine

We have a simple traffic that cycles through 3 states: green, yellow, and red. Since the sequence of these transitions is critical to the safety of motorists, we want to ensure that the order is always correct. To do so, let’s create a FSM to model this behavior.

State machines are represented by 2 definining characteristics:

  1. Their states
  2. Their transitions

To model our traffic light, let’s start by sketching out its states.

Possible States

We’ve listed out every possible state, represented by circles.

Let’s add some arrows to show the transitions. In our case, green can transition to yellow, yellow can transition to red, and red can transition to green.

States with transitions

This representation already tells us a great deal about our traffic light. Notice how it’s impossible to move from red to yellow, or from yellow to green. This is exactly what we want.

At this point, it’s also helpful to indicate what causes these transitions to occur. Let’s assume our traffic light uses timers to indicate when it should change states. Each light is associated with a timer, and when the timer expires, the state changes.

With this in mind, let’s label the state transitions.

States with named transitions

Now we have a state diagram we can work with! The next step is to translate this into Swift code.

Creating finite state machines in Swift

States in a state machine are mutually exclusive. This means that our state machine can only be in a single state at any point in time.

And whenever we talk about mutual exclusion in Swift, a good place to start is with an enum. We can represent the above state machine as an enum with 3 cases:

enum TrafficLightState {
	case green
	case yellow
	case red
}

Now let’s create a class that can hold the current state of our traffic light. We want to encapsulate both its current state and its transitions. This will ensure that the state can only change according to the diagram. We’ll also use this opportunity to define an initial state for our state machine.

struct TrafficLightStateMachine {
	enum TrafficLightState {
		case green
		case yellow
		case red
	}
	
	private var state: TrafficLightState = .red // Initial state
}

Notice that state is private. This is important. We absolutely do not want an enterprising developer to mess with our state. Since we can’t trust other developers (or ourselves) to be good stewards of the state machine’s internal state, it’s critical to make sure it remains private.

Now we can define the different state transitions as well, using the same names defined earlier in our state diagram:

class TrafficLightStateMachine {
	func greenTimerExpired() {
		switch state {
			case .green: 
				state = .yellow
			case .yellow, .red: 
				fatalError("Invalid state transition")
		}
	}
}

Let’s take a step back and explain what’s going. When the green timer expires on our traffic light, it will call greenTimerExpired() on our state machine. If the current state of our state machine is green (which it should be), we set the state to .yellow. However, if the current state isn’t .green, then we will have caught an error! We would have an invalid state transition, at which point we would we halt.

Is crashing really what we want to do here? We’ll have a look at that in a bit.

You can go ahead and fill in the state transitions for yellowTimerExpired() and redTimerExpired. Or you can download the complete implementation at the bottom of the article.

Communicating with the outside world

We have a state machine, but since its state is private, how will we know which state we’re in? Let’s create a change handler in order for our application to hook in, and make sure it’s called when our state changes.

class TrafficLightStateMachine {
	let stateDidChange: (TrafficLightState) -> Void
	
	private var state: TrafficLightState {
		didSet {
			stateDidChange(state)
		}
	}
}

And finally, let’s add a start() to kick things off. This will also communicate our initial state to the consumer of our state machine.

class TrafficLightStateMachine {
	//...
	
	func start() {
		stateDidChange(state)
	}
}

Handling invalid state transitions

If our state machine is handling a complex system like a traffic light, it’s probably a bad idea to crash on an invalid state transition. Imagine a one-in-a-million bug taking down our traffic light and leaving motorists with nothing to rely on.

This is why, in some cases, it’s helpful to define an error state. Lucky for us, this already exists in the world of traffic lights: the flashing red light.

Let’s update our state diagram to reflect this change.

States with error

In order to update our state machine, let’s add a case to our enum.

enum TrafficLightState {
	case green
	case yellow
	case red
	case flashingRed // Error state
}

When we do this, something wonderful happens. The compiler starts flagging errors in our state transitions, reminding us that we need to account for this new state.

This leads us to another key point: avoid using default switch statement handling.

If we had used default, we wouldn’t have gotten all these useful error messages reminding us to update our state transitions. This is part of the beauty of Swift; we can rely on the compiler to keep us honest and alert about updating our code. Wonderful!

Now let’s fix these errors by updating our state transitions from earlier:

class TrafficLightStateMachine {
	//...
	
	func greenTimerExpired() {
		switch state {
			case .green: 
				state = .yellow
			case .yellow, .red, .flashingRed:  // Updated
				state = .flashingRed           // Updated
		}
	}
	
	//...
}

What about the timers?

One natural question to ask ourselves is what do we do with the timers that govern the current state of our state machine.

The simple answer is that, well, it doesn’t matter! Our state machine class isn’t concerned with how the states are managed in code, it only cares about when the transition happens. Instead of thinking of our state machine class as a manager, it’s more like a bookkeeper. It makes sure that the state changes in ways we expect. In other words, it serves to enforce the rules that we’ve created for ourselves.

This means that we’re free, as consumers of this state machine, to implement the actual traffic light however we want. It can be a physical system with microcontrollers, or a playground with text output. As long as we’re informing the state machine of our transitions, it’s doing its job.

Wrapping up

State machines are a great way to both encapsulate and enforce state transitions. As we’ve seen here, they can be used to drive domain logic, but the uses of state machines extend far beyond that. We’ll explore those uses in upcoming articles.

Some key takeaways:

  1. When building FSMs, start by defining your states and transitions
  2. Make your FSM’s state private (or readonly) to enforce proper use
  3. Avoid using default in your switch handling.
  4. Consider using a dedicated error state instead of crashing on invalid transitions

Interested in seeing this state machine in action in a Swift Playground? If so, join our newsletter and we’ll send it straight to your email!

Meta
Originally published:
Saturday, Jan 9, 2021
Revised on:
  • resume_back_end_developer.pdf
  • coverletter_back_end_developer.pdf
Respond:
Any questions? Let us know on Twitter!
About
School of Swift is a training organization to help iOS teams with continuous education.