Fix iPad location selection entries not being unmarked as selected

2 files changed, 22 insertions, 8 deletions

diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
index 3272f0e65b..d6460343e0 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
@@ -136,6 +136,12 @@ final class LocationDataSource:
         ], reloadExisting: true)
     }
 
+    func scrollToSelectedRelay() {
+        indexPathForSelectedRelay().flatMap {
+            tableView.scrollToRow(at: $0, at: .middle, animated: false)
+        }
+    }
+
     // Called from `LocationDiffableDataSourceProtocol`.
     func nodeShowsChildren(_ node: LocationNode) -> Bool {
         node.showsChildren
@@ -277,12 +283,22 @@ extension LocationDataSource: UITableViewDelegate {
 
     func tableView(_ tableView: UITableView, willDisplay cell: UITableViewCell, forRowAt indexPath: IndexPath) {
         if let item = itemIdentifier(for: indexPath), item == selectedItem {
-            cell.setSelected(true, animated: false)
+            tableView.selectRow(at: indexPath, animated: false, scrollPosition: .none)
         }
     }
 
+    func tableView(_ tableView: UITableView, willSelectRowAt indexPath: IndexPath) -> IndexPath? {
+        if let indexPath = indexPathForSelectedRelay() {
+            tableView.deselectRow(at: indexPath, animated: false)
+            selectedItem = nil
+        }
+
+        return indexPath
+    }
+
     func tableView(_ tableView: UITableView, didSelectRowAt indexPath: IndexPath) {
         guard let item = itemIdentifier(for: indexPath) else { return }
+        selectedItem = item
 
         var customListSelection: UserSelectedRelays.CustomListSelection?
         if let topmostNode = item.node.root as? CustomListLocationNode {
@@ -303,12 +319,6 @@ extension LocationDataSource: UITableViewDelegate {
     private func scrollToTop(animated: Bool) {
         tableView.setContentOffset(.zero, animated: animated)
     }
-
-    private func scrollToSelectedRelay() {
-        indexPathForSelectedRelay().flatMap {
-            tableView.scrollToRow(at: $0, at: .middle, animated: false)
-        }
-    }
 }
 
 extension LocationDataSource: LocationCellDelegate {
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationViewController.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationViewController.swift
index 11b5c0bbb7..99f7f51488 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationViewController.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationViewController.swift
@@ -96,9 +96,13 @@ final class LocationViewController: UIViewController {
         }
     }
 
+    override func viewWillAppear(_ animated: Bool) {
+        super.viewWillAppear(animated)
+        dataSource?.scrollToSelectedRelay()
+    }
+
     override func viewDidAppear(_ animated: Bool) {
         super.viewDidAppear(animated)
-
         tableView.flashScrollIndicators()
     }