Add LocationDataSource

13 files changed, 774 insertions, 522 deletions

diff --git a/ios/MullvadVPN.xcodeproj/project.pbxproj b/ios/MullvadVPN.xcodeproj/project.pbxproj
index fc0441ccec..9e5371459d 100644
--- a/ios/MullvadVPN.xcodeproj/project.pbxproj
+++ b/ios/MullvadVPN.xcodeproj/project.pbxproj
@@ -65,6 +65,7 @@
 		5835B7CC233B76CB0096D79F /* TunnelManager.swift in Sources */ = {isa = PBXBuildFile; fileRef = 5835B7CB233B76CB0096D79F /* TunnelManager.swift */; };
 		583BC70724FE4DC500C9DE04 /* Optional+DispatchQueue.swift in Sources */ = {isa = PBXBuildFile; fileRef = 583BC70624FE4DC400C9DE04 /* Optional+DispatchQueue.swift */; };
 		583BC70824FE4DC500C9DE04 /* Optional+DispatchQueue.swift in Sources */ = {isa = PBXBuildFile; fileRef = 583BC70624FE4DC400C9DE04 /* Optional+DispatchQueue.swift */; };
+		583DA21425FA4B5C00318683 /* LocationDataSource.swift in Sources */ = {isa = PBXBuildFile; fileRef = 583DA21325FA4B5C00318683 /* LocationDataSource.swift */; };
 		5840250122B1124600E4CFEC /* IPAddress+Codable.swift in Sources */ = {isa = PBXBuildFile; fileRef = 5840250022B1124600E4CFEC /* IPAddress+Codable.swift */; };
 		5840250222B1124600E4CFEC /* IPAddress+Codable.swift in Sources */ = {isa = PBXBuildFile; fileRef = 5840250022B1124600E4CFEC /* IPAddress+Codable.swift */; };
 		5840250422B11AB700E4CFEC /* MullvadEndpoint.swift in Sources */ = {isa = PBXBuildFile; fileRef = 5840250322B11AB700E4CFEC /* MullvadEndpoint.swift */; };
@@ -200,7 +201,6 @@
 		58EF581125D69DB400AEBA94 /* StatusImageView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58EF581025D69DB400AEBA94 /* StatusImageView.swift */; };
 		58F19E35228C15BA00C7710B /* SpinnerActivityIndicatorView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58F19E34228C15BA00C7710B /* SpinnerActivityIndicatorView.swift */; };
 		58F3C0962492617E003E76BE /* AsyncOperation.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58E973DD24850EB600096F90 /* AsyncOperation.swift */; };
-		58F3C0A0249BBF1E003E76BE /* DiffableDataSources in Frameworks */ = {isa = PBXBuildFile; productRef = 58F3C09F249BBF1E003E76BE /* DiffableDataSources */; };
 		58F3C0A4249CB069003E76BE /* HeaderBarView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58F3C0A3249CB069003E76BE /* HeaderBarView.swift */; };
 		58F3C0A624A50157003E76BE /* relays.json in Resources */ = {isa = PBXBuildFile; fileRef = 58F3C0A524A50155003E76BE /* relays.json */; };
 		58F3C0A724A50C02003E76BE /* relays.json in Resources */ = {isa = PBXBuildFile; fileRef = 58F3C0A524A50155003E76BE /* relays.json */; };
@@ -309,6 +309,7 @@
 		582BB1B42295780F0055B6EF /* AccountExpiry.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AccountExpiry.swift; sourceTree = "<group>"; };
 		5835B7CB233B76CB0096D79F /* TunnelManager.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TunnelManager.swift; sourceTree = "<group>"; };
 		583BC70624FE4DC400C9DE04 /* Optional+DispatchQueue.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "Optional+DispatchQueue.swift"; sourceTree = "<group>"; };
+		583DA21325FA4B5C00318683 /* LocationDataSource.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = LocationDataSource.swift; sourceTree = "<group>"; };
 		5840250022B1124600E4CFEC /* IPAddress+Codable.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "IPAddress+Codable.swift"; sourceTree = "<group>"; };
 		5840250322B11AB700E4CFEC /* MullvadEndpoint.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = MullvadEndpoint.swift; sourceTree = "<group>"; };
 		5845F841236CBACD00B2D93C /* PacketTunnelIpc.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = PacketTunnelIpc.swift; sourceTree = "<group>"; };
@@ -432,7 +433,6 @@
 			buildActionMask = 2147483647;
 			files = (
 				585834F824D2BC1F00A8AF56 /* Logging in Frameworks */,
-				58F3C0A0249BBF1E003E76BE /* DiffableDataSources in Frameworks */,
 				58BA7947257901A5006FAEA0 /* WireGuardKit in Frameworks */,
 				586BD68422B7BBE400BB7F9F /* NetworkExtension.framework in Frameworks */,
 			);
@@ -644,6 +644,7 @@
 				5891BF5025E66B1E006D6FB0 /* UIBarButtonItem+KeyboardNavigation.swift */,
 				5891BF1B25E3E3EB006D6FB0 /* Bundle+ProductVersion.swift */,
 				585CA70E25F8C44600B47C62 /* UIMetrics.swift */,
+				583DA21325FA4B5C00318683 /* LocationDataSource.swift */,
 			);
 			path = MullvadVPN;
 			sourceTree = "<group>";
@@ -744,7 +745,6 @@
 			);
 			name = MullvadVPN;
 			packageProductDependencies = (
-				58F3C09F249BBF1E003E76BE /* DiffableDataSources */,
 				585834F724D2BC1F00A8AF56 /* Logging */,
 				58BA7946257901A5006FAEA0 /* WireGuardKit */,
 			);
@@ -845,7 +845,6 @@
 			);
 			mainGroup = 58CE5E57224146200008646E;
 			packageReferences = (
-				58F3C09E249BBF1E003E76BE /* XCRemoteSwiftPackageReference "DiffableDataSources" */,
 				585834F624D2BC1F00A8AF56 /* XCRemoteSwiftPackageReference "swift-log" */,
 				58BA79192578F092006FAEA0 /* XCRemoteSwiftPackageReference "wireguard-apple" */,
 			);
@@ -1006,6 +1005,7 @@
 				58B67B482602079E008EF58E /* RelaySelector.swift in Sources */,
 				58DF28A52417CB4B00E836B0 /* AppStorePaymentManager.swift in Sources */,
 				580EE22124B3240100F9D8A1 /* TransformOperationObserver.swift in Sources */,
+				583DA21425FA4B5C00318683 /* LocationDataSource.swift in Sources */,
 				582BB1AF229566420055B6EF /* SettingsCell.swift in Sources */,
 				5873884D239E6D7E00E96C4E /* EmbeddedViewContainerView.swift in Sources */,
 				583BC70724FE4DC500C9DE04 /* Optional+DispatchQueue.swift in Sources */,
@@ -1593,14 +1593,6 @@
 				version = "1.0.12-22";
 			};
 		};
-		58F3C09E249BBF1E003E76BE /* XCRemoteSwiftPackageReference "DiffableDataSources" */ = {
-			isa = XCRemoteSwiftPackageReference;
-			repositoryURL = "https://github.com/ra1028/DiffableDataSources.git";
-			requirement = {
-				kind = exactVersion;
-				version = 0.4.0;
-			};
-		};
 /* End XCRemoteSwiftPackageReference section */
 
 /* Begin XCSwiftPackageProductDependency section */
@@ -1629,11 +1621,6 @@
 			package = 58BA79192578F092006FAEA0 /* XCRemoteSwiftPackageReference "wireguard-apple" */;
 			productName = WireGuardKit;
 		};
-		58F3C09F249BBF1E003E76BE /* DiffableDataSources */ = {
-			isa = XCSwiftPackageProductDependency;
-			package = 58F3C09E249BBF1E003E76BE /* XCRemoteSwiftPackageReference "DiffableDataSources" */;
-			productName = DiffableDataSources;
-		};
 /* End XCSwiftPackageProductDependency section */
 	};
 	rootObject = 58CE5E58224146200008646E /* Project object */;
diff --git a/ios/MullvadVPN.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved b/ios/MullvadVPN.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved
index 16dde9781a..58780c69e1 100644
--- a/ios/MullvadVPN.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved
+++ b/ios/MullvadVPN.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved
@@ -2,24 +2,6 @@
   "object": {
     "pins": [
       {
-        "package": "DiffableDataSources",
-        "repositoryURL": "https://github.com/ra1028/DiffableDataSources.git",
-        "state": {
-          "branch": null,
-          "revision": "581b4f8d1634e83c6b33caaafdc6a115a74650c3",
-          "version": "0.4.0"
-        }
-      },
-      {
-        "package": "DifferenceKit",
-        "repositoryURL": "https://github.com/ra1028/DifferenceKit.git",
-        "state": {
-          "branch": null,
-          "revision": "14c66681e12a38b81045f44c6c29724a0d4b0e72",
-          "version": "1.1.5"
-        }
-      },
-      {
         "package": "swift-log",
         "repositoryURL": "https://github.com/apple/swift-log.git",
         "state": {
diff --git a/ios/MullvadVPN/ConnectViewController.swift b/ios/MullvadVPN/ConnectViewController.swift
index 4ec1d0e979..b5cbb08d05 100644
--- a/ios/MullvadVPN/ConnectViewController.swift
+++ b/ios/MullvadVPN/ConnectViewController.swift
@@ -10,8 +10,7 @@ import UIKit
 import NetworkExtension
 import Logging
 
-class ConnectViewController: UIViewController, RootContainment, TunnelObserver,
-    SelectLocationDelegate
+class ConnectViewController: UIViewController, RootContainment, TunnelObserver
 {
     @IBOutlet var secureLabel: UILabel!
     @IBOutlet var countryLabel: UILabel!
@@ -19,6 +18,8 @@ class ConnectViewController: UIViewController, RootContainment, TunnelObserver,
     @IBOutlet var connectionPanel: ConnectionPanelView!
     @IBOutlet var buttonsStackView: UIStackView!
 
+    private var relayConstraints: RelayConstraints?
+
     private let logger = Logger(label: "ConnectViewController")
 
     private let connectButton = AppButton(style: .success)
@@ -75,6 +76,8 @@ class ConnectViewController: UIViewController, RootContainment, TunnelObserver,
 
         TunnelManager.shared.addObserver(self)
         self.tunnelState = TunnelManager.shared.tunnelState
+
+        fetchRelayConstraints()
     }
 
     override func viewDidAppear(_ animated: Bool) {
@@ -95,33 +98,6 @@ class ConnectViewController: UIViewController, RootContainment, TunnelObserver,
         // no-op
     }
 
-    // MARK: - SelectLocationDelegate
-
-    func selectLocationViewController(_ controller: SelectLocationViewController, didSelectLocation location: RelayLocation) {
-        controller.dismiss(animated: true) {
-            let relayConstraints = RelayConstraints(location: .only(location))
-
-            TunnelManager.shared.setRelayConstraints(relayConstraints) { [weak self] (result) in
-                DispatchQueue.main.async {
-                    guard let self = self else { return }
-
-                    switch result {
-                    case .success:
-                        self.logger.debug("Updated relay constraints: \(relayConstraints)")
-                        self.connectTunnel()
-
-                    case .failure(let error):
-                        self.logger.error(chainedError: error, message: "Failed to update relay constraints")
-                    }
-                }
-            }
-        }
-    }
-
-    func selectLocationViewControllerDidCancel(_ controller: SelectLocationViewController) {
-        controller.dismiss(animated: true)
-    }
-
     // MARK: - Private
 
     private func updateButtons() {
@@ -252,18 +228,76 @@ class ConnectViewController: UIViewController, RootContainment, TunnelObserver,
         }
     }
 
-    private func showSelectLocation() {
-        let selectLocationController = SelectLocationNavigationController()
-        selectLocationController.selectLocationDelegate = self
+    private func showSelectLocationModal() {
+        let contentController = SelectLocationViewController()
+        contentController.navigationItem.title = NSLocalizedString("Select location", comment: "Navigation title")
+        contentController.navigationItem.largeTitleDisplayMode = .never
+        contentController.navigationItem.rightBarButtonItem = UIBarButtonItem(barButtonSystemItem: .done, target: self, action: #selector(handleDismissSelectLocationController(_:)))
+
+        contentController.didSelectRelayLocation = { [weak self] (controller, relayLocation) in
+            controller.view.isUserInteractionEnabled = false
+            DispatchQueue.main.asyncAfter(deadline: .now() + .milliseconds(250)) {
+                controller.view.isUserInteractionEnabled = true
+                controller.dismiss(animated: true) {
+                    self?.selectLocationControllerDidSelectRelayLocation(relayLocation)
+                }
+            }
+        }
+
+        let navController = SelectLocationNavigationController(contentController: contentController)
+
+        view.isUserInteractionEnabled = false
+        contentController.setSelectedRelayLocation(self.relayConstraints?.location.value, animated: false, scrollPosition: .none)
+        contentController.prefetchData { (error) in
+            if let error = error {
+                self.logger.error(chainedError: error, message: "Failed to prefetch the relays for SelectLocationViewController")
+            }
+
+            self.present(navController, animated: true) {
+                self.view.isUserInteractionEnabled = true
+            }
+        }
+    }
+
+    private func fetchRelayConstraints() {
+        TunnelManager.shared.getRelayConstraints { (result) in
+            DispatchQueue.main.async {
+                switch result {
+                case .success(let relayConstraints):
+                    self.relayConstraints = relayConstraints
+
+                case .failure(let error):
+                    self.logger.error(chainedError: error)
+                }
+            }
+        }
+    }
 
-        // Disable root controller interaction
-        rootContainerController?.view.isUserInteractionEnabled = false
+    private func selectLocationControllerDidSelectRelayLocation(_ relayLocation: RelayLocation) {
+        let relayConstraints = makeRelayConstraints(relayLocation)
 
-        selectLocationController.prefetchData {
-            self.present(selectLocationController, animated: true)
+        self.setTunnelRelayConstraints(relayConstraints)
+        self.relayConstraints = relayConstraints
+    }
 
-            // Re-enable root controller interaction
-            self.rootContainerController?.view.isUserInteractionEnabled = true
+    private func makeRelayConstraints(_ location: RelayLocation) -> RelayConstraints {
+        return RelayConstraints(location: .only(location))
+    }
+
+    private func setTunnelRelayConstraints(_ relayConstraints: RelayConstraints) {
+        TunnelManager.shared.setRelayConstraints(relayConstraints) { [weak self] (result) in
+            guard let self = self else { return }
+
+            DispatchQueue.main.async {
+                switch result {
+                case .success:
+                    self.logger.debug("Updated relay constraints: \(relayConstraints)")
+                    self.connectTunnel()
+
+                case .failure(let error):
+                    self.logger.error(chainedError: error, message: "Failed to update relay constraints")
+                }
+            }
         }
     }
 
@@ -286,7 +320,11 @@ class ConnectViewController: UIViewController, RootContainment, TunnelObserver,
     }
 
     @objc func handleSelectLocation(_ sender: Any) {
-        showSelectLocation()
+        showSelectLocationModal()
+    }
+
+    @objc func handleDismissSelectLocationController(_ sender: Any) {
+        self.presentedViewController?.dismiss(animated: true)
     }
 
 }
diff --git a/ios/MullvadVPN/CustomNavigationBar.swift b/ios/MullvadVPN/CustomNavigationBar.swift
index 4168d62bb0..3f5a91e118 100644
--- a/ios/MullvadVPN/CustomNavigationBar.swift
+++ b/ios/MullvadVPN/CustomNavigationBar.swift
@@ -10,30 +10,63 @@ import UIKit
 
 class CustomNavigationBar: UINavigationBar {
 
-    override init(frame: CGRect) {
-        super.init(frame: frame)
-
-        commonInit()
+    var prefersOpaqueBackground: Bool {
+        didSet {
+            setOpaqueBackgroundAppearance(prefersOpaqueBackground)
+        }
     }
 
-    required init?(coder aDecoder: NSCoder) {
-        super.init(coder: aDecoder)
+    // Returns the distance from the title label to the bottom of navigation bar
+    var titleLabelBottomInset: CGFloat {
+        // Go two levels deep only
+        let subviewsToExamine = subviews.flatMap { (view) -> [UIView] in
+            return [view] + view.subviews
+        }
+
+        let titleLabel = subviewsToExamine.first { (view) -> Bool in
+            return view is UILabel
+        }
 
-        commonInit()
+        if let titleLabel = titleLabel {
+            let titleFrame = titleLabel.convert(titleLabel.bounds, to: self)
+            return max(bounds.maxY - titleFrame.maxY, 0)
+        } else {
+            return 0
+        }
     }
 
-    private func commonInit() {
+    override init(frame: CGRect) {
+        if #available(iOS 13, *) {
+            prefersOpaqueBackground = false
+        } else {
+            prefersOpaqueBackground = true
+        }
+
+        super.init(frame: frame)
+
         var margins = layoutMargins
-        margins.left = 24
-        margins.right = 24
+        margins.left = UIMetrics.contentLayoutMargins.left
+        margins.right = UIMetrics.contentLayoutMargins.right
         layoutMargins = margins
 
-        if #available(iOS 13, *) {
-            // no-op
-        } else {
+        setOpaqueBackgroundAppearance(prefersOpaqueBackground)
+    }
+
+    required init?(coder: NSCoder) {
+        fatalError("init(coder:) has not been implemented")
+    }
+
+    private func setOpaqueBackgroundAppearance(_ flag: Bool) {
+        if flag {
             barTintColor = .secondaryColor
+            backgroundColor = .secondaryColor
             shadowImage = UIImage()
             isTranslucent = false
+        } else {
+            barTintColor = nil
+            backgroundColor = nil
+            shadowImage = nil
+            isTranslucent = true
         }
     }
 
diff --git a/ios/MullvadVPN/HeaderBarView.swift b/ios/MullvadVPN/HeaderBarView.swift
index ce6cad973f..0fd4b8f7e9 100644
--- a/ios/MullvadVPN/HeaderBarView.swift
+++ b/ios/MullvadVPN/HeaderBarView.swift
@@ -30,7 +30,12 @@ class HeaderBarView: UIView {
     override init(frame: CGRect) {
         super.init(frame: frame)
 
-        layoutMargins = UIEdgeInsets(top: 20, left: 12, bottom: 0, right: 16)
+        layoutMargins = UIEdgeInsets(
+            top: 0,
+            left: UIMetrics.contentLayoutMargins.left,
+            bottom: 0,
+            right: UIMetrics.contentLayoutMargins.right
+        )
 
         let constraints = [
             logoImageView.leadingAnchor.constraint(equalTo: layoutMarginsGuide.leadingAnchor),
diff --git a/ios/MullvadVPN/LocationDataSource.swift b/ios/MullvadVPN/LocationDataSource.swift
new file mode 100644
index 0000000000..9bf5ab5e9e
--- /dev/null
+++ b/ios/MullvadVPN/LocationDataSource.swift
@@ -0,0 +1,443 @@
+//
+//  LocationDataSource.swift
+//  MullvadVPN
+//
+//  Created by pronebird on 11/03/2021.
+//  Copyright © 2021 Mullvad VPN AB. All rights reserved.
+//
+
+import UIKit
+
+protocol LocationDataSourceItemProtocol {
+    var location: RelayLocation { get }
+    var displayName: String { get }
+    var showsChildren: Bool { get }
+    var isActive: Bool { get }
+
+    var isCollapsible: Bool { get }
+    var indentationLevel: Int { get }
+}
+
+class LocationDataSource: NSObject, UITableViewDataSource {
+
+    private var nodeByLocation = [RelayLocation: Node]()
+    private var locationList = [RelayLocation]()
+    private var rootNode = makeRootNode()
+
+    typealias CellProviderBlock = (UITableView, IndexPath, LocationDataSourceItemProtocol) -> UITableViewCell?
+
+    private let tableView: UITableView
+    private let cellProvider: CellProviderBlock
+
+    private(set) var selectedRelayLocation: RelayLocation?
+
+    private class func makeRootNode() -> Node {
+        return Node(
+            type: .root,
+            location: RelayLocation.country("#root"),
+            displayName: "",
+            showsChildren: true,
+            isActive: true,
+            children: []
+        )
+    }
+
+    init(tableView: UITableView, cellProvider: @escaping CellProviderBlock) {
+        self.tableView = tableView
+        self.cellProvider = cellProvider
+        super.init()
+
+        tableView.dataSource = self
+    }
+
+    func setSelectedRelayLocation(_ relayLocation: RelayLocation?, showHiddenParents: Bool, animated: Bool, scrollPosition: UITableView.ScrollPosition, completion: (() -> Void)? = nil) {
+        self.selectedRelayLocation = relayLocation
+
+        if relayLocation == nil {
+            if let indexPath = tableView.indexPathForSelectedRow {
+                tableView.deselectRow(at: indexPath, animated: animated)
+            }
+            completion?()
+        } else {
+            let setSelection = {
+                if let indexPath = self.indexPathForSelectedRelay() {
+                    self.tableView.selectRow(at: indexPath, animated: animated, scrollPosition: scrollPosition)
+                }
+                completion?()
+            }
+
+            if let relayLocation = relayLocation, showHiddenParents {
+                showParents(relayLocation, animated: animated, completion: setSelection)
+            } else {
+                setSelection()
+            }
+        }
+    }
+
+
+    func setRelays(_ response: ServerRelaysResponse) {
+        let rootNode = Self.makeRootNode()
+        var nodeByLocation = [RelayLocation: Node]()
+
+        for relay in response.wireguard.relays {
+            guard case .city(let countryCode, let cityCode) = RelayLocation(dashSeparatedString: relay.location),
+                  let serverLocation = response.locations[relay.location] else { continue }
+
+            let relayLocation = RelayLocation.hostname(countryCode, cityCode, relay.hostname)
+
+            for ascendantOrSelf in relayLocation.ascendants + [relayLocation] {
+                guard !nodeByLocation.keys.contains(ascendantOrSelf) else {
+                    continue
+                }
+
+                // Maintain the `showsChildren` state when transitioning between relay lists
+                let wasShowingChildren = self.nodeByLocation[ascendantOrSelf]?.showsChildren ?? false
+
+                let node: Node
+                switch ascendantOrSelf {
+                case .country:
+                    node = Node(
+                        type: .country,
+                        location: ascendantOrSelf,
+                        displayName: serverLocation.country,
+                        showsChildren: wasShowingChildren,
+                        isActive: true,
+                        children: []
+                    )
+                    rootNode.addChild(node)
+
+                case .city(let countryCode, _):
+                    node = Node(
+                        type: .city,
+                        location: ascendantOrSelf,
+                        displayName: serverLocation.city,
+                        showsChildren: wasShowingChildren,
+                        isActive: true,
+                        children: []
+                    )
+                    nodeByLocation[.country(countryCode)]!.addChild(node)
+
+                case .hostname(let countryCode, let cityCode, _):
+                    node = Node(
+                        type: .relay,
+                        location: ascendantOrSelf,
+                        displayName: relay.hostname,
+                        showsChildren: false,
+                        isActive: relay.active,
+                        children: []
+                    )
+                    nodeByLocation[.city(countryCode, cityCode)]!.addChild(node)
+                }
+
+                nodeByLocation[ascendantOrSelf] = node
+            }
+        }
+
+        rootNode.sortChildrenRecursive()
+        rootNode.computeActiveChildrenRecursive()
+        self.nodeByLocation = nodeByLocation
+        self.rootNode = rootNode
+        self.locationList = rootNode.flatRelayLocationList()
+
+        tableView.reloadData()
+        if let indexPath = self.indexPathForSelectedRelay() {
+            tableView.selectRow(at: indexPath, animated: false, scrollPosition: .none)
+        }
+    }
+
+    func showChildren(_ relayLocation: RelayLocation, showHiddenParents: Bool = false, animated: Bool, completion: (() -> Void)? = nil) {
+        toggleChildrenInternal(
+            relayLocation,
+            show: true,
+            showHiddenParents: showHiddenParents,
+            animated: animated,
+            completion: completion
+        )
+    }
+
+    func hideChildren(_ relayLocation: RelayLocation, animated: Bool, completion: (() -> Void)? = nil) {
+        toggleChildrenInternal(
+            relayLocation,
+            show: false,
+            showHiddenParents: false,
+            animated: animated,
+            completion: completion
+        )
+    }
+
+    func toggleChildren(_ relayLocation: RelayLocation, animated: Bool, completion: (() -> Void)? = nil) {
+        guard let node = self.nodeByLocation[relayLocation] else { return }
+
+        toggleChildrenInternal(relayLocation, show: !node.showsChildren, showHiddenParents: false, animated: animated, completion: completion)
+    }
+
+    private func showParents(_ relayLocation: RelayLocation, animated: Bool, completion: (() -> Void)? = nil) {
+        switch relayLocation {
+        case .country:
+            completion?()
+        case .city:
+            if let countryLocation = relayLocation.ascendants.first {
+                toggleChildrenInternal(countryLocation, show: true, showHiddenParents: false, animated: animated, completion: completion)
+            }
+        case .hostname:
+            if let cityLocation = relayLocation.ascendants.last {
+                toggleChildrenInternal(cityLocation, show: true, showHiddenParents: true, animated: animated, completion: completion)
+            }
+        }
+    }
+
+    private func toggleChildrenInternal(_ relayLocation: RelayLocation, show: Bool, showHiddenParents: Bool, animated: Bool, completion: (() -> Void)? = nil) {
+        let affectedRelayLocations: [RelayLocation]
+        if showHiddenParents {
+            affectedRelayLocations = relayLocation.ascendants + [relayLocation]
+        } else {
+            affectedRelayLocations = [relayLocation]
+        }
+
+        let affectedNodes = affectedRelayLocations.compactMap { (relayLocation) -> Node? in
+            return nodeByLocation[relayLocation]
+        }
+
+        // Pick the topmost node to expand or collapse
+        guard let topNode = affectedNodes.first(where: { (node) -> Bool in
+            return node.isCollapsible && node.showsChildren != show
+        }) else {
+            completion?()
+            return
+        }
+
+        let numAffectedChildren = topNode.countChildrenRecursive { (node) -> Bool in
+            if show {
+                return node.showsChildren || affectedNodes.contains(where: { (otherNode) -> Bool in
+                    return node === otherNode
+                })
+            } else {
+                return node.showsChildren
+            }
+        }
+
+        let applyChanges = { () -> ChangeSet? in
+            guard let topIndexPath = self.indexPath(for: topNode.location) else { return nil }
+
+            affectedNodes.forEach { (node) in
+                node.showsChildren = show
+            }
+
+            let affectedRange = (topIndexPath.row + 1 ... topIndexPath.row + numAffectedChildren)
+            let affectedIndexPaths = affectedRange.map { (row) -> IndexPath in
+                return IndexPath(row: row, section: 0)
+            }
+
+            if show {
+                self.locationList.insert(contentsOf: topNode.flatRelayLocationList(), at: topIndexPath.row + 1)
+
+                return ChangeSet(
+                    insertIndexPaths: affectedIndexPaths,
+                    deleteIndexPaths: [],
+                    updateIndexPaths: [topIndexPath]
+                )
+            } else {
+                self.locationList.removeSubrange(affectedRange)
+
+                return ChangeSet(
+                    insertIndexPaths: [],
+                    deleteIndexPaths: affectedIndexPaths,
+                    updateIndexPaths: [topIndexPath]
+                )
+            }
+        }
+
+        let restoreSelection = {
+            if let indexPath = self.indexPathForSelectedRelay() {
+                self.tableView.selectRow(at: indexPath, animated: false, scrollPosition: .none)
+            }
+        }
+
+        if animated {
+            tableView.performBatchUpdates {
+                if let changeSet = applyChanges() {
+                    tableView.insertRows(at: changeSet.insertIndexPaths, with: .fade)
+                    tableView.deleteRows(at: changeSet.deleteIndexPaths, with: .fade)
+                    tableView.reloadRows(at: changeSet.updateIndexPaths, with: .none)
+                }
+            } completion: { (finished) in
+                restoreSelection()
+                completion?()
+            }
+        } else {
+            _ = applyChanges()
+            tableView.reloadData()
+            restoreSelection()
+            completion?()
+        }
+    }
+
+    func relayLocation(for indexPath: IndexPath) -> RelayLocation? {
+        return locationList[indexPath.row]
+    }
+
+    func item(for indexPath: IndexPath) -> LocationDataSourceItemProtocol? {
+        return self.relayLocation(for: indexPath)
+            .flatMap { (relayLocation) -> Node? in
+                return nodeByLocation[relayLocation]
+            }
+    }
+
+    func indexPath(for location: RelayLocation) -> IndexPath? {
+        return locationList.firstIndex(of: location).map { (index) -> IndexPath in
+            return IndexPath(row: index, section: 0)
+        }
+    }
+
+    func indexPathForSelectedRelay() -> IndexPath? {
+        return selectedRelayLocation.flatMap { (relayLocation) -> IndexPath? in
+            return self.indexPath(for: relayLocation)
+        }
+    }
+
+    // MARK: - UITableViewDataSource
+
+    func numberOfSections(in tableView: UITableView) -> Int {
+        return 1
+    }
+
+    func tableView(_ tableView: UITableView, numberOfRowsInSection section: Int) -> Int {
+        assert(section == 0)
+        return locationList.count
+    }
+
+    func tableView(_ tableView: UITableView, cellForRowAt indexPath: IndexPath) -> UITableViewCell {
+        assert(indexPath.section == 0)
+        let item = self.item(for: indexPath)!
+        return cellProvider(tableView, indexPath, item)!
+    }
+}
+
+extension LocationDataSource {
+
+    private enum NodeType {
+        case root
+        case country
+        case city
+        case relay
+    }
+
+    private class Node: LocationDataSourceItemProtocol {
+        let nodeType: NodeType
+        var location: RelayLocation
+        var displayName: String
+        var showsChildren: Bool
+        var isActive: Bool
+        var children: [Node]
+
+        var isCollapsible: Bool {
+            switch nodeType {
+            case .country, .city:
+                return true
+            case .root, .relay:
+                return false
+            }
+        }
+
+        var indentationLevel: Int {
+            switch nodeType {
+            case .root, .country:
+                return 0
+            case .city:
+                return 1
+            case .relay:
+                return 2
+            }
+        }
+
+        init(type: NodeType, location: RelayLocation, displayName: String, showsChildren: Bool, isActive: Bool, children: [Node]) {
+            self.nodeType = type
+            self.location = location
+            self.displayName = displayName
+            self.showsChildren = showsChildren
+            self.isActive = isActive
+            self.children = children
+        }
+
+        func addChild(_ child: Node) {
+            children.append(child)
+        }
+
+        func sortChildrenRecursive() {
+            sortChildren()
+            children.forEach { (node) in
+                node.sortChildrenRecursive()
+            }
+        }
+
+        func computeActiveChildrenRecursive() {
+            switch nodeType {
+            case .root, .country:
+                for node in children {
+                    node.computeActiveChildrenRecursive()
+                }
+                fallthrough
+            case .city:
+                isActive = children.contains(where: { (node) -> Bool in
+                    return node.isActive
+                })
+            case .relay:
+                break
+            }
+        }
+
+        func countChildrenRecursive(where condition: @escaping (Node) -> Bool) -> Int {
+            return children.reduce(into: 0) { (numVisibleChildren, node) in
+                numVisibleChildren += 1
+                if condition(node) {
+                    numVisibleChildren += node.countChildrenRecursive(where: condition)
+                }
+            }
+        }
+
+        func flatRelayLocationList() -> [RelayLocation] {
+            return children.reduce(into: []) { (array, node) in
+                Self.flatten(node: node, into: &array)
+            }
+        }
+
+        private func sortChildren() {
+            switch nodeType {
+            case .root, .country:
+                children.sort { (a, b) -> Bool in
+                    return lexicalSortComparator(a.displayName, b.displayName)
+                }
+            case .city:
+                children.sort { (a, b) -> Bool in
+                    return fileSortComparator(a.location.stringRepresentation, b.location.stringRepresentation)
+                }
+            case .relay:
+                break
+            }
+        }
+
+        private class func flatten(node: Node, into array: inout [RelayLocation]) {
+            array.append(node.location)
+            if node.showsChildren {
+                for child in node.children {
+                    Self.flatten(node: child, into: &array)
+                }
+            }
+        }
+    }
+
+    private struct ChangeSet {
+        let insertIndexPaths: [IndexPath]
+        let deleteIndexPaths: [IndexPath]
+        let updateIndexPaths: [IndexPath]
+    }
+
+}
+
+private func lexicalSortComparator(_ a: String, _ b: String) -> Bool {
+    return a.localizedCaseInsensitiveCompare(b) == .orderedAscending
+}
+
+private func fileSortComparator(_ a: String, _ b: String) -> Bool {
+    return a.localizedStandardCompare(b) == .orderedAscending
+}
diff --git a/ios/MullvadVPN/ProblemReportSubmissionOverlayView.swift b/ios/MullvadVPN/ProblemReportSubmissionOverlayView.swift
index 1d5b66ed50..5bbd9a5077 100644
--- a/ios/MullvadVPN/ProblemReportSubmissionOverlayView.swift
+++ b/ios/MullvadVPN/ProblemReportSubmissionOverlayView.swift
@@ -123,7 +123,7 @@ class ProblemReportSubmissionOverlayView: UIView {
         addSubviews()
         transitionToState(state)
 
-        layoutMargins = UIEdgeInsets(top: 8, left: 24, bottom: 24, right: 24)
+        layoutMargins = UIMetrics.contentLayoutMargins
     }
 
     required init?(coder: NSCoder) {
diff --git a/ios/MullvadVPN/ProblemReportViewController.swift b/ios/MullvadVPN/ProblemReportViewController.swift
index 6780808e83..2c8d495909 100644
--- a/ios/MullvadVPN/ProblemReportViewController.swift
+++ b/ios/MullvadVPN/ProblemReportViewController.swift
@@ -39,7 +39,7 @@ class ProblemReportViewController: UIViewController, UITextFieldDelegate, Condit
     private lazy var containerView: UIView = {
         let containerView = UIView()
         containerView.translatesAutoresizingMaskIntoConstraints = false
-        containerView.layoutMargins = UIEdgeInsets(top: 8, left: 24, bottom: 24, right: 24)
+        containerView.layoutMargins = UIMetrics.contentLayoutMargins
         containerView.backgroundColor = .clear
         return containerView
     }()
diff --git a/ios/MullvadVPN/RelayConstraints.swift b/ios/MullvadVPN/RelayConstraints.swift
index f4f12f5c71..430cbc9645 100644
--- a/ios/MullvadVPN/RelayConstraints.swift
+++ b/ios/MullvadVPN/RelayConstraints.swift
@@ -69,6 +69,21 @@ enum RelayLocation: Codable, Hashable {
     case city(String, String)
     case hostname(String, String, String)
 
+    init?(dashSeparatedString: String) {
+        let components = dashSeparatedString.split(separator: "-", maxSplits: 2).map(String.init)
+
+        switch components.count {
+        case 1:
+            self = .country(components[0])
+        case 2:
+            self = .city(components[0], components[1])
+        case 3:
+            self = .hostname(components[0], components[1], components[2])
+        default:
+            return nil
+        }
+    }
+
     init(from decoder: Decoder) throws {
         let container = try decoder.singleValueContainer()
 
@@ -103,6 +118,20 @@ enum RelayLocation: Codable, Hashable {
         }
     }
 
+    /// A list of `RelayLocation` items preceding the given one in the relay tree
+    var ascendants: [RelayLocation] {
+        switch self {
+        case .hostname(let country, let city, _):
+            return [.country(country), .city(country, city)]
+
+        case .city(let country, _):
+            return [.country(country)]
+
+        case .country:
+            return []
+        }
+    }
+
 }
 
 extension RelayLocation: CustomDebugStringConvertible {
diff --git a/ios/MullvadVPN/SelectLocationCell.swift b/ios/MullvadVPN/SelectLocationCell.swift
index 735a42df9c..1aa60d93bf 100644
--- a/ios/MullvadVPN/SelectLocationCell.swift
+++ b/ios/MullvadVPN/SelectLocationCell.swift
@@ -8,6 +8,8 @@
 
 import UIKit
 
+private let kCollapseButtonWidth: CGFloat = 24
+
 class SelectLocationCell: BasicTableViewCell {
     typealias CollapseHandler = (SelectLocationCell) -> Void
 
@@ -117,7 +119,7 @@ class SelectLocationCell: BasicTableViewCell {
             locationLabel.topAnchor.constraint(equalTo: contentView.layoutMarginsGuide.topAnchor),
             locationLabel.bottomAnchor.constraint(equalTo: contentView.layoutMarginsGuide.bottomAnchor),
 
-            collapseButton.widthAnchor.constraint(equalToConstant: 70),
+            collapseButton.widthAnchor.constraint(equalToConstant: UIMetrics.contentLayoutMargins.left + UIMetrics.contentLayoutMargins.right + kCollapseButtonWidth),
             collapseButton.topAnchor.constraint(equalTo: contentView.topAnchor),
             collapseButton.trailingAnchor.constraint(equalTo: contentView.trailingAnchor),
             collapseButton.bottomAnchor.constraint(equalTo: contentView.bottomAnchor)
@@ -131,6 +133,7 @@ class SelectLocationCell: BasicTableViewCell {
 
     private func updateDisabled() {
         locationLabel.alpha = isDisabled ? 0.2 : 1
+        collapseButton.alpha = isDisabled ? 0.2 : 1
     }
 
     private func updateBackgroundColor() {
diff --git a/ios/MullvadVPN/SelectLocationHeaderView.swift b/ios/MullvadVPN/SelectLocationHeaderView.swift
index d5b0268008..f6a977ab2d 100644
--- a/ios/MullvadVPN/SelectLocationHeaderView.swift
+++ b/ios/MullvadVPN/SelectLocationHeaderView.swift
@@ -8,29 +8,46 @@
 
 import UIKit
 
-class SelectLocationHeaderView: UIView {
-
-    let textLabel = UILabel()
-
-    override init(frame: CGRect) {
-        super.init(frame: frame)
-
-        layoutMargins = UIEdgeInsets(top: 24, left: 24, bottom: 24, right: 24)
+class SelectLocationHeaderView: UITableViewHeaderFooterView {
 
+    lazy var textContentLabel: UILabel = {
+        let textLabel = UILabel()
         textLabel.translatesAutoresizingMaskIntoConstraints = false
         textLabel.font = UIFont.systemFont(ofSize: 17)
         textLabel.textColor = UIColor(white: 1, alpha: 0.6)
         textLabel.numberOfLines = 0
         textLabel.text = NSLocalizedString("While connected, your real location is masked with a private and secure location in the selected region", comment: "")
+        return textLabel
+    }()
+
+    var topLayoutMarginAdjustmentForNavigationBarTitle: CGFloat = 0 {
+        didSet {
+            let value = UIMetrics.contentLayoutMargins.top - topLayoutMarginAdjustmentForNavigationBarTitle
+            contentView.layoutMargins.top = max(value, 0)
+        }
+    }
+
+    override init(reuseIdentifier: String?) {
+        super.init(reuseIdentifier: reuseIdentifier)
+
+        layoutMargins = .zero
+        contentView.layoutMargins = UIMetrics.contentLayoutMargins
+        contentView.addSubview(textContentLabel)
+
+        backgroundView = UIView()
+        backgroundView?.backgroundColor = .secondaryColor
+
+        let trailingConstraint = textContentLabel.trailingAnchor.constraint(equalTo: contentView.layoutMarginsGuide.trailingAnchor)
+        trailingConstraint.priority = UILayoutPriority(800)
 
-        addSubview(textLabel)
+        let bottomConstraint = textContentLabel.bottomAnchor.constraint(equalTo: contentView.layoutMarginsGuide.bottomAnchor)
+        bottomConstraint.priority = UILayoutPriority(800)
 
         NSLayoutConstraint.activate([
-            textLabel.topAnchor.constraint(equalTo: layoutMarginsGuide.topAnchor),
-            textLabel.leadingAnchor.constraint(equalTo: layoutMarginsGuide.leadingAnchor),
-            textLabel.trailingAnchor.constraint(equalTo: layoutMarginsGuide.trailingAnchor),
-            textLabel.bottomAnchor.constraint(equalTo: layoutMarginsGuide.bottomAnchor),
-            textLabel.bottomAnchor.constraint(equalTo: layoutMarginsGuide.bottomAnchor)
+            textContentLabel.topAnchor.constraint(equalTo: contentView.layoutMarginsGuide.topAnchor),
+            textContentLabel.leadingAnchor.constraint(equalTo: contentView.layoutMarginsGuide.leadingAnchor),
+            trailingConstraint,
+            bottomConstraint
         ])
     }
 
diff --git a/ios/MullvadVPN/SelectLocationNavigationController.swift b/ios/MullvadVPN/SelectLocationNavigationController.swift
index be5c3098d3..13da20e934 100644
--- a/ios/MullvadVPN/SelectLocationNavigationController.swift
+++ b/ios/MullvadVPN/SelectLocationNavigationController.swift
@@ -9,35 +9,17 @@
 import Foundation
 import UIKit
 
-protocol SelectLocationDelegate: class {
-    func selectLocationViewController(_ controller: SelectLocationViewController, didSelectLocation location: RelayLocation)
-    func selectLocationViewControllerDidCancel(_ controller: SelectLocationViewController)
-}
-
 class SelectLocationNavigationController: UINavigationController {
-    private weak var contentController: SelectLocationViewController?
-
-    weak var selectLocationDelegate: SelectLocationDelegate?
 
-    init() {
+    init(contentController: SelectLocationViewController) {
         super.init(navigationBarClass: CustomNavigationBar.self, toolbarClass: nil)
 
-        navigationBar.prefersLargeTitles = true
         navigationBar.barStyle = .black
         navigationBar.tintColor = .white
+        navigationBar.prefersLargeTitles = false
 
-        let contentController = SelectLocationViewController()
-        contentController.navigationItem.title = NSLocalizedString("Select location", comment: "Navigation title")
-        contentController.navigationItem.largeTitleDisplayMode = .always
-        contentController.navigationItem.rightBarButtonItem = UIBarButtonItem(barButtonSystemItem: .done, target: self, action: #selector(handleDone(_:)))
-
-        contentController.didSelectLocationHandler = { [weak self] (location) in
-            guard let self = self, let contentController = self.contentController else { return }
-
-            self.selectLocationDelegate?.selectLocationViewController(contentController, didSelectLocation: location)
-        }
+        (navigationBar as? CustomNavigationBar)?.prefersOpaqueBackground = true
 
-        self.contentController = contentController
         self.viewControllers = [contentController]
     }
 
@@ -50,14 +32,4 @@ class SelectLocationNavigationController: UINavigationController {
     required init?(coder aDecoder: NSCoder) {
         fatalError("init(coder:) has not been implemented")
     }
-
-    func prefetchData(_ completionHandler: @escaping () -> Void) {
-        contentController?.prefetchData(completionHandler: completionHandler)
-    }
-
-    @objc func handleDone(_ sender: AnyObject) {
-        if let contentController = contentController {
-            selectLocationDelegate?.selectLocationViewControllerDidCancel(contentController)
-        }
-    }
 }
diff --git a/ios/MullvadVPN/SelectLocationViewController.swift b/ios/MullvadVPN/SelectLocationViewController.swift
index 47a6dab663..39464bce5c 100644
--- a/ios/MullvadVPN/SelectLocationViewController.swift
+++ b/ios/MullvadVPN/SelectLocationViewController.swift
@@ -6,35 +6,49 @@
 //  Copyright © 2019 Mullvad VPN AB. All rights reserved.
 //
 
-import DiffableDataSources
 import UIKit
 import Logging
 
-private let kCellIdentifier = "Cell"
+class SelectLocationViewController: UIViewController, RelayCacheObserver, UITableViewDelegate {
 
-class SelectLocationViewController: UITableViewController, RelayCacheObserver {
+    private enum ReuseIdentifiers: String {
+        case cell
+        case header
+    }
 
-    private enum Error: ChainedError {
-        case loadRelayList(RelayCacheError)
-        case getRelayConstraints(TunnelManager.Error)
+    private lazy var tableView: UITableView = {
+        let tableView = UITableView(frame: view.bounds, style: .plain)
+        tableView.translatesAutoresizingMaskIntoConstraints = true
+        tableView.autoresizingMask = [.flexibleWidth, .flexibleHeight]
+        tableView.backgroundColor = .clear
+        tableView.separatorColor = .secondaryColor
+        tableView.separatorInset = .zero
+        tableView.estimatedRowHeight = 53
+        tableView.estimatedSectionHeaderHeight = 109
+        tableView.indicatorStyle = .white
 
-        var errorDescription: String? {
-            switch self {
-            case .loadRelayList:
-                return "Failure to load a relay list"
-            case .getRelayConstraints:
-                return "Failure to get relay constraints"
-            }
-        }
-    }
+        tableView.register(SelectLocationHeaderView.self, forHeaderFooterViewReuseIdentifier: ReuseIdentifiers.header.rawValue)
+        tableView.register(SelectLocationCell.self, forCellReuseIdentifier: ReuseIdentifiers.cell.rawValue)
+
+        return tableView
+    }()
 
     private let logger = Logger(label: "SelectLocationController")
-    private var cachedRelays: CachedRelays?
-    private var relayConstraints: RelayConstraints?
-    private var expandedItems = [RelayLocation]()
-    private var dataSource: DataSource?
+    private var dataSource: LocationDataSource?
+    private var setCachedRelaysOnViewDidLoad: CachedRelays?
+    private var setRelayLocationOnViewDidLoad: RelayLocation?
+    private var isViewAppeared = false
 
-    var didSelectLocationHandler: ((RelayLocation) -> Void)?
+    var didSelectRelayLocation: ((SelectLocationViewController, RelayLocation) -> Void)?
+    var scrollToSelectedRelayOnViewWillAppear = true
+
+    init() {
+        super.init(nibName: nil, bundle: nil)
+    }
+
+    required init?(coder: NSCoder) {
+        fatalError("init(coder:) has not been implemented")
+    }
 
     // MARK: - View lifecycle
 
@@ -42,440 +56,169 @@ class SelectLocationViewController: UITableViewController, RelayCacheObserver {
         super.viewDidLoad()
 
         view.backgroundColor = .secondaryColor
-        tableView.tableHeaderView = SelectLocationHeaderView(frame: CGRect(x: 0, y: 0, width: 50, height: 50))
-        tableView.register(SelectLocationCell.self, forCellReuseIdentifier: kCellIdentifier)
-        tableView.separatorColor = .secondaryColor
-        tableView.separatorInset = .zero
+        view.addSubview(tableView)
 
-        dataSource = DataSource(
+        dataSource = LocationDataSource(
             tableView: self.tableView,
             cellProvider: { [weak self] (tableView, indexPath, item) -> UITableViewCell? in
                 guard let self = self else { return nil }
 
                 let cell = tableView.dequeueReusableCell(
-                    withIdentifier: kCellIdentifier, for: indexPath) as! SelectLocationCell
+                    withIdentifier: ReuseIdentifiers.cell.rawValue, for: indexPath) as! SelectLocationCell
 
-                cell.accessibilityIdentifier = item.relayLocation.stringRepresentation
-                cell.isDisabled = !item.hasActiveRelays()
-                cell.locationLabel.text = item.displayName()
-                cell.statusIndicator.isActive = item.hasActiveRelays()
-                cell.showsCollapseControl = item.isCollapsibleLevel()
-                cell.isExpanded = self.expandedItems.contains(item.relayLocation)
+                cell.accessibilityIdentifier = item.location.stringRepresentation
+                cell.isDisabled = !item.isActive
+                cell.locationLabel.text = item.displayName
+                cell.statusIndicator.isActive = item.isActive
+                cell.showsCollapseControl = item.isCollapsible
+                cell.isExpanded = item.showsChildren
                 cell.didCollapseHandler = { [weak self] (cell) in
                     self?.collapseCell(cell)
                 }
 
                 return cell
-        })
+            })
 
-        dataSource?.defaultRowAnimation = .top
+        tableView.delegate = self
         tableView.dataSource = dataSource
 
-        RelayCache.shared.addObserver(self)
-
-        updateDataSource(animateDifferences: false) {
-            self.updateTableViewSelection(scroll: true, animated: false)
+        if let setCachedRelaysOnViewDidLoad = self.setCachedRelaysOnViewDidLoad {
+            dataSource?.setRelays(setCachedRelaysOnViewDidLoad.relays)
         }
-    }
-
-    override func viewDidLayoutSubviews() {
-        super.viewDidLayoutSubviews()
-
-        updateTableHeaderViewSizeIfNeeded()
-    }
 
-    // MARK: - UITableViewDelegate
-
-    override func tableView(_ tableView: UITableView, shouldHighlightRowAt indexPath: IndexPath) -> Bool {
-        return dataSource?.itemIdentifier(for: indexPath)?.hasActiveRelays() ?? false
-    }
-
-    override func tableView(_ tableView: UITableView, indentationLevelForRowAt indexPath: IndexPath) -> Int {
-        return dataSource?.itemIdentifier(for: indexPath)?.indentationLevel() ?? 0
-    }
-
-    override func tableView(_ tableView: UITableView, didSelectRowAt indexPath: IndexPath) {
-        guard let item = dataSource?.itemIdentifier(for: indexPath) else { return }
-
-        // Disable interaction with the controller after selection
-        tableView.isUserInteractionEnabled = false
-
-        DispatchQueue.main.asyncAfter(deadline: .now() + .milliseconds(250)) {
-            self.didSelectLocationHandler?(item.relayLocation)
+        if let setRelayLocationOnViewDidLoad = setRelayLocationOnViewDidLoad {
+            dataSource?.setSelectedRelayLocation(
+                setRelayLocationOnViewDidLoad,
+                showHiddenParents: true,
+                animated: false,
+                scrollPosition: .none
+            )
         }
-    }
 
-    // MARK: - RelayCacheObserver
-
-    func relayCache(_ relayCache: RelayCache, didUpdateCachedRelays cachedRelays: CachedRelays) {
-        self.didReceiveCachedRelays(cachedRelays) { (result) in
-            DispatchQueue.main.async {
-                switch result {
-                case .success(let (cachedRelays, relayConstraints)):
-                    self.didReceiveCachedRelays(cachedRelays, relayConstraints: relayConstraints)
-
-                case .failure(let error):
-                    self.logger.error(chainedError: error)
-                }
-            }
-        }
-    }
-
-    // MARK: - Public
-
-    func prefetchData(completionHandler: @escaping () -> Void) {
-        fetchRelays { (result) in
-            DispatchQueue.main.async {
-                switch result {
-                case .success(let (cachedRelays, relayConstraints)):
-                    self.didReceiveCachedRelays(cachedRelays, relayConstraints: relayConstraints)
-
-                case .failure(let error):
-                    self.logger.error(chainedError: error)
-                }
-
-                completionHandler()
-            }
-        }
-    }
-
-    // MARK: - Relay list handling
-
-    private func fetchRelays(completionHandler: @escaping (Result<(CachedRelays, RelayConstraints), Error>) -> Void) {
-        RelayCache.shared.read { (result) in
-            switch result {
-            case .success(let cachedRelays):
-                self.didReceiveCachedRelays(cachedRelays, completionHandler: completionHandler)
-
-            case .failure(let error):
-                completionHandler(.failure(.loadRelayList(error)))
-            }
-        }
+        RelayCache.shared.addObserver(self)
     }
 
-    private func didReceiveCachedRelays(_ cachedRelays: CachedRelays, completionHandler: @escaping (Result<(CachedRelays, RelayConstraints), Error>) -> Void) {
-        TunnelManager.shared.getRelayConstraints { (result) in
-            let result = result
-                .map { (cachedRelays, $0) }
-                .mapError { Error.getRelayConstraints($0) }
+    override func viewWillAppear(_ animated: Bool) {
+        super.viewWillAppear(animated)
 
-            completionHandler(result)
+        if let indexPath = dataSource?.indexPathForSelectedRelay(), scrollToSelectedRelayOnViewWillAppear, !isViewAppeared {
+            self.tableView.scrollToRow(at: indexPath, at: .middle, animated: false)
         }
     }
 
-    private func didReceiveCachedRelays(_ cachedRelays: CachedRelays, relayConstraints: RelayConstraints) {
-        self.cachedRelays = cachedRelays
-        self.relayConstraints = relayConstraints
-
-        let relayLocation = relayConstraints.location.value
-        expandedItems = relayLocation?.ascendants ?? []
+    override func viewDidAppear(_ animated: Bool) {
+        super.viewDidAppear(animated)
 
-        updateDataSource(animateDifferences: false)
-        tableView.reloadData()
+        isViewAppeared = true
 
-        updateTableViewSelection(scroll: true, animated: false)
+        tableView.flashScrollIndicators()
     }
 
-    private func computeIndexPathForSelectedLocation(relayLocation: RelayLocation) -> IndexPath? {
-        guard let row = dataSource?.snapshot()
-            .itemIdentifiers
-            .firstIndex(where: { $0.relayLocation == relayLocation }) else {
-                return nil
-        }
+    override func viewDidDisappear(_ animated: Bool) {
+        super.viewDidDisappear(animated)
 
-        return IndexPath(row: row, section: 0)
+        isViewAppeared = false
     }
 
-    // MARK: - Collapsible cells
-
-    private func updateTableViewSelection(scroll: Bool, animated: Bool) {
-        guard let relayLocation = relayConstraints?.location.value else { return }
-
-        let indexPath = computeIndexPathForSelectedLocation(relayLocation: relayLocation)
+    // MARK: - UITableViewDelegate
 
-        let scrollPosition: UITableView.ScrollPosition = scroll ? .middle : .none
-        tableView.selectRow(at: indexPath, animated: animated, scrollPosition: scrollPosition)
+    func tableView(_ tableView: UITableView, shouldHighlightRowAt indexPath: IndexPath) -> Bool {
+        return dataSource?.item(for: indexPath)?.isActive ?? false
     }
 
-    private func updateDataSource(animateDifferences: Bool, completion: (() -> Void)? = nil) {
-        let items = self.cachedRelays.map { (cachedRelays) -> [DataSourceItem] in
-            return cachedRelays.relays.makeDataSource { (item) -> Bool in
-                return expandedItems.contains(item.relayLocation)
-            }
-        } ?? []
-
-        var snapshot = DataSourceSnapshot()
-        snapshot.appendSections([.locations])
-        snapshot.appendItems(items, toSection: .locations)
-
-        dataSource?.apply(
-            snapshot,
-            animatingDifferences: animateDifferences,
-            completion: completion
-        )
+    func tableView(_ tableView: UITableView, indentationLevelForRowAt indexPath: IndexPath) -> Int {
+        return dataSource?.item(for: indexPath)?.indentationLevel ?? 0
     }
 
-    private func collapseCell(_ cell: SelectLocationCell) {
-        guard let cellIndexPath = tableView.indexPath(for: cell),
-            let item = dataSource?.itemIdentifier(for: cellIndexPath) else {
-                return
-        }
-
-        let itemLocation = item.relayLocation
-
-        if let index = expandedItems.firstIndex(of: itemLocation) {
-            expandedItems.remove(at: index)
-            cell.isExpanded = false
-        } else {
-            expandedItems.append(itemLocation)
-            cell.isExpanded = true
-        }
-
-        updateDataSource(animateDifferences: true) {
-            self.updateTableViewSelection(scroll: false, animated: true)
+    func tableView(_ tableView: UITableView, willDisplay cell: UITableViewCell, forRowAt indexPath: IndexPath) {
+        if let item = dataSource?.item(for: indexPath), item.location == dataSource?.selectedRelayLocation {
+            cell.setSelected(true, animated: false)
         }
     }
 
-    // MARK: - UITableView header
-
-    private func updateTableHeaderViewSizeIfNeeded() {
-        guard let header = tableView.tableHeaderView else { return }
+    func tableView(_ tableView: UITableView, didSelectRowAt indexPath: IndexPath) {
+        guard let item = dataSource?.item(for: indexPath) else { return }
 
-        // measure the view size
-        let sizeConstraint = CGSize(
-            width: tableView.bounds.width,
-            height: UIView.layoutFittingCompressedSize.height
+        dataSource?.setSelectedRelayLocation(
+            item.location,
+            showHiddenParents: false,
+            animated: false,
+            scrollPosition: .none
         )
-
-        let newSize = header.systemLayoutSizeFitting(sizeConstraint)
-        let oldSize = header.frame.size
-
-        if oldSize.height != newSize.height {
-            header.frame.size.height = newSize.height
-
-            // reset the header view to force UITableView layout pass
-            tableView.tableHeaderView = header
-        }
+        didSelectRelayLocation?(self, item.location)
     }
-}
-
-private extension RelayLocation {
 
-    /// A list of `RelayLocation` items preceding the given one in the relay tree
-    var ascendants: [RelayLocation] {
-        switch self {
-        case .hostname(let country, let city, _):
-            return [.country(country), .city(country, city)]
+    func tableView(_ tableView: UITableView, viewForHeaderInSection section: Int) -> UIView? {
+        assert(section == 0)
 
-        case .city(let country, _):
-            return [.country(country)]
+        let view = tableView.dequeueReusableHeaderFooterView(withIdentifier: ReuseIdentifiers.header.rawValue) as! SelectLocationHeaderView
 
-        case .country:
-            return []
+        // When contained within the navigation controller, we want the distance between the navigation title
+        // and the table header label to be exactly 24pt.
+        if let navigationBar = navigationController?.navigationBar as? CustomNavigationBar {
+            view.topLayoutMarginAdjustmentForNavigationBarTitle = navigationBar.titleLabelBottomInset
         }
-    }
 
-}
-
-/// Enum describing the table view sections
-private enum DataSourceSection {
-    case locations
-}
-
-/// Data source type
-private typealias DataSource = TableViewDiffableDataSource<DataSourceSection, DataSourceItem>
-
-/// Data source snapshot type
-private typealias DataSourceSnapshot = DiffableDataSourceSnapshot<DataSourceSection, DataSourceItem>
-
-/// A wrapper type for RelayList to be able to represent it as a flat list
-private enum DataSourceItem: Hashable {
-
-    struct Country {
-        let location: String
-        let name: String
-        let hasActiveRelays: Bool
-    }
-
-    struct City {
-        let location: String
-        let name: String
-        let hasActiveRelays: Bool
+        return view
     }
 
-    struct Hostname {
-        let location: String
-        let hostname: String
-        let active: Bool
-    }
-
-    case country(Country)
-    case city(City)
-    case hostname(Hostname)
+    // MARK: - RelayCacheObserver
 
-    var relayLocation: RelayLocation {
-        switch self {
-        case .country(let country):
-            return .country(country.location)
-        case .city(let city):
-            let split = city.location.split(separator: "-", maxSplits: 2).map(String.init)
-            return .city(split[0], split[1])
-        case .hostname(let host):
-            let split = host.location.split(separator: "-", maxSplits: 2).map(String.init)
-            return .hostname(split[0], split[1], host.hostname)
+    func relayCache(_ relayCache: RelayCache, didUpdateCachedRelays cachedRelays: CachedRelays) {
+        DispatchQueue.main.async {
+            self.didReceiveCachedRelays(cachedRelays)
         }
     }
 
-    static func == (lhs: DataSourceItem, rhs: DataSourceItem) -> Bool {
-        lhs.relayLocation == rhs.relayLocation
-    }
-
-    func hash(into hasher: inout Hasher) {
-        hasher.combine(relayLocation)
-    }
-
-    func indentationLevel() -> Int {
-        switch self {
-        case .country:
-            return 0
-        case .city:
-            return 1
-        case .hostname:
-            return 2
-        }
-    }
+    // MARK: - Public
 
-    func displayName() -> String {
-        switch self {
-        case .country(let country):
-            return country.name
-        case .city(let city):
-            return city.name
-        case .hostname(let relay):
-            return relay.hostname
-        }
-    }
+    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)
 
-    func hasActiveRelays() -> Bool {
-        switch self {
-        case .country(let country):
-            return country.hasActiveRelays
-        case .city(let city):
-            return city.hasActiveRelays
-        case .hostname(let host):
-            return host.active
+                case .failure(let error):
+                    completionHandler(error)
+                }
+            }
         }
     }
 
-    func isCollapsibleLevel() -> Bool {
-        switch self {
-        case .country, .city:
-            return self.hasActiveRelays()
-        case .hostname:
-            return false
+    func setSelectedRelayLocation(_ relayLocation: RelayLocation?, animated: Bool, scrollPosition: UITableView.ScrollPosition) {
+        guard isViewLoaded else {
+            self.setRelayLocationOnViewDidLoad = relayLocation
+            return
         }
-    }
-
-}
-
-extension ServerRelaysResponse {
-    fileprivate static func lexicalSortComparator(_ a: String, _ b: String) -> Bool {
-        return a.localizedCaseInsensitiveCompare(b) == .orderedAscending
-    }
 
-    fileprivate static func fileSortComparator(_ a: String, _ b: String) -> Bool {
-        return a.localizedStandardCompare(b) == .orderedAscending
+        self.dataSource?.setSelectedRelayLocation(
+            relayLocation,
+            showHiddenParents: true,
+            animated: animated,
+            scrollPosition: scrollPosition
+        )
     }
 
-    fileprivate func makeDataSource(evaluator: (DataSourceItem) -> Bool) -> [DataSourceItem] {
-        let relaysByCountry = Dictionary(grouping: wireguard.relays) { (relay) -> String in
-            return relay.location.split(separator: "-").first.flatMap(String.init)!
-        }
-
-        var items = [DataSourceItem]()
-
-        var countryItems = [DataSourceItem.Country]()
-        var cityItems = [String: [DataSourceItem.City]]()
-        var relayItems = [String: [DataSourceItem.Hostname]]()
-
-        for (countryCode, relays) in relaysByCountry {
-            let relaysByCity = Dictionary(grouping: relays) { (relay) -> String in
-                return relay.location
-            }
-
-            if let (cityCode, relays) = relaysByCity.first {
-                guard let location = locations[cityCode] else {
-                    continue
-                }
-
-                let country = DataSourceItem.Country(
-                    location: countryCode,
-                    name: location.country,
-                    hasActiveRelays: relays.contains(where: { (serverRelay) -> Bool in
-                        return serverRelay.active
-                    }))
-
-                countryItems.append(country)
-                if !evaluator(.country(country)) {
-                    continue
-                }
-            }
-
-            for (cityCode, relays) in relaysByCity {
-                guard let location = locations[cityCode] else {
-                    // TODO: log to file?
-                    print("Location not found: \(cityCode)")
-                    continue
-                }
-
-                let city = DataSourceItem.City(
-                    location: cityCode,
-                    name: location.city,
-                    hasActiveRelays: relays.contains(where: { (serverRelay) -> Bool in
-                        return serverRelay.active
-                    }))
-
-                if var cities = cityItems[countryCode] {
-                    cities.append(city)
-                    cityItems[countryCode] = cities
-                } else {
-                    cityItems[countryCode] = [city]
-                }
-
-                if !evaluator(.city(city)) {
-                    continue
-                }
-
-                relayItems[cityCode] = relays.map { (relay) -> DataSourceItem.Hostname in
-                    return DataSourceItem.Hostname(location: relay.location, hostname: relay.hostname, active: relay.active)
-                }
-            }
-        }
+    // MARK: - Relay list handling
 
-        countryItems.sort { (a, b) -> Bool in
-            return Self.lexicalSortComparator(a.name, b.name)
+    private func didReceiveCachedRelays(_ cachedRelays: CachedRelays) {
+        guard isViewLoaded else {
+            self.setCachedRelaysOnViewDidLoad = cachedRelays
+            return
         }
+        self.dataSource?.setRelays(cachedRelays.relays)
+    }
 
-        for country in countryItems {
-            items.append(.country(country))
-
-            if var cities = cityItems[country.location] {
-                cities.sort { (a, b) -> Bool in
-                    return Self.lexicalSortComparator(a.name, b.name)
-                }
-                for city in cities {
-                    items.append(.city(city))
+    // MARK: - Collapsible cells
 
-                    if var relays = relayItems[city.location] {
-                        relays.sort { (a, b) -> Bool in
-                            return Self.fileSortComparator(a.hostname, b.hostname)
-                        }
-                        items.append(contentsOf: relays.map { DataSourceItem.hostname($0) })
-                    }
-                }
-            }
+    private func collapseCell(_ cell: SelectLocationCell) {
+        guard let cellIndexPath = tableView.indexPath(for: cell),
+              let dataSource = dataSource, let location = dataSource.relayLocation(for: cellIndexPath) else {
+            return
         }
 
-        return items
+        dataSource.toggleChildren(location, animated: true)
     }
 }