State Machines in Swift

Sun, 3 Jan 2021 14:20:00 +0100

Implementing a trivial state machine in Swift.

I recently needed to implement a state machine in a project and was pleasantly surprised to find how easy this is in Swift. Declaring state transitions is practically built into the syntax of the language!

Plenty of articles have been written about state machines in Swift, but I was a little disappointed that they all seemed to over-complicate matters with protocols, transition classes, sets of valid states, etc — even the ones touted as the “simplest” or “easiest” implementations.

While it’s great to have these open source frameworks, in my opinion they miss the point that state machines in Swift are so simple, with so little boilerplate code, that importing a dependency in most cases is just unnecessary. In this article, I’ll describe a truly minimal implementation of a state machine in Swift.

I won’t be putting any sample code on Github as usual, because the whole point is that state machines are easier to implement from scratch each time 😄

A really simple state machine

At their heart, state machines have just two things:

1. A set of states

2. A set of transitions between those states

The first one is almost always implemented as an enum, because that’s the best way to represent a set of mutually exclusive options.

Let’s model the states of a simple web browser (or anything that loads remote content):

enum BrowserState {
    case empty      // initial state, nothing loaded yet
    case loading    // request sent, waiting for content
    case loaded     // successful, displaying content
    case error      // request failed
    case cancelled  // user cancelled the request
}

Now how should we define the valid transitions between states? In other languages like Objective-C, we might define a mapping between each state and those it can transition to, probably using NSDictionary and NSSet. Then to test if a transition was valid, we’d have to look up the current state in the map, and check if its set contained the target state.

In Swift however, this becomes much simpler with enum case pattern matching. We can create a tuple of the current and next state, then match it against all the valid transitions:

extension BrowserState {
    func canTransition(to state: BrowserState) -> Bool {
        switch (self, state) {
            case (.empty, .loading),        // start load
                 (.loading, .loaded),       // success
                 (.loading, .error),        // failure
                 (.loading, .cancelled),    // user cancel
                 (.error, .loading),        // retry
                 (.cancelled, .loading):    // retry
                return true
            default:
                return false
        }
    }
}

That’s it! 🤩

There is almost no boilerplate code or common logic, which to my mind, is why capturing this in a generic state machine implementation doesn’t add much value.

The handling of transition events and invalid states is also specific to each use case. Here’s one example of how we might use our BrowserState in a browser class:

class Browser {
    var state: BrowserState = .empty {
        willSet {
            precondition(state.canTransition(to: newValue),
                "Invalid transition: \(state) -> \(newValue)")
        }
    }
}

In just a few lines of code, we’ve set the initial state of the browser and made it impossible to transition to an invalid state.

Of course, you might want to fail a little more gracefully than by crashing the app. We can use transition functions to define custom behaviour while changing states:

extension Browser {
    func load() {
        if state.canTransition(to: .loading) {
            state = .loading
            // send request, then set state
            // to either .loaded or .error
        }
    }

    func cancel() {
        if state.canTransition(to: .cancelled) {
            state = .cancelled
        }
    }
}

These can be tailored to your Browser class, depending on whether you want to handle invalid transitions with assertions, exceptions, returning a BOOL, printing a message, etc.

Attaching values to state

A very handy feature of Swift is being to attach values to enum cases, which means that our state machine can also capture more information if necessary.

In reality, our Browser class would need to keep track of all kinds of extra state: the URL being loaded, the content to display, or any network error that occurred. It would make a lot of sense to attach these details to the BrowserState itself, since each value only applies in certain states:

enum BrowserState {
    case empty
    case loading(URL)
    case loaded(content: String)
    case error(Error)
    case cancelled
}

The best part is that our canTransition function doesn’t need to change, because pattern matching can ignore attached values.

If you wanted to have code executed before entering or exiting certain states, you could even attach closures to those case values.

Simpler transition declarations

Another useful feature of Swift’s pattern matching is that you can match against wildcard cases. To demonstrate this, let’s use a more fun example of a state machine — the behaviour of an enemy in a game:

enum EnemyState {
    case idle       // 😐
    case patrolling // 🚶‍♂️
    case suspicious // 👀
    case hunting    // 😠
    case attacking  // 🔫
    case dead       // 😵
}

With more states, there are a lot more possible transitions (30 in this case), but we shouldn’t have to list them all. If we take the idle case, there are a lot of potential transitions:

Starting a patrol ➞ patrolling

Hearing something ➞ suspicious

Finding a body ➞ hunting

Seeing the player ➞ attacking

Stealth killed by the player ➞ dead

Instead of listing all of these, we can simply match against case (.idle, _), where _ matches any case.

It turns out there are a few states that can transition to any other state, and any state can transition straight to dead:

extension EnemyState {
    func canTransition(to state: EnemyState) -> Bool {
        switch (self, state) {
            case (.idle, _),        // can jump to any state
                 (.patrolling, _),  // can jump to any state
                 (.suspicious, _),  // can jump to any state
                 (.hunting, .suspicious),   // didn't find anything
                 (.hunting, .attacking),    // found the player
                 (.attacking, .hunting),    // lost the player
                 (_, .dead):    // can be killed during any action
                return true
            default:
                return false
        }
    }
}

Usually an existing state machine implementation would require you to list out each possible transition explicitly, which could be an unnecessary inconvenience.

My advice would be to roll your own state machine from scratch — you get the full power of Swift’s syntax, can customise it to your exact use case, use fewer dependencies, and better understand what you code is doing. What’s not to like? 😉