[feat]: Add support for automatic ViewEnvironment bridging in WorkflowUI

Samples/TicTacToe/Sources/Authentication/LoginScreen.swift

@@ -27,6 +27,7 @@ struct LoginScreen: Screen {
 
     func viewControllerDescription(environment: ViewEnvironment) -> ViewControllerDescription {
         return ViewControllerDescription(
+            environment: environment,
             build: { LoginViewController() },
             update: { $0.update(with: self) }
         )

WorkflowUI.podspec

@@ -20,6 +20,7 @@ Pod::Spec.new do |s|
 
     s.dependency 'Workflow', "#{s.version}"
     s.dependency 'ViewEnvironment', "#{s.version}"
+    s.dependency 'ViewEnvironmentUI', "#{s.version}"
 
     s.pod_target_xcconfig = { 'APPLICATION_EXTENSION_API_ONLY' => 'YES' }
 

WorkflowUI/Sources/Hosting/WorkflowHostingController.swift

@@ -18,10 +18,13 @@
 
 import ReactiveSwift
 import UIKit
+@_spi(ViewEnvironmentWiring) import ViewEnvironmentUI
 import Workflow
 
 /// Drives view controllers from a root Workflow.
 public final class WorkflowHostingController<ScreenType, Output>: UIViewController where ScreenType: Screen {
+    public typealias CustomizeEnvironment = (inout ViewEnvironment) -> Void
+
     /// Emits output events from the bound workflow.
     public var output: Signal<Output, Never> {
         return workflowHost.output
@@ -33,31 +36,38 @@ public final class WorkflowHostingController<ScreenType, Output>: UIViewControll
 
     private let (lifetime, token) = Lifetime.make()
 
-    public var rootViewEnvironment: ViewEnvironment {
-        didSet {
-            update(screen: workflowHost.rendering.value, environment: rootViewEnvironment)
-        }
+    public var customizeEnvironment: CustomizeEnvironment {
+        didSet { setNeedsEnvironmentUpdate() }
     }
 
     public init<W: AnyWorkflowConvertible>(
         workflow: W,
-        rootViewEnvironment: ViewEnvironment = .empty,
+        customizeEnvironment: @escaping CustomizeEnvironment = { _ in },
         observers: [WorkflowObserver] = []
     ) where W.Rendering == ScreenType, W.Output == Output {
         self.workflowHost = WorkflowHost(
             workflow: workflow.asAnyWorkflow(),
             observers: observers
         )
 
+        self.customizeEnvironment = customizeEnvironment
+
+        var customizedEnvironment: ViewEnvironment = .empty
+        customizeEnvironment(&customizedEnvironment)
+
         self.rootViewController = workflowHost
             .rendering
             .value
-            .buildViewController(in: rootViewEnvironment)
-
-        self.rootViewEnvironment = rootViewEnvironment
+            .viewControllerDescription(environment: customizedEnvironment)
+            .buildViewController()
 
         super.init(nibName: nil, bundle: nil)
 
+        // Do not automatically forward environment did change notifications to the rendered screen's backing view
+        // controller. Instead rely on `ViewControllerDescription` to call `setNeedsEnvironmentUpdate()` when updates
+        // occur.
+        environmentDescendantsOverride = { [] }
+
         addChild(rootViewController)
         rootViewController.didMove(toParent: self)
 
@@ -68,7 +78,7 @@ public final class WorkflowHostingController<ScreenType, Output>: UIViewControll
             .observeValues { [weak self] screen in
                 guard let self = self else { return }
 
-                self.update(screen: screen, environment: self.rootViewEnvironment)
+                self.update(screen: screen, environment: self.environment)
             }
     }
 
@@ -82,8 +92,17 @@ public final class WorkflowHostingController<ScreenType, Output>: UIViewControll
     }
 
     private func update(screen: ScreenType, environment: ViewEnvironment) {
+        let previousRoot = rootViewController
+
         update(child: \.rootViewController, with: screen, in: environment)
 
+        if previousRoot !== rootViewController {
+            // If a new view controller was instantiated and added as a child we need to inform it that the environment
+            // should be re-requested in order to respond to customizations in this WorkflowHostingController or any
+            // view controller above it in the UIViewController hierarchy.
+            setNeedsEnvironmentUpdate()
+        }
+
         updatePreferredContentSizeIfNeeded()
     }
 
@@ -98,6 +117,11 @@ public final class WorkflowHostingController<ScreenType, Output>: UIViewControll
         updatePreferredContentSizeIfNeeded()
     }
 
+    public override func viewWillLayoutSubviews() {
+        super.viewWillLayoutSubviews()
+        applyEnvironmentIfNeeded()
+    }
+
     override public func viewDidLayoutSubviews() {
         super.viewDidLayoutSubviews()
         rootViewController.view.frame = view.bounds
@@ -150,4 +174,14 @@ public final class WorkflowHostingController<ScreenType, Output>: UIViewControll
     }
 }
 
+extension WorkflowHostingController: ViewEnvironmentObserving {
+    public func customize(environment: inout ViewEnvironment) {
+        customizeEnvironment(&environment)
+    }
+
+    public func environmentDidChange() {
+        update(screen: workflowHost.rendering.value, environment: environment)
+    }
+}
+
 #endif

WorkflowUI/Sources/ModuleExports.swift

@@ -1 +1,2 @@
 @_exported import ViewEnvironment
+@_exported import ViewEnvironmentUI

WorkflowUI/Sources/Screen/ScreenViewController.swift

@@ -17,6 +17,8 @@
 #if canImport(UIKit)
 
 import UIKit
+import ViewEnvironment
+@_spi(ViewEnvironmentWiring) import ViewEnvironmentUI
 
 /// Generic base class that can be subclassed in order to to define a UI implementation that is powered by the
 /// given screen type.
@@ -25,7 +27,7 @@ import UIKit
 /// ```
 /// struct MyScreen: Screen {
 ///     func viewControllerDescription(environment: ViewEnvironment) -> ViewControllerDescription {
-///         return MyScreenViewController.description(for: self)
+///         return MyScreenViewController.description(for: self, environment: environment)
 ///     }
 /// }
 ///
@@ -42,11 +44,11 @@ open class ScreenViewController<ScreenType: Screen>: UIViewController {
         return ScreenType.self
     }
 
-    public private(set) final var environment: ViewEnvironment
+    private var previousEnvironment: ViewEnvironment
 
     public required init(screen: ScreenType, environment: ViewEnvironment) {
         self.screen = screen
-        self.environment = environment
+        self.previousEnvironment = environment
         super.init(nibName: nil, bundle: nil)
     }
 
@@ -55,11 +57,11 @@ open class ScreenViewController<ScreenType: Screen>: UIViewController {
         fatalError("init(coder:) has not been implemented")
     }
 
-    public final func update(screen: ScreenType, environment: ViewEnvironment) {
+    public final func update(screen: ScreenType) {
         let previousScreen = self.screen
         self.screen = screen
-        let previousEnvironment = self.environment
-        self.environment = environment
+        let previousEnvironment = self.previousEnvironment
+        self.previousEnvironment = environment
         screenDidChange(from: previousScreen, previousEnvironment: previousEnvironment)
     }
 
@@ -79,8 +81,9 @@ extension ScreenViewController {
         ViewControllerDescription(
             performInitialUpdate: performInitialUpdate,
             type: self,
+            environment: environment,
             build: { self.init(screen: screen, environment: environment) },
-            update: { $0.update(screen: screen, environment: environment) }
+            update: { $0.update(screen: screen) }
         )
     }
 }

WorkflowUI/Sources/ViewControllerDescription/ViewControllerDescription.swift

@@ -17,6 +17,8 @@
 #if canImport(UIKit)
 
 import UIKit
+import ViewEnvironment
+@_spi(ViewEnvironmentWiring) import ViewEnvironmentUI
 
 /// A ViewControllerDescription acts as a "recipe" for building and updating a specific `UIViewController`.
 /// It describes how to _create_ and later _update_ a given view controller instance, without creating one
@@ -49,6 +51,7 @@ public struct ViewControllerDescription {
     /// type changed.
     public let kind: KindIdentifier
 
+    private let environment: ViewEnvironment
     private let build: () -> UIViewController
     private let update: (UIViewController) -> Void
 
@@ -59,6 +62,11 @@ public struct ViewControllerDescription {
     ///   - performInitialUpdate: If an initial call to `update(viewController:)`
     ///     will be performed when the view controller is created. Defaults to `true`.
     ///
+    ///   - environment: The `ViewEnvironment` that should be injected above the
+    ///     described view controller for ViewEnvironmentUI environment propagation.
+    ///     This is typically passed in from a `Screen` in its
+    ///     `viewControllerDescription(environment:)` method.
+    ///
     ///   - type: The type of view controller produced by this description.
     ///     Typically, should should be able to omit this parameter, but
     ///     in cases where type inference has trouble, it’s offered as
@@ -70,13 +78,16 @@ public struct ViewControllerDescription {
     public init<VC: UIViewController>(
         performInitialUpdate: Bool = true,
         type: VC.Type = VC.self,
+        environment: ViewEnvironment,
         build: @escaping () -> VC,
         update: @escaping (VC) -> Void
     ) {
         self.performInitialUpdate = performInitialUpdate
 
         self.kind = .init(VC.self)
 
+        self.environment = environment
+
         self.build = build
 
         self.update = { untypedViewController in
@@ -95,7 +106,10 @@ public struct ViewControllerDescription {
 
         if performInitialUpdate {
             // Perform an initial update of the built view controller
+            // Note that this also configures the environment ancestor node.
             update(viewController: viewController)
+        } else {
+            configureAncestor(for: viewController)
         }
 
         return viewController
@@ -126,8 +140,43 @@ public struct ViewControllerDescription {
             """
         )
 
+        configureAncestor(for: viewController)
+
         update(viewController)
     }
+
+    private func configureAncestor(for viewController: UIViewController) {
+        guard let ancestorOverride = viewController.environmentAncestorOverride else {
+            // If no ancestor is currently present establish the initial ancestor override
+            establishAncestorOverride(for: viewController)
+            return
+        }
+
+        let currentAncestor = ancestorOverride()
+        // Check whether the VC's ancestor was overridden by a ViewControllerDescription.
+        guard currentAncestor is PropagationNode else {
+            // Do not override the VC's ancestor if it was overridden by something outside of the
+            // `ViewControllerDescription`'s management of this node.
+            // The view controller we're managing, or the container it's contained in, likely needs to manage this in a
+            // special way.
+            return
+        }
+
+        // We must nil this out first or we'll hit an assertion which protects against overriding the ancestor when
+        // some other system has already attempted to provide an override.
+        viewController.environmentAncestorOverride = nil
+        establishAncestorOverride(for: viewController)
+    }
+
+    private func establishAncestorOverride(for viewController: UIViewController) {
+        let ancestor = PropagationNode(
+            viewController: viewController,
+            environment: environment
+        )
+        viewController.environmentAncestorOverride = { ancestor }
+
+        ancestor.setNeedsEnvironmentUpdate()
+    }
 }
 
 extension ViewControllerDescription {
@@ -168,4 +217,39 @@ extension ViewControllerDescription {
     }
 }
 
+extension ViewControllerDescription {
+    fileprivate struct PropagationNode: ViewEnvironmentObserving {
+
+        weak var viewController: UIViewController?
+
+        let environment: ViewEnvironment
+
+        init(
+            viewController: UIViewController,
+            environment: ViewEnvironment
+        ) {
+            self.viewController = viewController
+            self.environment = environment
+        }
+
+        var environmentAncestor: ViewEnvironmentPropagating? { nil }
+
+        var environmentDescendants: [ViewEnvironmentPropagating] {
+            [viewController].compactMap { $0 }
+        }
+
+        func customize(environment: inout ViewEnvironment) {
+            environment = self.environment
+        }
+
+        func setNeedsEnvironmentUpdate() {
+            setNeedsEnvironmentUpdateOnAppropriateDescendants()
+        }
+
+        func applyEnvironmentIfNeeded() {
+            /// `apply(environment:)` is not implemented so do nothing.
+        }
+    }
+}
+
 #endif

WorkflowUI/Tests/DescribedViewControllerTests.swift

@@ -247,12 +247,14 @@ fileprivate enum TestScreen: Screen, Equatable {
         switch self {
         case .counter(let count):
             return ViewControllerDescription(
+                environment: environment,
                 build: { CounterViewController(count: count) },
                 update: { $0.count = count }
             )
 
         case .message(let message):
             return ViewControllerDescription(
+                environment: environment,
                 build: { MessageViewController(message: message) },
                 update: { $0.message = message }
             )

WorkflowUI/Tests/UIViewControllerExtensionTests.swift

@@ -155,6 +155,7 @@ private struct Screen1: Screen {
     func viewControllerDescription(environment: ViewEnvironment) -> ViewControllerDescription {
         ViewControllerDescription(
             type: VC1.self,
+            environment: environment,
             build: { VC1(identifier: "1", recordEvent: recordEvent) },
             update: { $0.recordEvent = recordEvent }
         )
@@ -167,6 +168,7 @@ private struct Screen2: Screen {
     func viewControllerDescription(environment: ViewEnvironment) -> ViewControllerDescription {
         ViewControllerDescription(
             type: VC2.self,
+            environment: environment,
             build: { VC2(identifier: "2", recordEvent: recordEvent) },
             update: { $0.recordEvent = recordEvent }
         )