Simple Swift FSM

StateMachine.swift

//: Playground - noun: a place where people can play

import UIKit

/// A finite state machine
final class FSM<State, Transition>: ObservableObject where State: Hashable, Transition: Hashable {
    /// current state
    @Published private(set) var state: State
    
    /// state transition graph
    private(set) var transitions = [State:[Transition:State]]()
    
    
    /// Adds a transition `transition` from state `first` to state `second`
    /// - Parameters:
    ///   - transition: a transition
    ///   - first: the starting state
    ///   - second: the ending state
    func addTransition(_ transition: Transition, from first: State, to second: State) {
        transitions[first, default: [:]][transition] = second
    }
    
    /// Returns `true` if `transition` will advance state
    /// - Parameter transition: a transition
    /// - Returns: `true` if the transition will advance state
    func canAdvance(_ transition: Transition) -> Bool {
        return transitions[state]?[transition] != nil
    }
    
    /// a transition observation function
    typealias Observer = (_ old: State, _ new: State) -> ()
    
    /// Advances state by performing `transition`.
    /// - Parameters:
    ///   - transition: a transition
    ///   - observe: an optional transition observer
    /// - Returns: the new state
    @discardableResult
    func transition(_ transition: Transition, observe: Observer? = nil) -> State {
        let prev = state
        if let next = transitions[prev]?[transition], next != prev {
            state = next
            observe?(prev, next)
        }
        
        return state
    }
    
    /// a transition builder function, used to add transitions to a newly created state machine
    typealias TransitionBuilder = (FSM<State,Transition>)->()
    
    /// Initializes the state machine with state `state` then optionally executes `builder` where
    /// transitions between states can be configured on the receiver. If `builder` is nil, then transitions
    /// can be added to the state machine after initialization
    /// - Parameters:
    ///   - state: a starting state
    ///   - builder: a transition builder
    init(state: State, transitions builder: TransitionBuilder? = nil) {
        self.state = state
        builder?(self)
    }
    
    fileprivate init(state: State, transitions: [State:[Transition:State]]) {
        self.state = state
        self.transitions = transitions
    }
}
extension FSM: Equatable {
    static func ==(lhs: FSM<State,Transition>, rhs: FSM<State,Transition>) -> Bool {
        return lhs === rhs || lhs.state == rhs.state && lhs.transitions == rhs.transitions
    }
}

extension FSM: Codable where State: Codable, Transition: Codable {
    enum CodingKeys: String, CodingKey {
        case state
        case transitions
    }
    
    func encode(to encoder: Encoder) throws {
        var container = encoder.container(keyedBy: CodingKeys.self)
        
        try container.encode(state, forKey: .state)
        try container.encode(transitions, forKey: .transitions)
    }

    convenience init(from decoder: Decoder) throws {
        let container = try decoder.container(keyedBy: CodingKeys.self)
        
        let state = try container.decode(State.self, forKey: .state)
        let transitions = try container.decode([State:[Transition:State]].self, forKey: .transitions)
        
        self.init(state: state, transitions: transitions)
    }
}

//
// Sample Usage
//

enum Switch {
    case on, off
}
enum Position {
    case up, down
}

var lights = FSM<Switch, Position>(state: .off) {
    fsm in

    fsm.addTransition(.up, from: .off, to: .on)
    fsm.addTransition(.down, from: .on, to: .off)
}

// you can observe state changes, since FSM is an `ObservableObject`
let obs = lights.$state.sink(receiveValue: {
    print($0)
})

// sending transitions to the FSM causes its state to change for defined state changes
lights.transition(.up)
lights.transition(.down)
lights.transition(.down)
lights.transition(.up)

Very nice code. This is short and describes well how a state machine works. I like that you have included transitions explicitly. I think it is odd that other code examples have omitted that. That seems like that natural place to register actions for.

How about splitting advance and the actions performed? I assume that in a non trivial system, different parts of the system will want to perform different actions depending on what transition was observed. So you should be able to register multiple observers for for a transition.

I don't mean this as criticism, but rather asking your opinion to see if I have understood the intended usage of this FSM correctly. I assume you have done your current design because it makes the working of it simpler and easier to demonstrate.

Yep - the intention was to demo the concept as an alternative take to this post: http://everythingel.se/blog/how-to-build-a-finite-state-machine-fsm-with-swift-pt-2/

Breaking out actions is a natural thing to do in this code if you intend to make a StateMachine a kind of shared, observed object. Alternatively the code can be used as-is if you opt to use the class as a kind of internal variable whose consumer broadcasts the observable behaviours. To separate concerns, you could declare an ObservableStateMachine class that encapsulates StateMachine and provides external observability separately.

This doesn't work if either the State or Transition enumerators have associated valuea, since it is not hashable. Is there a way to work around it?

Looks like enums with associated values either can now, or will shortly able to be made hashable automatically via the compiler.

How about splitting advance and the actions performed? I assume that in a non trivial system, different parts of the system will want to perform different actions depending on what transition was observed. So you should be able to register multiple observers for for a transition.

I don't mean this as criticism, but rather asking your opinion to see if I have understood the intended usage of this FSM correctly. I assume you have done your current design because it makes the working of it simpler and easier to demonstrate.

@ordovician

It's been a while, but Swift now makes this easy to do without bloating the code, so I've updated the gist to make StateMachine (now FSM) an ObservableObject with an @Published state.