State Machines in Swift using enums

Finite State Machines are one of those things you don’t use very often, but occasionally they can rescue a codebase from turning into a huge mess. But they don’t have to be difficult and complicated. In this post I will show how to leverage Swift’s powerful enums to build a simple and type-safe state machine, without needing to bring in an external library as a dependency to your app.

Recently I was working on an iOS app that could be controlled from multiple sources: from the UI (of course), remotely from the server, and via bluetooth. We needed a way to ensure that events from these sources couldn’t take the app into a strange state. Eg we can’t have the app presenting a login page if you’re already logged in. A FSM will also make your app more unit-test-friendly which is no bad thing.

There are two parts to a State Machine: State and Events. For this example, pretend we’re creating a simple game. State can be modelled neatly as a Swift enum:

enum State {
	case introduction // On the first page of the game
	case help // In the help screen
	case choosingEpisode // Choosing an episode
	case choosingLevel(Episode) // After choosing an episode, player is now choosing
		// a level
	case playingLevel(Episode, Level) // Now playing the chosen level
	case finishedLevel(Episode, Level, Int) // Finished a level, displaying their
		// score.
}

You’ll notice that the last few cases have associated values:

• choosingLevel(Episode) is used to store the previously-selected episode when on the level selection page.
• playingLevel(Episode, Level) is used to store the details of the current level whilst playing it.
• finishedLevel(Episode, Level, Int) stores the level that was just finished, and the score the player achieved.

If you’re unfamiliar with enums with associated values, you can read more about that here: https://docs.swift.org/swift-book/LanguageGuide/Enumerations.html#ID148

Events are similarly modelled using an enum:

enum Event {
	case displayHelp // Takes the user from the intro to the help screen.
	case dismissHelp // Takes the user back to the intro from help.
	case startPlaying // Takes the user from the intro to the episode chooser.
	case chooseEpisode(Episode) // The user selected an episode, so take
		// the user from the episode chooser to the level chooser.
	case chooseLevel(Level) // The user selected a level, so take the user
		// from the level chooser to the gameplay.
	case completeLevel(Int) // Takes the user from gameplay to the completed
		// level screen, displaying their achieved score.
	case dismissFinishedLevel // Takes the user from the completed level
		// screen back to the level chooser to play again.
}

Your current state needs to be stored somewhere. A singleton is a simple way to do this:

class FiniteStateMachine {
	static let shared = FiniteStateMachine()
	
	var state = State.introduction // Initial app-startup state.
	...

Next we need a mechanism for selecting appropriate transitions given the ‘before’ state. For example, a ‘return home’ event might perform a different state transition if you are in a game (eg many view controllers to pop) vs only in the help screen (eg only one modal VC to dismiss). I call this mechanism the transition selector.

An important feature of the transition selector is that it throws if the incoming event isn’t valid for the given state. This is key to the state machine’s ability to prevent you getting into an invalid state. Here is the transition selector:

	...
	/// A transition is a closure that, once run, leaves you in a new state.
	typealias Transition = () throws -> (State)

	/// This selects the appropriate transition to handle the event.
	/// The choice of transition may depend upon the old state.
	func transition(forEvent event: Event) throws -> Transition {
		switch (state, event) {
			case (.introduction, .displayHelp): return presentHelp()
			case (.help, .dismissHelp): return dismissHelp()
			case (.introduction, .startPlaying): return presentEpisodeChooser()

			// Here we pull a value out of the event using `let episode`:
			case (.choosingEpisode, .chooseEpisode(let episode)):
				return presentLevelChooser(episode: episode)

			// Here is how we can pull the episode out of the 'before' state.
			case (.choosingLevel(let episode), .chooseLevel(let level)):
				return presentGamePlay(level: level)

			... further transitions as per your app ...

			default: throw Errors.transitionNotFound
		}
	}

	enum MyErrors: Error {
		case transitionNotFound
	}
	...

You may notice that transition(forEvent:) returns a closure (of type Transition), but to return said closure it calls functions, rather than returning those functions themselves. This is because those functions are curried. Now currying is one of those CompSci terms that can be super-complicated, but don’t fear, the gist of it is that these functions return closures with the parameters baked in. It turns out to simply be some boilerplate that you don’t really need to worry too much about. Bear with me please!

Here’s what those transition functions look like:

	static func presentLevelChooser(episode: Episode) -> Transition {
		return {
			// Perform the transition.
			let levelChooser = ...create level chooser view controller...
			levelChooser.levels = episode.levels
			NavigationManager.shared.rootNavController
				.pushViewController(levelChooser, animated: true)

			// Return the new state
			return State.choosingLevel(episode)
		}
	}

Finally you need a mechanism for handling events. This is the function that the rest of your app will call:

	func handle(event: Event) throws {
		let transition = try transition(forEvent: event)
		state = try transition()
	}

And so, here’s how to use this state machine. In your various event-sources in your app (UI, network, perhaps bluetooth) you’d call the state machine like so:

	try? FiniteStateMachine.shared.handle(event: .startPlaying)
	try? FiniteStateMachine.shared.handle(event: .chooseEpisode(myEpisode))
	try? FiniteStateMachine.shared.handle(event: .chooseLevel(myLevel))
	try? FiniteStateMachine.shared.handle(event: .completeLevel(100))
	try? FiniteStateMachine.shared.handle(event: .dismissFinishedLevel)		

As always, feel free to use this framework as a starting point and modify until you get something that best suits your use-case.

PS: A handy trick I like - you can easily store a Countdown Timer in one of your state enum associated values, to create a timeout for that state that is automatically invalidated when leaving that state.

Thanks for reading, and have a great week!

Photo by Malcolm Lightbody on Unsplash

Chris Hulbert

(Comp Sci, Hons - UTS)

iOS Developer (Freelancer / Contractor) in Australia.

I have worked at places such as Google, Cochlear, Assembly Payments, News Corp, Fox Sports, NineMSN, FetchTV, Coles, Woolworths, Trust Bank, and Westpac, among others. If you're looking for help developing an iOS app, drop me a line!

Get in touch:
[email protected]
github.com/chrishulbert
linkedin

 Subscribe via RSS