[breaking]: update action apply API

Author: jamieQ

Workflow/Sources/AnyWorkflowConvertible.swift

@@ -115,10 +115,27 @@ extension AnyWorkflowConvertible {
     /// Process an `Output`
     ///
     /// - Parameter apply: On `Output`, mutate `State` as necessary and return new `Output` (or `nil`).
-    public func onOutput<Parent>(_ apply: @escaping ((inout Parent.State, Output) -> Parent.Output?)) -> AnyWorkflow<Rendering, AnyWorkflowAction<Parent>> {
+    public func onOutput<Parent>(
+        _ apply: @escaping (inout Parent.State, ApplyContext<Parent>, Output) -> Parent.Output?
+    ) -> AnyWorkflow<Rendering, AnyWorkflowAction<Parent>> {
         asAnyWorkflow()
             .mapOutput { output in
-                AnyWorkflowAction { state -> Parent.Output? in
+                AnyWorkflowAction { state, context -> Parent.Output? in
+                    apply(&state, context, output)
+                }
+            }
+    }
+
+    /// Process an `Output`
+    ///
+    /// - Parameter apply: On `Output`, mutate `State` as necessary and return new `Output` (or `nil`).
+    @_disfavoredOverload
+    public func onOutput<Parent>(
+        _ apply: @escaping (inout Parent.State, Output) -> Parent.Output?
+    ) -> AnyWorkflow<Rendering, AnyWorkflowAction<Parent>> {
+        asAnyWorkflow()
+            .mapOutput { output in
+                AnyWorkflowAction { state, _ -> Parent.Output? in
                     apply(&state, output)
                 }
             }

Workflow/Sources/ApplyContext.swift

@@ -0,0 +1,87 @@
+/*
+ * Copyright 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.
+ */
+
+/// Internal utility protocol so that `WorkflowTesting` can provide an alternate implementation
+protocol ApplyContextType<WorkflowType> {
+    associatedtype WorkflowType: Workflow
+
+    subscript<Value>(
+        workflowValue keyPath: KeyPath<WorkflowType, Value>
+    ) -> Value { get }
+}
+
+/// Runtime context passed as a parameter to `WorkflowAction`'s `apply()` method
+/// that provides an integration point with the runtime that can be used to read property values
+/// off of the current `Workflow` instance.
+public struct ApplyContext<WorkflowType: Workflow> {
+    let wrappedContext: any ApplyContextType<WorkflowType>
+
+    init<Impl: ApplyContextType>(_ context: Impl)
+        where Impl.WorkflowType == WorkflowType
+    {
+        self.wrappedContext = context
+    }
+
+    public subscript<Value>(
+        workflowValue keyPath: KeyPath<WorkflowType, Value>
+    ) -> Value {
+        wrappedContext[workflowValue: keyPath]
+    }
+}
+
+extension ApplyContext {
+    static func make<Wrapped: ApplyContextType>(
+        implementation: Wrapped
+    ) -> ApplyContext<Wrapped.WorkflowType>
+        where Wrapped.WorkflowType == WorkflowType
+    {
+        ApplyContext(implementation)
+    }
+}
+
+// FIXME: this is currently a class so that we can zero out the storage
+// when it's invalidated. it'd be nice to eventually make the `ApplyContext`
+// type itself `~Escapable` since that's really the behavior that we want
+// to enforce.
+
+/// The `ApplyContext` used by the Workflow runtime when applying actions.
+final class ConcreteApplyContext<WorkflowType: Workflow>: ApplyContextType {
+    private(set) var storage: WorkflowType?
+
+    private var validatedStorage: WorkflowType {
+        guard let storage else {
+            fatalError("Attempt to use an ApplyContext for Workflow of type '\(WorkflowType.self)' after it was invalidated. The context is only valid during a call to an apply(...) method.")
+        }
+
+        return storage
+    }
+
+    init(storage: WorkflowType) {
+        self.storage = storage
+    }
+
+    subscript<Value>(
+        workflowValue keyPath: KeyPath<WorkflowType, Value>
+    ) -> Value {
+        validatedStorage[keyPath: keyPath]
+    }
+
+    // MARK: -
+
+    func invalidate() {
+        storage = nil
+    }
+}

Workflow/Sources/RenderContext.swift

@@ -180,12 +180,17 @@ protocol RenderContextType: AnyObject {
 extension RenderContext {
     public func makeSink<Event>(
         of eventType: Event.Type,
-        onEvent: @escaping (Event, inout WorkflowType.State) -> WorkflowType.Output?
+        onEvent: @escaping (
+            Event,
+            inout WorkflowType.State,
+            ApplyContext<WorkflowType>
+        ) -> WorkflowType.Output?
+        // (Event, inout WorkflowType.State) -> WorkflowType.Output?
     ) -> Sink<Event> {
         makeSink(of: AnyWorkflowAction.self)
             .contraMap { event in
-                AnyWorkflowAction<WorkflowType> { state in
-                    onEvent(event, &state)
+                AnyWorkflowAction<WorkflowType> { state, context in
+                    onEvent(event, &state, context)
                 }
             }
     }

Workflow/Sources/StateMutationSink.swift

@@ -44,7 +44,7 @@ public struct StateMutationSink<WorkflowType: Workflow> {
     ///   - update: The `State` mutation to perform.
     public func send(_ update: @escaping (inout WorkflowType.State) -> Void) {
         sink.send(
-            AnyWorkflowAction<WorkflowType> { state in
+            AnyWorkflowAction<WorkflowType> { state, _ in
                 update(&state)
                 return nil
             }

Workflow/Sources/WorkflowAction.swift

@@ -27,14 +27,22 @@ public protocol WorkflowAction<WorkflowType> {
     ///
     /// - Returns: An optional output event for the workflow. If an output event is returned, it will be passed up
     ///            the workflow hierarchy to this workflow's parent.
-    func apply(toState state: inout WorkflowType.State) -> WorkflowType.Output?
+    func apply(
+        toState state: inout WorkflowType.State,
+        context: ApplyContext<WorkflowType>
+    ) -> WorkflowType.Output?
+}
+
+extension WorkflowAction {
+    /// Closure type signature matching `WorkflowAction`'s `apply()` method.
+    public typealias ActionApplyClosure = (inout WorkflowType.State, ApplyContext<WorkflowType>) -> WorkflowType.Output?
 }
 
 /// A type-erased workflow action.
 ///
 /// The `AnyWorkflowAction` type forwards `apply` to an underlying workflow action, hiding its specific underlying type.
 public struct AnyWorkflowAction<WorkflowType: Workflow>: WorkflowAction {
-    private let _apply: (inout WorkflowType.State) -> WorkflowType.Output?
+    private let _apply: ActionApplyClosure
 
     /// The underlying type-erased `WorkflowAction`
     public let base: Any
@@ -50,7 +58,9 @@ public struct AnyWorkflowAction<WorkflowType: Workflow>: WorkflowAction {
             self = anyEvent
             return
         }
-        self._apply = { base.apply(toState: &$0) }
+        self._apply = {
+            base.apply(toState: &$0, context: $1)
+        }
         self.base = base
         self.isClosureBased = false
     }
@@ -59,7 +69,7 @@ public struct AnyWorkflowAction<WorkflowType: Workflow>: WorkflowAction {
     ///
     /// - Parameter apply: the apply function for the resulting action.
     public init(
-        _ apply: @escaping (inout WorkflowType.State) -> WorkflowType.Output?,
+        _ apply: @escaping ActionApplyClosure,
         fileID: StaticString = #fileID,
         line: UInt = #line
     ) {
@@ -71,16 +81,35 @@ public struct AnyWorkflowAction<WorkflowType: Workflow>: WorkflowAction {
         self.init(closureAction: closureAction)
     }
 
+    /// Creates a type-erased workflow action with the given `apply` implementation.
+    ///
+    /// - Parameter apply: the apply function for the resulting action.
+    @_disfavoredOverload
+    public init(
+        _ apply: @escaping (inout WorkflowType.State) -> WorkflowType.Output?,
+        fileID: StaticString = #fileID,
+        line: UInt = #line
+    ) {
+        self.init(
+            { state, _ in apply(&state) },
+            fileID: fileID,
+            line: line
+        )
+    }
+
     /// Private initializer forwarded to via `init(_ apply:...)`
     /// - Parameter closureAction: The `ClosureAction` wrapping the underlying `apply` closure.
     fileprivate init(closureAction: ClosureAction<WorkflowType>) {
-        self._apply = closureAction.apply(toState:)
+        self._apply = closureAction.apply(toState:context:)
         self.base = closureAction
         self.isClosureBased = true
     }
 
-    public func apply(toState state: inout WorkflowType.State) -> WorkflowType.Output? {
-        _apply(&state)
+    public func apply(
+        toState state: inout WorkflowType.State,
+        context: ApplyContext<WorkflowType>
+    ) -> WorkflowType.Output? {
+        _apply(&state, context)
     }
 }
 
@@ -89,15 +118,15 @@ extension AnyWorkflowAction {
     ///
     /// - Parameter output: The output event to send when this action is applied.
     public init(sendingOutput output: WorkflowType.Output) {
-        self = AnyWorkflowAction { state in
+        self = AnyWorkflowAction { _, _ in
             output
         }
     }
 
     /// Creates a type-erased workflow action that does nothing (it leaves state unchanged and does not emit an output
     /// event).
     public static var noAction: AnyWorkflowAction<WorkflowType> {
-        AnyWorkflowAction { state in
+        AnyWorkflowAction { _, _ in
             nil
         }
     }
@@ -109,12 +138,12 @@ extension AnyWorkflowAction {
 /// Mainly used to provide more useful debugging/telemetry information for `AnyWorkflow` instances
 /// defined via a closure.
 struct ClosureAction<WorkflowType: Workflow>: WorkflowAction {
-    private let _apply: (inout WorkflowType.State) -> WorkflowType.Output?
+    private let _apply: ActionApplyClosure
     let fileID: StaticString
     let line: UInt
 
     init(
-        _apply: @escaping (inout WorkflowType.State) -> WorkflowType.Output?,
+        _apply: @escaping ActionApplyClosure,
         fileID: StaticString,
         line: UInt
     ) {
@@ -123,8 +152,11 @@ struct ClosureAction<WorkflowType: Workflow>: WorkflowAction {
         self.line = line
     }
 
-    func apply(toState state: inout WorkflowType.State) -> WorkflowType.Output? {
-        _apply(&state)
+    func apply(
+        toState state: inout WorkflowType.State,
+        context: ApplyContext<WorkflowType>
+    ) -> WorkflowType.Output? {
+        _apply(&state, context)
     }
 }
 

Workflow/Sources/WorkflowNode.swift

@@ -218,7 +218,14 @@ extension WorkflowNode {
         defer { observerCompletion?(state, output) }
 
         /// Apply the action to the current state
-        output = action.apply(toState: &state)
+        do {
+            // TODO: can we avoid instantiating a class here somehow?
+            let context = ConcreteApplyContext(storage: workflow)
+            defer { context.invalidate() }
+
+            let wrappedContext = ApplyContext.make(implementation: context)
+            output = action.apply(toState: &state, context: wrappedContext)
+        }
 
         return output
     }

Workflow/Tests/AnyWorkflowActionTests.swift

@@ -83,7 +83,8 @@ final class AnyWorkflowActionTests: XCTestCase {
         XCTAssertEqual(log, [])
 
         var state: Void = ()
-        _ = erased.apply(toState: &state)
+        let ctx = ApplyContext(ConcreteApplyContext(storage: ExampleWorkflow()))
+        _ = erased.apply(toState: &state, context: ctx)
 
         XCTAssertEqual(log, ["action invoked"])
     }

Workflow/Tests/ApplyContextTests.swift

@@ -0,0 +1,79 @@
+/*
+ * Copyright 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 Testing
+
+@testable import Workflow
+
+@MainActor
+struct ApplyContextTests {
+    @Test
+    func concreteApplyContextInvalidatedAfterUse() async throws {
+        var escapedContext: ApplyContext<EscapingContextWorkflow>?
+        let onApply = { (context: ApplyContext<EscapingContextWorkflow>) in
+            #expect(context[workflowValue: \.property] == 42)
+            #expect(context.concreteStorage != nil)
+            escapedContext = context
+        }
+
+        let workflow = EscapingContextWorkflow(
+            property: 42,
+            onApply: onApply
+        )
+        let node = WorkflowNode(workflow: workflow)
+
+        let emitEvent = node.render()
+        node.enableEvents()
+
+        emitEvent()
+
+        #expect(escapedContext != nil)
+        #expect(escapedContext?.concreteStorage == nil)
+    }
+}
+
+// MARK: -
+
+private struct EscapingContextWorkflow: Workflow {
+    typealias Rendering = () -> Void
+    typealias State = Void
+
+    var property: Int
+    var onApply: ((ApplyContext<Self>) -> Void)?
+
+    func render(
+        state: State,
+        context: RenderContext<EscapingContextWorkflow>
+    ) -> Rendering {
+        let sink = context.makeSink(of: EscapingAction.self)
+        let action = EscapingAction(onApply: onApply)
+        return { sink.send(action) }
+    }
+
+    struct EscapingAction: WorkflowAction {
+        typealias WorkflowType = EscapingContextWorkflow
+
+        var onApply: ((ApplyContext<WorkflowType>) -> Void)?
+
+        func apply(
+            toState state: inout WorkflowType.State,
+            context: ApplyContext<WorkflowType>
+        ) -> WorkflowType.Output? {
+            onApply?(context)
+            return nil
+        }
+    }
+}