Combine parts of LocationDataSource and AddLocationDataSource to one

22 files changed, 288 insertions, 312 deletions

diff --git a/ios/MullvadVPN.xcodeproj/project.pbxproj b/ios/MullvadVPN.xcodeproj/project.pbxproj
index 632600e08f..270ecdabc9 100644
--- a/ios/MullvadVPN.xcodeproj/project.pbxproj
+++ b/ios/MullvadVPN.xcodeproj/project.pbxproj
@@ -531,6 +531,7 @@
 		7A6389EB2B7FAD7A008E77E1 /* SettingsFieldValidationErrorContentView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7A6389EA2B7FAD7A008E77E1 /* SettingsFieldValidationErrorContentView.swift */; };
 		7A6389ED2B7FADA1008E77E1 /* SettingsFieldValidationErrorConfiguration.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7A6389EC2B7FADA1008E77E1 /* SettingsFieldValidationErrorConfiguration.swift */; };
 		7A6389F82B864CDF008E77E1 /* LocationNode.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7A6389F72B864CDF008E77E1 /* LocationNode.swift */; };
+		7A6652B82BB44C3E0042D848 /* LocationDiffableDataSourceProtocol.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7A6652B62BB44B120042D848 /* LocationDiffableDataSourceProtocol.swift */; };
 		7A6B4F592AB8412E00123853 /* TunnelMonitorTimings.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7A6B4F582AB8412E00123853 /* TunnelMonitorTimings.swift */; };
 		7A6F2FA52AFA3CB2006D0856 /* AccountExpiryTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7A6F2FA42AFA3CB2006D0856 /* AccountExpiryTests.swift */; };
 		7A6F2FA72AFBB9AE006D0856 /* AccountExpiry.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7A6F2FA62AFBB9AE006D0856 /* AccountExpiry.swift */; };
@@ -1795,6 +1796,7 @@
 		7A6389EA2B7FAD7A008E77E1 /* SettingsFieldValidationErrorContentView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SettingsFieldValidationErrorContentView.swift; sourceTree = "<group>"; };
 		7A6389EC2B7FADA1008E77E1 /* SettingsFieldValidationErrorConfiguration.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SettingsFieldValidationErrorConfiguration.swift; sourceTree = "<group>"; };
 		7A6389F72B864CDF008E77E1 /* LocationNode.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = LocationNode.swift; sourceTree = "<group>"; };
+		7A6652B62BB44B120042D848 /* LocationDiffableDataSourceProtocol.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = LocationDiffableDataSourceProtocol.swift; sourceTree = "<group>"; };
 		7A6B4F582AB8412E00123853 /* TunnelMonitorTimings.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TunnelMonitorTimings.swift; sourceTree = "<group>"; };
 		7A6F2FA42AFA3CB2006D0856 /* AccountExpiryTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AccountExpiryTests.swift; sourceTree = "<group>"; };
 		7A6F2FA62AFBB9AE006D0856 /* AccountExpiry.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AccountExpiry.swift; sourceTree = "<group>"; };
@@ -2448,6 +2450,7 @@
 				F050AE4D2B70D7F8003F4EDB /* LocationCellViewModel.swift */,
 				583DA21325FA4B5C00318683 /* LocationDataSource.swift */,
 				F050AE5D2B739A73003F4EDB /* LocationDataSourceProtocol.swift */,
+				7A6652B62BB44B120042D848 /* LocationDiffableDataSourceProtocol.swift */,
 				7A6389F72B864CDF008E77E1 /* LocationNode.swift */,
 				F050AE512B70DFC0003F4EDB /* LocationSection.swift */,
 				F0BE65362B9F136A005CC385 /* LocationSectionHeaderView.swift */,
@@ -5349,6 +5352,7 @@
 				58607A4D2947287800BC467D /* AccountExpiryInAppNotificationProvider.swift in Sources */,
 				58C8191829FAA2C400DEB1B4 /* NotificationConfiguration.swift in Sources */,
 				58FF9FE82B07650A00E4C97D /* ButtonCellContentConfiguration.swift in Sources */,
+				7A6652B82BB44C3E0042D848 /* LocationDiffableDataSourceProtocol.swift in Sources */,
 				5827B0A82B0F49EF00CCBBA1 /* ProxyConfigurationInteractorProtocol.swift in Sources */,
 				7A5869B92B56E7F000640D27 /* IPOverrideViewControllerDelegate.swift in Sources */,
 				586C0D7A2B039CE300E7CDD7 /* ShadowsocksCipherPicker.swift in Sources */,
diff --git a/ios/MullvadVPN.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved b/ios/MullvadVPN.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved
index eaa4f1187c..02691892fe 100644
--- a/ios/MullvadVPN.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved
+++ b/ios/MullvadVPN.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved
@@ -1,5 +1,4 @@
 {
-  "originHash" : "c15149b2d59d9e9c72375f65339c04f41a19943e1117e682df27fc9f943fdc56",
   "pins" : [
     {
       "identity" : "swift-log",
@@ -19,5 +18,5 @@
       }
     }
   ],
-  "version" : 3
+  "version" : 2
 }
diff --git a/ios/MullvadVPN/Coordinators/ApplicationCoordinator.swift b/ios/MullvadVPN/Coordinators/ApplicationCoordinator.swift
index 627183b1f7..0c2b8678f8 100644
--- a/ios/MullvadVPN/Coordinators/ApplicationCoordinator.swift
+++ b/ios/MullvadVPN/Coordinators/ApplicationCoordinator.swift
@@ -43,14 +43,6 @@ final class ApplicationCoordinator: Coordinator, Presenting, RootContainerViewCo
      */
     private let secondaryNavigationContainer = RootContainerViewController()
 
-    private var customListRepository: CustomListRepositoryProtocol {
-        #if DEBUG
-        InMemoryCustomListRepository()
-        #else
-        CustomListRepository()
-        #endif
-    }
-
     /// Posts `preferredAccountNumber` notification when user inputs the account number instead of voucher code
     private let preferredAccountNumberSubject = PassthroughSubject<String, Never>()
 
@@ -719,7 +711,7 @@ final class ApplicationCoordinator: Coordinator, Presenting, RootContainerViewCo
             navigationController: navigationController,
             tunnelManager: tunnelManager,
             relayCacheTracker: relayCacheTracker,
-            customListRepository: customListRepository
+            customListRepository: CustomListRepository()
         )
 
         locationCoordinator.didFinish = { [weak self] _ in
diff --git a/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsCoordinator.swift b/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsCoordinator.swift
index d7b306d978..feb5bd415e 100644
--- a/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsCoordinator.swift
+++ b/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsCoordinator.swift
@@ -6,8 +6,6 @@
 //  Copyright © 2024 Mullvad VPN AB. All rights reserved.
 //
 
-import Combine
-import Foundation
 import MullvadSettings
 import MullvadTypes
 import Routing
diff --git a/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsDataSource.swift b/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsDataSource.swift
index 18564b8a12..048d3e51fa 100644
--- a/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsDataSource.swift
+++ b/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsDataSource.swift
@@ -6,36 +6,40 @@
 //  Copyright © 2024 Mullvad VPN AB. All rights reserved.
 //
 
-import Foundation
 import MullvadSettings
 import MullvadTypes
 import UIKit
 
-class AddLocationsDataSource: UITableViewDiffableDataSource<LocationSection, LocationCellViewModel> {
-    private let tableView: UITableView
-    private let nodes: [LocationNode]
+class AddLocationsDataSource:
+    UITableViewDiffableDataSource<LocationSection, LocationCellViewModel>,
+    LocationDiffableDataSourceProtocol {
     private var customListLocationNode: CustomListLocationNode
+    private let nodes: [LocationNode]
     var didUpdateCustomList: ((CustomListLocationNode) -> Void)?
+    let tableView: UITableView
+    let sections: [LocationSection]
 
     init(
         tableView: UITableView,
-        allLocations: [LocationNode],
+        allLocationNodes: [LocationNode],
         customList: CustomList
     ) {
         self.tableView = tableView
-        self.nodes = allLocations
+        self.nodes = allLocationNodes
 
         self.customListLocationNode = CustomListLocationNodeBuilder(
             customList: customList,
             allLocations: self.nodes
         ).customListLocationNode
 
+        let sections: [LocationSection] = [.customLists]
+        self.sections = sections
+
         super.init(tableView: tableView) { _, indexPath, itemIdentifier in
             let cell = tableView.dequeueReusableView(
-                withIdentifier: LocationSection.allCases[indexPath.section],
+                withIdentifier: sections[indexPath.section],
                 for: indexPath
-                // swiftlint:disable:next force_cast
-            ) as! LocationCell
+            ) as! LocationCell // swiftlint:disable:this force_cast
             cell.configure(item: itemIdentifier, behavior: .add)
             cell.selectionStyle = .none
             return cell
@@ -47,6 +51,14 @@ class AddLocationsDataSource: UITableViewDiffableDataSource<LocationSection, Loc
         reloadWithSelectedLocations()
     }
 
+    func nodeShowsChildren(_ node: LocationNode) -> Bool {
+        isLocationInCustomList(node: node)
+    }
+
+    func nodeShouldBeSelected(_ node: LocationNode) -> Bool {
+        customListLocationNode.children.contains(node)
+    }
+
     private func reloadWithSelectedLocations() {
         var locationsList: [LocationCellViewModel] = []
         nodes.forEach { node in
@@ -62,9 +74,13 @@ class AddLocationsDataSource: UITableViewDiffableDataSource<LocationSection, Loc
                 return
             }
 
-            // Walk tree backwards to determine which nodes should be expanded.
-            node.forEachAncestor { node in
-                node.showsChildren = true
+            // Only parents with partially selected children should be expanded.
+            node.forEachDescendant { descendantNode in
+                if customListLocationNode.children.contains(descendantNode) {
+                    descendantNode.forEachAncestor { descendantParentNode in
+                        descendantParentNode.showsChildren = true
+                    }
+                }
             }
 
             locationsList.append(contentsOf: recursivelyCreateCellViewModelTree(
@@ -73,53 +89,7 @@ class AddLocationsDataSource: UITableViewDiffableDataSource<LocationSection, Loc
                 indentationLevel: 1
             ))
         }
-        updateDataSnapshot(with: locationsList)
-    }
-
-    private func updateDataSnapshot(
-        with list: [LocationCellViewModel],
-        animated: Bool = false,
-        completion: (() -> Void)? = nil
-    ) {
-        var snapshot = NSDiffableDataSourceSnapshot<LocationSection, LocationCellViewModel>()
-
-        snapshot.appendSections([.customLists])
-        snapshot.appendItems(list, toSection: .customLists)
-
-        apply(snapshot, animatingDifferences: animated, completion: completion)
-    }
-
-    private func recursivelyCreateCellViewModelTree(
-        for node: LocationNode,
-        in section: LocationSection,
-        indentationLevel: Int
-    ) -> [LocationCellViewModel] {
-        var viewModels = [LocationCellViewModel]()
-        for childNode in node.children {
-            viewModels.append(
-                LocationCellViewModel(
-                    section: .customLists,
-                    node: childNode,
-                    indentationLevel: indentationLevel,
-                    isSelected: customListLocationNode.children.contains(childNode)
-                )
-            )
-
-            let indentationLevel = indentationLevel + 1
-
-            // Walk tree forward to determine which nodes should be expanded.
-            if isLocationInCustomList(node: childNode) {
-                viewModels.append(
-                    contentsOf: recursivelyCreateCellViewModelTree(
-                        for: childNode,
-                        in: section,
-                        indentationLevel: indentationLevel
-                    )
-                )
-            }
-        }
-
-        return viewModels
+        updateDataSnapshot(with: [locationsList])
     }
 
     private func isLocationInCustomList(node: LocationNode) -> Bool {
@@ -146,26 +116,23 @@ extension AddLocationsDataSource: UITableViewDelegate {
 
 extension AddLocationsDataSource: LocationCellDelegate {
     func toggleExpanding(cell: LocationCell) {
-        guard let indexPath = tableView.indexPath(for: cell),
-              let item = itemIdentifier(for: indexPath) else { return }
-        let isExpanded = item.node.showsChildren
-
-        item.node.showsChildren = !isExpanded
-
-        var locationList = snapshot().itemIdentifiers
-
-        if !isExpanded {
-            locationList.addSubNodes(from: item, at: indexPath)
-        } else {
-            locationList.removeSubNodes(from: item.node)
+        let items = toggledItems(for: cell).first!.map { item in
+            var item = item
+            if containsChild(parent: customListLocationNode, child: item.node) {
+                item.isSelected = true
+            }
+            return item
         }
 
-        updateDataSnapshot(with: locationList, animated: true, completion: {
-            self.scroll(to: item, animated: true)
+        updateDataSnapshot(with: [items], reloadExisting: true, completion: {
+            if let indexPath = self.tableView.indexPath(for: cell),
+               let item = self.itemIdentifier(for: indexPath) {
+                self.scroll(to: item, animated: true)
+            }
         })
     }
 
-    func toggleSelection(cell: LocationCell) {
+    func toggleSelecting(cell: LocationCell) {
         guard let index = tableView.indexPath(for: cell)?.row else { return }
 
         var locationList = snapshot().itemIdentifiers
@@ -181,36 +148,12 @@ extension AddLocationsDataSource: LocationCellDelegate {
         } else {
             customListLocationNode.remove(selectedLocation: item.node, with: locationList)
         }
-        updateDataSnapshot(with: locationList, completion: {
+        updateDataSnapshot(with: [locationList], completion: {
             self.didUpdateCustomList?(self.customListLocationNode)
         })
     }
 }
 
-extension AddLocationsDataSource {
-    private func scroll(to item: LocationCellViewModel, animated: Bool) {
-        guard
-            let visibleIndexPaths = tableView.indexPathsForVisibleRows,
-            let indexPath = indexPath(for: item)
-        else { return }
-
-        if item.node.children.count > visibleIndexPaths.count {
-            tableView.scrollToRow(at: indexPath, at: .top, animated: animated)
-        } else {
-            if let last = item.node.children.last {
-                if let lastInsertedIndexPath = self.indexPath(for: LocationCellViewModel(
-                    section: .customLists,
-                    node: last
-                )),
-                    let lastVisibleIndexPath = visibleIndexPaths.last,
-                    lastInsertedIndexPath >= lastVisibleIndexPath {
-                    tableView.scrollToRow(at: lastInsertedIndexPath, at: .bottom, animated: animated)
-                }
-            }
-        }
-    }
-}
-
 // MARK: - Toggle selection in table view
 
 fileprivate extension [LocationCellViewModel] {
diff --git a/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsViewController.swift b/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsViewController.swift
index 8fc9928d8c..c728982fdb 100644
--- a/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsViewController.swift
+++ b/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsViewController.swift
@@ -28,28 +28,9 @@ class AddLocationsViewController: UIViewController {
         tableView.rowHeight = 56
         tableView.indicatorStyle = .white
         tableView.accessibilityIdentifier = .addLocationsView
-        tableView.allowsMultipleSelection = true
-        tableView.tableHeaderView = nil
-        tableView.sectionHeaderHeight = .zero
         return tableView
     }()
 
-    private lazy var backBarButton: UIBarButtonItem = {
-        let backBarButton = UIBarButtonItem(
-            primaryAction: UIAction(
-                image: UIImage(resource: .iconBack),
-                handler: { [weak self] _ in
-                    guard let self else { return }
-                    delegate?.didBack()
-                    navigationController?.popViewController(animated: true)
-                }
-            )
-        )
-        backBarButton.style = .done
-
-        return backBarButton
-    }()
-
     init(
         allLocationsNodes: [LocationNode],
         customList: CustomList
@@ -67,11 +48,18 @@ class AddLocationsViewController: UIViewController {
         super.viewDidLoad()
         tableView.backgroundColor = view.backgroundColor
         view.backgroundColor = .secondaryColor
-        navigationItem.leftBarButtonItem = backBarButton
         addConstraints()
         setUpDataSource()
     }
 
+    override func didMove(toParent parent: UIViewController?) {
+        super.didMove(toParent: parent)
+
+        if parent == nil {
+            delegate?.didBack()
+        }
+    }
+
     private func addConstraints() {
         view.addConstrainedSubviews([tableView]) {
             tableView.pinEdgesToSuperview()
@@ -81,7 +69,7 @@ class AddLocationsViewController: UIViewController {
     private func setUpDataSource() {
         dataSource = AddLocationsDataSource(
             tableView: tableView,
-            allLocations: nodes.copy(),
+            allLocationNodes: nodes.copy(),
             customList: customList
         )
 
diff --git a/ios/MullvadVPN/Coordinators/CustomLists/CustomListDataSourceConfiguration.swift b/ios/MullvadVPN/Coordinators/CustomLists/CustomListDataSourceConfiguration.swift
index faf4e17764..77bbefd8bc 100644
--- a/ios/MullvadVPN/Coordinators/CustomLists/CustomListDataSourceConfiguration.swift
+++ b/ios/MullvadVPN/Coordinators/CustomLists/CustomListDataSourceConfiguration.swift
@@ -63,6 +63,21 @@ class CustomListDataSourceConfiguration: NSObject {
 }
 
 extension CustomListDataSourceConfiguration: UITableViewDelegate {
+    func tableView(_ tableView: UITableView, viewForHeaderInSection section: Int) -> UIView? {
+        return nil
+    }
+
+    func tableView(_ tableView: UITableView, heightForHeaderInSection section: Int) -> CGFloat {
+        let sectionIdentifier = dataSource.snapshot().sectionIdentifiers[section]
+
+        return switch sectionIdentifier {
+        case .name:
+            16
+        default:
+            UITableView.automaticDimension
+        }
+    }
+
     func tableView(_ tableView: UITableView, heightForRowAt indexPath: IndexPath) -> CGFloat {
         UIMetrics.SettingsCell.customListsCellHeight
     }
diff --git a/ios/MullvadVPN/Coordinators/CustomLists/EditCustomListCoordinator.swift b/ios/MullvadVPN/Coordinators/CustomLists/EditCustomListCoordinator.swift
index 9e753f22d6..5545f1bc95 100644
--- a/ios/MullvadVPN/Coordinators/CustomLists/EditCustomListCoordinator.swift
+++ b/ios/MullvadVPN/Coordinators/CustomLists/EditCustomListCoordinator.swift
@@ -94,6 +94,6 @@ extension EditCustomListCoordinator: CustomListViewControllerDelegate {
 
         coordinator.start()
 
-        coordinator.addChild(coordinator)
+        addChild(coordinator)
     }
 }
diff --git a/ios/MullvadVPN/Coordinators/CustomLists/EditLocationsCoordinator.swift b/ios/MullvadVPN/Coordinators/CustomLists/EditLocationsCoordinator.swift
index c9687c31ea..9255a2bc29 100644
--- a/ios/MullvadVPN/Coordinators/CustomLists/EditLocationsCoordinator.swift
+++ b/ios/MullvadVPN/Coordinators/CustomLists/EditLocationsCoordinator.swift
@@ -6,7 +6,6 @@
 //  Copyright © 2024 Mullvad VPN AB. All rights reserved.
 //
 
-import Foundation
 import MullvadSettings
 import MullvadTypes
 import Routing
diff --git a/ios/MullvadVPN/Coordinators/CustomLists/ListCustomListCoordinator.swift b/ios/MullvadVPN/Coordinators/CustomLists/ListCustomListCoordinator.swift
index 4753c74f72..fbdab2fba8 100644
--- a/ios/MullvadVPN/Coordinators/CustomLists/ListCustomListCoordinator.swift
+++ b/ios/MullvadVPN/Coordinators/CustomLists/ListCustomListCoordinator.swift
@@ -92,7 +92,7 @@ class ListCustomListCoordinator: Coordinator, Presentable, Presenting {
         }
 
         tunnelManager.updateSettings([.relayConstraints(relayConstraints)]) { [weak self] in
-            self?.tunnelManager.startTunnel()
+            self?.tunnelManager.reconnectTunnel(selectNewRelay: true)
         }
     }
 
diff --git a/ios/MullvadVPN/Coordinators/CustomLists/ListCustomListViewController.swift b/ios/MullvadVPN/Coordinators/CustomLists/ListCustomListViewController.swift
index 460cc5e2f7..78a14a298c 100644
--- a/ios/MullvadVPN/Coordinators/CustomLists/ListCustomListViewController.swift
+++ b/ios/MullvadVPN/Coordinators/CustomLists/ListCustomListViewController.swift
@@ -103,7 +103,6 @@ class ListCustomListViewController: UIViewController {
         tableView.separatorColor = .secondaryColor
         tableView.separatorInset = .zero
         tableView.separatorStyle = .singleLine
-        tableView.contentInset.top = 16
         tableView.rowHeight = UIMetrics.SettingsCell.customListsCellHeight
         tableView.registerReusableViews(from: CellReuseIdentifier.self)
     }
diff --git a/ios/MullvadVPN/Coordinators/LocationCoordinator.swift b/ios/MullvadVPN/Coordinators/LocationCoordinator.swift
index 53e164274e..0f808438a3 100644
--- a/ios/MullvadVPN/Coordinators/LocationCoordinator.swift
+++ b/ios/MullvadVPN/Coordinators/LocationCoordinator.swift
@@ -15,8 +15,8 @@ import UIKit
 class LocationCoordinator: Coordinator, Presentable, Presenting {
     private let tunnelManager: TunnelManager
     private let relayCacheTracker: RelayCacheTracker
+    private let customListRepository: CustomListRepositoryProtocol
     private var cachedRelays: CachedRelays?
-    private var customListRepository: CustomListRepositoryProtocol
 
     let navigationController: UINavigationController
 
diff --git a/ios/MullvadVPN/UI appearance/UIMetrics.swift b/ios/MullvadVPN/UI appearance/UIMetrics.swift
index f4d306590c..c3fbe48cff 100644
--- a/ios/MullvadVPN/UI appearance/UIMetrics.swift
+++ b/ios/MullvadVPN/UI appearance/UIMetrics.swift
@@ -173,7 +173,7 @@ extension UIMetrics {
     static let contentInsets = UIEdgeInsets(top: 24, left: 24, bottom: 24, right: 24)
 
     /// Common layout margins for location cell presentation
-    static let locationCellLayoutMargins = NSDirectionalEdgeInsets(top: 16, leading: 28, bottom: 16, trailing: 12)
+    static let locationCellLayoutMargins = NSDirectionalEdgeInsets(top: 16, leading: 16, bottom: 16, trailing: 12)
 
     /// Layout margins used by content heading displayed below the large navigation title.
     static let contentHeadingLayoutMargins = NSDirectionalEdgeInsets(top: 8, leading: 24, bottom: 24, trailing: 24)
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/AllLocationDataSource.swift b/ios/MullvadVPN/View controllers/SelectLocation/AllLocationDataSource.swift
index a6e9e1bab0..f2f955968f 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/AllLocationDataSource.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/AllLocationDataSource.swift
@@ -63,7 +63,8 @@ class AllLocationDataSource: LocationDataSourceProtocol {
             let countryNode = CountryLocationNode(
                 name: serverLocation.country,
                 code: LocationNode.combineNodeCodes([countryCode]),
-                locations: [location]
+                locations: [location],
+                isActive: relay.active
             )
 
             if !rootNode.children.contains(countryNode) {
@@ -75,7 +76,8 @@ class AllLocationDataSource: LocationDataSourceProtocol {
             let cityNode = CityLocationNode(
                 name: serverLocation.city,
                 code: LocationNode.combineNodeCodes([countryCode, cityCode]),
-                locations: [location]
+                locations: [location],
+                isActive: relay.active
             )
 
             if let countryNode = rootNode.countryFor(code: countryCode),
@@ -89,7 +91,8 @@ class AllLocationDataSource: LocationDataSourceProtocol {
             let hostNode = HostLocationNode(
                 name: relay.hostname,
                 code: LocationNode.combineNodeCodes([hostCode]),
-                locations: [location]
+                locations: [location],
+                isActive: relay.active
             )
 
             if let countryNode = rootNode.countryFor(code: countryCode),
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/CustomListsDataSource.swift b/ios/MullvadVPN/View controllers/SelectLocation/CustomListsDataSource.swift
index e9cad7cf1c..6d8f8989fa 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/CustomListsDataSource.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/CustomListsDataSource.swift
@@ -26,17 +26,9 @@ class CustomListsDataSource: LocationDataSourceProtocol {
     /// Constructs a collection of node trees by copying each matching counterpart
     /// from the complete list of nodes created in ``AllLocationDataSource``.
     func reload(allLocationNodes: [LocationNode], isFiltered: Bool) {
-        nodes = repository.fetchAll().compactMap { customList in
-            let listNode = CustomListLocationNode(
-                name: customList.name,
-                code: customList.name.lowercased(),
-                locations: customList.locations,
-                customList: customList
-            )
-
-            listNode.children = customList.locations.compactMap { location in
-                copy(location, from: allLocationNodes, withParent: listNode)
-            }
+        nodes = repository.fetchAll().compactMap { list in
+            let customListWrapper = CustomListLocationNodeBuilder(customList: list, allLocations: allLocationNodes)
+            let listNode = customListWrapper.customListLocationNode
 
             listNode.forEachDescendant { node in
                 // Each item in a section in a diffable data source needs to be unique.
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/InMemoryCustomListRepository.swift b/ios/MullvadVPN/View controllers/SelectLocation/InMemoryCustomListRepository.swift
index 600692b8d5..4a72fa61e1 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/InMemoryCustomListRepository.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/InMemoryCustomListRepository.swift
@@ -16,7 +16,7 @@ class InMemoryCustomListRepository: CustomListRepositoryProtocol {
         CustomList(
             id: UUID(uuidString: "F17948CB-18E2-4F84-82CD-5780F94216DB")!,
             name: "Netflix",
-            locations: [.city("al", "tia")]
+            locations: [.hostname("al", "tia", "al-tia-wg-001")]
         ),
         CustomList(
             id: UUID(uuidString: "4104C603-B35D-4A64-8865-96C0BF33D57F")!,
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationCell.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationCell.swift
index 79f847095f..24a1ce6fac 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationCell.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationCell.swift
@@ -8,23 +8,14 @@
 
 import UIKit
 
-//protocol LocationCellDelegate: AnyObject {
-//    func toggle(cell: LocationCell)
-//}
-
 protocol LocationCellDelegate: AnyObject {
     func toggleExpanding(cell: LocationCell)
-    func toggleSelection(cell: LocationCell)
+    func toggleSelecting(cell: LocationCell)
 }
 
 class LocationCell: UITableViewCell {
     weak var delegate: LocationCellDelegate?
 
-    private let container: UIStackView = {
-        let view = UIStackView()
-        return view
-    }()
-
     private let locationLabel: UILabel = {
         let label = UILabel()
         label.font = UIFont.systemFont(ofSize: 16)
@@ -64,8 +55,16 @@ class LocationCell: UITableViewCell {
         return button
     }()
 
-    private var behavior: LocationCellBehavior = .select
+    private var locationLabelLeadingMargin: CGFloat {
+        switch behavior {
+        case .add:
+            0
+        case .select:
+            12
+        }
+    }
 
+    private var behavior: LocationCellBehavior = .select
     private let chevronDown = UIImage(resource: .iconChevronDown)
     private let chevronUp = UIImage(resource: .iconChevronUp)
 
@@ -152,7 +151,7 @@ class LocationCell: UITableViewCell {
 
         updateCollapseImage()
         updateAccessibilityCustomActions()
-//        updateDisabled()
+        updateDisabled()
         updateBackgroundColor()
         setLayoutMargins()
 
@@ -161,9 +160,9 @@ class LocationCell: UITableViewCell {
             statusIndicator,
             locationLabel,
             collapseButton,
-            checkboxButton
+            checkboxButton,
         ]) {
-            tickImageView.leadingAnchor.constraint(equalTo: contentView.layoutMarginsGuide.leadingAnchor)
+            tickImageView.pinEdgesToSuperviewMargins(PinnableEdges([.leading(0)]))
             tickImageView.centerYAnchor.constraint(equalTo: contentView.centerYAnchor)
 
             statusIndicator.widthAnchor.constraint(equalToConstant: 16)
@@ -171,29 +170,24 @@ class LocationCell: UITableViewCell {
             statusIndicator.centerXAnchor.constraint(equalTo: tickImageView.centerXAnchor)
             statusIndicator.centerYAnchor.constraint(equalTo: tickImageView.centerYAnchor)
 
-            checkboxButton.pinEdgeToSuperviewMargin(.leading(.zero))
-            checkboxButton.centerYAnchor.constraint(equalTo: contentView.centerYAnchor)
-            checkboxButton.widthAnchor
-                .constraint(
-                    equalToConstant: 44.0
-                )
-            checkboxButton.heightAnchor.constraint(equalTo: checkboxButton.widthAnchor, multiplier: 1, constant: 0)
+            checkboxButton.pinEdgesToSuperview(PinnableEdges([.top(0), .bottom(0)]))
+            checkboxButton.trailingAnchor.constraint(equalTo: locationLabel.leadingAnchor, constant: 14)
+            checkboxButton.widthAnchor.constraint(
+                equalToConstant: UIMetrics.contentLayoutMargins.leading + UIMetrics.contentLayoutMargins.trailing + 24
+            )
 
+            locationLabel.pinEdgesToSuperviewMargins(PinnableEdges([.top(0), .bottom(0)]))
             locationLabel.leadingAnchor.constraint(
                 equalTo: statusIndicator.trailingAnchor,
-                constant: 12
+                constant: locationLabelLeadingMargin
             )
-
             locationLabel.trailingAnchor.constraint(lessThanOrEqualTo: collapseButton.leadingAnchor)
                 .withPriority(.defaultHigh)
-            locationLabel.pinEdgesToSuperviewMargins(PinnableEdges([.top(.zero), .bottom(.zero)]))
 
-            collapseButton.widthAnchor
-                .constraint(
-                    equalToConstant: UIMetrics.contentLayoutMargins.leading + UIMetrics
-                        .contentLayoutMargins.trailing + 24.0
-                )
-            collapseButton.pinEdgesToSuperviewMargins(.all().excluding(.leading))
+            collapseButton.widthAnchor.constraint(
+                equalToConstant: UIMetrics.contentLayoutMargins.leading + UIMetrics.contentLayoutMargins.trailing + 24
+            )
+            collapseButton.pinEdgesToSuperview(.all().excluding(.leading))
         }
     }
 
@@ -305,18 +299,19 @@ class LocationCell: UITableViewCell {
     }
 
     @objc private func toggleCheckboxButton(_ sender: UIControl) {
-        delegate?.toggleSelection(cell: self)
+        delegate?.toggleSelecting(cell: self)
     }
 }
 
-enum LocationCellBehavior {
-    case add
-    case select
-}
-
 extension LocationCell {
+    enum LocationCellBehavior {
+        case add
+        case select
+    }
+
     func configure(item: LocationCellViewModel, behavior: LocationCellBehavior) {
         accessibilityIdentifier = item.node.code
+        isDisabled = !item.node.isActive
         locationLabel.text = item.node.name
         showsCollapseControl = !item.node.children.isEmpty
         isExpanded = item.node.showsChildren
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationCellViewModel.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationCellViewModel.swift
index 50418e21d8..df0a3ba62c 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationCellViewModel.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationCellViewModel.swift
@@ -12,7 +12,7 @@ struct LocationCellViewModel: Hashable {
     let section: LocationSection
     let node: LocationNode
     var indentationLevel = 0
-    var isSelected: Bool = false
+    var isSelected = false
 
     func hash(into hasher: inout Hasher) {
         hasher.combine(section)
@@ -21,8 +21,8 @@ struct LocationCellViewModel: Hashable {
 
     static func == (lhs: Self, rhs: Self) -> Bool {
         lhs.node == rhs.node &&
-        lhs.section == rhs.section &&
-        lhs.isSelected == rhs.isSelected
+            lhs.section == rhs.section &&
+            lhs.isSelected == rhs.isSelected
     }
 }
 
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
index 9c1ca8ca06..8340dbf40e 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
@@ -12,12 +12,15 @@ import MullvadSettings
 import MullvadTypes
 import UIKit
 
-final class LocationDataSource: UITableViewDiffableDataSource<LocationSection, LocationCellViewModel> {
+final class LocationDataSource:
+    UITableViewDiffableDataSource<LocationSection, LocationCellViewModel>,
+    LocationDiffableDataSourceProtocol {
     private var currentSearchString = ""
-    private let tableView: UITableView
     private var dataSources: [LocationDataSourceProtocol] = []
     private var selectedItem: LocationCellViewModel?
     private var hasFilter = false
+    let tableView: UITableView
+    let sections: [LocationSection]
 
     var didSelectRelayLocations: ((UserSelectedRelays) -> Void)?
     var didTapEditCustomLists: (() -> Void)?
@@ -28,6 +31,10 @@ final class LocationDataSource: UITableViewDiffableDataSource<LocationSection, L
         customLists: LocationDataSourceProtocol
     ) {
         self.tableView = tableView
+
+        let sections: [LocationSection] = LocationSection.allCases
+        self.sections = sections
+
         #if DEBUG
         self.dataSources.append(customLists)
         #endif
@@ -35,10 +42,9 @@ final class LocationDataSource: UITableViewDiffableDataSource<LocationSection, L
 
         super.init(tableView: tableView) { _, indexPath, itemIdentifier in
             let cell = tableView.dequeueReusableView(
-                withIdentifier: LocationSection.allCases[indexPath.section],
+                withIdentifier: sections[indexPath.section],
                 for: indexPath
-                // swiftlint:disable:next force_cast
-            ) as! LocationCell
+            ) as! LocationCell // swiftlint:disable:this force_cast
             cell.configure(item: itemIdentifier, behavior: .select)
             return cell
         }
@@ -71,7 +77,7 @@ final class LocationDataSource: UITableViewDiffableDataSource<LocationSection, L
     func filterRelays(by searchString: String, scrollToSelected: Bool = true) {
         currentSearchString = searchString
 
-        let list = LocationSection.allCases.enumerated().map { index, section in
+        let list = sections.enumerated().map { index, section in
             dataSources[index]
                 .search(by: searchString)
                 .flatMap { node in
@@ -109,33 +115,16 @@ final class LocationDataSource: UITableViewDiffableDataSource<LocationSection, L
         filterRelays(by: currentSearchString, scrollToSelected: false)
     }
 
-    private func indexPathForSelectedRelay() -> IndexPath? {
-        selectedItem.flatMap { indexPath(for: $0) }
+    func nodeShowsChildren(_ node: LocationNode) -> Bool {
+        node.showsChildren
     }
 
-    private func updateDataSnapshot(
-        with list: [[LocationCellViewModel]],
-        reloadExisting: Bool = false,
-        animated: Bool = false,
-        completion: (() -> Void)? = nil
-    ) {
-        var snapshot = NSDiffableDataSourceSnapshot<LocationSection, LocationCellViewModel>()
-        let sections = LocationSection.allCases
-
-        snapshot.appendSections(sections)
-        for (index, section) in sections.enumerated() {
-            let items = list[index]
-
-            snapshot.appendItems(items, toSection: section)
-
-            if reloadExisting {
-                snapshot.reconfigureOrReloadItems(items)
-            }
-        }
+    func nodeShouldBeSelected(_ node: LocationNode) -> Bool {
+        false
+    }
 
-        DispatchQueue.main.async {
-            self.apply(snapshot, animatingDifferences: animated, completion: completion)
-        }
+    private func indexPathForSelectedRelay() -> IndexPath? {
+        selectedItem.flatMap { indexPath(for: $0) }
     }
 
     private func mapSelectedItem(from selectedRelays: UserSelectedRelays?) {
@@ -165,11 +154,13 @@ final class LocationDataSource: UITableViewDiffableDataSource<LocationSection, L
 
         let rootNode = selectedItem.node.root
 
+        // Exit early if no changes to the node tree are necessary.
         guard selectedItem.node != rootNode else {
             completion?()
             return
         }
 
+        // Make sure we have an index path for the selected item.
         guard let indexPath = indexPath(for: LocationCellViewModel(
             section: selectedItem.section,
             node: rootNode
@@ -180,16 +171,18 @@ final class LocationDataSource: UITableViewDiffableDataSource<LocationSection, L
             node.showsChildren = true
         }
 
+        // Construct node tree.
         let nodesToAdd = recursivelyCreateCellViewModelTree(
             for: rootNode,
             in: selectedItem.section,
             indentationLevel: 1
         )
 
+        // Insert the new node tree below the select item.
         var snapshotItems = snapshot().itemIdentifiers(inSection: selectedItem.section)
         snapshotItems.insert(contentsOf: nodesToAdd, at: indexPath.row + 1)
 
-        let list = LocationSection.allCases.enumerated().map { index, section in
+        let list = sections.enumerated().map { index, section in
             index == indexPath.section
                 ? snapshotItems
                 : snapshot().itemIdentifiers(inSection: section)
@@ -203,36 +196,6 @@ final class LocationDataSource: UITableViewDiffableDataSource<LocationSection, L
         )
     }
 
-    private func recursivelyCreateCellViewModelTree(
-        for node: LocationNode,
-        in section: LocationSection,
-        indentationLevel: Int
-    ) -> [LocationCellViewModel] {
-        var viewModels = [LocationCellViewModel]()
-
-        for childNode in node.children where !childNode.isHiddenFromSearch {
-            viewModels.append(
-                LocationCellViewModel(
-                    section: section,
-                    node: childNode,
-                    indentationLevel: indentationLevel
-                )
-            )
-
-            if childNode.showsChildren {
-                viewModels.append(
-                    contentsOf: recursivelyCreateCellViewModelTree(
-                        for: childNode,
-                        in: section,
-                        indentationLevel: indentationLevel + 1
-                    )
-                )
-            }
-        }
-
-        return viewModels
-    }
-
     override func tableView(_ tableView: UITableView, cellForRowAt indexPath: IndexPath) -> UITableViewCell {
         // swiftlint:disable:next force_cast
         let cell = super.tableView(tableView, cellForRowAt: indexPath) as! LocationCell
@@ -243,7 +206,7 @@ final class LocationDataSource: UITableViewDiffableDataSource<LocationSection, L
 
 extension LocationDataSource: UITableViewDelegate {
     func tableView(_ tableView: UITableView, viewForHeaderInSection section: Int) -> UIView? {
-        switch LocationSection.allCases[section] {
+        switch sections[section] {
         case .allLocations:
             return LocationSectionHeaderView(
                 configuration: LocationSectionHeaderView.Configuration(name: LocationSection.allLocations.description)
@@ -265,7 +228,7 @@ extension LocationDataSource: UITableViewDelegate {
     }
 
     func tableView(_ tableView: UITableView, heightForFooterInSection section: Int) -> CGFloat {
-        switch LocationSection.allCases[section] {
+        switch sections[section] {
         case .allLocations:
             return .zero
         case .customLists:
@@ -304,72 +267,31 @@ extension LocationDataSource: UITableViewDelegate {
 
         didSelectRelayLocations?(relayLocations)
     }
-}
 
-extension LocationDataSource: LocationCellDelegate {
-    func toggleSelection(cell: LocationCell) {
-        print("Just put a print statement // Marco")
+    private func scrollToTop(animated: Bool) {
+        tableView.setContentOffset(.zero, animated: animated)
     }
 
+    private func scrollToSelectedRelay() {
+        indexPathForSelectedRelay().flatMap {
+            tableView.scrollToRow(at: $0, at: .middle, animated: false)
+        }
+    }
+}
+
+extension LocationDataSource: LocationCellDelegate {
     func toggleExpanding(cell: LocationCell) {
         guard let indexPath = tableView.indexPath(for: cell),
               let item = itemIdentifier(for: indexPath) else { return }
 
-        let sections = LocationSection.allCases
-        let section = sections[indexPath.section]
-        let isExpanded = item.node.showsChildren
-        var locationList = snapshot().itemIdentifiers(inSection: section)
-
-        item.node.showsChildren = !isExpanded
-
-        if !isExpanded {
-            locationList.addSubNodes(from: item, at: indexPath)
-        } else {
-            locationList.removeSubNodes(from: item.node)
-        }
-
-        let list = sections.enumerated().map { index, section in
-            index == indexPath.section
-                ? locationList
-                : snapshot().itemIdentifiers(inSection: section)
-        }
+        let items = toggledItems(for: cell)
 
-        updateDataSnapshot(with: list, reloadExisting: true, completion: {
+        updateDataSnapshot(with: items, reloadExisting: true, completion: {
             self.scroll(to: item, animated: true)
         })
     }
-}
-
-extension LocationDataSource {
-    private func scroll(to item: LocationCellViewModel, animated: Bool) {
-        guard
-            let visibleIndexPaths = tableView.indexPathsForVisibleRows,
-            let indexPath = indexPath(for: item)
-        else { return }
-
-        if item.node.children.count > visibleIndexPaths.count {
-            tableView.scrollToRow(at: indexPath, at: .top, animated: animated)
-        } else {
-            if let last = item.node.children.last {
-                if let lastInsertedIndexPath = self.indexPath(for: LocationCellViewModel(
-                    section: LocationSection.allCases[indexPath.section],
-                    node: last
-                )),
-                    let lastVisibleIndexPath = visibleIndexPaths.last,
-                    lastInsertedIndexPath >= lastVisibleIndexPath {
-                    tableView.scrollToRow(at: lastInsertedIndexPath, at: .bottom, animated: animated)
-                }
-            }
-        }
-    }
 
-    private func scrollToTop(animated: Bool) {
-        tableView.setContentOffset(.zero, animated: animated)
-    }
-
-    private func scrollToSelectedRelay() {
-        indexPathForSelectedRelay().flatMap {
-            tableView.scrollToRow(at: $0, at: .middle, animated: false)
-        }
+    func toggleSelecting(cell: LocationCell) {
+        // No op.
     }
 }
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationDiffableDataSourceProtocol.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationDiffableDataSourceProtocol.swift
new file mode 100644
index 0000000000..0450be0a81
--- /dev/null
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationDiffableDataSourceProtocol.swift
@@ -0,0 +1,119 @@
+//
+//  LocationDiffableDataSourceProtocol.swift
+//  MullvadVPNUITests
+//
+//  Created by Jon Petersson on 2024-03-27.
+//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
+//
+
+import MullvadTypes
+import UIKit
+
+protocol LocationDiffableDataSourceProtocol: UITableViewDiffableDataSource<LocationSection, LocationCellViewModel> {
+    var tableView: UITableView { get }
+    var sections: [LocationSection] { get }
+    func nodeShowsChildren(_ node: LocationNode) -> Bool
+    func nodeShouldBeSelected(_ node: LocationNode) -> Bool
+}
+
+extension LocationDiffableDataSourceProtocol {
+    func scroll(to item: LocationCellViewModel, animated: Bool) {
+        guard
+            let visibleIndexPaths = tableView.indexPathsForVisibleRows,
+            let indexPath = indexPath(for: item)
+        else { return }
+
+        if item.node.children.count > visibleIndexPaths.count {
+            tableView.scrollToRow(at: indexPath, at: .top, animated: animated)
+        } else {
+            if let last = item.node.children.last {
+                if let lastInsertedIndexPath = self.indexPath(for: LocationCellViewModel(
+                    section: sections[indexPath.section],
+                    node: last
+                )),
+                    let lastVisibleIndexPath = visibleIndexPaths.last,
+                    lastInsertedIndexPath >= lastVisibleIndexPath {
+                    tableView.scrollToRow(at: lastInsertedIndexPath, at: .bottom, animated: animated)
+                }
+            }
+        }
+    }
+
+    func toggledItems(for cell: LocationCell) -> [[LocationCellViewModel]] {
+        guard let indexPath = tableView.indexPath(for: cell),
+              let item = itemIdentifier(for: indexPath) else { return [[]] }
+
+        let section = sections[indexPath.section]
+        let isExpanded = item.node.showsChildren
+        var locationList = snapshot().itemIdentifiers(inSection: section)
+
+        item.node.showsChildren = !isExpanded
+
+        if !isExpanded {
+            locationList.addSubNodes(from: item, at: indexPath)
+        } else {
+            locationList.removeSubNodes(from: item.node)
+        }
+
+        return sections.enumerated().map { index, section in
+            index == indexPath.section
+                ? locationList
+                : snapshot().itemIdentifiers(inSection: section)
+        }
+    }
+
+    func updateDataSnapshot(
+        with list: [[LocationCellViewModel]],
+        reloadExisting: Bool = false,
+        animated: Bool = false,
+        completion: (() -> Void)? = nil
+    ) {
+        var snapshot = NSDiffableDataSourceSnapshot<LocationSection, LocationCellViewModel>()
+
+        snapshot.appendSections(sections)
+        for (index, section) in sections.enumerated() {
+            let items = list[index]
+
+            snapshot.appendItems(items, toSection: section)
+
+            if reloadExisting {
+                snapshot.reconfigureOrReloadItems(items)
+            }
+        }
+
+        DispatchQueue.main.async {
+            self.apply(snapshot, animatingDifferences: animated, completion: completion)
+        }
+    }
+
+    func recursivelyCreateCellViewModelTree(
+        for node: LocationNode,
+        in section: LocationSection,
+        indentationLevel: Int
+    ) -> [LocationCellViewModel] {
+        var viewModels = [LocationCellViewModel]()
+
+        for childNode in node.children where !childNode.isHiddenFromSearch {
+            viewModels.append(
+                LocationCellViewModel(
+                    section: section,
+                    node: childNode,
+                    indentationLevel: indentationLevel,
+                    isSelected: nodeShouldBeSelected(childNode)
+                )
+            )
+
+            if nodeShowsChildren(childNode) {
+                viewModels.append(
+                    contentsOf: recursivelyCreateCellViewModelTree(
+                        for: childNode,
+                        in: section,
+                        indentationLevel: indentationLevel + 1
+                    )
+                )
+            }
+        }
+
+        return viewModels
+    }
+}
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationNode.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationNode.swift
index ac9277d63a..fbf2fbf8fb 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationNode.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationNode.swift
@@ -13,6 +13,7 @@ class LocationNode {
     let name: String
     var code: String
     var locations: [RelayLocation]
+    var isActive: Bool
     weak var parent: LocationNode?
     var children: [LocationNode]
     var showsChildren: Bool
@@ -22,6 +23,7 @@ class LocationNode {
         name: String,
         code: String,
         locations: [RelayLocation] = [],
+        isActive: Bool = true,
         parent: LocationNode? = nil,
         children: [LocationNode] = [],
         showsChildren: Bool = false,
@@ -30,6 +32,7 @@ class LocationNode {
         self.name = name
         self.code = code
         self.locations = locations
+        self.isActive = isActive
         self.parent = parent
         self.children = children
         self.showsChildren = showsChildren
@@ -91,6 +94,7 @@ extension LocationNode {
             name: name,
             code: code,
             locations: locations,
+            isActive: isActive,
             parent: parent,
             children: [],
             showsChildren: showsChildren,
@@ -137,6 +141,7 @@ class CustomListLocationNode: LocationNode {
         name: String,
         code: String,
         locations: [RelayLocation] = [],
+        isActive: Bool = true,
         parent: LocationNode? = nil,
         children: [LocationNode] = [],
         showsChildren: Bool = false,
@@ -149,6 +154,7 @@ class CustomListLocationNode: LocationNode {
             name: name,
             code: code,
             locations: locations,
+            isActive: isActive,
             parent: parent,
             children: children,
             showsChildren: showsChildren,
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationSectionHeaderView.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationSectionHeaderView.swift
index 49c9cbce20..4a137d9cc1 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationSectionHeaderView.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationSectionHeaderView.swift
@@ -74,7 +74,9 @@ class LocationSectionHeaderView: UIView, UIContentView {
 
     private func applyAppearance() {
         backgroundColor = .primaryColor
-        directionalLayoutMargins = NSDirectionalEdgeInsets(top: 8, leading: 16, bottom: 8, trailing: 24)
+
+        let leadingInset = UIMetrics.locationCellLayoutMargins.leading + 6
+        directionalLayoutMargins = NSDirectionalEdgeInsets(top: 8, leading: leadingInset, bottom: 8, trailing: 24)
     }
 }