From 7a8f3198651ec48ae76aa2b5c011edd20cf16fcb Mon Sep 17 00:00:00 2001
From: Westin Newell <westin@squareup.com>
Date: Thu, 30 Mar 2023 10:37:34 -0700
Subject: [PATCH] Add bridging to ScreenViewController

---
 .../Sources/Screen/ScreenViewController.swift | 28 +++++++++++++------
 1 file changed, 20 insertions(+), 8 deletions(-)

diff --git a/WorkflowUI/Sources/Screen/ScreenViewController.swift b/WorkflowUI/Sources/Screen/ScreenViewController.swift
index 143c10cb8..52c7e9925 100644
--- a/WorkflowUI/Sources/Screen/ScreenViewController.swift
+++ b/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.
@@ -39,31 +41,41 @@ open class ScreenViewController<ScreenType: Screen>: UIViewController {
     public private(set) final var screen: ScreenType
 
     public final var screenType: Screen.Type {
-        return ScreenType.self
+        ScreenType.self
     }
 
-    public private(set) final var environment: ViewEnvironment
+    private var _environment: ViewEnvironment
 
     public required init(screen: ScreenType, environment: ViewEnvironment) {
         self.screen = screen
-        self.environment = environment
+        self._environment = environment
         super.init(nibName: nil, bundle: nil)
+
+        let ancestor = ViewEnvironmentPropagationNode(
+            environmentDescendants: { [weak self] in
+                [self].compactMap { $0 }
+            },
+            customizeEnvironment: { [weak self] environment in
+                guard let self else { return }
+                environment = self._environment
+            }
+        )
+
+        environmentAncestorOverride = { ancestor }
     }
 
     @available(*, unavailable)
-    public required init?(coder aDecoder: NSCoder) {
-        fatalError("init(coder:) has not been implemented")
-    }
+    public required init?(coder aDecoder: NSCoder) { fatalError() }
 
     public final func update(screen: ScreenType, environment: ViewEnvironment) {
         let previousScreen = self.screen
         self.screen = screen
         let previousEnvironment = self.environment
-        self.environment = environment
+        _environment = environment
+        setNeedsEnvironmentUpdate()
         screenDidChange(from: previousScreen, previousEnvironment: previousEnvironment)
     }
 
-    /// Subclasses should override this method in order to update any relevant UI bits when the screen model changes.
     open func screenDidChange(from previousScreen: ScreenType, previousEnvironment: ViewEnvironment) {}
 }