Merge branch 'select-location-delegate'

2 files changed, 51 insertions, 38 deletions

diff --git a/ios/MullvadVPN/LocationDataSource.swift b/ios/MullvadVPN/LocationDataSource.swift
index 9bf5ab5e9e..9d33f09b3e 100644
--- a/ios/MullvadVPN/LocationDataSource.swift
+++ b/ios/MullvadVPN/LocationDataSource.swift
@@ -30,6 +30,8 @@ class LocationDataSource: NSObject, UITableViewDataSource {
     private let cellProvider: CellProviderBlock
 
     private(set) var selectedRelayLocation: RelayLocation?
+    private var lastShowHiddenParents = false
+    private var lastScrollPosition: UITableView.ScrollPosition = .none
 
     private class func makeRootNode() -> Node {
         return Node(
@@ -52,6 +54,8 @@ class LocationDataSource: NSObject, UITableViewDataSource {
 
     func setSelectedRelayLocation(_ relayLocation: RelayLocation?, showHiddenParents: Bool, animated: Bool, scrollPosition: UITableView.ScrollPosition, completion: (() -> Void)? = nil) {
         self.selectedRelayLocation = relayLocation
+        self.lastShowHiddenParents = showHiddenParents
+        self.lastScrollPosition = scrollPosition
 
         if relayLocation == nil {
             if let indexPath = tableView.indexPathForSelectedRow {
@@ -78,6 +82,7 @@ class LocationDataSource: NSObject, UITableViewDataSource {
     func setRelays(_ response: ServerRelaysResponse) {
         let rootNode = Self.makeRootNode()
         var nodeByLocation = [RelayLocation: Node]()
+        let dataSourceWasEmpty = locationList.isEmpty
 
         for relay in response.wireguard.relays {
             guard case .city(let countryCode, let cityCode) = RelayLocation(dashSeparatedString: relay.location),
@@ -140,8 +145,26 @@ class LocationDataSource: NSObject, UITableViewDataSource {
         self.locationList = rootNode.flatRelayLocationList()
 
         tableView.reloadData()
-        if let indexPath = self.indexPathForSelectedRelay() {
-            tableView.selectRow(at: indexPath, animated: false, scrollPosition: .none)
+
+        let setSelection = { (_ scrollPosition: UITableView.ScrollPosition) in
+            if let indexPath = self.indexPathForSelectedRelay() {
+                self.tableView.selectRow(at: indexPath, animated: false, scrollPosition: scrollPosition)
+            }
+        }
+
+        // Sometimes the selected relay may be set before the data source is populated with relays.
+        // In that case restore the selection using cached parameters from the last call to
+        // `setSelectedRelayLocation`.
+        if let selectedRelayLocation = self.selectedRelayLocation, dataSourceWasEmpty {
+            if lastShowHiddenParents {
+                showParents(selectedRelayLocation, animated: false) {
+                    setSelection(self.lastScrollPosition)
+                }
+            } else {
+                setSelection(self.lastScrollPosition)
+            }
+        } else {
+            setSelection(.none)
         }
     }
 
diff --git a/ios/MullvadVPN/SelectLocationViewController.swift b/ios/MullvadVPN/SelectLocationViewController.swift
index 39464bce5c..e308f0d26b 100644
--- a/ios/MullvadVPN/SelectLocationViewController.swift
+++ b/ios/MullvadVPN/SelectLocationViewController.swift
@@ -9,7 +9,11 @@
 import UIKit
 import Logging
 
-class SelectLocationViewController: UIViewController, RelayCacheObserver, UITableViewDelegate {
+protocol SelectLocationViewControllerDelegate: class {
+    func selectLocationViewController(_ controller: SelectLocationViewController, didSelectRelayLocation relayLocation: RelayLocation)
+}
+
+class SelectLocationViewController: UIViewController, UITableViewDelegate {
 
     private enum ReuseIdentifiers: String {
         case cell
@@ -37,9 +41,10 @@ class SelectLocationViewController: UIViewController, RelayCacheObserver, UITabl
     private var dataSource: LocationDataSource?
     private var setCachedRelaysOnViewDidLoad: CachedRelays?
     private var setRelayLocationOnViewDidLoad: RelayLocation?
+    private var setScrollPositionOnViewDidLoad: UITableView.ScrollPosition = .none
     private var isViewAppeared = false
 
-    var didSelectRelayLocation: ((SelectLocationViewController, RelayLocation) -> Void)?
+    weak var delegate: SelectLocationViewControllerDelegate?
     var scrollToSelectedRelayOnViewWillAppear = true
 
     init() {
@@ -91,11 +96,9 @@ class SelectLocationViewController: UIViewController, RelayCacheObserver, UITabl
                 setRelayLocationOnViewDidLoad,
                 showHiddenParents: true,
                 animated: false,
-                scrollPosition: .none
+                scrollPosition: setScrollPositionOnViewDidLoad
             )
         }
-
-        RelayCache.shared.addObserver(self)
     }
 
     override func viewWillAppear(_ animated: Bool) {
@@ -120,6 +123,16 @@ class SelectLocationViewController: UIViewController, RelayCacheObserver, UITabl
         isViewAppeared = false
     }
 
+    override func viewWillTransition(to size: CGSize, with coordinator: UIViewControllerTransitionCoordinator) {
+        super.viewWillTransition(to: size, with: coordinator)
+
+        coordinator.animate { (context) in
+            if let indexPath = self.dataSource?.indexPathForSelectedRelay() {
+                self.tableView.scrollToRow(at: indexPath, at: .middle, animated: false)
+            }
+        }
+    }
+
     // MARK: - UITableViewDelegate
 
     func tableView(_ tableView: UITableView, shouldHighlightRowAt indexPath: IndexPath) -> Bool {
@@ -145,7 +158,8 @@ class SelectLocationViewController: UIViewController, RelayCacheObserver, UITabl
             animated: false,
             scrollPosition: .none
         )
-        didSelectRelayLocation?(self, item.location)
+
+        self.delegate?.selectLocationViewController(self, didSelectRelayLocation: item.location)
     }
 
     func tableView(_ tableView: UITableView, viewForHeaderInSection section: Int) -> UIView? {
@@ -162,34 +176,20 @@ class SelectLocationViewController: UIViewController, RelayCacheObserver, UITabl
         return view
     }
 
-    // MARK: - RelayCacheObserver
-
-    func relayCache(_ relayCache: RelayCache, didUpdateCachedRelays cachedRelays: CachedRelays) {
-        DispatchQueue.main.async {
-            self.didReceiveCachedRelays(cachedRelays)
-        }
-    }
-
     // MARK: - Public
 
-    func prefetchData(completionHandler: @escaping (RelayCacheError?) -> Void) {
-        RelayCache.shared.read { (result) in
-            DispatchQueue.main.async {
-                switch result {
-                case .success(let cachedRelays):
-                    self.didReceiveCachedRelays(cachedRelays)
-                    completionHandler(nil)
-
-                case .failure(let error):
-                    completionHandler(error)
-                }
-            }
+    func setCachedRelays(_ cachedRelays: CachedRelays) {
+        guard isViewLoaded else {
+            self.setCachedRelaysOnViewDidLoad = cachedRelays
+            return
         }
+        self.dataSource?.setRelays(cachedRelays.relays)
     }
 
     func setSelectedRelayLocation(_ relayLocation: RelayLocation?, animated: Bool, scrollPosition: UITableView.ScrollPosition) {
         guard isViewLoaded else {
             self.setRelayLocationOnViewDidLoad = relayLocation
+            self.setScrollPositionOnViewDidLoad = scrollPosition
             return
         }
 
@@ -201,16 +201,6 @@ class SelectLocationViewController: UIViewController, RelayCacheObserver, UITabl
         )
     }
 
-    // MARK: - Relay list handling
-
-    private func didReceiveCachedRelays(_ cachedRelays: CachedRelays) {
-        guard isViewLoaded else {
-            self.setCachedRelaysOnViewDidLoad = cachedRelays
-            return
-        }
-        self.dataSource?.setRelays(cachedRelays.relays)
-    }
-
     // MARK: - Collapsible cells
 
     private func collapseCell(_ cell: SelectLocationCell) {