[Prototype] Update WorkflowHostinController on layout if ancestor hierarchy has changed

Author: n8chur

ViewEnvironmentUI/Sources/ViewEnvironmentPropagating.swift

@@ -272,6 +272,27 @@ extension ViewEnvironmentPropagating {
             objc_setAssociatedObject(self, &AssociatedKeys.descendantsOverride, newValue, .OBJC_ASSOCIATION_RETAIN)
         }
     }
+
+    @_spi(ViewEnvironmentWiring)
+    public typealias EnvironmentAncestorPath = ViewEnvironmentPropagatingAncestorPath
+
+    @_spi(ViewEnvironmentWiring)
+    public var environmentAncestorPath: EnvironmentAncestorPath {
+        var path = EnvironmentAncestorPath()
+
+        if let first = environmentAncestor {
+            for node in sequence(first: first, next: \.environmentAncestor) {
+                path.nodes.append(ObjectIdentifier(node))
+            }
+        }
+
+        return path
+    }
+}
+
+@_spi(ViewEnvironmentWiring)
+public struct ViewEnvironmentPropagatingAncestorPath: Equatable {
+    var nodes: [ObjectIdentifier] = []
 }
 
 /// 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 ancestorPath: EnvironmentAncestorPath?
+
     public init<W: AnyWorkflowConvertible>(
         workflow: W,
         customizeEnvironment: @escaping CustomizeEnvironment = { _ in },
@@ -79,7 +81,7 @@ 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, ancestorPath: self.environmentAncestorPath)
             }
     }
 
@@ -92,7 +94,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, ancestorPath: EnvironmentAncestorPath) {
+        self.ancestorPath = ancestorPath
+
+        let environment = environment
         let previousRoot = rootViewController
 
         update(child: \.rootViewController, with: screen, in: environment)
@@ -120,7 +125,11 @@ public final class WorkflowHostingController<ScreenType, Output>: WorkflowUIView
 
     override public func viewWillLayoutSubviews() {
         super.viewWillLayoutSubviews()
-        applyEnvironmentIfNeeded()
+
+        let ancestorPath = environmentAncestorPath
+        if ancestorPath != self.ancestorPath {
+            update(screen: workflowHost.rendering.value, ancestorPath: ancestorPath)
+        }
     }
 
     override public func viewDidLayoutSubviews() {
@@ -181,7 +190,7 @@ extension WorkflowHostingController: ViewEnvironmentObserving {
     }
 
     public func environmentDidChange() {
-        update(screen: workflowHost.rendering.value, environment: environment)
+        update(screen: workflowHost.rendering.value, ancestorPath: environmentAncestorPath)
     }
 }
 

WorkflowUI/Tests/WorkflowHostingControllerTests.swift

@@ -224,6 +224,57 @@ 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) }
+            )
+        )
+
+        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)
+
+        hostingController.view.setNeedsLayout()
+        hostingController.view.layoutIfNeeded()
+
+        XCTAssertEqual(changedEnvironments.count, 2)
+
+        hostingController.view.setNeedsLayout()
+        hostingController.view.layoutIfNeeded()
+
+        XCTAssertEqual(changedEnvironments.count, 2)
+
+        container.willMove(toParent: nil)
+        container.view.removeFromSuperview()
+        container.removeFromParent()
+
+        let root2 = UIViewController()
+        root2.addChild(container)
+        root2.view.addSubview(container.view)
+        container.didMove(toParent: root2)
+
+        hostingController.view.setNeedsLayout()
+        hostingController.view.layoutIfNeeded()
+
+        XCTAssertEqual(changedEnvironments.count, 3)
+
+        hostingController.view.setNeedsLayout()
+        hostingController.view.layoutIfNeeded()
+
+        XCTAssertEqual(changedEnvironments.count, 3)
+    }
 }
 
 fileprivate struct SubscribingWorkflow: Workflow {
@@ -303,7 +354,6 @@ fileprivate struct EnvironmentObservingWorkflow: Workflow {
 }
 
 fileprivate final class EnvironmentCustomizingViewController: UIViewController, ViewEnvironmentObserving {
-
     var customizeEnvironment: (inout ViewEnvironment) -> Void
 
     init(customizeEnvironment: @escaping (inout ViewEnvironment) -> Void) {