Fix lazy lodaing for custom listsimprove-select-location-lazy-expansion-ios

4 files changed, 88 insertions, 64 deletions

diff --git a/ios/MullvadVPN/View controllers/SelectLocation/DataSource/AllLocationDataSource.swift b/ios/MullvadVPN/View controllers/SelectLocation/DataSource/AllLocationDataSource.swift
index bb24c2f759..3af6b4a9de 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/DataSource/AllLocationDataSource.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/DataSource/AllLocationDataSource.swift
@@ -13,7 +13,7 @@ import MullvadTypes
 /// Relay data grouped for lazy child node creation
 struct RelaysByLocation {
     var countryName: String
-    var cities: [String: CityRelays] // keyed by cityCode
+    var cities: [String: CityRelays]  // keyed by cityCode
 
     struct CityRelays {
         var cityName: String
@@ -28,19 +28,23 @@ class AllLocationDataSource: SearchableLocationDataSource {
     private var relaysByCountry: [String: RelaysByLocation] = [:]
 
     /// Constructs a collection of node trees from relays fetched from the API.
-    /// Only creates country nodes initially - children are created lazily when expanded.
+    /// Creates country and city nodes eagerly, but relay/host nodes are created lazily when expanded.
+    /// This hybrid approach maintains performance (relays are ~40,000) while allowing custom lists
+    /// and selection to work correctly (they need to traverse countries and cities).
     func reload(_ relays: LocationRelays) {
         let expandedCodes = collectExpandedCodes()
 
-        // Group relays by country and city for lazy loading
+        // Group relays by country and city for lazy loading of hosts
         relaysByCountry = [:]
         var countryActiveState: [String: Bool] = [:]
+        var cityActiveState: [String: Bool] = [:]
 
         for relay in relays.relays {
             guard let serverLocation = relays.locations[relay.location.rawValue] else { continue }
 
             let countryCode = String(relay.location.country)
             let cityCode = String(relay.location.city)
+            let countryCityCode = LocationNode.combineNodeCodes([countryCode, cityCode])
 
             // Initialize country entry if needed
             if relaysByCountry[countryCode] == nil {
@@ -63,10 +67,11 @@ class AllLocationDataSource: SearchableLocationDataSource {
             // Track active state
             if relay.active {
                 countryActiveState[countryCode] = true
+                cityActiveState[countryCityCode] = true
             }
         }
 
-        // Create only country nodes initially
+        // Create country and city nodes eagerly, but relay nodes lazily
         var countryNodes: [LocationNode] = []
         for (countryCode, countryData) in relaysByCountry {
             let countryLocation = RelayLocation.country(countryCode)
@@ -77,31 +82,47 @@ class AllLocationDataSource: SearchableLocationDataSource {
                 isActive: countryActiveState[countryCode] ?? false,
                 showsChildren: expandedCodes.contains(countryCode)
             )
-            countryNode.hasLazyChildren = true
-            countryNodes.append(countryNode)
 
-            // If this country was expanded, populate its children immediately
-            if expandedCodes.contains(countryCode) {
-                populateChildren(for: countryNode, expandedCodes: expandedCodes)
+            // Create city nodes eagerly (needed for custom lists and selection)
+            var cityNodes: [LocationNode] = []
+            for (cityCode, cityData) in countryData.cities {
+                let countryCityCode = LocationNode.combineNodeCodes([countryCode, cityCode])
+                let cityLocation = RelayLocation.city(countryCode, cityCode)
+
+                let cityNode = LocationNode(
+                    name: NSLocalizedString(cityData.cityName, comment: ""),
+                    code: countryCityCode,
+                    locations: [cityLocation],
+                    isActive: cityActiveState[countryCityCode] ?? false,
+                    parent: countryNode,
+                    showsChildren: expandedCodes.contains(countryCityCode)
+                )
+                // Mark city as having lazy children (relays)
+                cityNode.hasLazyChildren = true
+                cityNodes.append(cityNode)
+
+                // If this city was expanded, populate its relay children immediately
+                if expandedCodes.contains(countryCityCode) {
+                    populateCityChildren(
+                        node: cityNode, countryCode: countryCode, cityCode: cityCode, cityData: cityData)
+                }
             }
+
+            cityNodes.sort()
+            countryNode.children = cityNodes
+            countryNodes.append(countryNode)
         }
 
         countryNodes.sort()
         nodes = countryNodes
     }
 
-    /// Populates children for a node if they haven't been created yet.
-    /// Call this when a node is about to be expanded.
+    /// Populates relay children for a city node if they haven't been created yet.
+    /// Call this when a city node is about to be expanded.
     func populateChildren(for node: LocationNode, expandedCodes: Set<String> = []) {
         guard node.hasLazyChildren, node.children.isEmpty else { return }
 
-        // Check if this is a country node
-        if let countryData = relaysByCountry[node.code] {
-            populateCountryChildren(node: node, countryData: countryData, expandedCodes: expandedCodes)
-            return
-        }
-
-        // Check if this is a city node (code format: "countryCode-cityCode")
+        // City node code format: "countryCode-cityCode"
         let codeParts = node.code.split(separator: "-")
         if codeParts.count >= 2 {
             let countryCode = String(codeParts[0])
@@ -112,42 +133,6 @@ class AllLocationDataSource: SearchableLocationDataSource {
         }
     }
 
-    private func populateCountryChildren(
-        node: LocationNode,
-        countryData: RelaysByLocation,
-        expandedCodes: Set<String>
-    ) {
-        var cityNodes: [LocationNode] = []
-        let countryCode = node.code
-
-        for (cityCode, cityData) in countryData.cities {
-            let countryCityCode = LocationNode.combineNodeCodes([countryCode, cityCode])
-            let cityLocation = RelayLocation.city(countryCode, cityCode)
-
-            // Check if city has any active relays
-            let hasActiveRelay = cityData.relays.contains { $0.active }
-
-            let cityNode = LocationNode(
-                name: NSLocalizedString(cityData.cityName, comment: ""),
-                code: countryCityCode,
-                locations: [cityLocation],
-                isActive: hasActiveRelay,
-                parent: node,
-                showsChildren: expandedCodes.contains(countryCityCode)
-            )
-            cityNode.hasLazyChildren = true
-            cityNodes.append(cityNode)
-
-            // If this city was expanded, populate its children immediately
-            if expandedCodes.contains(countryCityCode) {
-                populateCityChildren(node: cityNode, countryCode: countryCode, cityCode: cityCode, cityData: cityData)
-            }
-        }
-
-        cityNodes.sort()
-        node.children = cityNodes
-    }
-
     private func populateCityChildren(
         node: LocationNode,
         countryCode: String,
@@ -194,6 +179,18 @@ class AllLocationDataSource: SearchableLocationDataSource {
         guard let location = selectedRelays.locations.first else {
             return nil
         }
+
+        // For hostname lookups, ensure the city's relay children are populated first
+        if case let .hostname(countryCode, cityCode, _) = location {
+            if let cityNode =
+                rootNode
+                .countryFor(code: countryCode)?
+                .cityFor(codes: [countryCode, cityCode])
+            {
+                populateChildren(for: cityNode)
+            }
+        }
+
         return descendantNode(in: rootNode, for: location, baseCodes: [])
     }
 }
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/DataSource/CustomListLocationNodeBuilder.swift b/ios/MullvadVPN/View controllers/SelectLocation/DataSource/CustomListLocationNodeBuilder.swift
index 56ffb14d17..596fb7d729 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/DataSource/CustomListLocationNodeBuilder.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/DataSource/CustomListLocationNodeBuilder.swift
@@ -13,6 +13,14 @@ import MullvadTypes
 struct CustomListLocationNodeBuilder {
     let customList: CustomList
     let allLocations: [LocationNode]
+    /// Optional data source to populate lazy children when looking up hostnames
+    let dataSource: AllLocationDataSource?
+
+    init(customList: CustomList, allLocations: [LocationNode], dataSource: AllLocationDataSource? = nil) {
+        self.customList = customList
+        self.allLocations = allLocations
+        self.dataSource = dataSource
+    }
 
     var customListLocationNode: CustomListLocationNode {
         let listNode = CustomListLocationNode(
@@ -29,22 +37,33 @@ struct CustomListLocationNodeBuilder {
         listNode.children = listNode.locations.compactMap { location in
             switch location {
             case let .country(countryCode):
-                rootNode
+                return
+                    rootNode
                     .countryFor(code: countryCode)?
                     .copy(withParent: listNode)
 
             case let .city(countryCode, cityCode):
-                rootNode
+                return
+                    rootNode
                     .countryFor(code: countryCode)?
                     .cityFor(codes: [countryCode, cityCode])?
                     .copy(withParent: listNode)
 
             case let .hostname(countryCode, cityCode, hostCode):
-                rootNode
+                // For hostname lookups, we need to ensure the city's relay children are populated
+                if let cityNode =
+                    rootNode
                     .countryFor(code: countryCode)?
-                    .cityFor(codes: [countryCode, cityCode])?
-                    .hostFor(code: hostCode)?
-                    .copy(withParent: listNode)
+                    .cityFor(codes: [countryCode, cityCode])
+                {
+                    // Populate lazy children if needed
+                    dataSource?.populateChildren(for: cityNode)
+                    return
+                        cityNode
+                        .hostFor(code: hostCode)?
+                        .copy(withParent: listNode)
+                }
+                return nil
             }
         }
 
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/DataSource/CustomListsDataSource.swift b/ios/MullvadVPN/View controllers/SelectLocation/DataSource/CustomListsDataSource.swift
index ba5d67b3c4..d9a78d3c61 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/DataSource/CustomListsDataSource.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/DataSource/CustomListsDataSource.swift
@@ -21,10 +21,17 @@ class CustomListsDataSource: SearchableLocationDataSource {
 
     /// Constructs a collection of node trees by copying each matching counterpart
     /// from the complete list of nodes created in ``AllLocationDataSource``.
-    func reload(allLocationNodes: [LocationNode]) {
+    /// - Parameters:
+    ///   - allLocationNodes: The location nodes from AllLocationDataSource
+    ///   - dataSource: Optional data source to populate lazy children when looking up hostnames
+    func reload(allLocationNodes: [LocationNode], dataSource: AllLocationDataSource? = nil) {
         let expandedCodes = collectExpandedCodes()
         nodes = repository.fetchAll().map { list in
-            let customListWrapper = CustomListLocationNodeBuilder(customList: list, allLocations: allLocationNodes)
+            let customListWrapper = CustomListLocationNodeBuilder(
+                customList: list,
+                allLocations: allLocationNodes,
+                dataSource: dataSource
+            )
             let listNode = customListWrapper.customListLocationNode
             listNode.showsChildren = expandedCodes.contains(listNode.code)
 
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/SelectLocationViewModel.swift b/ios/MullvadVPN/View controllers/SelectLocation/SelectLocationViewModel.swift
index a07be1e391..035ca809ae 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/SelectLocationViewModel.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/SelectLocationViewModel.swift
@@ -299,8 +299,9 @@ class SelectLocationViewModelImpl: SelectLocationViewModel {
     }
 
     private func refreshCustomLists() {
-        exitCustomListsDataSource.reload(allLocationNodes: exitContext.locations)
-        entryCustomListsDataSource.reload(allLocationNodes: entryContext.locations)
+        exitCustomListsDataSource.reload(allLocationNodes: exitContext.locations, dataSource: exitLocationsDataSource)
+        entryCustomListsDataSource.reload(
+            allLocationNodes: entryContext.locations, dataSource: entryLocationsDataSource)
 
         exitContext.customLists = exitCustomListsDataSource.nodes
         entryContext.customLists = entryCustomListsDataSource.nodes