GKStateMachine in Swift

StateMachines.swift

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

import Cocoa

class State {
    weak var stateMachine: StateMachine?

    func isValidNextState(stateClass: AnyClass) -> Bool {
        return true
    }

    func didEnterWithPreviousState(previousState: State?) {}

    func updateWithDeltaTime(seconds: NSTimeInterval) {}

    func willExitWithNextState(nextState: State) {}
}

class StateMachine {
    // Can't use a map of [AnyClass: State] because AnyClass isn't hashable
    let states: [State]

    var currentState: State?

    init(states: [State]) {
        self.states = states
        for state in self.states {
            state.stateMachine = self
        }
    }

    func stateForClass(stateClass: AnyClass) -> State? {
        // We can't use an easy map of [AnyClass : State], so just loop over the states
        // array. We're not likely to have enough states that it'll be a problem anyway.
        for state in self.states {
            if let c: AnyClass = object_getClass(state) {
                // This feels weird: === explicitly compares memory locations, which maybe works?
                if c === stateClass {
                    return state
                }
            }
        }

        return nil
    }

    func canEnterState(stateClass: AnyClass) -> Bool {
        if let state = stateForClass(stateClass) {
            if self.currentState == nil {
                return true
            }

            return state.isValidNextState(stateClass)
        }

        return false
    }

    func enterState(stateClass: AnyClass) -> Bool {
        if let nextState = stateForClass(stateClass) {
            if let previousState = self.currentState {
                if !previousState.isValidNextState(stateClass) {
                    return false
                }

                previousState.willExitWithNextState(nextState)
            }

            self.currentState = nextState
            nextState.didEnterWithPreviousState(self.currentState)
            return true
        }
        return false
    }

    func updateWithDeltaTime(seconds: NSTimeInterval) {
        if let state = self.currentState {
            state.updateWithDeltaTime(seconds)
        }
    }
}

//// Tests
/********************************************
* The state graph used in these tests
*
* StateOne -----> StateTwo -----> State Three
*     ^            |  ^                |
*     |            |  |                |
*     -------------|  |----------------|
*
* StateThree will set properties when it receives any of the lifecycle messages
* so that we can make sure they were called.
*
* StateFour is a valid JLFGKStateClass, but isn't part of the state machine.
*
* And then a fake one that's just not a state: NotAState
*/

class StateOne : State {
    override func isValidNextState(stateClass: AnyClass) -> Bool {
        return stateClass === StateTwo.self
    }
}

class StateTwo : State {
    override func isValidNextState(stateClass: AnyClass) -> Bool {
        return stateClass === StateOne.self || stateClass === StateThree.self
    }
}

class StateThree : State {
    var didEnterStateCalled = false
    var willLeaveStateCalled = false
    var deltaTimeAtLastUpdate: NSTimeInterval = 0.0

    override func isValidNextState(stateClass: AnyClass) -> Bool {
        return stateClass === StateTwo.self
    }

    override func didEnterWithPreviousState(previousState: State?) {
        self.didEnterStateCalled = true
    }

    override func willExitWithNextState(nextState: State) {
        self.willLeaveStateCalled = true
    }

    override func updateWithDeltaTime(seconds: NSTimeInterval) {
        self.deltaTimeAtLastUpdate = seconds
    }
}

class StateFour : State {
    override func isValidNextState(stateClass: AnyClass) -> Bool {
        assertionFailure("Should never reach this")
        return true
    }
}

var stateMachine = StateMachine(states: [StateOne(), StateTwo(), StateThree()])

// Can we pull the proper states back out?
var stateOne = stateMachine.stateForClass(StateOne)
var stateTwo = stateMachine.stateForClass(StateTwo)
var stateThree = stateMachine.stateForClass(StateThree)
var stateFour = stateMachine.stateForClass(StateFour)

if let stateOne = stateOne {
    stateOne.isValidNextState(StateTwo)
}

// This should work, since the state machine doesn't have a current state yet.
stateMachine.enterState(StateOne)

// These should work: the transitions are allowed
stateMachine.enterState(StateTwo)
stateMachine.enterState(StateThree)

// Make sure the lifecycle methods are being called
stateMachine.updateWithDeltaTime(0.12)
stateMachine.enterState(StateTwo)

if let stateThree = stateThree {
    if let state = stateThree as? StateThree {
        state.didEnterStateCalled
        state.willLeaveStateCalled
        state.deltaTimeAtLastUpdate
    }
}

// That all looks good!