Simple State Machine in Swift

Our programs often use State Machines even though we might not even realise it.

How often do we have something called State and associated rules how states can change? Well there we most likely have a State Machine.

Jira Ticket example

Let’s take a Ticket in issue tracking tool Jira as an example of a State Machine. Each Ticket can have multiple states and states can change only following a set of rules according to Default workflow diagram.

Ticket can transition states following these rules:

• Open ticket cannot be Reopened.
• In Progress ticket cannot be Closed without first being Resolved.
• Resolved ticket cannot be put In Progress without first being Reopened.
• and so on…

Instead of writing all conditions here, it can be better expressed using the following diagram:

Image taken from Jira: Working with workflows

Show me the code

There are existing libraries for State Machine implementation in Swift, some of them very sophisticated and advanced. But often a simple solution is just enough…

StateMachine protocol

protocol StateMachine {
    associatedtype State: Hashable
    var state: State { get }
    /// States transitions definition.
    /// Must contain set of allowed next states for each state.
    static var stateTransitions: [State: Set<State>] { get }
    func canTransition(to nextState: State) -> Bool
    mutating func transition(to nextState: State)
}

A state machine is defined using StateMachine protocol. Let’s have a closer look…

associatedtype State: Hashable

This is a type representing a State of a state machine. This will most likely be an enum. It is Hashable so it can be stored in Set of allowed state transitions.

var state: State

Accessor to current state of a state machine. Every state machine must have a state at all times.

static var stateTransitions: [State: Set<State>]

This is the main magic behind our state machine: its state transition table. For each possible state, it defines what are the next states the state machine is allowed to transition to. It is basically code definition of the state diagram above.

func canTransition(to nextState: State) -> Bool

Determines whether the transition to next state candidate is allowed from current state.

mutating func transition(to nextState: State)

Performs the transition of a state machine to next state. It is marked as mutating because it mutates the state.

Extension with default implementation

extension StateMachine {
    static func isStateTransitionAllowed(from state: State, to nextState: State) -> Bool {
        guard let validNextStates = self.stateTransitions[state] else {
            fatalError("No transitions for \(state) defined in stateTransitions")
        }
        return validNextStates.contains(nextState)
    }
    func canTransition(to nextState: State) -> Bool {
        return type(of: self).isStateTransitionAllowed(from: self.state, to: nextState)
    }
}

The extension provides default implementation of canTransition(to:). By default does the expected: state transition is allowed when stateTransitions allows it.

Example implementation

struct JiraTicket: StateMachine {
    let identifier: String
    init(identifier: String) {
        self.identifier = identifier
    }
    // MARK: StateMachine
    enum State: CaseIterable {
        case open
        case reopened
        case inProgress
        case resolved
        case closed
    }
    private(set) var state: State = .open
    static let stateTransitions: [State: Set<State>] = [
        .open: [.inProgress, .resolved, .closed],
        .reopened: [.inProgress, .resolved, .closed],
        .inProgress: [.open, .resolved],
        .resolved: [.reopened, .closed],
        .closed: [.reopened]
    ]
    mutating func transition(to nextState: State) {
        precondition(self.canTransition(to: nextState), "Invalid state transition (\(self.state) -> \(nextState))!")
        self.state = nextState
    }
}

The example implementation is pretty straightforward.

• enum State represents all possible states of a ticket.
• stateTransitions is defined with all possible transitions.
• transition(to:) performs the state transition.
  • It changes value of a private variable.
  • There is a precondition to check if the transition is valid. This implementation considers invalid state transition being a programmer’s error. canTransition(to:) should always be consulted first before changing the state.

Usage

var ticket = JiraTicket(identifier: "XX-123")
ticket.transition(to: .inProgress)
// do something
ticket.transition(to: .resolved)

Ticket can transition states when mutable (defined as var).

let ticket = JiraTicket(identifier: "XX-123")
ticket.transition(to: .inProgress) // ERROR: Cannot use mutating member on immutable value

Because the transition(to:) is mutating, it cannot be called on immutable ticket defined using let.