[fix] View environment customizations on hosting VCs now propagate to…

… SwiftUI environment (#297)

* ModeledHostingController implementations now conform to ViewEnvironmentObserving

* Make StateAccessor init public

WorkflowSwiftUI/Sources/ObservableScreen.swift

@@ -56,7 +56,7 @@ public extension ObservableScreen {
             },
             update: { hostingController in
                 hostingController.setModel(model)
-                hostingController.setViewEnvironment(environment)
+                // ViewEnvironment updates are handled by the ModeledHostingController internally
             }
         )
     }
@@ -89,9 +89,10 @@ private final class ViewEnvironmentHolder: ObservableObject {
     }
 }
 
-private final class ModeledHostingController<Model, Content: View>: UIHostingController<ModifiedContent<Content, ViewEnvironmentModifier>> {
+private final class ModeledHostingController<Model, Content: View>: UIHostingController<ModifiedContent<Content, ViewEnvironmentModifier>>, ViewEnvironmentObserving {
     let setModel: (Model) -> Void
-    let setViewEnvironment: (ViewEnvironment) -> Void
+
+    private let viewEnvironmentHolder: ViewEnvironmentHolder
 
     var swiftUIScreenSizingOptions: SwiftUIScreenSizingOptions {
         didSet {
@@ -111,10 +112,8 @@ private final class ModeledHostingController<Model, Content: View>: UIHostingCon
         rootView: Content,
         sizingOptions swiftUIScreenSizingOptions: SwiftUIScreenSizingOptions
     ) {
-        let viewEnvironmentHolder = ViewEnvironmentHolder(viewEnvironment: viewEnvironment)
-
         self.setModel = setModel
-        self.setViewEnvironment = { viewEnvironmentHolder.viewEnvironment = $0 }
+        self.viewEnvironmentHolder = ViewEnvironmentHolder(viewEnvironment: viewEnvironment)
         self.swiftUIScreenSizingOptions = swiftUIScreenSizingOptions
 
         super.init(
@@ -169,6 +168,12 @@ private final class ModeledHostingController<Model, Content: View>: UIHostingCon
         }
     }
 
+    override func viewWillLayoutSubviews() {
+        super.viewWillLayoutSubviews()
+
+        applyEnvironmentIfNeeded()
+    }
+
     private func updateSizingOptionsIfNeeded() {
         if #available(iOS 16.0, *) {
             self.sizingOptions = swiftUIScreenSizingOptions.uiHostingControllerSizingOptions
@@ -190,6 +195,12 @@ private final class ModeledHostingController<Model, Content: View>: UIHostingCon
             view.setNeedsLayout()
         }
     }
+
+    // MARK: ViewEnvironmentObserving
+
+    func apply(environment: ViewEnvironment) {
+        viewEnvironmentHolder.viewEnvironment = environment
+    }
 }
 
 fileprivate extension SwiftUIScreenSizingOptions {

WorkflowSwiftUI/Sources/StateAccessor.swift

@@ -12,6 +12,14 @@
 public struct StateAccessor<State: ObservableState> {
     let state: State
     let sendValue: (@escaping (inout State) -> Void) -> Void
+
+    public init(
+        state: State,
+        sendValue: @escaping (@escaping (inout State) -> Void) -> Void
+    ) {
+        self.state = state
+        self.sendValue = sendValue
+    }
 }
 
 extension StateAccessor: ObservableModel {

WorkflowSwiftUI/Tests/ObservableScreenTests.swift

@@ -0,0 +1,80 @@
+#if canImport(UIKit)
+
+import SwiftUI
+import ViewEnvironment
+@_spi(ViewEnvironmentWiring) import ViewEnvironmentUI
+import WorkflowSwiftUI
+import XCTest
+
+final class ObservableScreenTests: XCTestCase {
+    func test_viewEnvironmentObservation() {
+        // Ensure that environment customizations made on the view controller
+        // are propagated to the SwiftUI view environment.
+
+        var state = MyState()
+
+        let viewController = TestKeyEmittingScreen(
+            model: MyModel(
+                accessor: StateAccessor(
+                    state: state,
+                    sendValue: { $0(&state) }
+                )
+            )
+        )
+        .buildViewController(in: .empty)
+
+        let lifetime = viewController.addEnvironmentCustomization { environment in
+            environment[TestKey.self] = 1
+        }
+
+        viewController.view.layoutIfNeeded()
+
+        XCTAssertEqual(state.emittedValue, 1)
+
+        withExtendedLifetime(lifetime) {}
+    }
+}
+
+private struct TestKey: ViewEnvironmentKey {
+    static var defaultValue: Int = 0
+}
+
+@ObservableState
+private struct MyState {
+    var emittedValue: TestKey.Value?
+}
+
+private struct MyModel: ObservableModel {
+    typealias State = MyState
+
+    let accessor: StateAccessor<State>
+}
+
+private struct TestKeyEmittingScreen: ObservableScreen {
+    typealias Model = MyModel
+
+    var model: Model
+
+    let sizingOptions: WorkflowSwiftUI.SwiftUIScreenSizingOptions = [.preferredContentSize]
+
+    static func makeView(store: Store<Model>) -> some View {
+        ContentView(store: store)
+    }
+
+    struct ContentView: View {
+        @Environment(\.viewEnvironment)
+        var viewEnvironment: ViewEnvironment
+
+        var store: Store<Model>
+
+        var body: some View {
+            WithPerceptionTracking {
+                let _ = { store.emittedValue = viewEnvironment[TestKey.self] }()
+                Color.clear
+                    .frame(width: 1, height: 1)
+            }
+        }
+    }
+}
+
+#endif

WorkflowSwiftUIExperimental/Sources/SwiftUIScreen.swift

@@ -65,8 +65,8 @@ public extension SwiftUIScreen {
             },
             update: {
                 $0.modelSink.send(self)
-                $0.viewEnvironmentSink.send(environment)
                 $0.swiftUIScreenSizingOptions = sizingOptions
+                // ViewEnvironment updates are handled by the ModeledHostingController internally
             }
         )
     }
@@ -92,7 +92,7 @@ private struct EnvironmentInjectingView<Content: View>: View {
     }
 }
 
-private final class ModeledHostingController<Model, Content: View>: UIHostingController<Content> {
+private final class ModeledHostingController<Model, Content: View>: UIHostingController<Content>, ViewEnvironmentObserving {
     let modelSink: Sink<Model>
     let viewEnvironmentSink: Sink<ViewEnvironment>
     var swiftUIScreenSizingOptions: SwiftUIScreenSizingOptions {
@@ -122,6 +122,7 @@ private final class ModeledHostingController<Model, Content: View>: UIHostingCon
         updateSizingOptionsIfNeeded()
     }
 
+    @available(*, unavailable)
     required init?(coder aDecoder: NSCoder) {
         fatalError("not implemented")
     }
@@ -148,7 +149,8 @@ private final class ModeledHostingController<Model, Content: View>: UIHostingCon
             // not updated appropriately after the first layout.
             // UI-5797
             if !hasLaidOutOnce,
-                swiftUIScreenSizingOptions.contains(.preferredContentSize) {
+               swiftUIScreenSizingOptions.contains(.preferredContentSize)
+            {
                 let size = view.sizeThatFits(view.frame.size)
 
                 if preferredContentSize != size {
@@ -164,13 +166,20 @@ private final class ModeledHostingController<Model, Content: View>: UIHostingCon
         }
     }
 
+    override func viewWillLayoutSubviews() {
+        super.viewWillLayoutSubviews()
+
+        applyEnvironmentIfNeeded()
+    }
+
     private func updateSizingOptionsIfNeeded() {
         if #available(iOS 16.0, *) {
             self.sizingOptions = swiftUIScreenSizingOptions.uiHostingControllerSizingOptions
         }
 
         if !swiftUIScreenSizingOptions.contains(.preferredContentSize),
-            preferredContentSize != .zero {
+           preferredContentSize != .zero
+        {
             preferredContentSize = .zero
         }
     }
@@ -184,11 +193,17 @@ private final class ModeledHostingController<Model, Content: View>: UIHostingCon
             view.setNeedsLayout()
         }
     }
+
+    // MARK: ViewEnvironmentObserving
+
+    func apply(environment: ViewEnvironment) {
+        viewEnvironmentSink.send(environment)
+    }
 }
 
-extension SwiftUIScreenSizingOptions {
+fileprivate extension SwiftUIScreenSizingOptions {
     @available(iOS 16.0, *)
-    fileprivate var uiHostingControllerSizingOptions: UIHostingControllerSizingOptions {
+    var uiHostingControllerSizingOptions: UIHostingControllerSizingOptions {
         var options = UIHostingControllerSizingOptions()
 
         if contains(.preferredContentSize) {

WorkflowSwiftUIExperimental/Tests/SwiftUIScreenTests.swift

@@ -2,6 +2,8 @@
 
 import SwiftUI
 import UIKit
+import ViewEnvironment
+@_spi(ViewEnvironmentWiring) import ViewEnvironmentUI
 import WorkflowSwiftUIExperimental
 import XCTest
 
@@ -54,6 +56,32 @@ final class SwiftUIScreenTests: XCTestCase {
 
         XCTAssertEqual(viewController.preferredContentSize, .zero)
     }
+
+    func test_viewEnvironmentObservation() {
+        // Ensure that environment customizations made on the view controller
+        // are propagated to the SwiftUI view environment.
+
+        var emittedValue: Int?
+
+        let viewController = TestKeyEmittingScreen(onTestKeyEmission: { value in
+            emittedValue = value
+        })
+        .buildViewController(in: .empty)
+
+        let lifetime = viewController.addEnvironmentCustomization { environment in
+            environment[TestKey.self] = 1
+        }
+
+        viewController.view.layoutIfNeeded()
+
+        XCTAssertEqual(emittedValue, 1)
+
+        withExtendedLifetime(lifetime) {}
+    }
+}
+
+private struct TestKey: ViewEnvironmentKey {
+    static var defaultValue: Int = 0
 }
 
 private struct ContentScreen: SwiftUIScreen {
@@ -65,4 +93,27 @@ private struct ContentScreen: SwiftUIScreen {
     }
 }
 
+private struct TestKeyEmittingScreen: SwiftUIScreen {
+    var onTestKeyEmission: (TestKey.Value) -> Void
+
+    let sizingOptions: SwiftUIScreenSizingOptions = [.preferredContentSize]
+
+    static func makeView(model: ObservableValue<Self>) -> some View {
+        ContentView(onTestKeyEmission: model.onTestKeyEmission)
+    }
+
+    struct ContentView: View {
+        @Environment(\.viewEnvironment)
+        var viewEnvironment: ViewEnvironment
+
+        var onTestKeyEmission: (TestKey.Value) -> Void
+
+        var body: some View {
+            let _ = onTestKeyEmission(viewEnvironment[TestKey.self])
+            Color.clear
+                .frame(width: 1, height: 1)
+        }
+    }
+}
+
 #endif