FSM

FSM.swift

//
//  StateMachine.swift
//
//
//  Created by Christian Treffs on 24.06.20.
//
public typealias FSM<State, StateInstance> = StateMachine<State, StateInstance> where State: Equatable

public final class StateMachine<State, StateInstance> where State: Equatable {
    /// Current state of this finite-state machine.
    public private(set) lazy var current: State = { fatalError("no current state set!") }()
    private lazy var lazyCurrentInstance: StateInstance? = { fatalError("no current instance set!") }()

    private var actionsOnEnter: [TransitionAction<State>]
    private var actionsOnExit: [TransitionAction<State>]

    private var stateInstances: [StateId: StateInstanceBox<State, StateInstance>]

    private init() {
        actionsOnEnter = []
        actionsOnExit = []
        stateInstances = [:]
    }

    private var subcriptionsId: UInt = 0

    private func nextSubscriptionId() -> UInt {
        defer { subcriptionsId += 1 }
        return subcriptionsId
    }

    @discardableResult
    public func onEnterState(_ state: State, execute closure: @escaping () -> Void) -> TransitionActionSubscription {
        let subscription = TransitionActionSubscription(id: nextSubscriptionId(), isOnEnter: true)
        let action = TransitionAction(id: subscription.id, source: nil, destination: state, closure: closure)
        actionsOnEnter.append(action)
        return subscription
    }

    @discardableResult
    public func onEnterState(_ state: State, from sourceState: State, execute closure: @escaping () -> Void) -> TransitionActionSubscription {
        let subscription = TransitionActionSubscription(id: nextSubscriptionId(), isOnEnter: true)
        let action = TransitionAction(id: subscription.id, source: sourceState, destination: state, closure: closure)
        actionsOnEnter.append(action)
        return subscription
    }

    @discardableResult
    public func onExitState(_ state: State, execute closure: @escaping () -> Void) -> TransitionActionSubscription {
        let subscription = TransitionActionSubscription(id: nextSubscriptionId(), isOnEnter: false)
        let action = TransitionAction(id: subscription.id, source: state, destination: nil, closure: closure)
        actionsOnExit.append(action)
        return subscription
    }

    @discardableResult
    public func onExitState(_ state: State, to destinationState: State, execute closure: @escaping () -> Void) -> TransitionActionSubscription {
        let subscription = TransitionActionSubscription(id: nextSubscriptionId(), isOnEnter: false)
        let action = TransitionAction(id: subscription.id, source: state, destination: destinationState, closure: closure)
        actionsOnExit.append(action)
        return subscription
    }

    @discardableResult
    public func unsubscribe(_ subscription: TransitionActionSubscription) -> Bool {
        switch subscription.isOnEnter {
        case true:
            if let index = actionsOnEnter.firstIndex(where: { $0.id == subscription.id }) {
                actionsOnEnter.swapAt(index, actionsOnEnter.endIndex - 1)
                actionsOnEnter.removeLast()
                return true
            }

        case false:
            if let index = actionsOnExit.firstIndex(where: { $0.id == subscription.id }) {
                actionsOnExit.swapAt(index, actionsOnExit.endIndex - 1)
                actionsOnExit.removeLast()
                return true
            }
        }
        return false
    }

    private func shouldPerformTransition(_ transition: Transition<State>) -> Bool {
        transition.isTransient
    }
}

extension StateMachine: Equatable where State: Equatable {
    public static func == (lhs: StateMachine<State, StateInstance>, rhs: StateMachine<State, StateInstance>) -> Bool {
        lhs.current == rhs.current
    }
}

extension StateMachine where StateInstance == Void {
    public convenience init(initial initialState: State) {
        self.init()
        current = initialState
    }
}

extension StateMachine where State: Hashable {
    public var currentInstance: StateInstance? {
        self.lazyCurrentInstance
    }

    public convenience init(initial initialState: State, initialInstance: @autoclosure @escaping () -> StateInstance) {
        self.init()
        addInstance(for: initialState, instance: initialInstance())
        performTransition(Transition(from: nil, to: initialState))
    }

    @discardableResult
    public func addInstance(for state: State, instance: @autoclosure @escaping () -> StateInstance) -> Bool {
        let stateId = StateId.make(from: state)
        guard stateInstances[stateId] == nil else {
            return false
        }

        let box = StateInstanceBox<State, StateInstance> {  [weak self] in
            let newInstance = instance()
            if let stateObserving = newInstance as? StateChangeObserving {
                self?.onEnterState(state, execute: stateObserving.onEnter)
                self?.onExitState(state, execute: stateObserving.onExit)
            }
            return newInstance
        }
        stateInstances[stateId] = box
        return true
    }

    @discardableResult
    public func transition(to newState: State) -> Bool {
        let transition = Transition(from: current, to: newState)
        guard shouldPerformTransition(transition) else {
            return false
        }
        performTransition(transition)
        return true
    }

    private func performTransition(_ transition: Transition<State>) {
        let newState = transition.destination
        let newInstance = getInstance(for: newState)

        performExit(transition)
        current = newState
        lazyCurrentInstance = newInstance
        performEnter(transition)
    }

    private func getInstance(for state: State) -> StateInstance? {
        stateInstances[StateId.make(from: state)]?.instance
    }

    /// an entry action: performed when entering the state, and
    private func performEnter(_ transition: Transition<State>) {
        actionsOnEnter
            .forEach { if $0.isApplicable(on: transition) { $0.closure() } }
    }

    /// an exit action: performed when exiting the state.
    private func performExit(_ transition: Transition<State>) {
        actionsOnExit
            .forEach { if $0.isApplicable(on: transition) { $0.closure() } }
    }
}

public struct Transition<State> where State: Equatable {
    public let source: State?
    public let destination: State

    public init(from source: State?, to destination: State) {
        self.source = source
        self.destination = destination
    }

    public var isTransient: Bool {
        source != destination
    }
}

extension Transition: Equatable where State: Equatable { }

public struct TransitionActionSubscription {
    let id: UInt
    let isOnEnter: Bool
}

private struct TransitionAction<State: Equatable> {
    let id: UInt
    let source: State?
    let destination: State?
    let closure: () -> Void

    func isApplicable(on transition: Transition<State>) -> Bool {
        switch (source, destination) {
        case let (.some(source), .some(destination)):
            return transition.source == source &&
                transition.destination == destination

        case let (.some(source), .none) :
            return transition.source == source

        case let (.none, .some(destination)):
            return transition.destination == destination

        case (.none, .none):
            return true
        }
    }
}

private enum StateId: Equatable, Hashable {
    case hashValue(Int)

    static func make<State>(from state: State) -> StateId where State: Hashable {
        .hashValue(state.hashValue)
    }
}

private class StateInstanceBox<State, StateInstance> {
    let constructor: () -> StateInstance
    lazy var instance: StateInstance = constructor()

    init(_ stateConstructor: @escaping () -> StateInstance) {
        constructor = stateConstructor
    }
}

public protocol StateChangeObserving {
    func onEnter()
    func onExit()
}