SubtreeManager.swift

/*
 * Copyright 2020 Square Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import Dispatch

extension WorkflowNode {
    /// Manages the subtree of a workflow. Specifically, this type encapsulates the logic required to update and manage
    /// the lifecycle of nested workflows across multiple render passes.
    final class SubtreeManager {
        var onUpdate: ((Output) -> Void)?

        /// Sinks from the outside world (i.e. UI)
        private(set) var eventPipes: [EventPipe] = []

        /// Reusable sinks from the previous render pass
        private var previousSinks: [ObjectIdentifier: AnyReusableSink] = [:]

        /// The current array of children
        private(set) var childWorkflows: [ChildKey: AnyChildWorkflow] = [:]

        /// The current array of side-effects
        private(set) var sideEffectLifetimes: [AnyHashable: SideEffectLifetime] = [:]

        private let session: WorkflowSession

        /// Reference to the context object for the entity hosting the corresponding node.
        private let hostContext: HostContext

        private var observer: WorkflowObserver? {
            hostContext.observer
        }

        init(
            session: WorkflowSession,
            hostContext: HostContext
        ) {
            self.session = session
            self.hostContext = hostContext
        }

        /// Performs an update pass using the given closure.
        func render<Rendering>(
            _ actions: (RenderContext<WorkflowType>) -> Rendering
        ) -> Rendering {
            /// Invalidate the previous action handlers.
            for eventPipe in eventPipes {
                eventPipe.invalidate()
            }

            /// Create a workflow context containing the existing children
            let context = Context(
                previousSinks: previousSinks,
                originalChildWorkflows: childWorkflows,
                originalSideEffectLifetimes: sideEffectLifetimes,
                hostContext: hostContext,
                session: session
            )

            let wrapped = RenderContext.make(implementation: context)

            /// Pass the context into the closure to allow a render to take place
            let rendering = actions(wrapped)

            wrapped.invalidate()

            /// After the render is complete, assign children using *only the children that were used during the render
            /// pass.* This means that any pre-existing children that were *not* used during the render pass are removed
            /// as a result of this call to `render`.
            childWorkflows = context.usedChildWorkflows
            sideEffectLifetimes = context.usedSideEffectLifetimes

            /// Captured the reusable sinks from this render pass.
            previousSinks = context.sinkStore.usedSinks

            /// Capture all the pipes to be enabled after render completes.
            eventPipes = context.eventPipes
            eventPipes.append(contentsOf: context.sinkStore.eventPipes)

            /// Set all event pipes to `pending`.
            eventPipes.forEach { $0.setPending() }

            /// Return the rendered result
            return rendering
        }

        /// Enable the eventPipes for the previous rendering. The eventPipes are not valid until this has
        /// be called. It is an error to call this twice without generating a new rendering.
        func enableEvents() {
            /// Enable all action pipes.
            for eventPipe in eventPipes {
                eventPipe.enable { [weak self] output in
                    self?.handle(output: output)
                }
            }

            /// Enable all child workflows.
            for child in childWorkflows {
                child.value.enableEvents()
            }
        }

        func makeDebugSnapshot() -> [WorkflowHierarchyDebugSnapshot.Child] {
            childWorkflows
                .sorted(by: { lhs, rhs -> Bool in
                    lhs.key.key < rhs.key.key
                })
                .map {
                    WorkflowHierarchyDebugSnapshot.Child(
                        key: $0.key.key,
                        snapshot: $0.value.makeDebugSnapshot()
                    )
                }
        }

        private func handle(output: Output) {
            onUpdate?(output)
        }
    }
}

extension WorkflowNode.SubtreeManager {
    enum Output {
        case update(any WorkflowAction<WorkflowType>, source: WorkflowUpdateDebugInfo.Source)
        case childDidUpdate(WorkflowUpdateDebugInfo)
    }
}

// MARK: - Render Context

extension WorkflowNode.SubtreeManager {
    /// The workflow context implementation used by the subtree manager.
    fileprivate final class Context: RenderContextType {
        private(set) var eventPipes: [EventPipe]

        private(set) var sinkStore: SinkStore

        private let originalChildWorkflows: [ChildKey: AnyChildWorkflow]
        private(set) var usedChildWorkflows: [ChildKey: AnyChildWorkflow]

        private let originalSideEffectLifetimes: [AnyHashable: SideEffectLifetime]
        private(set) var usedSideEffectLifetimes: [AnyHashable: SideEffectLifetime]

        private let hostContext: HostContext
        private let session: WorkflowSession

        private var observer: WorkflowObserver? {
            hostContext.observer
        }

        init(
            previousSinks: [ObjectIdentifier: AnyReusableSink],
            originalChildWorkflows: [ChildKey: AnyChildWorkflow],
            originalSideEffectLifetimes: [AnyHashable: SideEffectLifetime],
            hostContext: HostContext,
            session: WorkflowSession
        ) {
            self.eventPipes = []
            self.sinkStore = SinkStore(previousSinks: previousSinks)

            self.originalChildWorkflows = originalChildWorkflows
            self.usedChildWorkflows = [:]

            self.originalSideEffectLifetimes = originalSideEffectLifetimes
            self.usedSideEffectLifetimes = [:]

            self.hostContext = hostContext
            self.session = session
        }

        func render<Child, Action>(
            workflow: Child,
            key: String,
            outputMap: @escaping (Child.Output) -> Action
        ) -> Child.Rendering
            where Child: Workflow,
            Action: WorkflowAction,
            WorkflowType == Action.WorkflowType
        {
            /// A unique key used to identify this child workflow
            let childKey = ChildKey(childType: Child.self, key: key)

            /// If the key already exists in `used`, then a workflow of the same type has been rendered multiple times
            /// during this render pass with the same key. This is not allowed.
            guard usedChildWorkflows[childKey] == nil else {
                fatalError("Child workflows of the same type must be given unique keys. Duplicate workflows of type \(Child.self) were encountered with the key \"\(key)\" in \(WorkflowType.self)")
            }

            let child: ChildWorkflow<Child>
            let eventPipe = EventPipe()
            eventPipes.append(eventPipe)

            /// See if we can reuse an existing child node for the given key.
            if let existing = originalChildWorkflows[childKey] {
                /// Cast the untyped child into a specific typed child. Because our children are keyed by their workflow
                /// type, this should never fail.
                guard let existing = existing as? ChildWorkflow<Child> else {
                    fatalError("ChildKey class type does not match the underlying workflow type.")
                }

                /// Update the existing child
                existing.update(
                    workflow: workflow,
                    outputMap: { outputMap($0) },
                    eventPipe: eventPipe
                )
                child = existing
            } else {
                /// We could not find an existing child matching the given child key, so we will generate a new child.
                /// This spins up a new workflow node, etc to host the newly created child.
                child = ChildWorkflow<Child>(
                    workflow: workflow,
                    outputMap: { outputMap($0) },
                    eventPipe: eventPipe,
                    key: key,
                    hostContext: hostContext,
                    parentSession: session
                )
            }

            /// Store the resolved child in `used`. This allows us to a) hold on to any used children after this render
            /// pass, and b) ensure that we never allow the use of a given workflow type with identical keys.
            usedChildWorkflows[childKey] = child
            return child.render()
        }

        func makeSink<Action>(of actionType: Action.Type) -> Sink<Action> where Action: WorkflowAction, WorkflowType == Action.WorkflowType {
            let reusableSink = sinkStore.findOrCreate(actionType: Action.self)

            let sink = Sink<Action> { [weak reusableSink] action in
                WorkflowLogger.logSinkEvent(ref: SignpostRef(), action: action)

                // use a weak reference as we'd like control over the lifetime
                reusableSink?.handle(action: action)
            }

            return sink
        }

        func runSideEffect(key: AnyHashable, action: (Lifetime) -> Void) {
            if let existingSideEffect = originalSideEffectLifetimes[key] {
                usedSideEffectLifetimes[key] = existingSideEffect
            } else {
                let sideEffectLifetime = SideEffectLifetime()
                action(sideEffectLifetime.lifetime)
                usedSideEffectLifetimes[key] = sideEffectLifetime
            }
        }
    }
}

// MARK: - Reusable Sink

extension WorkflowNode.SubtreeManager {
    fileprivate struct SinkStore {
        var eventPipes: [EventPipe] {
            usedSinks.values.map { reusableSink -> EventPipe in
                reusableSink.eventPipe
            }
        }

        private var previousSinks: [ObjectIdentifier: AnyReusableSink]
        private(set) var usedSinks: [ObjectIdentifier: AnyReusableSink]

        init(previousSinks: [ObjectIdentifier: AnyReusableSink]) {
            self.previousSinks = previousSinks
            self.usedSinks = [:]
        }

        mutating func findOrCreate<Action: WorkflowAction>(actionType: Action.Type) -> ReusableSink<Action> {
            let key = ObjectIdentifier(actionType)

            let reusableSink: ReusableSink<Action>

            if let previousSink = previousSinks.removeValue(forKey: key) as? ReusableSink<Action> {
                // Reuse a previous sink, creating a new event pipe to send the action through.
                previousSink.eventPipe = EventPipe()
                reusableSink = previousSink
            } else if let usedSink = usedSinks[key] as? ReusableSink<Action> {
                // Multiple sinks using the same backing sink.
                reusableSink = usedSink
            } else {
                // Create a new reusable sink.
                reusableSink = ReusableSink<Action>()
            }

            usedSinks[key] = reusableSink

            return reusableSink
        }
    }

    /// Type-erased base class for reusable sinks.
    fileprivate class AnyReusableSink {
        var eventPipe: EventPipe

        init() {
            self.eventPipe = EventPipe()
        }
    }

    fileprivate final class ReusableSink<Action: WorkflowAction>: AnyReusableSink where Action.WorkflowType == WorkflowType {
        func handle(action: Action) {
            let output = Output.update(action, source: .external)

            if case .pending = eventPipe.validationState {
                // Workflow is currently processing an `event`.
                // Scheduling it to be processed after.
                DispatchQueue.workflowExecution.async { [weak self] in
                    self?.eventPipe.handle(event: output)
                }
                return
            }
            eventPipe.handle(event: output)
        }
    }
}

// MARK: - EventPipe

extension WorkflowNode.SubtreeManager {
    final class EventPipe {
        var validationState: ValidationState
        enum ValidationState {
            case preparing
            case pending
            case valid(handler: (Output) -> Void)
            case invalid
        }

        /// Utility to detect reentrancy in `handle()`
        private var isHandlingEvent: Bool = false

        init() {
            self.validationState = .preparing
        }

        func handle(event: Output) {
            dispatchPrecondition(condition: .onQueue(DispatchQueue.workflowExecution))

            let isReentrantCall = isHandlingEvent
            isHandlingEvent = true
            defer { isHandlingEvent = isReentrantCall }

            switch validationState {
            case .preparing:
                fatalError("[\(WorkflowType.self)] Sink sent an action inside `render`. Sinks are not valid until `render` has completed.")

            case .pending:
                fatalError("[\(WorkflowType.self)] Action sent to pipe while in the `pending` state.")

            case .valid(let handler):
                handler(event)

            case .invalid:
                #if DEBUG
                // Reentrancy seems to often be due to UIKit behaviors over
                // which we have little control (e.g. synchronous resignation
                // of first responder after a new Rendering is assigned). Emit
                // some debug info in these cases.
                if isReentrantCall {
                    print("[\(WorkflowType.self)]: ℹ️ Sink sent another action after it was invalidated but before its original action handling was resolved. This new action will be ignored. If this is unexpected, set a Swift error breakpoint on `\(InvalidSinkSentAction.self)` to debug.")
                }

                do {
                    throw InvalidSinkSentAction()
                } catch {}
                #endif

                // If we're invalid and this is the first time `handle()` has
                // been called, then it's likely we've somehow been inadvertently
                // retained from the 'outside world'. Fail more loudly in this case.
                assert(isReentrantCall, """
                    [\(WorkflowType.self)]: Sink sent an action after it was invalidated. This action will be ignored.
                """)
            }
        }

        func setPending() {
            guard case .preparing = validationState else {
                fatalError("Attempted to `setPending` an EventPipe that was not in the preparing state.")
            }
            validationState = .pending
        }

        func enable(with handler: @escaping (Output) -> Void) {
            guard case .pending = validationState else {
                fatalError("EventPipe can only be enabled from the `pending` state")
            }
            validationState = .valid(handler: handler)
        }

        func invalidate() {
            validationState = .invalid
        }
    }
}

// MARK: - ChildKey

extension WorkflowNode.SubtreeManager {
    struct ChildKey: Hashable {
        var childType: Any.Type
        var key: String

        init(childType: (some Any).Type, key: String) {
            self.childType = childType
            self.key = key
        }

        func hash(into hasher: inout Hasher) {
            hasher.combine(ObjectIdentifier(childType))
            hasher.combine(key)
        }

        static func == (lhs: ChildKey, rhs: ChildKey) -> Bool {
            lhs.childType == rhs.childType
                && lhs.key == rhs.key
        }
    }
}

// MARK: - Child Workflows

extension WorkflowNode.SubtreeManager {
    /// Abstract base class for running children in the subtree.
    class AnyChildWorkflow {
        fileprivate var eventPipe: EventPipe

        fileprivate init(eventPipe: EventPipe) {
            self.eventPipe = eventPipe
        }

        func enableEvents() {
            fatalError()
        }

        func makeDebugSnapshot() -> WorkflowHierarchyDebugSnapshot {
            fatalError()
        }
    }

    fileprivate final class ChildWorkflow<W: Workflow>: AnyChildWorkflow {
        private let node: WorkflowNode<W>
        private var outputMap: (W.Output) -> any WorkflowAction<WorkflowType>

        init(
            workflow: W,
            outputMap: @escaping (W.Output) -> any WorkflowAction<WorkflowType>,
            eventPipe: EventPipe,
            key: String,
            hostContext: HostContext,
            parentSession: WorkflowSession?
        ) {
            self.outputMap = outputMap
            self.node = WorkflowNode<W>(
                workflow: workflow,
                key: key,
                hostContext: hostContext,
                parentSession: parentSession
            )

            super.init(eventPipe: eventPipe)

            node.onOutput = { [weak self] output in
                self?.handle(workflowOutput: output)
            }
        }

        override func enableEvents() {
            node.enableEvents()
        }

        func render() -> W.Rendering {
            node.render()
        }

        func update(
            workflow: W,
            outputMap: @escaping (W.Output) -> any WorkflowAction<WorkflowType>,
            eventPipe: EventPipe
        ) {
            self.outputMap = outputMap
            self.eventPipe = eventPipe
            node.update(workflow: workflow)
        }

        private func handle(workflowOutput: WorkflowNode<W>.Output) {
            let output = if let outputEvent = workflowOutput.outputEvent {
                Output.update(
                    outputMap(outputEvent),
                    source: .subtree(workflowOutput.debugInfo)
                )
            } else {
                Output.childDidUpdate(workflowOutput.debugInfo)
            }

            eventPipe.handle(event: output)
        }

        override func makeDebugSnapshot() -> WorkflowHierarchyDebugSnapshot {
            node.makeDebugSnapshot()
        }
    }
}

// MARK: - Side Effects

extension WorkflowNode.SubtreeManager {
    class SideEffectLifetime {
        fileprivate let lifetime: Lifetime

        fileprivate init() {
            self.lifetime = Lifetime()
        }

        deinit {
            // Explicitly end the lifetime in case someone retained it from outside
            lifetime.end()
        }
    }
}

// MARK: - Debugging Utilities

#if DEBUG
private struct InvalidSinkSentAction: Error {}
#endif