refactoring location datasource

12 files changed, 608 insertions, 432 deletions

diff --git a/ios/MullvadVPN.xcodeproj/project.pbxproj b/ios/MullvadVPN.xcodeproj/project.pbxproj
index 17a4b9ae7f..888d2782ac 100644
--- a/ios/MullvadVPN.xcodeproj/project.pbxproj
+++ b/ios/MullvadVPN.xcodeproj/project.pbxproj
@@ -789,6 +789,13 @@
 		F050AE582B7376C6003F4EDB /* CustomListRepository.swift in Sources */ = {isa = PBXBuildFile; fileRef = F050AE562B7376C6003F4EDB /* CustomListRepository.swift */; };
 		F050AE5A2B7376F4003F4EDB /* CustomList.swift in Sources */ = {isa = PBXBuildFile; fileRef = F050AE592B7376F4003F4EDB /* CustomList.swift */; };
 		F050AE5C2B73797D003F4EDB /* CustomListRepositoryTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = F050AE5B2B73797D003F4EDB /* CustomListRepositoryTests.swift */; };
+		F050AE4C2B70D5A7003F4EDB /* SelectLocationNode.swift in Sources */ = {isa = PBXBuildFile; fileRef = F050AE4B2B70D5A7003F4EDB /* SelectLocationNode.swift */; };
+		F050AE4E2B70D7F8003F4EDB /* LocationCellViewModel.swift in Sources */ = {isa = PBXBuildFile; fileRef = F050AE4D2B70D7F8003F4EDB /* LocationCellViewModel.swift */; };
+		F050AE502B70DC4F003F4EDB /* SelectLocationNodeProtocol.swift in Sources */ = {isa = PBXBuildFile; fileRef = F050AE4F2B70DC4F003F4EDB /* SelectLocationNodeProtocol.swift */; };
+		F050AE522B70DFC0003F4EDB /* SelectLocationGroup.swift in Sources */ = {isa = PBXBuildFile; fileRef = F050AE512B70DFC0003F4EDB /* SelectLocationGroup.swift */; };
+		F050AE5E2B739A73003F4EDB /* LocationDataSourceProtocol.swift in Sources */ = {isa = PBXBuildFile; fileRef = F050AE5D2B739A73003F4EDB /* LocationDataSourceProtocol.swift */; };
+		F050AE602B73A41E003F4EDB /* AllLocationDataSource.swift in Sources */ = {isa = PBXBuildFile; fileRef = F050AE5F2B73A41E003F4EDB /* AllLocationDataSource.swift */; };
+		F050AE622B74DBAC003F4EDB /* CustomListsDataSource.swift in Sources */ = {isa = PBXBuildFile; fileRef = F050AE612B74DBAC003F4EDB /* CustomListsDataSource.swift */; };
 		F05F39942B21C6C6006E60A7 /* relays.json in Resources */ = {isa = PBXBuildFile; fileRef = 58F3C0A524A50155003E76BE /* relays.json */; };
 		F05F39972B21C735006E60A7 /* RelayCache.swift in Sources */ = {isa = PBXBuildFile; fileRef = 5820675A26E6576800655B05 /* RelayCache.swift */; };
 		F05F39982B21C73C006E60A7 /* CachedRelays.swift in Sources */ = {isa = PBXBuildFile; fileRef = 585DA87626B024A600B8C587 /* CachedRelays.swift */; };
@@ -1898,6 +1905,13 @@
 		F050AE562B7376C6003F4EDB /* CustomListRepository.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = CustomListRepository.swift; sourceTree = "<group>"; };
 		F050AE592B7376F4003F4EDB /* CustomList.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = CustomList.swift; sourceTree = "<group>"; };
 		F050AE5B2B73797D003F4EDB /* CustomListRepositoryTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = CustomListRepositoryTests.swift; sourceTree = "<group>"; };
+		F050AE4B2B70D5A7003F4EDB /* SelectLocationNode.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SelectLocationNode.swift; sourceTree = "<group>"; };
+		F050AE4D2B70D7F8003F4EDB /* LocationCellViewModel.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = LocationCellViewModel.swift; sourceTree = "<group>"; };
+		F050AE4F2B70DC4F003F4EDB /* SelectLocationNodeProtocol.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SelectLocationNodeProtocol.swift; sourceTree = "<group>"; };
+		F050AE512B70DFC0003F4EDB /* SelectLocationGroup.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SelectLocationGroup.swift; sourceTree = "<group>"; };
+		F050AE5D2B739A73003F4EDB /* LocationDataSourceProtocol.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = LocationDataSourceProtocol.swift; sourceTree = "<group>"; };
+		F050AE5F2B73A41E003F4EDB /* AllLocationDataSource.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AllLocationDataSource.swift; sourceTree = "<group>"; };
+		F050AE612B74DBAC003F4EDB /* CustomListsDataSource.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = CustomListsDataSource.swift; sourceTree = "<group>"; };
 		F06045E52B231EB700B2D37A /* URLSessionTransport.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = URLSessionTransport.swift; sourceTree = "<group>"; };
 		F06045E92B23217E00B2D37A /* ShadowsocksTransport.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = ShadowsocksTransport.swift; sourceTree = "<group>"; };
 		F06045EB2B2322A500B2D37A /* Jittered.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Jittered.swift; sourceTree = "<group>"; };
@@ -2324,9 +2338,16 @@
 		583FE01729C196F3006E85F9 /* SelectLocation */ = {
 			isa = PBXGroup;
 			children = (
+				F050AE5F2B73A41E003F4EDB /* AllLocationDataSource.swift */,
+				F050AE612B74DBAC003F4EDB /* CustomListsDataSource.swift */,
 				58435AC129CB2A350099C71B /* LocationCellFactory.swift */,
+				F050AE4D2B70D7F8003F4EDB /* LocationCellViewModel.swift */,
 				583DA21325FA4B5C00318683 /* LocationDataSource.swift */,
+				F050AE5D2B739A73003F4EDB /* LocationDataSourceProtocol.swift */,
 				5888AD82227B11080051EB06 /* SelectLocationCell.swift */,
+				F050AE512B70DFC0003F4EDB /* SelectLocationGroup.swift */,
+				F050AE4B2B70D5A7003F4EDB /* SelectLocationNode.swift */,
+				F050AE4F2B70DC4F003F4EDB /* SelectLocationNodeProtocol.swift */,
 				5888AD86227B17950051EB06 /* SelectLocationViewController.swift */,
 			);
 			path = SelectLocation;
@@ -4995,6 +5016,7 @@
 				58968FAE28743E2000B799DC /* TunnelInteractor.swift in Sources */,
 				7A1A26472A29CF0800B978AA /* RelayFilterDataSource.swift in Sources */,
 				5864AF0929C78850005B0CD9 /* PreferencesCellFactory.swift in Sources */,
+				F050AE4E2B70D7F8003F4EDB /* LocationCellViewModel.swift in Sources */,
 				58CEB30C2AFD586600E6E088 /* DynamicBackgroundConfiguration.swift in Sources */,
 				587B7536266528A200DEF7E9 /* NotificationManager.swift in Sources */,
 				5820EDA9288FE064006BF4E4 /* DeviceManagementInteractor.swift in Sources */,
@@ -5008,6 +5030,7 @@
 				5871FB96254ADE4E0051A0A4 /* ConsolidatedApplicationLog.swift in Sources */,
 				F0E8E4C52A60499100ED26A3 /* AccountDeletionViewController.swift in Sources */,
 				7A9CCCC12A96302800DD6A34 /* AccountCoordinator.swift in Sources */,
+				F050AE502B70DC4F003F4EDB /* SelectLocationNodeProtocol.swift in Sources */,
 				58FEEB58260B662E00A621A8 /* AutomaticKeyboardResponder.swift in Sources */,
 				5846227326E22A160035F7C2 /* StorePaymentObserver.swift in Sources */,
 				F0E3618B2A4ADD2F00AEEF2B /* WelcomeContentView.swift in Sources */,
@@ -5049,6 +5072,7 @@
 				7A9CCCBE2A96302800DD6A34 /* AccountDeletionCoordinator.swift in Sources */,
 				588527B4276B4F2F00BAA373 /* SetAccountOperation.swift in Sources */,
 				58FF9FE02B075ABC00E4C97D /* EditAccessMethodViewController.swift in Sources */,
+				F050AE622B74DBAC003F4EDB /* CustomListsDataSource.swift in Sources */,
 				F0DA87472A9CB9A2006044F1 /* AccountExpiryRow.swift in Sources */,
 				585CA70F25F8C44600B47C62 /* UIMetrics.swift in Sources */,
 				E1187ABD289BBB850024E748 /* OutOfTimeContentView.swift in Sources */,
@@ -5081,6 +5105,7 @@
 				063F026628FFE11C001FA09F /* RESTCreateApplePaymentResponse+Localization.swift in Sources */,
 				58DF28A52417CB4B00E836B0 /* StorePaymentManager.swift in Sources */,
 				583DA21425FA4B5C00318683 /* LocationDataSource.swift in Sources */,
+				F050AE602B73A41E003F4EDB /* AllLocationDataSource.swift in Sources */,
 				587EB6742714520600123C75 /* PreferencesDataSourceDelegate.swift in Sources */,
 				582BB1AF229566420055B6EF /* SettingsCell.swift in Sources */,
 				7AF9BE8E2A331C7B00DBFEDB /* RelayFilterViewModel.swift in Sources */,
@@ -5186,7 +5211,9 @@
 				584592612639B4A200EF967F /* TermsOfServiceContentView.swift in Sources */,
 				5875960A26F371FC00BF6711 /* Tunnel+Messaging.swift in Sources */,
 				586C0D912B03D8A400E7CDD7 /* AccessMethodHeaderFooterReuseIdentifier.swift in Sources */,
+				F050AE4C2B70D5A7003F4EDB /* SelectLocationNode.swift in Sources */,
 				7A2960F62A963F7500389B82 /* AlertCoordinator.swift in Sources */,
+				F050AE522B70DFC0003F4EDB /* SelectLocationGroup.swift in Sources */,
 				063687BA28EB234F00BE7161 /* PacketTunnelTransport.swift in Sources */,
 				A9C342C12ACC37E30045F00E /* TunnelStatusBlockObserver.swift in Sources */,
 				587425C12299833500CA2045 /* RootContainerViewController.swift in Sources */,
@@ -5206,6 +5233,7 @@
 				5827B0BF2B14B37D00CCBBA1 /* Publisher+PreviousValue.swift in Sources */,
 				7A9CCCB62A96302800DD6A34 /* OutOfTimeCoordinator.swift in Sources */,
 				5827B0AA2B0F4C9100CCBBA1 /* EditAccessMethodViewControllerDelegate.swift in Sources */,
+				F050AE5E2B739A73003F4EDB /* LocationDataSourceProtocol.swift in Sources */,
 				7A5869A82B5140C200640D27 /* MethodSettingsValidationErrorContentView.swift in Sources */,
 				A99E5EE22B762ED30033F241 /* ProblemReportViewController+ViewManagement.swift in Sources */,
 				7A5869A22B502EA800640D27 /* MethodSettingsSectionIdentifier.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
deleted file mode 100644
index 02691892fe..0000000000
--- a/ios/MullvadVPN.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved
+++ /dev/null
@@ -1,22 +0,0 @@
-{
-  "pins" : [
-    {
-      "identity" : "swift-log",
-      "kind" : "remoteSourceControl",
-      "location" : "https://github.com/apple/swift-log.git",
-      "state" : {
-        "revision" : "173f567a2dfec11d74588eea82cecea555bdc0bc",
-        "version" : "1.4.0"
-      }
-    },
-    {
-      "identity" : "wireguard-apple",
-      "kind" : "remoteSourceControl",
-      "location" : "https://github.com/mullvad/wireguard-apple.git",
-      "state" : {
-        "revision" : "11a00c20dc03f2751db47e94f585c0778c7bde82"
-      }
-    }
-  ],
-  "version" : 2
-}
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/AllLocationDataSource.swift b/ios/MullvadVPN/View controllers/SelectLocation/AllLocationDataSource.swift
new file mode 100644
index 0000000000..0d4c6b9bce
--- /dev/null
+++ b/ios/MullvadVPN/View controllers/SelectLocation/AllLocationDataSource.swift
@@ -0,0 +1,95 @@
+//
+//  AllLocationDataSource.swift
+//  MullvadVPN
+//
+//  Created by Mojgan on 2024-02-07.
+//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
+//
+
+import Foundation
+import MullvadREST
+import MullvadTypes
+
+class AllLocationDataSource: LocationDataSourceProtocol {
+    var nodeByLocation = [RelayLocation: SelectLocationNode]()
+    private var locationList = [RelayLocation]()
+
+    func search(by text: String) -> [RelayLocation] {
+        if text.isEmpty {
+            return locationList
+        } else {
+            var filteredLocations: [RelayLocation] = []
+            locationList.forEach { location in
+                guard let countryNode = nodeByLocation[location] else { return }
+                countryNode.showsChildren = false
+
+                if text.isEmpty || countryNode.displayName.fuzzyMatch(text) {
+                    filteredLocations.append(countryNode.location)
+                }
+
+                for cityNode in countryNode.children {
+                    cityNode.showsChildren = false
+
+                    let relaysContainSearchString = cityNode.children
+                        .contains(where: { $0.displayName.fuzzyMatch(text) })
+
+                    if cityNode.displayName.fuzzyMatch(text) || relaysContainSearchString {
+                        if !filteredLocations.contains(where: { $0 == countryNode.location }) {
+                            filteredLocations.append(countryNode.location)
+                        }
+
+                        filteredLocations.append(cityNode.location)
+                        countryNode.showsChildren = true
+
+                        if relaysContainSearchString {
+                            cityNode.children.map { $0.location }.forEach {
+                                filteredLocations.append($0)
+                            }
+                            cityNode.showsChildren = true
+                        }
+                    }
+                }
+            }
+
+            return filteredLocations
+        }
+    }
+
+    func reload(
+        _ response: MullvadREST.REST.ServerRelaysResponse,
+        relays: [MullvadREST.REST.ServerRelay]
+    ) -> [RelayLocation] {
+        nodeByLocation.removeAll()
+        let rootNode = self.makeRootNode(name: SelectLocationGroup.allLocations.description)
+
+        for relay in relays {
+            guard case let .city(countryCode, 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 = nodeByLocation[ascendantOrSelf]?.showsChildren ?? false
+
+                let node = createNode(
+                    root: rootNode,
+                    ascendantOrSelf: ascendantOrSelf,
+                    serverLocation: serverLocation,
+                    relay: relay,
+                    wasShowingChildren: wasShowingChildren
+                )
+                nodeByLocation[ascendantOrSelf] = node
+            }
+        }
+
+        rootNode.sortChildrenRecursive()
+        rootNode.computeActiveChildrenRecursive()
+        locationList = rootNode.flatRelayLocationList()
+        return locationList
+    }
+}
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/CustomListsDataSource.swift b/ios/MullvadVPN/View controllers/SelectLocation/CustomListsDataSource.swift
new file mode 100644
index 0000000000..25df241bef
--- /dev/null
+++ b/ios/MullvadVPN/View controllers/SelectLocation/CustomListsDataSource.swift
@@ -0,0 +1,27 @@
+//
+//  CustomListsDataSource.swift
+//  MullvadVPN
+//
+//  Created by Mojgan on 2024-02-08.
+//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
+//
+
+import Foundation
+import MullvadREST
+import MullvadTypes
+
+class CustomListsDataSource: LocationDataSourceProtocol {
+    var nodeByLocation = [RelayLocation: SelectLocationNode]()
+    private var locationList = [RelayLocation]()
+
+    func search(by text: String) -> [RelayLocation] {
+        return []
+    }
+
+    func reload(
+        _ response: MullvadREST.REST.ServerRelaysResponse,
+        relays: [MullvadREST.REST.ServerRelay]
+    ) -> [RelayLocation] {
+        locationList
+    }
+}
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationCellFactory.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationCellFactory.swift
index 46e7d97aac..6e98f1f729 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationCellFactory.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationCellFactory.swift
@@ -10,22 +10,26 @@ import MullvadTypes
 import UIKit
 
 protocol LocationCellEventHandler {
-    func collapseCell(for item: RelayLocation)
+    func collapseCell(for item: LocationCellViewModel)
+    func node(for item: LocationCellViewModel) -> SelectLocationNode?
 }
 
 final class LocationCellFactory: CellFactoryProtocol {
-    var nodeByLocation = [RelayLocation: LocationDataSource.Node]()
     var delegate: LocationCellEventHandler?
     let tableView: UITableView
+    let reuseIdentifier: String
 
-    init(tableView: UITableView, nodeByLocation: [RelayLocation: LocationDataSource.Node]) {
+    init(
+        tableView: UITableView,
+        reuseIdentifier: String
+    ) {
         self.tableView = tableView
-        self.nodeByLocation = nodeByLocation
+        self.reuseIdentifier = reuseIdentifier
     }
 
-    func makeCell(for item: RelayLocation, indexPath: IndexPath) -> UITableViewCell {
+    func makeCell(for item: LocationCellViewModel, indexPath: IndexPath) -> UITableViewCell {
         let cell = tableView.dequeueReusableCell(
-            withIdentifier: LocationDataSource.CellReuseIdentifiers.locationCell.rawValue,
+            withIdentifier: reuseIdentifier,
             for: indexPath
         )
 
@@ -34,9 +38,9 @@ final class LocationCellFactory: CellFactoryProtocol {
         return cell
     }
 
-    func configureCell(_ cell: UITableViewCell, item: RelayLocation, indexPath: IndexPath) {
+    func configureCell(_ cell: UITableViewCell, item: LocationCellViewModel, indexPath: IndexPath) {
         guard let cell = cell as? SelectLocationCell,
-              let node = nodeByLocation[item] else { return }
+              let node = delegate?.node(for: item) else { return }
 
         cell.accessibilityIdentifier = node.location.stringRepresentation
         cell.isDisabled = !node.isActive
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationCellViewModel.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationCellViewModel.swift
new file mode 100644
index 0000000000..c90e504442
--- /dev/null
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationCellViewModel.swift
@@ -0,0 +1,14 @@
+//
+//  LocationCellViewModel.swift
+//  MullvadVPN
+//
+//  Created by Mojgan on 2024-02-05.
+//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
+//
+
+import MullvadTypes
+
+struct LocationCellViewModel: Hashable {
+    let group: SelectLocationGroup
+    let location: RelayLocation
+}
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
index e3891d214d..334c051f81 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
@@ -6,59 +6,32 @@
 //  Copyright © 2021 Mullvad VPN AB. All rights reserved.
 //
 
+import Combine
 import MullvadREST
 import MullvadTypes
 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 }
-}
-
-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]()
+final class LocationDataSource: UITableViewDiffableDataSource<SelectLocationGroup, LocationCellViewModel> {
     private var currentSearchString = ""
-
     private let tableView: UITableView
     private let locationCellFactory: LocationCellFactory
+    private var datasources: [LocationDataSourceProtocol] = []
 
-    private class func makeRootNode() -> Node {
-        Node(
-            type: .root,
-            location: RelayLocation.country("#root"),
-            displayName: "",
-            showsChildren: true,
-            isActive: true,
-            children: []
-        )
-    }
-
-    var selectedRelayLocation: RelayLocation?
+    var selectedRelayLocation: LocationCellViewModel?
     var didSelectRelayLocation: ((RelayLocation) -> Void)?
 
-    init(tableView: UITableView) {
+    init(
+        tableView: UITableView,
+        allLocations: LocationDataSourceProtocol,
+        customLists: LocationDataSourceProtocol
+    ) {
         self.tableView = tableView
+        self.datasources.append(customLists)
+        self.datasources.append(allLocations)
 
         let locationCellFactory = LocationCellFactory(
             tableView: tableView,
-            nodeByLocation: nodeByLocation
+            reuseIdentifier: SelectLocationGroup.Cell.locationCell.reuseIdentifier
         )
         self.locationCellFactory = locationCellFactory
 
@@ -77,282 +50,133 @@ final class LocationDataSource: UITableViewDiffableDataSource<Int, RelayLocation
         let relays = response.wireguard.relays.filter { relay in
             return RelaySelector.relayMatchesFilter(relay, filter: filter)
         }
-
-        let rootNode = Self.makeRootNode()
-        nodeByLocation.removeAll()
-
-        for relay in relays {
-            guard case let .city(countryCode, 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 = nodeByLocation[ascendantOrSelf]?.showsChildren ?? false
-
-                let node = createNode(
-                    ascendantOrSelf: ascendantOrSelf,
-                    serverLocation: serverLocation,
-                    relay: relay,
-                    rootNode: rootNode,
-                    wasShowingChildren: wasShowingChildren
-                )
-                nodeByLocation[ascendantOrSelf] = node
-            }
+        var list: [[LocationCellViewModel]] = []
+        for section in 0 ..< datasources.count {
+            list.append(
+                datasources[section]
+                    .reload(response, relays: relays)
+                    .map { LocationCellViewModel(group: SelectLocationGroup.allCases[section], location: $0) }
+            )
         }
-
-        rootNode.sortChildrenRecursive()
-        rootNode.computeActiveChildrenRecursive()
-        locationList = rootNode.flatRelayLocationList()
-
         filterRelays(by: currentSearchString)
     }
 
     func indexPathForSelectedRelay() -> IndexPath? {
-        selectedRelayLocation.flatMap { indexPath(for: $0) }
+        selectedRelayLocation.flatMap {
+            return indexPath(for: $0)
+        }
     }
 
     func filterRelays(by searchString: String) {
         currentSearchString = searchString
-
-        if currentSearchString.isEmpty {
-            return resetLocationList()
-        }
-
-        var filteredLocations = [RelayLocation]()
-
-        locationList.forEach { location in
-            guard let countryNode = nodeByLocation[location] else { return }
-            countryNode.showsChildren = false
-
-            if searchString.isEmpty || countryNode.displayName.fuzzyMatch(searchString) {
-                filteredLocations.append(countryNode.location)
-            }
-
-            for cityNode in countryNode.children {
-                cityNode.showsChildren = false
-
-                let relaysContainSearchString = cityNode.children.contains(where: { node in
-                    node.displayName.fuzzyMatch(searchString)
-                })
-
-                if cityNode.displayName.fuzzyMatch(searchString) || relaysContainSearchString {
-                    if !filteredLocations.contains(countryNode.location) {
-                        filteredLocations.append(countryNode.location)
-                    }
-
-                    filteredLocations.append(cityNode.location)
-                    countryNode.showsChildren = true
-
-                    if relaysContainSearchString {
-                        filteredLocations.append(contentsOf: cityNode.children.map { $0.location })
-                        cityNode.showsChildren = true
-                    }
-                }
-            }
-        }
-
-        updateDataSnapshot(with: filteredLocations, reloadExisting: true) { [weak self] in
-            self?.scrollToTop(animated: false)
-        }
-    }
-
-    private func createNode(
-        ascendantOrSelf: RelayLocation,
-        serverLocation: REST.ServerLocation,
-        relay: REST.ServerRelay,
-        rootNode: Node,
-        wasShowingChildren: Bool
-    ) -> Node {
-        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 let .city(countryCode, _):
-            node = Node(
-                type: .city,
-                location: ascendantOrSelf,
-                displayName: serverLocation.city,
-                showsChildren: wasShowingChildren,
-                isActive: true,
-                children: []
+        var list: [[LocationCellViewModel]] = []
+        for section in 0 ..< SelectLocationGroup.allCases.count {
+            list.append(
+                datasources[section]
+                    .search(by: searchString)
+                    .map { LocationCellViewModel(group: SelectLocationGroup.allCases[section], location: $0) }
             )
-            nodeByLocation[.country(countryCode)]!.addChild(node)
-
-        case let .hostname(countryCode, cityCode, _):
-            node = Node(
-                type: .relay,
-                location: ascendantOrSelf,
-                displayName: relay.hostname,
-                showsChildren: false,
-                isActive: relay.active,
-                children: []
-            )
-            nodeByLocation[.city(countryCode, cityCode)]!.addChild(node)
         }
-
-        return node
+        updateDataSnapshot(with: list, reloadExisting: !searchString.isEmpty)
+        if searchString.isEmpty {
+            self.setSelectedRelayLocation(self.selectedRelayLocation, animated: false, completion: {
+                self.scrollToSelectedRelay()
+            })
+        } else {
+            self.scrollToTop(animated: false)
+        }
     }
 
     private func updateDataSnapshot(
-        with locations: [RelayLocation],
+        with list: [[LocationCellViewModel]],
         reloadExisting: Bool = false,
         animated: Bool = false,
         completion: (() -> Void)? = nil
     ) {
-        updateCellFactory(with: nodeByLocation)
-
-        var snapshot = NSDiffableDataSourceSnapshot<Int, RelayLocation>()
+        var snapshot = NSDiffableDataSourceSnapshot<SelectLocationGroup, LocationCellViewModel>()
 
-        snapshot.appendSections([0])
-        snapshot.appendItems(locations)
+        let sections = Array(SelectLocationGroup.allCases)
+        snapshot.appendSections(sections)
+        for (index, section) in sections.enumerated() {
+            snapshot.appendItems(list[index], toSection: section)
+        }
 
         if reloadExisting {
-            snapshot.reloadItems(locations)
+            snapshot.reloadSections(SelectLocationGroup.allCases)
         }
 
         apply(snapshot, animatingDifferences: animated, completion: completion)
     }
 
     private func registerClasses() {
-        CellReuseIdentifiers.allCases.forEach { enumCase in
+        SelectLocationGroup.allCases.forEach {
             tableView.register(
-                enumCase.reusableViewClass,
-                forCellReuseIdentifier: enumCase.rawValue
+                $0.cell.reusableViewClass,
+                forCellReuseIdentifier: $0.cell.reuseIdentifier
             )
         }
     }
 
-    private func updateCellFactory(with nodeByLocation: [RelayLocation: Node]) {
-        locationCellFactory.nodeByLocation = nodeByLocation
+    private func item(for indexPath: IndexPath) -> LocationCellViewModel? {
+        itemIdentifier(for: indexPath)
     }
 
     private func setSelectedRelayLocation(
-        _ relayLocation: RelayLocation?,
+        _ relayLocation: LocationCellViewModel?,
         animated: Bool,
         completion: (() -> Void)? = nil
     ) {
         selectedRelayLocation = relayLocation
-        var locationList = snapshot().itemIdentifiers
-
-        guard let selectedRelayLocation,
-              !locationList.contains(selectedRelayLocation) else { return }
-
-        let selectedLocationTree = selectedRelayLocation.ascendants + [selectedRelayLocation]
-
-        guard let topLocation = selectedLocationTree.first,
-              let topNode = nodeByLocation[topLocation],
-              let indexPath = indexPath(for: topLocation)
-        else {
-            return
-        }
-
-        selectedLocationTree.forEach { location in
-            nodeByLocation[location]?.showsChildren = true
-        }
-
-        locationList.addLocations(topNode.flatRelayLocationList(), at: indexPath.row + 1)
-        updateDataSnapshot(with: locationList, reloadExisting: true, animated: animated, completion: completion)
-    }
-
-    private func toggleChildren(
-        _ relayLocation: RelayLocation,
-        show: Bool,
-        animated: Bool
-    ) {
-        guard let node = nodeByLocation[relayLocation],
-              let indexPath = indexPath(for: node.location),
-              let cell = tableView.cellForRow(at: indexPath) else { return }
-
-        node.showsChildren = show
-        locationCellFactory.configureCell(cell, item: node.location, indexPath: indexPath)
+        selectedRelayLocation
+            .flatMap { item in
+                let group = item.group
+                var locationList = snapshot().itemIdentifiers(inSection: group)
+                guard !locationList.contains(item) else {
+                    completion?()
+                    return
+                }
+                let selectedLocationTree = item.location.ascendants + [item.location]
 
-        var locationList = snapshot().itemIdentifiers
-        let locationsToEdit = node.flatRelayLocationList()
+                guard let first = selectedLocationTree.first else { return }
+                let topLocation = LocationCellViewModel(group: group, location: first)
 
-        if show {
-            locationList.addLocations(locationsToEdit, at: indexPath.row + 1)
-        } else {
-            locationsToEdit.forEach { nodeByLocation[$0]?.showsChildren = false }
-            locationList.removeLocations(locationsToEdit)
-        }
-
-        updateDataSnapshot(with: locationList, animated: animated) { [weak self] in
-            guard let visibleIndexPaths = self?.tableView.indexPathsForVisibleRows else { return }
+                guard let indexPath = indexPath(for: topLocation),
+                      let topNode = node(for: topLocation) else {
+                    return
+                }
 
-            let scrollToNodeTop = {
-                if let firstInsertedIndexPath = self?.indexPath(for: node.location) {
-                    self?.tableView.scrollToRow(
-                        at: firstInsertedIndexPath,
-                        at: .top,
-                        animated: animated
-                    )
+                selectedLocationTree.forEach { location in
+                    node(for: LocationCellViewModel(group: group, location: location))?.showsChildren = true
                 }
-            }
 
-            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
-                    )
+                locationList.addLocations(
+                    topNode.flatRelayLocationList().map { LocationCellViewModel(group: group, location: $0) },
+                    at: indexPath.row + 1
+                )
+                var list: [[LocationCellViewModel]] = Array(repeating: [], count: datasources.count)
+                for index in 0 ..< list.count {
+                    list[index] = (index == indexPath.section)
+                        ? locationList
+                        : snapshot().itemIdentifiers(inSection: SelectLocationGroup.allCases[index])
                 }
-            }
 
-            if node.children.count > visibleIndexPaths.count {
-                scrollToNodeTop()
-            } else {
-                scrollToNodeBottom()
+                updateDataSnapshot(
+                    with: list,
+                    reloadExisting: true,
+                    animated: animated,
+                    completion: completion
+                )
             }
-        }
-    }
-
-    private func resetLocationList() {
-        nodeByLocation.values.forEach { $0.showsChildren = false }
-
-        updateDataSnapshot(with: locationList, reloadExisting: true)
-        setSelectedRelayLocation(selectedRelayLocation, animated: false)
-
-        if let indexPath = indexPathForSelectedRelay() {
-            tableView.scrollToRow(at: indexPath, at: .middle, animated: false)
-        }
-    }
-
-    private func item(for indexPath: IndexPath) -> LocationDataSourceItemProtocol? {
-        itemIdentifier(for: indexPath).flatMap { nodeByLocation[$0] }
-    }
-
-    private func scrollToTop(animated: Bool) {
-        tableView.setContentOffset(.zero, animated: animated)
     }
 }
 
 extension LocationDataSource: UITableViewDelegate {
     func tableView(_ tableView: UITableView, shouldHighlightRowAt indexPath: IndexPath) -> Bool {
-        item(for: indexPath)?.isActive ?? false
+        item(for: indexPath).flatMap { node(for: $0) }?.isActive ?? false
     }
 
     func tableView(_ tableView: UITableView, indentationLevelForRowAt indexPath: IndexPath) -> Int {
-        item(for: indexPath)?.indentationLevel ?? 0
+        item(for: indexPath).flatMap { node(for: $0) }?.indentationLevel ?? 0
     }
 
     func tableView(
@@ -361,168 +185,101 @@ extension LocationDataSource: UITableViewDelegate {
         forRowAt indexPath: IndexPath
     ) {
         if let item = item(for: indexPath),
-           item.location == selectedRelayLocation {
+           item == selectedRelayLocation {
             cell.setSelected(true, animated: false)
         }
     }
 
     func tableView(_ tableView: UITableView, didSelectRowAt indexPath: IndexPath) {
-        guard let item = item(for: indexPath),
-              item.location != selectedRelayLocation
-        else {
-            return
-        }
-
-        if let indexPath = indexPathForSelectedRelay(),
-           let cell = tableView.cellForRow(at: indexPath) {
-            cell.setSelected(false, animated: false)
-        }
-
-        setSelectedRelayLocation(
-            item.location,
-            animated: false
-        )
-
-        didSelectRelayLocation?(item.location)
+        item(for: indexPath)
+            .flatMap { item in
+                guard item.location != selectedRelayLocation?.location else { return }
+                didSelectRelayLocation?(item.location)
+                setSelectedRelayLocation(item, animated: false)
+                indexPathForSelectedRelay().flatMap {
+                    let cell = tableView.cellForRow(at: $0)
+                    cell?.setSelected(false, animated: false)
+                }
+            }
     }
 }
 
 extension LocationDataSource: LocationCellEventHandler {
-    func collapseCell(for item: RelayLocation) {
-        guard let node = nodeByLocation[item] else { return }
-
-        toggleChildren(
-            item,
-            show: !node.showsChildren,
-            animated: true
-        )
-    }
-}
-
-extension LocationDataSource {
-    enum NodeType {
-        case root
-        case country
-        case city
-        case relay
-    }
-
-    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
+    func collapseCell(for item: LocationCellViewModel) {
+        indexPath(for: item).flatMap { indexPath in
+            guard let node = self.node(for: item),
+                  let cell = tableView.cellForRow(at: indexPath) else { return }
+            let isExpanded = node.showsChildren
+            let group = SelectLocationGroup.allCases[indexPath.section]
+            node.showsChildren = !isExpanded
+            locationCellFactory.configureCell(
+                cell,
+                item: LocationCellViewModel(group: group, location: node.location),
+                indexPath: indexPath
+            )
+            var locationList = snapshot().itemIdentifiers(inSection: group)
+            let locationsToEdit = node.flatRelayLocationList().map { LocationCellViewModel(group: group, location: $0) }
+            if !isExpanded {
+                locationList.addLocations(locationsToEdit, at: indexPath.row + 1)
+            } else {
+                locationsToEdit.forEach { self.node(for: $0)?.showsChildren = false }
+                locationList.removeLocations(locationsToEdit)
             }
-        }
-
-        var indentationLevel: Int {
-            switch nodeType {
-            case .root, .country:
-                return 0
-            case .city:
-                return 1
-            case .relay:
-                return 2
+            var list: [[LocationCellViewModel]] = Array(repeating: [], count: datasources.count)
+            for index in 0 ..< list.count {
+                list[index] = (index == indexPath.section)
+                    ? locationList
+                    : snapshot().itemIdentifiers(inSection: SelectLocationGroup.allCases[index])
             }
+            self.updateDataSnapshot(with: list, completion: {
+                self.scroll(to: item, animated: true)
+            })
         }
+    }
 
-        init(
-            type: NodeType,
-            location: RelayLocation,
-            displayName: String,
-            showsChildren: Bool,
-            isActive: Bool,
-            children: [Node]
-        ) {
-            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 node(for item: LocationCellViewModel) -> SelectLocationNode? {
+        guard let sectionIndex = SelectLocationGroup.allCases.firstIndex(of: item.group) else {
+            return nil
         }
+        return datasources[sectionIndex].nodeByLocation[item.location]
+    }
+}
 
-        func sortChildrenRecursive() {
-            sortChildren()
-            children.forEach { node in
-                node.sortChildrenRecursive()
-            }
-        }
+extension LocationDataSource {
+    private func scroll(to location: LocationCellViewModel, animated: Bool) {
+        guard let visibleIndexPaths = tableView.indexPathsForVisibleRows,
+              let indexPath = indexPath(for: location),
+              let node = node(for: location) else { return }
 
-        func computeActiveChildrenRecursive() {
-            switch nodeType {
-            case .root, .country:
-                for node in children {
-                    node.computeActiveChildrenRecursive()
+        if node.children.count > visibleIndexPaths.count {
+            tableView.scrollToRow(at: indexPath, at: .top, animated: animated)
+        } else {
+            node.children.last.flatMap { last in
+                if let lastInsertedIndexPath = self.indexPath(for: LocationCellViewModel(
+                    group: SelectLocationGroup.allCases[indexPath.section],
+                    location: last.location
+                )),
+                    let lastVisibleIndexPath = visibleIndexPaths.last,
+                    lastInsertedIndexPath >= lastVisibleIndexPath {
+                    tableView.scrollToRow(at: lastInsertedIndexPath, at: .bottom, animated: animated)
                 }
-                fallthrough
-            case .city:
-                isActive = children.contains(where: { node -> Bool in
-                    node.isActive
-                })
-            case .relay:
-                break
-            }
-        }
-
-        func flatRelayLocationList(includeHiddenChildren: Bool = false) -> [RelayLocation] {
-            children.reduce(into: []) { array, node in
-                Self.flatten(node: node, into: &array, includeHiddenChildren: includeHiddenChildren)
             }
         }
+    }
 
-        private func sortChildren() {
-            switch nodeType {
-            case .root, .country:
-                children.sort { a, b -> Bool in
-                    lexicalSortComparator(a.displayName, b.displayName)
-                }
-            case .city:
-                children.sort { a, b -> Bool in
-                    fileSortComparator(
-                        a.location.stringRepresentation,
-                        b.location.stringRepresentation
-                    )
-                }
-            case .relay:
-                break
-            }
-        }
+    private func scrollToTop(animated: Bool) {
+        tableView.setContentOffset(.zero, animated: animated)
+    }
 
-        private class func flatten(node: Node, into array: inout [RelayLocation], includeHiddenChildren: Bool) {
-            array.append(node.location)
-            if includeHiddenChildren || node.showsChildren {
-                for child in node.children {
-                    Self.flatten(node: child, into: &array, includeHiddenChildren: includeHiddenChildren)
-                }
-            }
+    private func scrollToSelectedRelay() {
+        indexPathForSelectedRelay().flatMap {
+            tableView.scrollToRow(at: $0, at: .middle, animated: false)
         }
     }
 }
 
-private func lexicalSortComparator(_ a: String, _ b: String) -> Bool {
-    a.localizedCaseInsensitiveCompare(b) == .orderedAscending
-}
-
-private func fileSortComparator(_ a: String, _ b: String) -> Bool {
-    a.localizedStandardCompare(b) == .orderedAscending
-}
-
-private extension [RelayLocation] {
-    mutating func addLocations(_ locations: [RelayLocation], at index: Int) {
+private extension [LocationCellViewModel] {
+    mutating func addLocations(_ locations: [LocationCellViewModel], at index: Int) {
         if index < count {
             insert(contentsOf: locations, at: index)
         } else {
@@ -530,11 +287,9 @@ private extension [RelayLocation] {
         }
     }
 
-    mutating func removeLocations(_ locations: [RelayLocation]) {
+    mutating func removeLocations(_ locations: [LocationCellViewModel]) {
         removeAll(where: { location in
             locations.contains(location)
         })
     }
-
-    // swiftlint:disable:next file_length
 }
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSourceProtocol.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSourceProtocol.swift
new file mode 100644
index 0000000000..f00cdf3408
--- /dev/null
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSourceProtocol.swift
@@ -0,0 +1,68 @@
+//
+//  LocationDataSourceProtocol.swift
+//  MullvadVPN
+//
+//  Created by Mojgan on 2024-02-07.
+//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
+//
+
+import Foundation
+import MullvadREST
+import MullvadTypes
+import UIKit
+
+protocol LocationDataSourceProtocol {
+    var nodeByLocation: [RelayLocation: SelectLocationNode] { get }
+
+    func search(by text: String) -> [RelayLocation]
+
+    func reload(
+        _ response: REST.ServerRelaysResponse,
+        relays: [REST.ServerRelay]
+    ) -> [RelayLocation]
+}
+
+extension LocationDataSourceProtocol {
+    func makeRootNode(name: String) -> SelectLocationNode {
+        SelectLocationNode(nodeType: .root, location: .country("#root"), displayName: name)
+    }
+
+    func createNode(
+        root: SelectLocationNode,
+        ascendantOrSelf: RelayLocation,
+        serverLocation: REST.ServerLocation,
+        relay: REST.ServerRelay,
+        wasShowingChildren: Bool
+    ) -> SelectLocationNode {
+        let node: SelectLocationNode
+
+        switch ascendantOrSelf {
+        case .country:
+            node = SelectLocationNode(
+                nodeType: .country,
+                location: ascendantOrSelf,
+                displayName: serverLocation.country,
+                showsChildren: wasShowingChildren
+            )
+            root.addChild(node)
+        case let .city(countryCode, _):
+            node = SelectLocationNode(
+                nodeType: .city,
+                location: ascendantOrSelf,
+                displayName: serverLocation.city,
+                showsChildren: wasShowingChildren
+            )
+            nodeByLocation[.country(countryCode)]!.addChild(node)
+
+        case let .hostname(countryCode, cityCode, _):
+            node = SelectLocationNode(
+                nodeType: .relay,
+                location: ascendantOrSelf,
+                displayName: relay.hostname,
+                isActive: relay.active
+            )
+            nodeByLocation[.city(countryCode, cityCode)]!.addChild(node)
+        }
+        return node
+    }
+}
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/SelectLocationGroup.swift b/ios/MullvadVPN/View controllers/SelectLocation/SelectLocationGroup.swift
new file mode 100644
index 0000000000..1d5b66dcf2
--- /dev/null
+++ b/ios/MullvadVPN/View controllers/SelectLocation/SelectLocationGroup.swift
@@ -0,0 +1,59 @@
+//
+//  SelectLocationSectionGroup.swift
+//  MullvadVPN
+//
+//  Created by Mojgan on 2024-02-05.
+//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
+//
+
+import Foundation
+enum SelectLocationGroup: Hashable, CustomStringConvertible, CaseIterable {
+    case customLists
+    case allLocations
+
+    var description: String {
+        switch self {
+        case .customLists:
+            return NSLocalizedString(
+                "SELECT_LOCATION_ADD_CUSTOM_LISTS",
+                value: "Custom lists",
+                comment: ""
+            )
+        case .allLocations:
+            return NSLocalizedString(
+                "SELECT_LOCATION_ALL_LOCATIONS",
+                value: "All locations",
+                comment: ""
+            )
+        }
+    }
+
+    var cell: Cell {
+        Cell.locationCell
+    }
+
+    static var allCases: [SelectLocationGroup] {
+        #if DEBUG
+        return [.customLists, .allLocations]
+        #else
+        return [.allLocations]
+        #endif
+    }
+}
+
+extension SelectLocationGroup {
+    enum Cell: String, CaseIterable {
+        case locationCell
+
+        var reusableViewClass: AnyClass {
+            switch self {
+            case .locationCell:
+                return SelectLocationCell.self
+            }
+        }
+
+        var reuseIdentifier: String {
+            self.rawValue
+        }
+    }
+}
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/SelectLocationNode.swift b/ios/MullvadVPN/View controllers/SelectLocation/SelectLocationNode.swift
new file mode 100644
index 0000000000..789075d15f
--- /dev/null
+++ b/ios/MullvadVPN/View controllers/SelectLocation/SelectLocationNode.swift
@@ -0,0 +1,123 @@
+//
+//  SelectLocationNode.swift
+//  MullvadVPN
+//
+//  Created by Mojgan on 2024-02-05.
+//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
+//
+
+import MullvadTypes
+
+enum LocationNodeType {
+    case root
+    case country
+    case city
+    case relay
+}
+
+class SelectLocationNode: SelectLocationNodeProtocol {
+    var children: [SelectLocationNode]
+    var showsChildren: Bool
+    var nodeType: LocationNodeType
+    var location: RelayLocation
+    var displayName: String
+    var isActive: Bool
+
+    init(
+        nodeType: LocationNodeType,
+        location: RelayLocation,
+        displayName: String = "",
+        isActive: Bool = true,
+        showsChildren: Bool = false,
+        children: [SelectLocationNode] = []
+    ) {
+        self.showsChildren = showsChildren
+        self.nodeType = nodeType
+        self.location = location
+        self.displayName = displayName
+        self.isActive = isActive
+        self.children = children
+    }
+
+    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
+        }
+    }
+
+    func addChild(_ child: SelectLocationNode) {
+        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
+                node.isActive
+            })
+        case .relay:
+            break
+        }
+    }
+
+    func flatRelayLocationList(includeHiddenChildren: Bool = false) -> [RelayLocation] {
+        children.reduce(into: []) { array, node in
+            Self.flatten(node: node, into: &array, includeHiddenChildren: includeHiddenChildren)
+        }
+    }
+
+    private func sortChildren() {
+        switch nodeType {
+        case .root, .country:
+            children.sort { a, b -> Bool in
+                a.displayName.localizedCaseInsensitiveCompare(b.displayName) == .orderedAscending
+            }
+        case .city:
+            children.sort { a, b -> Bool in
+                a.location.stringRepresentation
+                    .localizedStandardCompare(b.location.stringRepresentation) == .orderedAscending
+            }
+        case .relay:
+            break
+        }
+    }
+
+    private class func flatten(
+        node: SelectLocationNode,
+        into array: inout [RelayLocation],
+        includeHiddenChildren: Bool
+    ) {
+        array.append(node.location)
+        if includeHiddenChildren || node.showsChildren {
+            for child in node.children {
+                Self.flatten(node: child, into: &array, includeHiddenChildren: includeHiddenChildren)
+            }
+        }
+    }
+}
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/SelectLocationNodeProtocol.swift b/ios/MullvadVPN/View controllers/SelectLocation/SelectLocationNodeProtocol.swift
new file mode 100644
index 0000000000..2d45f3f2d8
--- /dev/null
+++ b/ios/MullvadVPN/View controllers/SelectLocation/SelectLocationNodeProtocol.swift
@@ -0,0 +1,18 @@
+//
+//  SelectLocationNodeProtocol.swift
+//  MullvadVPN
+//
+//  Created by Mojgan on 2024-02-05.
+//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
+//
+
+import MullvadTypes
+
+protocol SelectLocationNodeProtocol {
+    var location: RelayLocation { get }
+    var displayName: String { get }
+    var showsChildren: Bool { get }
+    var isActive: Bool { get }
+    var isCollapsible: Bool { get }
+    var indentationLevel: Int { get }
+}
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/SelectLocationViewController.swift b/ios/MullvadVPN/View controllers/SelectLocation/SelectLocationViewController.swift
index 47e9a4309f..5f8f5d2145 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/SelectLocationViewController.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/SelectLocationViewController.swift
@@ -114,12 +114,19 @@ final class SelectLocationViewController: UIViewController {
     // MARK: - Private
 
     private func setUpDataSource() {
-        dataSource = LocationDataSource(tableView: tableView)
+        dataSource = LocationDataSource(
+            tableView: tableView,
+            allLocations: AllLocationDataSource(),
+            customLists: CustomListsDataSource()
+        )
         dataSource?.didSelectRelayLocation = { [weak self] location in
             self?.didSelectRelay?(location)
         }
 
-        dataSource?.selectedRelayLocation = relayLocation
+        dataSource?.selectedRelayLocation = relayLocation.flatMap { LocationCellViewModel(
+            group: .allLocations,
+            location: $0
+        ) }
 
         if let cachedRelays {
             dataSource?.setRelays(cachedRelays.relays, filter: filter)