Apply filtering on custom lists

3 files changed, 14 insertions, 11 deletions

diff --git a/ios/MullvadVPN/View controllers/SelectLocation/CustomListsDataSource.swift b/ios/MullvadVPN/View controllers/SelectLocation/CustomListsDataSource.swift
index 99ad974f20..dd041642cd 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/CustomListsDataSource.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/CustomListsDataSource.swift
@@ -25,16 +25,16 @@ 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]) {
-        nodes = repository.fetchAll().map { list in
+    func reload(allLocationNodes: [LocationNode], isFiltered: Bool) {
+        nodes = repository.fetchAll().compactMap { customList in
             let listNode = CustomListLocationNode(
-                name: list.name,
-                code: list.name.lowercased(),
-                locations: list.locations,
-                customList: list
+                name: customList.name,
+                code: customList.name.lowercased(),
+                locations: customList.locations,
+                customList: customList
             )
 
-            listNode.children = list.locations.compactMap { location in
+            listNode.children = customList.locations.compactMap { location in
                 copy(location, from: allLocationNodes, withParent: listNode)
             }
 
@@ -46,7 +46,7 @@ class CustomListsDataSource: LocationDataSourceProtocol {
                 node.code = LocationNode.combineNodeCodes([listNode.code, node.code])
             }
 
-            return listNode
+            return (isFiltered && listNode.children.isEmpty) ? nil : listNode
         }
     }
 
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
index 4e95c9f592..c8c626320b 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
@@ -17,6 +17,7 @@ final class LocationDataSource: UITableViewDiffableDataSource<LocationSection, L
     private let tableView: UITableView
     private var dataSources: [LocationDataSourceProtocol] = []
     private var selectedItem: LocationCellViewModel?
+    private var hasFilter = false
 
     var didSelectRelayLocations: ((UserSelectedRelays) -> Void)?
     var didTapEditCustomLists: (() -> Void)?
@@ -47,6 +48,8 @@ final class LocationDataSource: UITableViewDiffableDataSource<LocationSection, L
     }
 
     func setRelays(_ response: REST.ServerRelaysResponse, selectedRelays: UserSelectedRelays?, filter: RelayFilter) {
+        hasFilter = filter.providers != .any || filter.ownership != .any
+
         let allLocationsDataSource =
             dataSources.first(where: { $0 is AllLocationDataSource }) as? AllLocationDataSource
 
@@ -58,7 +61,7 @@ final class LocationDataSource: UITableViewDiffableDataSource<LocationSection, L
         }
 
         allLocationsDataSource?.reload(response, relays: relays)
-        customListsDataSource?.reload(allLocationNodes: allLocationsDataSource?.nodes ?? [])
+        customListsDataSource?.reload(allLocationNodes: allLocationsDataSource?.nodes ?? [], isFiltered: hasFilter)
 
         mapSelectedItem(from: selectedRelays)
         filterRelays(by: currentSearchString)
@@ -101,7 +104,7 @@ final class LocationDataSource: UITableViewDiffableDataSource<LocationSection, L
         let customListsDataSource =
             dataSources.first(where: { $0 is CustomListsDataSource }) as? CustomListsDataSource
 
-        customListsDataSource?.reload(allLocationNodes: allLocationsDataSource?.nodes ?? [])
+        customListsDataSource?.reload(allLocationNodes: allLocationsDataSource?.nodes ?? [], isFiltered: hasFilter)
 
         mapSelectedItem(from: selectedRelays)
         filterRelays(by: currentSearchString, scrollToSelected: false)
diff --git a/ios/MullvadVPNTests/Location/CustomListsDataSourceTests.swift b/ios/MullvadVPNTests/Location/CustomListsDataSourceTests.swift
index 4431498b7d..a14e2296f8 100644
--- a/ios/MullvadVPNTests/Location/CustomListsDataSourceTests.swift
+++ b/ios/MullvadVPNTests/Location/CustomListsDataSourceTests.swift
@@ -75,7 +75,7 @@ class CustomListsDataSourceTests: XCTestCase {
 extension CustomListsDataSourceTests {
     private func setUpDataSource() {
         dataSource = CustomListsDataSource(repository: CustomListsRepositoryStub(customLists: customLists))
-        dataSource.reload(allLocationNodes: allLocationNodes)
+        dataSource.reload(allLocationNodes: allLocationNodes, isFiltered: false)
     }
 
     private func createAllLocationNodes() {