State Machine

state-machine.swift

// MARK: - Read-Write Queue

class ReadWriteQueue {
    private let specificKey = DispatchSpecificKey<String>()
    private let queue: DispatchQueue

    private var isAlreadyInQueue: Bool {
        return DispatchQueue.getSpecific(key: specificKey) == queue.label
    }

    init(label: String = "read-write.queue") {
        queue = DispatchQueue(label: label, attributes: .concurrent)
        queue.setSpecific(key: specificKey, value: label)
    }

    deinit {
        queue.setSpecific(key: specificKey, value: nil)
    }

    // solving readers-writers problem: any amount of readers can access data at a time, but only one writer is allowed at a time
    // - reads are executed concurrently on executing queue, but are executed synchronously on a calling queue
    // - write blocks executing queue, but is executed asynchronously on a calling queue so it doesn't have to wait
    // note:
    //  it's fine to have async write, and sync reads, because write blocks queue and reads are executed synchronously;
    //  so if we want ro read after writing, we'll still be waiting (reads are sync) for write to finish and allow reads to execute;
    func write(_ work: @escaping () -> Void) {
        if isAlreadyInQueue {
            work()
        } else {
            queue.async(flags: .barrier, execute: work)
        }
    }

    // if we're already executing inside queue, then no need to add task there synchronosuly since it can lead to a deadlock
    func read<T>(_  work: () throws -> T) rethrows -> T {
        if isAlreadyInQueue {
            return try work()
        } else {
            return try queue.sync(execute: work)
        }
    }
}

// MARK: - State Machine

class StateMachine<State, Event> {
    typealias Schema = (State, Event) -> State?

    private let queue = ReadWriteQueue(label: "state-machine.queue")
    private let schema: Schema
    private var _state: State
    var state: State {
        return queue.read { _state }
    }

    // Handlers
    var shouldTransit: ((State, Event, State) -> Bool)?
    var didTransit: ((State, Event, State) -> Void)?

    // Initialize
    init(initialState: State, schema: @escaping Schema) {
        _state = initialState
        self.schema = schema
    }

    // Handle Event
    func handle(_ event: Event, callback: ((State?) -> Void)? = nil) {
        queue.write { [weak self] in
            guard let strongSelf = self else { return }
            guard let newState = strongSelf.schema(strongSelf._state, event) else {
                callback?(nil)
                return
            }
            if let shouldTransit = strongSelf.shouldTransit {
                if shouldTransit(strongSelf._state, event, newState) {
                    strongSelf.changeState(to: newState, withEvent: event)
                } else {
                    callback?(nil)
                    return
                }
            } else {
                strongSelf.changeState(to: newState, withEvent: event)
            }
            callback?(strongSelf._state)
        }
    }

    private func changeState(to newState: State, withEvent event: Event) {
        let fromState = _state
        _state = newState
        didTransit?(fromState, event, newState)
    }

    // Force reset to State
    func reset(to state: State) {
        queue.write { [weak self] in
            self?._state = state
        }
    }
}

Usage:

enum Number {
    case one
    case two
    case three
}

enum Operation {
    case increment
    case decrement
}

let sm = StateMachine<Number, Operation>(initialState: .one, schema: { (state, event) in
    switch (state, event) {
    case (.one, .increment): return .two

    case (.two, .increment): return .three
    case (.two, .decrement): return .one

    case (.three, .decrement): return .two

    default: return nil
    }
})

sm.handle(.increment)
sm.state // .two

sm.reset(to: .one)

sm.handle(.decrement)
sm.state // .one