Merge branch 'IOS-53'

4 files changed, 216 insertions, 369 deletions

diff --git a/ios/MullvadVPN.xcodeproj/project.pbxproj b/ios/MullvadVPN.xcodeproj/project.pbxproj
index 0872f8f484..35fa04a23d 100644
--- a/ios/MullvadVPN.xcodeproj/project.pbxproj
+++ b/ios/MullvadVPN.xcodeproj/project.pbxproj
@@ -343,6 +343,7 @@
 		7AD8490D29BA1EC500878E53 /* SettingsCellFactory.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7AD8490C29BA1EC500878E53 /* SettingsCellFactory.swift */; };
 		7AD8490F29BA26B000878E53 /* CellFactoryProtocol.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7AD8490E29BA26B000878E53 /* CellFactoryProtocol.swift */; };
 		7AD8491129BA316500878E53 /* PreferencesCellFactory.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7AD8491029BA316500878E53 /* PreferencesCellFactory.swift */; };
+		7AF1E73A29C47727002C6633 /* LocationCellFactory.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7AF1E73929C47727002C6633 /* LocationCellFactory.swift */; };
 		E1187ABC289BBB850024E748 /* OutOfTimeViewController.swift in Sources */ = {isa = PBXBuildFile; fileRef = E1187ABA289BBB850024E748 /* OutOfTimeViewController.swift */; };
 		E1187ABD289BBB850024E748 /* OutOfTimeContentView.swift in Sources */ = {isa = PBXBuildFile; fileRef = E1187ABB289BBB850024E748 /* OutOfTimeContentView.swift */; };
 		E158B360285381C60002F069 /* StringFormatter.swift in Sources */ = {isa = PBXBuildFile; fileRef = E158B35F285381C60002F069 /* StringFormatter.swift */; };
@@ -909,6 +910,7 @@
 		7AD8490C29BA1EC500878E53 /* SettingsCellFactory.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SettingsCellFactory.swift; sourceTree = "<group>"; };
 		7AD8490E29BA26B000878E53 /* CellFactoryProtocol.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = CellFactoryProtocol.swift; sourceTree = "<group>"; };
 		7AD8491029BA316500878E53 /* PreferencesCellFactory.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = PreferencesCellFactory.swift; sourceTree = "<group>"; };
+		7AF1E73929C47727002C6633 /* LocationCellFactory.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = LocationCellFactory.swift; sourceTree = "<group>"; };
 		E1187ABA289BBB850024E748 /* OutOfTimeViewController.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = OutOfTimeViewController.swift; sourceTree = "<group>"; };
 		E1187ABB289BBB850024E748 /* OutOfTimeContentView.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = OutOfTimeContentView.swift; sourceTree = "<group>"; };
 		E158B35F285381C60002F069 /* StringFormatter.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = StringFormatter.swift; sourceTree = "<group>"; };
@@ -1423,6 +1425,7 @@
 				58727282265D173C00F315B2 /* LaunchScreen.storyboard */,
 				58E20770274672CA00DE5D77 /* LaunchViewController.swift */,
 				583DA21325FA4B5C00318683 /* LocationDataSource.swift */,
+				7AF1E73929C47727002C6633 /* LocationCellFactory.swift */,
 				58B993B02608A34500BA7811 /* LoginContentView.swift */,
 				58CE5E65224146200008646E /* LoginViewController.swift */,
 				58C3F4F82964B08300D72515 /* MapViewController.swift */,
@@ -2386,6 +2389,7 @@
 				587EB6742714520600123C75 /* PreferencesDataSourceDelegate.swift in Sources */,
 				584555F42991176200DD0657 /* UIPresentationController+Private.swift in Sources */,
 				582BB1AF229566420055B6EF /* SettingsCell.swift in Sources */,
+				7AF1E73A29C47727002C6633 /* LocationCellFactory.swift in Sources */,
 				58F3C0A4249CB069003E76BE /* HeaderBarView.swift in Sources */,
 				587A01FC23F1F0BE00B68763 /* SimulatorTunnelProviderHost.swift in Sources */,
 				5819C2172729595500D6EC38 /* SettingsAddDNSEntryCell.swift in Sources */,
diff --git a/ios/MullvadVPN/LocationCellFactory.swift b/ios/MullvadVPN/LocationCellFactory.swift
new file mode 100644
index 0000000000..fa3915f175
--- /dev/null
+++ b/ios/MullvadVPN/LocationCellFactory.swift
@@ -0,0 +1,54 @@
+//
+//  LocationCellFactory.swift
+//  MullvadVPN
+//
+//  Created by Jon Petersson on 2023-03-17.
+//  Copyright © 2023 Mullvad VPN AB. All rights reserved.
+//
+
+import MullvadTypes
+import UIKit
+
+protocol LocationCellEventHandler {
+    func collapseCell(for item: RelayLocation)
+}
+
+final class LocationCellFactory: CellFactoryProtocol {
+    var nodeByLocation = [RelayLocation: LocationDataSource.Node]()
+    var delegate: LocationCellEventHandler?
+    let tableView: UITableView
+
+    init(tableView: UITableView, nodeByLocation: [RelayLocation: LocationDataSource.Node]) {
+        self.tableView = tableView
+        self.nodeByLocation = nodeByLocation
+    }
+
+    func makeCell(for item: RelayLocation, indexPath: IndexPath) -> UITableViewCell {
+        let cell = tableView.dequeueReusableCell(
+            withIdentifier: LocationDataSource.CellReuseIdentifiers.locationCell.rawValue,
+            for: indexPath
+        )
+
+        configureCell(cell, item: item, indexPath: indexPath)
+
+        return cell
+    }
+
+    func configureCell(
+        _ cell: UITableViewCell,
+        item: RelayLocation,
+        indexPath: IndexPath
+    ) {
+        guard let cell = cell as? SelectLocationCell,
+              let node = nodeByLocation[item] else { return }
+
+        cell.accessibilityIdentifier = node.location.stringRepresentation
+        cell.isDisabled = !node.isActive
+        cell.locationLabel.text = node.displayName
+        cell.showsCollapseControl = node.isCollapsible
+        cell.isExpanded = node.showsChildren
+        cell.didCollapseHandler = { [weak self] cell in
+            self?.delegate?.collapseCell(for: item)
+        }
+    }
+}
diff --git a/ios/MullvadVPN/LocationDataSource.swift b/ios/MullvadVPN/LocationDataSource.swift
index 0f58953882..62d9b3a4e8 100644
--- a/ios/MullvadVPN/LocationDataSource.swift
+++ b/ios/MullvadVPN/LocationDataSource.swift
@@ -20,23 +20,25 @@ protocol LocationDataSourceItemProtocol {
     var indentationLevel: Int { get }
 }
 
-class LocationDataSource: NSObject, UITableViewDataSource {
-    typealias CellProviderBlock = (UITableView, IndexPath, LocationDataSourceItemProtocol)
-        -> UITableViewCell
-    typealias CellConfiguratorBlock = (UITableViewCell, IndexPath, LocationDataSourceItemProtocol)
-        -> Void
+final class LocationDataSource: UITableViewDiffableDataSource<Int, RelayLocation> {
+    enum CellReuseIdentifiers: String, CaseIterable {
+        case locationCell
+
+        var reusableViewClass: AnyClass {
+            switch self {
+            case .locationCell:
+                return SelectLocationCell.self
+            }
+        }
+    }
 
     private var nodeByLocation = [RelayLocation: Node]()
     private var locationList = [RelayLocation]()
     private var rootNode = makeRootNode()
+    private(set) var selectedRelayLocation: RelayLocation?
 
     private let tableView: UITableView
-    private let cellProvider: CellProviderBlock
-    private let cellConfigurator: CellConfiguratorBlock
-
-    private(set) var selectedRelayLocation: RelayLocation?
-    private var lastShowHiddenParents = false
-    private var lastScrollPosition: UITableView.ScrollPosition = .none
+    private let locationCellFactory: LocationCellFactory
 
     private class func makeRootNode() -> Node {
         return Node(
@@ -49,60 +51,46 @@ class LocationDataSource: NSObject, UITableViewDataSource {
         )
     }
 
-    init(
-        tableView: UITableView,
-        cellProvider: @escaping CellProviderBlock,
-        cellConfigurator: @escaping CellConfiguratorBlock
-    ) {
+    var didSelectRelayLocation: ((RelayLocation) -> Void)?
+
+    init(tableView: UITableView) {
         self.tableView = tableView
-        self.cellProvider = cellProvider
-        self.cellConfigurator = cellConfigurator
 
-        super.init()
+        let locationCellFactory = LocationCellFactory(
+            tableView: tableView,
+            nodeByLocation: nodeByLocation
+        )
+        self.locationCellFactory = locationCellFactory
+
+        super.init(tableView: tableView) { tableView, indexPath, itemIdentifier in
+            locationCellFactory.makeCell(for: itemIdentifier, indexPath: indexPath)
+        }
+
+        tableView.delegate = self
+        locationCellFactory.delegate = self
 
-        tableView.dataSource = self
+        defaultRowAnimation = .fade
+        registerClasses()
     }
 
     func setSelectedRelayLocation(
         _ relayLocation: RelayLocation?,
-        showHiddenParents: Bool,
-        animated: Bool,
-        scrollPosition: UITableView.ScrollPosition,
-        completion: (() -> Void)? = nil
+        animated: Bool
     ) {
         selectedRelayLocation = relayLocation
-        lastShowHiddenParents = showHiddenParents
-        lastScrollPosition = scrollPosition
 
-        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()
-            }
+        let selectedLocationTree = selectedRelayLocation?.ascendants ?? []
+        selectedLocationTree.forEach { location in
+            nodeByLocation[location]?.showsChildren = true
         }
+
+        updateCellFactory(with: nodeByLocation)
+        updateDataSnapshot(with: locationList, animated: animated)
     }
 
     func setRelays(_ response: REST.ServerRelaysResponse) {
         let rootNode = Self.makeRootNode()
         var nodeByLocation = [RelayLocation: Node]()
-        let dataSourceWasEmpty = locationList.isEmpty
 
         for relay in response.wireguard.relays {
             guard case let .city(
@@ -119,7 +107,7 @@ class LocationDataSource: NSObject, UITableViewDataSource {
                 }
 
                 // Maintain the `showsChildren` state when transitioning between relay lists
-                let wasShowingChildren = self.nodeByLocation[ascendantOrSelf]?
+                let wasShowingChildren = nodeByLocation[ascendantOrSelf]?
                     .showsChildren ?? false
 
                 let node: Node
@@ -168,288 +156,172 @@ class LocationDataSource: NSObject, UITableViewDataSource {
         self.rootNode = rootNode
         locationList = rootNode.flatRelayLocationList()
 
-        tableView.reloadData()
-
-        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 = selectedRelayLocation, dataSourceWasEmpty {
-            if lastShowHiddenParents {
-                showParents(selectedRelayLocation, animated: false) {
-                    setSelection(self.lastScrollPosition)
-                }
-            } else {
-                setSelection(lastScrollPosition)
-            }
-        } else {
-            setSelection(.none)
-        }
+        updateCellFactory(with: nodeByLocation)
+        updateDataSnapshot(with: locationList)
     }
 
-    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 indexPathForSelectedRelay() -> IndexPath? {
+        return selectedRelayLocation.flatMap { indexPath(for: $0) }
     }
 
-    func toggleChildren(
-        _ relayLocation: RelayLocation,
-        animated: Bool,
+    private func updateDataSnapshot(
+        with locations: [RelayLocation],
+        animated: Bool = false,
         completion: (() -> Void)? = nil
     ) {
-        guard let node = nodeByLocation[relayLocation] else { return }
+        var snapshot = NSDiffableDataSourceSnapshot<Int, RelayLocation>()
+        snapshot.appendSections([0])
 
-        toggleChildrenInternal(
-            relayLocation,
-            show: !node.showsChildren,
-            showHiddenParents: false,
-            animated: animated,
-            completion: completion
-        )
-    }
+        for location in locations {
+            snapshot.appendItems([location])
 
-    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
-                )
+            guard let countryNode = nodeByLocation[location], countryNode.showsChildren else {
+                continue
             }
-        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]
-        }
+            for cityNode in countryNode.children {
+                snapshot.appendItems([cityNode.location])
 
-        // 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
-        }
+                guard cityNode.showsChildren else {
+                    continue
+                }
 
-        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
+                snapshot.appendItems(cityNode.children.map { $0.location })
             }
         }
 
-        let applyChanges = { () -> ChangeSet? in
-            guard let topIndexPath = self.indexPath(for: topNode.location) else { return nil }
+        apply(snapshot, animatingDifferences: animated, completion: completion)
+    }
 
-            affectedNodes.forEach { node in
-                node.showsChildren = show
-            }
+    private func registerClasses() {
+        CellReuseIdentifiers.allCases.forEach { enumCase in
+            tableView.register(
+                enumCase.reusableViewClass,
+                forCellReuseIdentifier: enumCase.rawValue
+            )
+        }
+    }
 
-            let affectedRange = (topIndexPath.row + 1 ... topIndexPath.row + numAffectedChildren)
-            let affectedIndexPaths = affectedRange.map { row -> IndexPath in
-                return IndexPath(row: row, section: 0)
-            }
+    private func updateCellFactory(with nodeByLocation: [RelayLocation: Node]) {
+        locationCellFactory.nodeByLocation = nodeByLocation
+    }
 
-            if show {
-                self.locationList.insert(
-                    contentsOf: topNode.flatRelayLocationList(),
-                    at: topIndexPath.row + 1
-                )
+    private func toggleChildren(
+        _ relayLocation: RelayLocation,
+        show: Bool,
+        animated: Bool
+    ) {
+        guard let node = nodeByLocation[relayLocation] else { return }
 
-                return ChangeSet(
-                    insertIndexPaths: affectedIndexPaths,
-                    deleteIndexPaths: [],
-                    updateIndexPaths: [topIndexPath]
-                )
-            } else {
-                self.locationList.removeSubrange(affectedRange)
+        node.showsChildren = show
 
-                return ChangeSet(
-                    insertIndexPaths: [],
-                    deleteIndexPaths: affectedIndexPaths,
-                    updateIndexPaths: [topIndexPath]
-                )
-            }
+        if let indexPath = indexPath(for: node.location),
+           let cell = tableView.cellForRow(at: indexPath)
+        {
+            locationCellFactory.configureCell(cell, item: node.location, indexPath: indexPath)
         }
 
-        let restoreSelection = {
-            if let indexPath = self.indexPathForSelectedRelay() {
-                self.tableView.selectRow(at: indexPath, animated: false, scrollPosition: .none)
-            }
-        }
+        updateCellFactory(with: nodeByLocation)
+        updateDataSnapshot(with: locationList, animated: animated) { [weak self] in
+            guard let visibleIndexPaths = self?.tableView.indexPathsForVisibleRows else { return }
 
-        let scrollToInsertedIndexPaths = { [weak tableView] (changeSet: ChangeSet) in
-            guard let lastInsertedIndexPath = changeSet.insertIndexPaths.last,
-                  let lastUpdatedIndexPath = changeSet.updateIndexPaths.last,
-                  let visibleIndexPaths = tableView?.indexPathsForVisibleRows,
-                  let lastVisibleIndexPath = visibleIndexPaths.last,
-                  lastInsertedIndexPath >= lastVisibleIndexPath
-            else {
-                return
-            }
-            if changeSet.insertIndexPaths.count >= visibleIndexPaths.count {
-                tableView?.scrollToRow(at: lastUpdatedIndexPath, at: .top, animated: animated)
-            } else {
-                tableView?.scrollToRow(at: lastInsertedIndexPath, at: .bottom, animated: animated)
+            let scrollToNodeTop = {
+                if let firstInsertedIndexPath = self?.indexPath(for: node.location) {
+                    self?.tableView.scrollToRow(
+                        at: firstInsertedIndexPath,
+                        at: .top,
+                        animated: animated
+                    )
+                }
             }
-        }
 
-        if animated {
-            guard let changeSet = applyChanges() else {
-                completion?()
-                return
+            let scrollToNodeBottom = {
+                if let location = node.children.last?.location,
+                   let lastInsertedIndexPath = self?.indexPath(for: location),
+                   let lastVisibleIndexPath = visibleIndexPaths.last,
+                   lastInsertedIndexPath >= lastVisibleIndexPath
+                {
+                    self?.tableView.scrollToRow(
+                        at: lastInsertedIndexPath,
+                        at: .bottom,
+                        animated: animated
+                    )
+                }
             }
 
-            tableView.performBatchUpdates {
-                tableView.insertRows(at: changeSet.insertIndexPaths, with: .fade)
-                tableView.deleteRows(at: changeSet.deleteIndexPaths, with: .fade)
-                changeSet.updateIndexPaths.forEach { indexPath in
-                    guard let item = item(for: indexPath) else {
-                        assertionFailure()
-                        return
-                    }
-
-                    if let cell = tableView.cellForRow(at: indexPath) {
-                        cellConfigurator(cell, indexPath, item)
-                    }
-                }
-            } completion: { finished in
-                scrollToInsertedIndexPaths(changeSet)
-                restoreSelection()
-                completion?()
+            if node.children.count > visibleIndexPaths.count {
+                scrollToNodeTop()
+            } else {
+                scrollToNodeBottom()
             }
-        } else {
-            _ = applyChanges()
-            tableView.reloadData()
-            restoreSelection()
-            completion?()
         }
     }
 
-    func relayLocation(for indexPath: IndexPath) -> RelayLocation? {
-        return locationList[indexPath.row]
+    private func item(for indexPath: IndexPath) -> LocationDataSourceItemProtocol? {
+        return itemIdentifier(for: indexPath).flatMap { nodeByLocation[$0] }
     }
+}
 
-    func item(for indexPath: IndexPath) -> LocationDataSourceItemProtocol? {
-        return relayLocation(for: indexPath)
-            .flatMap { relayLocation -> Node? in
-                return nodeByLocation[relayLocation]
-            }
+extension LocationDataSource: UITableViewDelegate {
+    func tableView(_ tableView: UITableView, shouldHighlightRowAt indexPath: IndexPath) -> Bool {
+        return item(for: indexPath)?.isActive ?? false
     }
 
-    func indexPath(for location: RelayLocation) -> IndexPath? {
-        return locationList.firstIndex(of: location).map { index -> IndexPath in
-            return IndexPath(row: index, section: 0)
-        }
+    func tableView(_ tableView: UITableView, indentationLevelForRowAt indexPath: IndexPath) -> Int {
+        return item(for: indexPath)?.indentationLevel ?? 0
     }
 
-    func indexPathForSelectedRelay() -> IndexPath? {
-        return selectedRelayLocation.flatMap { relayLocation -> IndexPath? in
-            return self.indexPath(for: relayLocation)
+    func tableView(
+        _ tableView: UITableView,
+        willDisplay cell: UITableViewCell,
+        forRowAt indexPath: IndexPath
+    ) {
+        if let item = item(for: indexPath),
+           item.location == selectedRelayLocation
+        {
+            cell.setSelected(true, animated: false)
         }
     }
 
-    // MARK: - UITableViewDataSource
+    func tableView(_ tableView: UITableView, didSelectRowAt indexPath: IndexPath) {
+        guard let item = item(for: indexPath) else { return }
 
-    func numberOfSections(in tableView: UITableView) -> Int {
-        return 1
-    }
+        if let indexPath = indexPathForSelectedRelay(),
+           let cell = tableView.cellForRow(at: indexPath)
+        {
+            cell.setSelected(false, animated: false)
+        }
 
-    func tableView(_ tableView: UITableView, numberOfRowsInSection section: Int) -> Int {
-        assert(section == 0)
-        return locationList.count
-    }
+        setSelectedRelayLocation(
+            item.location,
+            animated: false
+        )
 
-    func tableView(_ tableView: UITableView, cellForRowAt indexPath: IndexPath) -> UITableViewCell {
-        assert(indexPath.section == 0)
-        let item = item(for: indexPath)!
-        let cell = cellProvider(tableView, indexPath, item)
+        didSelectRelayLocation?(item.location)
+    }
+}
 
-        cellConfigurator(cell, indexPath, item)
+extension LocationDataSource: LocationCellEventHandler {
+    func collapseCell(for item: RelayLocation) {
+        guard let node = nodeByLocation[item] else { return }
 
-        return cell
+        toggleChildren(
+            item,
+            show: !node.showsChildren,
+            animated: true
+        )
     }
 }
 
 extension LocationDataSource {
-    private enum NodeType {
+    enum NodeType {
         case root
         case country
         case city
         case relay
     }
 
-    private class Node: LocationDataSourceItemProtocol {
+    class Node: LocationDataSourceItemProtocol {
         let nodeType: NodeType
         var location: RelayLocation
         var displayName: String
@@ -562,12 +434,6 @@ extension LocationDataSource {
             }
         }
     }
-
-    private struct ChangeSet {
-        let insertIndexPaths: [IndexPath]
-        let deleteIndexPaths: [IndexPath]
-        let updateIndexPaths: [IndexPath]
-    }
 }
 
 private func lexicalSortComparator(_ a: String, _ b: String) -> Bool {
diff --git a/ios/MullvadVPN/SelectLocationViewController.swift b/ios/MullvadVPN/SelectLocationViewController.swift
index 17ea45c6fc..6b27797a07 100644
--- a/ios/MullvadVPN/SelectLocationViewController.swift
+++ b/ios/MullvadVPN/SelectLocationViewController.swift
@@ -19,8 +19,6 @@ protocol SelectLocationViewControllerDelegate: AnyObject {
 }
 
 class SelectLocationViewController: UIViewController, UITableViewDelegate {
-    static let cellReuseIdentifier = "Cell"
-
     private var tableView: UITableView?
 
     private let tableHeaderFooterView = SelectLocationHeaderView()
@@ -75,12 +73,14 @@ class SelectLocationViewController: UIViewController, UITableViewDelegate {
         tableView.estimatedRowHeight = 53
         tableView.indicatorStyle = .white
 
-        tableView.register(
-            SelectLocationCell.self,
-            forCellReuseIdentifier: Self.cellReuseIdentifier
-        )
-
         self.tableView = tableView
+        dataSource = LocationDataSource(tableView: tableView)
+        tableView.dataSource = dataSource
+
+        dataSource?.didSelectRelayLocation = { [weak self] location in
+            guard let self = self else { return }
+            self.delegate?.selectLocationViewController(self, didSelectRelayLocation: location)
+        }
 
         view.accessibilityElements = [tableHeaderFooterView, tableView]
         view.backgroundColor = .secondaryColor
@@ -89,31 +89,6 @@ class SelectLocationViewController: UIViewController, UITableViewDelegate {
         tableHeaderFooterView.translatesAutoresizingMaskIntoConstraints = false
         view.addSubview(tableHeaderFooterView)
 
-        dataSource = LocationDataSource(
-            tableView: tableView,
-            cellProvider: { tableView, indexPath, item in
-                return tableView.dequeueReusableCell(
-                    withIdentifier: Self.cellReuseIdentifier,
-                    for: indexPath
-                )
-            },
-            cellConfigurator: { [weak self] cell, indexPath, item in
-                guard let cell = cell as? SelectLocationCell else { return }
-
-                cell.accessibilityIdentifier = item.location.stringRepresentation
-                cell.isDisabled = !item.isActive
-                cell.locationLabel.text = item.displayName
-                cell.showsCollapseControl = item.isCollapsible
-                cell.isExpanded = item.showsChildren
-                cell.didCollapseHandler = { [weak self] cell in
-                    self?.collapseCell(cell)
-                }
-            }
-        )
-
-        tableView.delegate = self
-        tableView.dataSource = dataSource
-
         tableHeaderFooterViewTopConstraints = [
             tableHeaderFooterView.topAnchor.constraint(equalTo: view.topAnchor),
             tableView.topAnchor.constraint(equalTo: tableHeaderFooterView.bottomAnchor),
@@ -140,9 +115,7 @@ class SelectLocationViewController: UIViewController, UITableViewDelegate {
         if let setRelayLocationOnViewDidLoad = setRelayLocationOnViewDidLoad {
             dataSource?.setSelectedRelayLocation(
                 setRelayLocationOnViewDidLoad,
-                showHiddenParents: true,
-                animated: false,
-                scrollPosition: setScrollPositionOnViewDidLoad
+                animated: false
             )
         }
     }
@@ -194,41 +167,6 @@ class SelectLocationViewController: UIViewController, UITableViewDelegate {
         updateTableHeaderTopLayoutMargin()
     }
 
-    // MARK: - UITableViewDelegate
-
-    func tableView(_ tableView: UITableView, shouldHighlightRowAt indexPath: IndexPath) -> Bool {
-        return dataSource?.item(for: indexPath)?.isActive ?? false
-    }
-
-    func tableView(_ tableView: UITableView, indentationLevelForRowAt indexPath: IndexPath) -> Int {
-        return dataSource?.item(for: indexPath)?.indentationLevel ?? 0
-    }
-
-    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)
-        }
-    }
-
-    func tableView(_ tableView: UITableView, didSelectRowAt indexPath: IndexPath) {
-        guard let item = dataSource?.item(for: indexPath) else { return }
-
-        dataSource?.setSelectedRelayLocation(
-            item.location,
-            showHiddenParents: false,
-            animated: false,
-            scrollPosition: .none
-        )
-
-        delegate?.selectLocationViewController(self, didSelectRelayLocation: item.location)
-    }
-
     // MARK: - Public
 
     func setCachedRelays(_ cachedRelays: CachedRelays) {
@@ -252,25 +190,10 @@ class SelectLocationViewController: UIViewController, UITableViewDelegate {
 
         dataSource?.setSelectedRelayLocation(
             relayLocation,
-            showHiddenParents: true,
-            animated: animated,
-            scrollPosition: scrollPosition
+            animated: animated
         )
     }
 
-    // MARK: - Collapsible cells
-
-    private func collapseCell(_ cell: SelectLocationCell) {
-        guard let cellIndexPath = tableView?.indexPath(for: cell),
-              let dataSource = dataSource,
-              let location = dataSource.relayLocation(for: cellIndexPath)
-        else {
-            return
-        }
-
-        dataSource.toggleChildren(location, animated: true)
-    }
-
     // MARK: - Private
 
     private func updateTableHeaderTopLayoutMargin() {