Embedded update pt 3 (#4168)

## Summary
Previous PR: #4150
Makes EmbeddedPaymentElement view a view that contains embedded view so
we can swap to the updated embedded view with an animation.

Makes EmbeddedPaymentElement restore previous customer in **_payment
method list selection_ (1)** on `update`, call `didUpdateHeight` and
`didUpdatePaymentOption` if necessary.

Still to come:
- Restore previous customer input in ***form (2)*** + E2E tests (e.g.
load PI -> fill out card form -> update to SI -> expect form to be
preserved but w/o checkbox)
- Make confirm handle in-flight and failed update calls.
- (Bonus) Cancel network calls etc. from previous update to reduce
battery/network usage. Can apply this to FC.update as well.

## Motivation
https://jira.corp.stripe.com/browse/MOBILESDK-2583

## Testing
See test

## Changelog
Not user facing

StripePaymentSheet/StripePaymentSheet/Source/Helpers/DownloadManager.swift

@@ -10,7 +10,7 @@ import UIKit
 
 /// For internal SDK use only.
 @objc(STP_Internal_DownloadManager)
-// TODO: Refactor this API shape! https://github.com/stripe/stripe-ios/pull/3487#discussion_r1561337866
+// TODO: https://jira.corp.stripe.com/browse/MOBILESDK-2604 Refactor this!
 @_spi(STP) public class DownloadManager: NSObject {
     public typealias UpdateImageHandler = (UIImage) -> Void
 

StripePaymentSheet/StripePaymentSheet/Source/PaymentSheet/Embedded/EmbeddedPaymentElement+Internal.swift

@@ -4,15 +4,27 @@
 //
 //  Created by Yuki Tokuhiro on 10/10/24.
 //
+@_spi(STP) import StripeCore
+@_spi(STP) import StripePayments
+@_spi(STP) import StripePaymentsUI
+@_spi(STP) import StripeUICore
+import UIKit
 
 extension EmbeddedPaymentElement {
     @MainActor
     static func makeView(
         configuration: Configuration,
         loadResult: PaymentSheetLoader.LoadResult,
         analyticsHelper: PaymentSheetAnalyticsHelper,
+        previousPaymentOption: PaymentOption? = nil,
         delegate: EmbeddedPaymentMethodsViewDelegate? = nil
     ) -> EmbeddedPaymentMethodsView {
+        // Restore the customer's previous payment method.
+        // Caveats:
+        // - Only payment method details (including checkbox state) and billing details are restored
+        // - Only restored if the previous input resulted in a completed form i.e. partial or invalid input is still discarded
+        // TODO: Restore the form, if any
+
         let shouldShowApplePay = PaymentSheet.isApplePayEnabled(elementsSession: loadResult.elementsSession, configuration: configuration)
         let shouldShowLink = PaymentSheet.isLinkEnabled(elementsSession: loadResult.elementsSession, configuration: configuration)
         let savedPaymentMethodAccessoryType = RowButton.RightAccessoryButton.getAccessoryButtonType(
@@ -23,7 +35,23 @@ extension EmbeddedPaymentElement {
             allowsPaymentMethodRemoval: loadResult.elementsSession.allowsRemovalOfPaymentMethodsForPaymentSheet()
         )
         let initialSelection: EmbeddedPaymentMethodsView.Selection? = {
-            // Default to the customer's default or the first saved payment method, if any
+            // Select the previous payment option
+            switch previousPaymentOption {
+            case .applePay:
+                return .applePay
+            case .link:
+                return .link
+            case .external(paymentMethod: let paymentMethod, billingDetails: _):
+                return .new(paymentMethodType: .external(paymentMethod))
+            case .saved(paymentMethod: let paymentMethod, confirmParams: _):
+                return .saved(paymentMethod: paymentMethod)
+            case .new(confirmParams: let confirmParams):
+                return .new(paymentMethodType: confirmParams.paymentMethodType)
+            case nil:
+                break
+            }
+
+            // If there's no previous customer input, default to the customer's default or the first saved payment method, if any
             let customerDefault = CustomerPaymentOption.defaultPaymentMethod(for: configuration.customer?.id)
             switch customerDefault {
             case .applePay:
@@ -66,3 +94,34 @@ extension EmbeddedPaymentElement: EmbeddedPaymentMethodsViewDelegate {
         delegate?.embeddedPaymentElementDidUpdatePaymentOption(embeddedPaymentElement: self)
     }
 }
+
+// MARK: - EmbeddedPaymentElement.PaymentOptionDisplayData
+
+extension EmbeddedPaymentElement.PaymentOptionDisplayData {
+    init(paymentOption: PaymentOption, mandateText: NSAttributedString?) {
+        self.mandateText = mandateText
+        self.image = paymentOption.makeIcon(updateImageHandler: nil) // ☠️ This can make a blocking network request TODO: https://jira.corp.stripe.com/browse/MOBILESDK-2604 Refactor this!
+        switch paymentOption {
+        case .applePay:
+            label = String.Localized.apple_pay
+            paymentMethodType = "apple_pay"
+            billingDetails = nil
+        case .saved(let paymentMethod, _):
+            label = paymentMethod.paymentSheetLabel
+            paymentMethodType = paymentMethod.type.identifier
+            billingDetails = paymentMethod.billingDetails?.toPaymentSheetBillingDetails()
+        case .new(let confirmParams):
+            label = confirmParams.paymentSheetLabel
+            paymentMethodType = confirmParams.paymentMethodType.identifier
+            billingDetails = confirmParams.paymentMethodParams.billingDetails?.toPaymentSheetBillingDetails()
+        case .link(let option):
+            label = option.paymentSheetLabel
+            paymentMethodType = STPPaymentMethodType.link.identifier
+            billingDetails = option.billingDetails?.toPaymentSheetBillingDetails()
+        case .external(let paymentMethod, let stpBillingDetails):
+            label = paymentMethod.label
+            paymentMethodType = paymentMethod.type
+            billingDetails = stpBillingDetails.toPaymentSheetBillingDetails()
+        }
+    }
+}

StripePaymentSheet/StripePaymentSheet/Source/PaymentSheet/Embedded/EmbeddedPaymentElement.swift

@@ -44,12 +44,16 @@ public final class EmbeddedPaymentElement {
         public let paymentMethodType: String
         /// If you set `configuration.embeddedViewDisplaysMandateText = false`, this text must be displayed to the customer near your “Buy” button to comply with regulations.
         public let mandateText: NSAttributedString?
+
     }
 
     /// Contains information about the customer's selected payment option.
     /// Use this to display the payment option in your own UI
     public var paymentOption: PaymentOptionDisplayData? {
-        return embeddedPaymentMethodsView.displayData
+        guard let _paymentOption else {
+            return nil
+        }
+        return .init(paymentOption: _paymentOption, mandateText: embeddedPaymentMethodsView.mandateText)
     }
 
     /// An asynchronous failable initializer
@@ -103,7 +107,8 @@ public final class EmbeddedPaymentElement {
         currentUpdateTask?.cancel()
         _ = await currentUpdateTask?.value
         // Start the new update task
-        let currentUpdateTask = Task { [weak self, configuration, paymentOption, analyticsHelper] in
+        let currentUpdateTask = Task { @MainActor [weak self, configuration, analyticsHelper] in
+            // ⚠️ Don't modify `self` until the end to avoid being canceled halfway through and leaving self in a partially updated state.
             // 1. Reload v1/elements/session.
             let loadResult: PaymentSheetLoader.LoadResult
             do {
@@ -126,26 +131,28 @@ public final class EmbeddedPaymentElement {
                 configuration: configuration,
                 loadResult: loadResult,
                 analyticsHelper: analyticsHelper,
+                previousPaymentOption: self?._paymentOption,
                 delegate: self
-                // TODO: https://jira.corp.stripe.com/browse/MOBILESDK-2583 Restore previous payment option
             )
 
-            // 2. Pre-load image into cache
-            // Hack: Accessing paymentOption has the side-effect of ensuring its `image` property is loaded (from the internet instead of disk) before we call the completion handler.
+            // 3. Pre-load image into cache
             // Call this on a detached Task b/c this synchronously (!) loads the image from network and we don't want to block the main actor
             let fetchPaymentOption = Task.detached(priority: .userInitiated) {
-                return await embeddedPaymentMethodsView.displayData
+                // This has the nasty side effect of synchronously downloading the image (see https://jira.corp.stripe.com/browse/MOBILESDK-2604)
+                // This caches it so that DownloadManager doesn't block the main thread when the merchant tries to access the image
+                return await embeddedPaymentMethodsView.selection?.paymentMethodType?.makeImage(updateHandler: nil)
             }
             _ = await fetchPaymentOption.value
 
             guard let self, !Task.isCancelled else {
                 return .canceled
             }
-            // At this point, we're the latest update - update self properties and inform our delegate.
+            // At this point, we're still the latest update and update is successful - update self properties and inform our delegate.
+            let oldPaymentOption = self.paymentOption
             self.loadResult = loadResult
             self.embeddedPaymentMethodsView = embeddedPaymentMethodsView
             self.containerView.updateEmbeddedPaymentMethodsView(embeddedPaymentMethodsView)
-            if paymentOption != embeddedPaymentMethodsView.displayData {
+            if oldPaymentOption != self.paymentOption {
                 self.delegate?.embeddedPaymentElementDidUpdatePaymentOption(embeddedPaymentElement: self)
             }
             return .succeeded
@@ -169,6 +176,31 @@ public final class EmbeddedPaymentElement {
     internal private(set) var loadResult: PaymentSheetLoader.LoadResult
     internal private(set) var currentUpdateTask: Task<UpdateResult, Never>?
     private let analyticsHelper: PaymentSheetAnalyticsHelper
+    internal var _paymentOption: PaymentOption? {
+        // TODO: Handle forms. See `PaymentSheetVerticalViewController.selectedPaymentOption`.
+        // TODO: Handle CVC recollection
+        switch embeddedPaymentMethodsView.selection {
+        case .applePay:
+            return .applePay
+        case .link:
+            return .link(option: .wallet)
+        case let .new(paymentMethodType: paymentMethodType):
+            let params = IntentConfirmParams(type: paymentMethodType)
+            params.setDefaultBillingDetailsIfNecessary(for: configuration)
+            switch paymentMethodType {
+            case .stripe:
+                return .new(confirmParams: params)
+            case .external(let type):
+                return .external(paymentMethod: type, billingDetails: params.paymentMethodParams.nonnil_billingDetails)
+            case .instantDebits, .linkCardBrand:
+                return .new(confirmParams: params)
+            }
+        case .saved(paymentMethod: let paymentMethod):
+            return .saved(paymentMethod: paymentMethod, confirmParams: nil)
+        case .none:
+            return nil
+        }
+    }
 
     private init(
         configuration: Configuration,
@@ -278,3 +310,12 @@ extension EmbeddedPaymentElement {
     public typealias BillingDetailsCollectionConfiguration = PaymentSheet.BillingDetailsCollectionConfiguration
     public typealias ExternalPaymentMethodConfiguration = PaymentSheet.ExternalPaymentMethodConfiguration
 }
+
+// MARK: - EmbeddedPaymentElement.PaymentOptionDisplayData
+
+extension EmbeddedPaymentElement.PaymentOptionDisplayData {
+    public static func == (lhs: Self, rhs: Self) -> Bool {
+        // Unfortunately, we need to manually define this because the implementation of Equatable on UIImage does not work
+        return lhs.image.pngData() == rhs.image.pngData() && rhs.label == lhs.label && lhs.billingDetails == rhs.billingDetails && lhs.paymentMethodType == rhs.paymentMethodType && lhs.mandateText == rhs.mandateText
+    }
+}

StripePaymentSheet/StripePaymentSheet/Source/PaymentSheet/Embedded/EmbeddedPaymentMethodsView.swift

@@ -6,7 +6,6 @@
 //
 
 import Foundation
-@_spi(STP) import StripeCore
 @_spi(STP) import StripeUICore
 import UIKit
 
@@ -21,11 +20,6 @@ class EmbeddedPaymentMethodsView: UIView {
 
     typealias Selection = VerticalPaymentMethodListSelection // TODO(porter) Maybe define our own later
 
-    var displayData: EmbeddedPaymentElement.PaymentOptionDisplayData? {
-        guard let selection else { return nil }
-        return .init(selection: selection, mandateText: mandateView.attributedText)
-    }
-
     private let appearance: PaymentSheet.Appearance
     private(set) var selection: Selection? {
         didSet {
@@ -37,8 +31,13 @@ class EmbeddedPaymentMethodsView: UIView {
     }
     private let mandateProvider: MandateTextProvider
     private let shouldShowMandate: Bool
+    /// A bit hacky; this is the mandate text for the given payment method, *regardless* of whether it is shown in the view.
+    /// It'd be better if the source of truth of mandate text was not the view and instead an independent `func mandateText(...) -> NSAttributedString` function, but this is hard b/c US Bank Account doesn't show mandate in certain states.
+    var mandateText: NSAttributedString? {
+        mandateView.attributedText
+    }
 
-    lazy var stackView: UIStackView = {
+    private(set) lazy var stackView: UIStackView = {
         let stackView = UIStackView()
         stackView.axis = .vertical
         stackView.spacing = appearance.embeddedPaymentElement.style == .floatingButton ? appearance.embeddedPaymentElement.row.floating.spacing : 0
@@ -136,16 +135,24 @@ class EmbeddedPaymentMethodsView: UIView {
             stackView.addArrangedSubview(linkRowButton)
         }
 
+        // Add all non-card PMs (card is added above)
         for paymentMethodType in paymentMethodTypes where paymentMethodType != .stripe(.card) {
-            stackView.addArrangedSubview(RowButton.makeForPaymentMethodType(paymentMethodType: paymentMethodType,
-                                                                            subtitle: VerticalPaymentMethodListViewController.subtitleText(for: paymentMethodType),
-                                                                            savedPaymentMethodType: savedPaymentMethod?.type,
-                                                                            appearance: rowButtonAppearance,
-                                                                            shouldAnimateOnPress: true,
-                                                                            isEmbedded: true,
-                                                                            didTap: { [weak self] rowButton in
-                self?.didTap(selectedRowButton: rowButton, selection: .new(paymentMethodType: paymentMethodType))
-            }))
+            let selection: Selection = .new(paymentMethodType: paymentMethodType)
+            let rowButton = RowButton.makeForPaymentMethodType(
+                paymentMethodType: paymentMethodType,
+                subtitle: VerticalPaymentMethodListViewController.subtitleText(for: paymentMethodType),
+                savedPaymentMethodType: savedPaymentMethod?.type,
+                appearance: rowButtonAppearance,
+                shouldAnimateOnPress: true,
+                isEmbedded: true,
+                didTap: { [weak self] rowButton in
+                    self?.didTap(selectedRowButton: rowButton, selection: selection)
+                }
+            )
+            if initialSelection == selection {
+                rowButton.isSelected = true
+            }
+            stackView.addArrangedSubview(rowButton)
         }
 
         if appearance.embeddedPaymentElement.style != .floatingButton {
@@ -249,37 +256,3 @@ extension PaymentSheet.Appearance.EmbeddedPaymentElement.Style {
         }
     }
 }
-@_spi(STP) import StripePayments
-@_spi(STP) import StripePaymentsUI
-
-extension EmbeddedPaymentElement.PaymentOptionDisplayData {
-    init(selection: EmbeddedPaymentMethodsView.Selection, mandateText: NSAttributedString?) {
-        self.mandateText = mandateText
-
-        switch selection {
-        case .new(paymentMethodType: let paymentMethodType):
-            image = paymentMethodType.makeImage(
-                forDarkBackground: UITraitCollection.current.isDarkMode,
-                updateHandler: nil
-            )
-            label = paymentMethodType.displayName
-            self.paymentMethodType = paymentMethodType.identifier
-            billingDetails = nil // TODO(porter) Handle billing details when we present forms (maybe set this to defaultBillingDetails) if billingDetailsConfiguration.attachDefaultsToPaymentMethod is true
-        case .saved(paymentMethod: let paymentMethod):
-            image = paymentMethod.makeIcon()
-            label = paymentMethod.paymentSheetLabel
-            paymentMethodType = paymentMethod.type.identifier
-            billingDetails = paymentMethod.billingDetails?.toPaymentSheetBillingDetails()
-        case .applePay:
-            image = Image.apple_pay_mark.makeImage().withRenderingMode(.alwaysOriginal)
-            label = .Localized.apple_pay
-            paymentMethodType = "apple_pay"
-            billingDetails = nil // TODO(porter) Handle billing details when we present forms
-        case .link:
-            image = Image.link_logo.makeImage()
-            label = STPPaymentMethodType.link.displayName
-            paymentMethodType = STPPaymentMethodType.link.identifier
-            billingDetails = nil // TODO(porter) Handle billing details when we present forms
-        }
-    }
-}