[fix]: update WorkflowHostingController on layout if ancestor hierarchy has changed (#227)

Author: n8chur

ViewEnvironmentUI/Sources/ViewEnvironmentPropagating.swift

@@ -272,6 +272,65 @@ extension ViewEnvironmentPropagating {
             objc_setAssociatedObject(self, &AssociatedKeys.descendantsOverride, newValue, .OBJC_ASSOCIATION_RETAIN)
         }
     }
+
+    /// Returns an `Equatable` representation of the `ViewEnvironmentPropagating` environment
+    /// ancestor tree path.
+    ///
+    /// This can be useful, for example, if you need to determine if any ancestor was inserted or
+    /// removed above this node.
+    ///
+    /// The `Equatable` implementation of this type compares the tree as an array of weak
+    /// references.
+    ///
+    @_spi(ViewEnvironmentWiring)
+    public var environmentAncestorPath: EnvironmentAncestorPath {
+        var path = EnvironmentAncestorPath()
+
+        if let first = environmentAncestor {
+            for node in sequence(first: first, next: \.environmentAncestor) {
+                path.append(node)
+            }
+        }
+
+        return path
+    }
+
+    @_spi(ViewEnvironmentWiring)
+    public typealias EnvironmentAncestorPath = ViewEnvironmentPropagatingAncestorPath
+}
+
+/// An `Equatable` representation of the `ViewEnvironmentPropagating` environment ancestor tree
+/// path.
+///
+/// This can be useful, for example, if you need to determine if any ancestor was inserted or
+/// removed above this node.
+///
+/// The `Equatable` implementation of this type compares the tree as an array of weak references.
+///
+@_spi(ViewEnvironmentWiring)
+public struct ViewEnvironmentPropagatingAncestorPath: Equatable {
+    private var nodes: [WeakBox] = []
+
+    fileprivate mutating func append(_ node: ViewEnvironmentPropagating) {
+        nodes.append(WeakBox(node))
+    }
+
+    // Use a weak box to avoid retaining the node.
+    //
+    // We do this instead of `ObjectIdentifier` because `ObjectIdentifier`s are only valid for the
+    // lifetime of the object being identified—the value of the pointer could be re-used if it is
+    // deallocated.
+    private struct WeakBox: Equatable {
+        weak var node: ViewEnvironmentPropagating?
+
+        init(_ node: ViewEnvironmentPropagating) {
+            self.node = node
+        }
+
+        static func == (lhs: WeakBox, rhs: WeakBox) -> Bool {
+            lhs.node === rhs.node
+        }
+    }
 }
 
 /// A closure that is called when the `ViewEnvironment` needs to be updated.

WorkflowUI/Sources/Hosting/WorkflowHostingController.swift

@@ -41,6 +41,8 @@ public final class WorkflowHostingController<ScreenType, Output>: WorkflowUIView
 
     private let (lifetime, token) = Lifetime.make()
 
+    private var lastEnvironmentAncestorPath: EnvironmentAncestorPath?
+
     public init<W: AnyWorkflowConvertible>(
         workflow: W,
         customizeEnvironment: @escaping CustomizeEnvironment = { _ in },
@@ -79,7 +81,10 @@ public final class WorkflowHostingController<ScreenType, Output>: WorkflowUIView
             .observeValues { [weak self] screen in
                 guard let self = self else { return }
 
-                self.update(screen: screen, environment: self.environment)
+                self.update(
+                    screen: screen,
+                    environmentAncestorPath: self.environmentAncestorPath
+                )
             }
     }
 
@@ -92,7 +97,10 @@ public final class WorkflowHostingController<ScreenType, Output>: WorkflowUIView
         fatalError("init(coder:) has not been implemented")
     }
 
-    private func update(screen: ScreenType, environment: ViewEnvironment) {
+    private func update(screen: ScreenType, environmentAncestorPath: EnvironmentAncestorPath) {
+        lastEnvironmentAncestorPath = environmentAncestorPath
+
+        let environment = environment
         let previousRoot = rootViewController
 
         update(child: \.rootViewController, with: screen, in: environment)
@@ -120,7 +128,14 @@ public final class WorkflowHostingController<ScreenType, Output>: WorkflowUIView
 
     override public func viewWillLayoutSubviews() {
         super.viewWillLayoutSubviews()
-        applyEnvironmentIfNeeded()
+
+        let environmentAncestorPath = environmentAncestorPath
+        if environmentAncestorPath != lastEnvironmentAncestorPath {
+            update(
+                screen: workflowHost.rendering.value,
+                environmentAncestorPath: environmentAncestorPath
+            )
+        }
     }
 
     override public func viewDidLayoutSubviews() {
@@ -181,7 +196,10 @@ extension WorkflowHostingController: ViewEnvironmentObserving {
     }
 
     public func environmentDidChange() {
-        update(screen: workflowHost.rendering.value, environment: environment)
+        update(
+            screen: workflowHost.rendering.value,
+            environmentAncestorPath: environmentAncestorPath
+        )
     }
 }
 

WorkflowUI/Tests/WorkflowHostingControllerTests.swift

@@ -224,6 +224,63 @@ class WorkflowHostingControllerTests: XCTestCase {
             XCTAssertEqual(environment[ScreenKey.self], true)
         }
     }
+
+    func test_environment_updates_on_layout_in_new_hierarchy() {
+        var changedEnvironments: [ViewEnvironment] = []
+        let hostingController = WorkflowHostingController(
+            workflow: EnvironmentObservingWorkflow(
+                value: "first",
+                onEnvironmentDidChange: { changedEnvironments.append($0) }
+            )
+        )
+
+        // Setup the initial hierarchy
+        let root1 = UIViewController()
+        let container = UIViewController()
+        root1.addChild(container)
+        root1.view.addSubview(container.view)
+        container.didMove(toParent: root1)
+
+        container.addChild(hostingController)
+        container.view.addSubview(hostingController.view)
+        hostingController.didMove(toParent: container)
+
+        XCTAssertEqual(changedEnvironments.count, 1)
+
+        // Triggering a layout should cause an update to the workflow's rendering since the
+        // ancestor path has changed since the `WorkflowHostingController` was initialized
+        hostingController.view.setNeedsLayout()
+        hostingController.view.layoutIfNeeded()
+        XCTAssertEqual(changedEnvironments.count, 2)
+
+        // There should be no environment update if the Workflow state and ancestor tree path has
+        // not changed
+        hostingController.view.setNeedsLayout()
+        hostingController.view.layoutIfNeeded()
+        XCTAssertEqual(changedEnvironments.count, 2)
+
+        // Change the environment ancestor path
+        container.willMove(toParent: nil)
+        container.view.removeFromSuperview()
+        container.removeFromParent()
+
+        let root2 = UIViewController()
+        root2.addChild(container)
+        root2.view.addSubview(container.view)
+        container.didMove(toParent: root2)
+
+        // An environment update should occur since the ancestor path has changed since the last
+        // update and/or layout.
+        hostingController.view.setNeedsLayout()
+        hostingController.view.layoutIfNeeded()
+        XCTAssertEqual(changedEnvironments.count, 3)
+
+        // There should be no environment update if the Workflow state and ancestor tree path has
+        // not changed
+        hostingController.view.setNeedsLayout()
+        hostingController.view.layoutIfNeeded()
+        XCTAssertEqual(changedEnvironments.count, 3)
+    }
 }
 
 fileprivate struct SubscribingWorkflow: Workflow {
@@ -303,7 +360,6 @@ fileprivate struct EnvironmentObservingWorkflow: Workflow {
 }
 
 fileprivate final class EnvironmentCustomizingViewController: UIViewController, ViewEnvironmentObserving {
-
     var customizeEnvironment: (inout ViewEnvironment) -> Void
 
     init(customizeEnvironment: @escaping (inout ViewEnvironment) -> Void) {