Skip to content

Commit

Permalink
Fix crash on UITableViewDiffableDataSource
Browse files Browse the repository at this point in the history
  • Loading branch information
mojganii committed Nov 8, 2024
1 parent e914beb commit a1267a3
Show file tree
Hide file tree
Showing 6 changed files with 139 additions and 94 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,19 @@ import MullvadSettings
import MullvadTypes
import UIKit

//
// AddLocationsDataSource.swift
// MullvadVPN
//
// Created by Mojgan on 2024-02-29.
// Copyright © 2024 Mullvad VPN AB. All rights reserved.
//

import Combine
import MullvadSettings
import MullvadTypes
import UIKit

class AddLocationsDataSource:
UITableViewDiffableDataSource<LocationSection, LocationCellViewModel>,
LocationDiffableDataSourceProtocol {
Expand Down Expand Up @@ -83,7 +96,7 @@ class AddLocationsDataSource:
indentationLevel: 1
))
}
updateDataSnapshot(with: [locationsList])
reloadDataSnapshot(with: [locationsList])
}

private func isLocationInCustomList(node: LocationNode) -> Bool {
Expand All @@ -110,20 +123,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) {
Expand All @@ -142,7 +147,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
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ extension [LocationCellViewModel] {
section: section,
node: $0,
indentationLevel: item.indentationLevel + 1,
isSelected: false
isSelected: item.isSelected
)
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,6 @@ final class LocationDataSource:

allLocationsDataSource?.reload(relaysWithLocation)
customListsDataSource?.reload(allLocationNodes: allLocationsDataSource?.nodes ?? [])

setSelectedRelays(selectedRelays)
filterRelays(by: currentSearchString)
}
Expand All @@ -80,11 +79,9 @@ final class LocationDataSource:
}
}

updateDataSnapshot(with: list, reloadExisting: !searchString.isEmpty) {
self.tableView.reloadData()

reloadDataSnapshot(with: list) {
if searchString.isEmpty {
self.updateSelection(self.selectedLocation, animated: false, completion: {
self.updateSelection(completion: {
self.scrollToSelectedRelay()
})
} else {
Expand Down Expand Up @@ -129,22 +126,24 @@ final class LocationDataSource:
return recursivelyCreateCellViewModelTree(for: rootNode, in: .customLists, indentationLevel: 0)
}

updateDataSnapshot(with: [
reloadDataSnapshot(with: [
list,
snapshot().itemIdentifiers(inSection: .allLocations),
], reloadExisting: true)
])
}

func setSelectedRelays(_ selectedRelays: RelaySelection) {
selectedLocation = mapSelection(from: selectedRelays.selected)

excludedLocation = mapSelection(from: selectedRelays.excluded)
excludedLocation?.excludedRelayTitle = selectedRelays.excludedTitle

tableView.reloadData()
updateSelection(completion: {
self.scrollToSelectedRelay()
})
}

func scrollToSelectedRelay() {
// MARK: - Private functions

private func scrollToSelectedRelay() {
indexPathForSelectedRelay().flatMap {
tableView.scrollToRow(at: $0, at: .middle, animated: false)
}
Expand Down Expand Up @@ -185,11 +184,11 @@ final class LocationDataSource:
return nil
}

private func updateSelection(_ item: LocationCellViewModel?, animated: Bool, completion: (() -> Void)? = nil) {
selectedLocation = item
private func updateSelection(completion: (() -> Void)? = nil) {
guard let selectedLocation else { return }

let rootNode = selectedLocation.node.root
var snapshot = snapshot()

// Exit early if no changes to the node tree should be made.
guard selectedLocation.node != rootNode else {
Expand All @@ -215,22 +214,12 @@ final class LocationDataSource:
indentationLevel: 1
)

// Insert the new node tree below the selected item.
var snapshotItems = snapshot().itemIdentifiers(inSection: selectedLocation.section)
snapshotItems.insert(contentsOf: nodesToAdd, at: indexPath.row + 1)

let list = sections.enumerated().map { index, section in
index == indexPath.section
? snapshotItems
: snapshot().itemIdentifiers(inSection: section)
}
let existingItems = snapshot.itemIdentifiers(inSection: selectedLocation.section)
snapshot.deleteItems(nodesToAdd)
snapshot.insertItems(nodesToAdd, afterItem: existingItems[indexPath.row])

updateDataSnapshot(
with: list,
reloadExisting: true,
animated: animated,
completion: completion
)
// Apply the updated snapshot
self.applySnapshotUsingReloadData(snapshot, completion: completion)
}

override func tableView(_ tableView: UITableView, cellForRowAt indexPath: IndexPath) -> UITableViewCell {
Expand Down Expand Up @@ -310,25 +299,21 @@ extension LocationDataSource: UITableViewDelegate {
}

func tableView(_ tableView: UITableView, willDisplay cell: UITableViewCell, forRowAt indexPath: IndexPath) {
if let item = itemIdentifier(for: indexPath), item == selectedLocation {
cell.setSelected(true, animated: false)
if let item = itemIdentifier(for: indexPath) {
cell.setSelected(item == selectedLocation, animated: false)
}
}

func tableView(_ tableView: UITableView, willSelectRowAt indexPath: IndexPath) -> IndexPath? {
if let indexPath = indexPathForSelectedRelay() {
if let cell = tableView.cellForRow(at: indexPath) {
cell.setSelected(false, animated: false)
}
tableView.deselectRow(at: indexPath, animated: false)
}

return indexPath
}

func tableView(_ tableView: UITableView, didSelectRowAt indexPath: IndexPath) {
guard let item = itemIdentifier(for: indexPath) else { return }
selectedLocation = item

var customListSelection: UserSelectedRelays.CustomListSelection?
if let topmostNode = item.node.root as? CustomListLocationNode {
customListSelection = UserSelectedRelays.CustomListSelection(
Expand All @@ -341,7 +326,6 @@ extension LocationDataSource: UITableViewDelegate {
locations: item.node.locations,
customListSelection: customListSelection
)

didSelectRelayLocations?(relayLocations)
}

Expand All @@ -354,12 +338,9 @@ extension LocationDataSource: LocationCellDelegate {
func toggleExpanding(cell: LocationCell) {
guard let indexPath = tableView.indexPath(for: cell),
let item = itemIdentifier(for: indexPath) else { return }

let items = toggledItems(for: cell)

updateDataSnapshot(with: items, reloadExisting: true, completion: {
toggledItems(for: cell) {
self.scroll(to: item, animated: true)
})
}
}

func toggleSelecting(cell: LocationCell) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -39,49 +39,109 @@ extension LocationDiffableDataSourceProtocol {
}
}

func toggledItems(for cell: LocationCell) -> [[LocationCellViewModel]] {
func toggledItems(for cell: LocationCell, completion: (() -> Void)? = nil) {
guard let indexPath = tableView.indexPath(for: cell),
let item = itemIdentifier(for: indexPath) else { return [[]] }

let item = itemIdentifier(for: indexPath) else { return }
let snapshot = snapshot()
let section = sections[indexPath.section]
let isExpanded = item.node.showsChildren
var locationList = snapshot().itemIdentifiers(inSection: section)
var locationList = snapshot.itemIdentifiers(inSection: section)

item.node.showsChildren = !isExpanded

if !isExpanded {
locationList.addSubNodes(from: item, at: indexPath)
addItem(locationList, toSection: section, index: indexPath.row, completion: completion)
} else {
locationList.removeSubNodes(from: item.node)
updateSnapshotRetainingOnly(locationList, toSection: section, completion: completion)
}
}

private func addItem(
_ items: [LocationCellViewModel],
toSection section: LocationSection,
index: Int,
completion: (() -> Void)? = nil
) {
var snapshot = snapshot()
let existingItems = snapshot.itemIdentifiers(inSection: section)

return sections.enumerated().map { index, section in
index == indexPath.section
? locationList
: snapshot().itemIdentifiers(inSection: section)
// Filter itemsToAdd to only include items not already in the section
let uniqueItems = items.filter { item in
existingItems.firstIndex(where: { $0 == item }) == nil
}

// Insert unique items at the specified index
if index < existingItems.count {
snapshot.insertItems(uniqueItems, afterItem: existingItems[index])
} else {
// If the index is beyond bounds, append to the end
snapshot.appendItems(uniqueItems, toSection: section)
}
applyAndReconfigureSnapshot(snapshot, in: section, completion: completion)
}

func updateDataSnapshot(
with list: [[LocationCellViewModel]],
reloadExisting: Bool = false,
animated: Bool = false,
private func updateSnapshotRetainingOnly(
_ itemsToKeep: [LocationCellViewModel],
toSection section: LocationSection,
completion: (() -> Void)? = nil
) {
var snapshot = NSDiffableDataSourceSnapshot<LocationSection, LocationCellViewModel>()
var snapshot = snapshot()

snapshot.appendSections(sections)
for (index, section) in sections.enumerated() {
let items = list[index]
// Ensure the section exists in the snapshot
guard snapshot.sectionIdentifiers.contains(section) else { return }

snapshot.appendItems(items, toSection: section)
// Get the current items in the section
let currentItems = snapshot.itemIdentifiers(inSection: section)

if reloadExisting {
snapshot.reconfigureItems(items)
// Determine the items that should be deleted
let itemsToDelete = currentItems.filter { !itemsToKeep.contains($0) }
snapshot.deleteItems(itemsToDelete)

// Apply the updated snapshot
applyAndReconfigureSnapshot(snapshot, in: section, completion: completion)
}

private func applyAndReconfigureSnapshot(
_ snapshot: NSDiffableDataSourceSnapshot<LocationSection, LocationCellViewModel>,
in section: LocationSection,
completion: (() -> Void)? = nil
) {
DispatchQueue.main.async {
self.apply(snapshot, animatingDifferences: true) {
// After adding, reconfigure specified items to update their content
var updatedSnapshot = self.snapshot()

// Ensure the items exist in the snapshot before attempting to reconfigure
let existingItems = updatedSnapshot.itemIdentifiers(inSection: section)

// Reconfigure the specified items
updatedSnapshot.reconfigureItems(existingItems)

// Apply the reconfigured snapshot without animations to avoid any flickering
self.apply(updatedSnapshot, animatingDifferences: false)
}
}
}

func reloadDataSnapshot(
with list: [[LocationCellViewModel]],
animated: Bool = false,
completion: (() -> Void)? = nil
) {
var snapshot = self.snapshot()
snapshot.deleteAllItems()
for (index, section) in sections.enumerated() {
// Check if the section already exists in the snapshot, and add it if not
if !snapshot.sectionIdentifiers.contains(section) {
snapshot.appendSections([section])
}
let items = list[index]
snapshot.appendItems(items, toSection: section)
}
DispatchQueue.main.async {
// Apply the reconfiguration snapshot without animation to avoid unwanted animations
self.apply(snapshot, animatingDifferences: animated, completion: completion)
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -72,11 +72,6 @@ final class LocationViewController: UIViewController {
}
}

override func viewWillAppear(_ animated: Bool) {
super.viewWillAppear(animated)
dataSource?.scrollToSelectedRelay()
}

override func viewDidAppear(_ animated: Bool) {
super.viewDidAppear(animated)
tableView.flashScrollIndicators()
Expand Down
Loading

0 comments on commit a1267a3

Please sign in to comment.