From 4c21ed7b6f47584b737df072173ef943eba4e15c Mon Sep 17 00:00:00 2001
From: Nick Porter <porter@stripe.com>
Date: Mon, 3 Feb 2025 13:23:37 -0700
Subject: [PATCH] Consolidate and improve UIViewController traversing

---
 .../Embedded/EmbeddedViewRepresentable.swift  | 23 +++----------------
 .../BottomSheetViewController.swift           |  2 +-
 .../UIViewController+StripeUICore.swift       | 23 +++++++++++--------
 3 files changed, 18 insertions(+), 30 deletions(-)

diff --git a/StripePaymentSheet/StripePaymentSheet/Source/PaymentSheet/Embedded/EmbeddedViewRepresentable.swift b/StripePaymentSheet/StripePaymentSheet/Source/PaymentSheet/Embedded/EmbeddedViewRepresentable.swift
index 5804d3af31c..549780cad77 100644
--- a/StripePaymentSheet/StripePaymentSheet/Source/PaymentSheet/Embedded/EmbeddedViewRepresentable.swift
+++ b/StripePaymentSheet/StripePaymentSheet/Source/PaymentSheet/Embedded/EmbeddedViewRepresentable.swift
@@ -7,6 +7,7 @@
 
 import SwiftUI
 @_spi(STP) import StripeCore
+@_spi(STP) import StripeUICore
 
 struct EmbeddedViewRepresentable: UIViewRepresentable {
     @ObservedObject var viewModel: EmbeddedPaymentElementViewModel
@@ -43,26 +44,8 @@ struct EmbeddedViewRepresentable: UIViewRepresentable {
 
 // MARK: UIWindow and UIViewController helpers
 
-extension UIWindow {    
+extension UIWindow {
     static var visibleViewController: UIViewController? {
-        UIApplication.shared.stp_hackilyFumbleAroundUntilYouFindAKeyWindow()?.rootViewController?.topMostViewController
+        UIApplication.shared.stp_hackilyFumbleAroundUntilYouFindAKeyWindow()?.rootViewController?.findTopMostPresentedViewController()
     }
 }
-
-extension UIViewController {
-    var topMostViewController: UIViewController {
-        if let nav = self as? UINavigationController {
-            // Use visibleViewController for navigation stacks
-            return nav.visibleViewController?.topMostViewController ?? nav
-        } else if let tab = self as? UITabBarController {
-            // Use selectedViewController for tab controllers
-            return tab.selectedViewController?.topMostViewController ?? tab
-        } else if let presented = presentedViewController {
-            // Recurse for any presented controllers
-            return presented.topMostViewController
-        }
-        
-        return self
-    }
-}
-
diff --git a/StripePaymentSheet/StripePaymentSheet/Source/PaymentSheet/ViewControllers/BottomSheetViewController.swift b/StripePaymentSheet/StripePaymentSheet/Source/PaymentSheet/ViewControllers/BottomSheetViewController.swift
index d98b16f8c17..dd38707483b 100644
--- a/StripePaymentSheet/StripePaymentSheet/Source/PaymentSheet/ViewControllers/BottomSheetViewController.swift
+++ b/StripePaymentSheet/StripePaymentSheet/Source/PaymentSheet/ViewControllers/BottomSheetViewController.swift
@@ -478,7 +478,7 @@ extension BottomSheetViewController: UIScrollViewDelegate {
 extension BottomSheetViewController: PaymentSheetAuthenticationContext {
 
     func authenticationPresentingViewController() -> UIViewController {
-        return findTopMostPresentedViewController() ?? self
+        return findTopMostPresentedViewController()
     }
 
     func configureSafariViewController(_ viewController: SFSafariViewController) {
diff --git a/StripeUICore/StripeUICore/Source/Categories/UIViewController+StripeUICore.swift b/StripeUICore/StripeUICore/Source/Categories/UIViewController+StripeUICore.swift
index 0b83259d59e..ca031dc1d4b 100644
--- a/StripeUICore/StripeUICore/Source/Categories/UIViewController+StripeUICore.swift
+++ b/StripeUICore/StripeUICore/Source/Categories/UIViewController+StripeUICore.swift
@@ -40,15 +40,20 @@ import UIKit
         return self
     }
 
-    /// Walks the presented view controller hierarchy and return the top most presented controller.
-    /// - Returns: Returns the top most presented view controller, or `nil` if this view controller is not presenting another controller.
-    func findTopMostPresentedViewController() -> UIViewController? {
-        var topMostController = self.presentedViewController
-
-        while let presented = topMostController?.presentedViewController {
-            topMostController = presented
+    /// Returns the topmost view controller in the hierarchy.
+    /// - Returns: The topmost `UIViewController`, or `self` if no higher controller is found.
+    func findTopMostPresentedViewController() -> UIViewController {
+        if let nav = self as? UINavigationController {
+            // Use visibleViewController for navigation stacks
+            return nav.visibleViewController?.findTopMostPresentedViewController() ?? nav
+        } else if let tab = self as? UITabBarController {
+            // Use selectedViewController for tab controllers
+            return tab.selectedViewController?.findTopMostPresentedViewController() ?? tab
+        } else if let presented = presentedViewController {
+            // Recurse for any presented controllers
+            return presented.findTopMostPresentedViewController()
         }
-
-        return topMostController
+        
+        return self
     }
 }