Fix filters in LocationCoordinator

5 files changed, 54 insertions, 68 deletions

diff --git a/ios/MullvadVPN/Containers/Root/RootContainerViewController.swift b/ios/MullvadVPN/Containers/Root/RootContainerViewController.swift
index 0d005d66b1..2c405395bc 100644
--- a/ios/MullvadVPN/Containers/Root/RootContainerViewController.swift
+++ b/ios/MullvadVPN/Containers/Root/RootContainerViewController.swift
@@ -308,40 +308,6 @@ class RootContainerViewController: UIViewController {
         )
     }
 
-    /// Add account and settings bar buttons into the presentation container to make them accessible even
-    /// when the root container is covered with a modal.
-    func addTrailingButtonsToPresentationContainer(_ presentationContainer: UIView) {
-        let accountButton = getPresentationContainerAccountButton()
-        let settingsButton = getPresentationContainerSettingsButton()
-
-        presentationContainerAccountButton = accountButton
-        presentationContainerSettingsButton = settingsButton
-
-        // Hide the account button inside the header bar to avoid color blending issues
-        headerBarView.accountButton.alpha = 0
-        headerBarView.settingsButton.alpha = 0
-
-        presentationContainer.addConstrainedSubviews([accountButton, settingsButton]) {
-            accountButton.centerXAnchor
-                .constraint(equalTo: headerBarView.accountButton.centerXAnchor)
-            accountButton.centerYAnchor
-                .constraint(equalTo: headerBarView.accountButton.centerYAnchor)
-
-            settingsButton.centerXAnchor
-                .constraint(equalTo: headerBarView.settingsButton.centerXAnchor)
-            settingsButton.centerYAnchor
-                .constraint(equalTo: headerBarView.settingsButton.centerYAnchor)
-        }
-    }
-
-    func removeTrailingButtonsFromPresentationContainer() {
-        presentationContainerAccountButton?.removeFromSuperview()
-        presentationContainerSettingsButton?.removeFromSuperview()
-
-        headerBarView.accountButton.alpha = 1
-        headerBarView.settingsButton.alpha = 1
-    }
-
     func setOverrideHeaderBarHidden(_ isHidden: Bool?, animated: Bool) {
         overrideHeaderBarHidden = isHidden
 
diff --git a/ios/MullvadVPN/Coordinators/LocationCoordinator.swift b/ios/MullvadVPN/Coordinators/LocationCoordinator.swift
index 222a20db1b..2f5b9f6472 100644
--- a/ios/MullvadVPN/Coordinators/LocationCoordinator.swift
+++ b/ios/MullvadVPN/Coordinators/LocationCoordinator.swift
@@ -16,7 +16,7 @@ class LocationCoordinator: Coordinator, Presentable, Presenting {
     private let tunnelManager: TunnelManager
     private let relayCacheTracker: RelayCacheTracker
     private let customListRepository: CustomListRepositoryProtocol
-    private var cachedRelays: LocationRelays?
+    private var locationRelays: LocationRelays?
 
     let navigationController: UINavigationController
 
@@ -71,22 +71,37 @@ class LocationCoordinator: Coordinator, Presentable, Presenting {
             guard let self else { return }
             didFinish?(self)
         }
-
         relayCacheTracker.addObserver(self)
 
         if let cachedRelays = try? relayCacheTracker.getCachedRelays() {
-            let locationRelays = LocationRelays(
-                relays: cachedRelays.relays.wireguard.relays,
-                locations: cachedRelays.relays.locations
+            updateRelaysWithLocationFrom(
+                cachedRelays: cachedRelays,
+                filter: relayFilter,
+                controllerWrapper: locationViewControllerWrapper
             )
-            self.cachedRelays = locationRelays
-
-            locationViewControllerWrapper.setCachedRelays(locationRelays, filter: relayFilter)
         }
 
         navigationController.pushViewController(locationViewControllerWrapper, animated: false)
     }
 
+    private func updateRelaysWithLocationFrom(
+        cachedRelays: CachedRelays,
+        filter: RelayFilter,
+        controllerWrapper: LocationViewControllerWrapper
+    ) {
+        var relaysWithLocation = LocationRelays(
+            relays: cachedRelays.relays.wireguard.relays,
+            locations: cachedRelays.relays.locations
+        )
+        relaysWithLocation.relays = relaysWithLocation.relays.filter { relay in
+            RelaySelector.relayMatchesFilter(relay, filter: filter)
+        }
+
+        self.locationRelays = relaysWithLocation
+
+        controllerWrapper.setRelaysWithLocation(relaysWithLocation, filter: filter)
+    }
+
     private func makeRelayFilterCoordinator(forModalPresentation isModalPresentation: Bool)
         -> RelayFilterCoordinator {
         let navigationController = CustomNavigationController()
@@ -100,12 +115,13 @@ class LocationCoordinator: Coordinator, Presentable, Presenting {
         relayFilterCoordinator.didFinish = { [weak self] coordinator, filter in
             guard let self else { return }
 
-            if var cachedRelays, let filter {
-                cachedRelays.relays = cachedRelays.relays.filter { relay in
-                    RelaySelector.relayMatchesFilter(relay, filter: filter)
-                }
-
-                locationViewControllerWrapper?.setCachedRelays(cachedRelays, filter: filter)
+            if let cachedRelays = try? relayCacheTracker.getCachedRelays(), let locationViewControllerWrapper,
+               let filter {
+                updateRelaysWithLocationFrom(
+                    cachedRelays: cachedRelays,
+                    filter: filter,
+                    controllerWrapper: locationViewControllerWrapper
+                )
             }
 
             coordinator.dismiss(animated: true)
@@ -169,9 +185,9 @@ extension LocationCoordinator: RelayCacheTrackerObserver {
             relays: cachedRelays.relays.wireguard.relays,
             locations: cachedRelays.relays.locations
         )
-        self.cachedRelays = locationRelays
+        self.locationRelays = locationRelays
 
-        locationViewControllerWrapper?.setCachedRelays(locationRelays, filter: relayFilter)
+        locationViewControllerWrapper?.setRelaysWithLocation(locationRelays, filter: relayFilter)
     }
 }
 
@@ -200,8 +216,12 @@ extension LocationCoordinator: LocationViewControllerWrapperDelegate {
 
         tunnelManager.updateSettings([.relayConstraints(relayConstraints)])
 
-        if let cachedRelays {
-            locationViewControllerWrapper?.setCachedRelays(cachedRelays, filter: filter)
+        if let cachedRelays = try? relayCacheTracker.getCachedRelays(), let locationViewControllerWrapper {
+            updateRelaysWithLocationFrom(
+                cachedRelays: cachedRelays,
+                filter: filter,
+                controllerWrapper: locationViewControllerWrapper
+            )
         }
     }
 
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
index 9fe7490d5f..60cd693a3c 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
@@ -54,14 +54,14 @@ final class LocationDataSource:
         defaultRowAnimation = .fade
     }
 
-    func setRelays(_ cachedRelays: LocationRelays, selectedRelays: RelaySelection) {
+    func setRelays(_ relaysWithLocation: LocationRelays, selectedRelays: RelaySelection) {
         let allLocationsDataSource =
             dataSources.first(where: { $0 is AllLocationDataSource }) as? AllLocationDataSource
 
         let customListsDataSource =
             dataSources.first(where: { $0 is CustomListsDataSource }) as? CustomListsDataSource
 
-        allLocationsDataSource?.reload(cachedRelays)
+        allLocationsDataSource?.reload(relaysWithLocation)
         customListsDataSource?.reload(allLocationNodes: allLocationsDataSource?.nodes ?? [])
 
         setSelectedRelays(selectedRelays)
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationViewController.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationViewController.swift
index 425dc954c5..587f5ab146 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationViewController.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationViewController.swift
@@ -23,7 +23,7 @@ final class LocationViewController: UIViewController {
     private let topContentView = UIStackView()
     private let filterView = RelayFilterView()
     private var dataSource: LocationDataSource?
-    private var cachedRelays: LocationRelays?
+    private var relaysWithLocation: LocationRelays?
     private var filter = RelayFilter()
     private var selectedRelays: RelaySelection
     private var daitaEnabled: Bool
@@ -86,8 +86,8 @@ final class LocationViewController: UIViewController {
 
     // MARK: - Public
 
-    func setCachedRelays(_ cachedRelays: LocationRelays, filter: RelayFilter) {
-        self.cachedRelays = cachedRelays
+    func setRelaysWithLocation(_ relaysWithLocation: LocationRelays, filter: RelayFilter) {
+        self.relaysWithLocation = relaysWithLocation
         self.filter = filter
 
         if filterViewShouldBeHidden {
@@ -97,7 +97,7 @@ final class LocationViewController: UIViewController {
             filterView.setFilter(filter)
         }
 
-        dataSource?.setRelays(cachedRelays, selectedRelays: selectedRelays)
+        dataSource?.setRelays(relaysWithLocation, selectedRelays: selectedRelays)
     }
 
     func refreshCustomLists() {
@@ -125,15 +125,15 @@ final class LocationViewController: UIViewController {
         dataSource?.didTapEditCustomLists = { [weak self] in
             guard let self else { return }
 
-            if let cachedRelays {
+            if let relaysWithLocation {
                 let allLocationDataSource = AllLocationDataSource()
-                allLocationDataSource.reload(cachedRelays)
+                allLocationDataSource.reload(relaysWithLocation)
                 delegate?.navigateToCustomLists(nodes: allLocationDataSource.nodes)
             }
         }
 
-        if let cachedRelays {
-            dataSource?.setRelays(cachedRelays, selectedRelays: selectedRelays)
+        if let relaysWithLocation {
+            dataSource?.setRelays(relaysWithLocation, selectedRelays: selectedRelays)
         }
     }
 
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationViewControllerWrapper.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationViewControllerWrapper.swift
index ae1f12aead..4f8147f746 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationViewControllerWrapper.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationViewControllerWrapper.swift
@@ -107,19 +107,19 @@ final class LocationViewControllerWrapper: UIViewController {
         swapViewController()
     }
 
-    func setCachedRelays(_ cachedRelays: LocationRelays, filter: RelayFilter) {
-        var daitaFilteredRelays = cachedRelays
+    func setRelaysWithLocation(_ relaysWithLocation: LocationRelays, filter: RelayFilter) {
+        var daitaFilteredRelays = relaysWithLocation
         if daitaEnabled {
-            daitaFilteredRelays.relays = cachedRelays.relays.filter { relay in
+            daitaFilteredRelays.relays = relaysWithLocation.relays.filter { relay in
                 relay.daita == true
             }
         }
 
         if multihopEnabled {
-            entryLocationViewController.setCachedRelays(daitaFilteredRelays, filter: filter)
-            exitLocationViewController.setCachedRelays(cachedRelays, filter: filter)
+            entryLocationViewController.setRelaysWithLocation(daitaFilteredRelays, filter: filter)
+            exitLocationViewController.setRelaysWithLocation(relaysWithLocation, filter: filter)
         } else {
-            exitLocationViewController.setCachedRelays(daitaFilteredRelays, filter: filter)
+            exitLocationViewController.setRelaysWithLocation(daitaFilteredRelays, filter: filter)
         }
     }