How To Build a Finite State Machine (FSM) with Swift (Pt. 2)

In the previous post, the FSM we designed was based on actually two distinct protocols. One for the states and the other one for the delegate. I did that, because I thought it would be the most flexible way. However, fooling a bit more around with it, it turns out that it is not and actually is even more complicated 😜!

So, here we go with the shiny new improved implementation of a super fancy FSM. We still have two protocols, but everything is combined in a single delegate.

The protocol for the states looks like this:

protocol StateProtocol : Hashable {}  

The one for the delegate looks like this:

protocol FSMDelegate  
{
    typealias State

    func getNextState(state: State) -> State
    func switchState(state: State, toState: State) -> State
}

What have we done here? First, the protocol for the states does not pre-define anything but that it must be hashable (which is good for comparison). Of course, the real states have to inherit from this. Second, FSMDelegate does now define typealias State, which is a kind of template parameter (Swift slang: generic, see Apple1 and SO2). Furthermore, we dropped equalsState, because StateProtocol is now Hashable.

No comes some magic, because the definition of the StateMachine class is not as straightforward as one could think:

class StateMachine<  
    StateT: StateProtocol,
    DelegateT: FSMDelegate where DelegateT.State == StateT>
{
    var state    : StateT
    var delegate : DelegateT?

    init(initialState: StateT)
    {
        self.state = initialState
    }

    func updateFSM()
    {
        if let delegate = delegate
        {
            let nextState = delegate.getNextState(state)
            if nextState != state
            {
                state = delegate.switchState(state, toState: nextState)
            }
        }
    }   
}

Starting from bottom, updateFSM is only slightly changed, but now captures the delegate first, if it is not nil. The most interesting part here is the definition of the class, because this must be a generic and also define the typealias in StateProtocol. Therefore, we first declare a generic type StateT which must comply with the StateProtocol type. Then, in order to define State in FSMDelegate, we use the somewhat cryptic clause DelegateT: FSMDelegate where DelegateT.State == StateT. In words: "define a type DelegateT which is of type FSMDelegate with the type State in DelegateT - i.e. Statein FSMDelegate - being StateT, which in turn is a StateProtocol. Then, the types DelegateT and StateT become usable for the implementation.

Last point to clarify: how to use this new version? I'll show this again with the example of applying states to some figure. E.g. the enum becomes:

enum FigureStates : StateProtocol  
{
    case Sitting, Walking, Standing

    func getNextState() -> FigureStates
    {
        switch self
        {
            case Sitting  : return Standing
            case Walking  : return Standing
            case Standing : return Sitting
        }
    }
}

Note, that we could omit the equalsState function defined in the previous version. The class which instantiates the delegate becomes this:

class Figure : FSMDelegate  
{
    private let fsm : StateMachine<FigureStates, FigureAI>

    init(initialState: FigureStates)
    {
        self.fsm = StateMachine<FigureStates, Figure>(initialState: initialState)
        fsm.delegate = self
    }

    func getNextState(state: FigureStates) -> FigureStates
    {
        return state.getNextState()
    }

    func switchState(state: FigureStates, toState: FigureStates) -> FigureStates
    {
        switch (state, toState)
        {
            case (.Sitting, .Standing):
                makeFigureStandup();
                return .Standing

            //... more cases

            default: return state
        }
    }
}

Here we can see, how the StateMachine class is instantiated, i.e. by supplying the actual types for the generics in this class. Another positive side-effect is, that we can omit all the (nasty) casts we had to use before.

IMO this version is now much cleaner than the previous one. Enjoy!