Fix crash on UITableViewDiffableDataSource

mullvad · Nov 12, 2024 · 66c756b · 66c756b
1 parent c838835
commit 66c756b
Show file tree

Hide file tree

Showing 9 changed files with 217 additions and 173 deletions.
diff --git a/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsDataSource.swift b/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsDataSource.swift
@@ -83,7 +83,7 @@ class AddLocationsDataSource:
                 indentationLevel: 1
             ))
         }
-        updateDataSnapshot(with: [locationsList])
+        reloadDataSnapshot(with: [locationsList])
     }
 
     private func isLocationInCustomList(node: LocationNode) -> Bool {
@@ -110,20 +110,12 @@ extension AddLocationsDataSource: UITableViewDelegate {
 
 extension AddLocationsDataSource: LocationCellDelegate {
     func toggleExpanding(cell: LocationCell) {
-        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: [items], reloadExisting: true, completion: {
+        toggledItems(for: cell) {
             if let indexPath = self.tableView.indexPath(for: cell),
                let item = self.itemIdentifier(for: indexPath) {
                 self.scroll(to: item, animated: true)
             }
-        })
+        }
     }
 
     func toggleSelecting(cell: LocationCell) {
@@ -142,7 +134,7 @@ extension AddLocationsDataSource: LocationCellDelegate {
         } else {
             customListLocationNode.remove(selectedLocation: item.node, with: locationList)
         }
-        updateDataSnapshot(with: [locationList], completion: {
+        reloadDataSnapshot(with: [locationList], completion: {
             let locations = self.customListLocationNode.children.reduce([]) { partialResult, locationNode in
                 partialResult + locationNode.locations
             }

diff --git a/ios/MullvadVPN/View controllers/SelectLocation/AllLocationDataSource.swift b/ios/MullvadVPN/View controllers/SelectLocation/AllLocationDataSource.swift
@@ -22,6 +22,7 @@ class AllLocationDataSource: LocationDataSourceProtocol {
     /// and city names.
     func reload(_ relays: LocationRelays) {
         let rootNode = RootLocationNode()
+        let expandedRelays = nodes.flatMap { [$0] + $0.flattened }.filter { $0.showsChildren }.map { $0.code }
 
         for relay in relays.relays {
             guard case
@@ -32,7 +33,13 @@ class AllLocationDataSource: LocationDataSourceProtocol {
             let relayLocation = RelayLocation.hostname(countryCode, cityCode, relay.hostname)
 
             for ancestorOrSelf in relayLocation.ancestors + [relayLocation] {
-                addLocation(ancestorOrSelf, rootNode: rootNode, serverLocation: serverLocation, relay: relay)
+                addLocation(
+                    ancestorOrSelf,
+                    rootNode: rootNode,
+                    serverLocation: serverLocation,
+                    relay: relay,
+                    showsChildren: expandedRelays.contains(ancestorOrSelf.stringRepresentation)
+                )
             }
         }
 
@@ -56,15 +63,17 @@ class AllLocationDataSource: LocationDataSourceProtocol {
         _ location: RelayLocation,
         rootNode: LocationNode,
         serverLocation: REST.ServerLocation,
-        relay: REST.ServerRelay
+        relay: REST.ServerRelay,
+        showsChildren: Bool
     ) {
         switch location {
         case let .country(countryCode):
             let countryNode = LocationNode(
                 name: serverLocation.country,
                 code: LocationNode.combineNodeCodes([countryCode]),
                 locations: [location],
-                isActive: true // Defaults to true, updated when children are populated.
+                isActive: true, // Defaults to true, updated when children are populated.
+                showsChildren: showsChildren
             )
 
             if !rootNode.children.contains(countryNode) {
@@ -77,7 +86,8 @@ class AllLocationDataSource: LocationDataSourceProtocol {
                 name: serverLocation.city,
                 code: LocationNode.combineNodeCodes([countryCode, cityCode]),
                 locations: [location],
-                isActive: true // Defaults to true, updated when children are populated.
+                isActive: true, // Defaults to true, updated when children are populated.
+                showsChildren: showsChildren
             )
 
             if let countryNode = rootNode.countryFor(code: countryCode),
@@ -92,7 +102,8 @@ class AllLocationDataSource: LocationDataSourceProtocol {
                 name: relay.hostname,
                 code: LocationNode.combineNodeCodes([hostCode]),
                 locations: [location],
-                isActive: relay.active
+                isActive: relay.active,
+                showsChildren: showsChildren
             )
 
             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
@@ -26,16 +26,19 @@ 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]) {
+        let expandedRelays = nodes.flatMap { [$0] + $0.flattened }.filter { $0.showsChildren }.map { $0.code }
         nodes = repository.fetchAll().map { list in
             let customListWrapper = CustomListLocationNodeBuilder(customList: list, allLocations: allLocationNodes)
             let listNode = customListWrapper.customListLocationNode
+            listNode.showsChildren = expandedRelays.contains(listNode.code)
 
             listNode.forEachDescendant { node in
                 // Each item in a section in a diffable data source needs to be unique.
                 // Since LocationCellViewModel partly depends on LocationNode.code for
                 // equality, each node code needs to be prefixed with the code of its
                 // parent custom list to uphold this.
                 node.code = LocationNode.combineNodeCodes([listNode.code, node.code])
+                node.showsChildren = expandedRelays.contains(node.code)
             }
 
             return listNode

diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationCellViewModel.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationCellViewModel.swift
@@ -34,13 +34,14 @@ extension [LocationCellViewModel] {
     mutating func addSubNodes(from item: LocationCellViewModel, at indexPath: IndexPath) {
         let section = LocationSection.allCases[indexPath.section]
         let row = indexPath.row + 1
+        item.node.showsChildren = true
 
         let locations = item.node.children.map {
             LocationCellViewModel(
                 section: section,
                 node: $0,
                 indentationLevel: item.indentationLevel + 1,
-                isSelected: false
+                isSelected: item.isSelected
             )
         }