Improve selection handling

Demo/ASCollectionViewDemo/Screens/PhotoGrid/PhotoGridScreen.swift

@@ -7,7 +7,7 @@ import UIKit
 struct PhotoGridScreen: View
 {
 	@State var data: [Post] = DataSource.postsForGridSection(1, number: 1000)
-	@State var selectedItems: Set<Int> = []
+	@State var selectedIndexes: Set<Int> = []
 
 	@Environment(\.editMode) private var editMode
 	var isEditing: Bool
@@ -22,7 +22,7 @@ struct PhotoGridScreen: View
 		ASCollectionViewSection(
 			id: 0,
 			data: data,
-			selectedItems: $selectedItems,
+			selectedIndexes: $selectedIndexes,
 			onCellEvent: onCellEvent,
 			dragDropConfig: dragDropConfig,
 			contextMenuProvider: contextMenuProvider)
@@ -65,6 +65,8 @@ struct PhotoGridScreen: View
 		ASCollectionView(
 			section: section)
 			.layout(self.layout)
+			.allowsSelection(self.isEditing)
+			.allowsMultipleSelection(self.isEditing)
 			.edgesIgnoringSafeArea(.all)
 			.navigationBarTitle("Explore", displayMode: .large)
 			.navigationBarItems(
@@ -76,7 +78,7 @@ struct PhotoGridScreen: View
 						Button(action: {
 							withAnimation {
 								// We want the cell removal to be animated, so explicitly specify `withAnimation`
-								self.data.remove(atOffsets: IndexSet(self.selectedItems))
+								self.data.remove(atOffsets: IndexSet(self.selectedIndexes))
 							}
 						})
 						{

Demo/ASCollectionViewDemo/Screens/Waterfall/WaterfallScreen.swift

@@ -8,7 +8,7 @@ import UIKit
 struct WaterfallScreen: View
 {
 	@State var data: [[Post]] = (0 ... 10).map { DataSource.postsForWaterfallSection($0, number: 100) }
-	@State var selectedItems: [SectionID: Set<Int>] = [:]
+	@State var selectedIndexes: [SectionID: Set<Int>] = [:]
 	@State var columnMinSize: CGFloat = 150
 
 	@Environment(\.editMode) private var editMode
@@ -25,7 +25,7 @@ struct WaterfallScreen: View
 			ASCollectionViewSection(
 				id: offset,
 				data: sectionData,
-				selectedItems: $selectedItems[offset],
+				selectedIndexes: $selectedIndexes[offset],
 				onCellEvent: onCellEvent)
 			{ item, state in
 				GeometryReader
@@ -91,6 +91,8 @@ struct WaterfallScreen: View
 			ASCollectionView(
 				sections: sections)
 				.layout(self.layout)
+				.allowsSelection(self.isEditing)
+				.allowsMultipleSelection(self.isEditing)
 				.customDelegate(WaterfallScreenLayoutDelegate.init)
 				.contentInsets(.init(top: 0, left: 10, bottom: 10, right: 10))
 				.navigationBarTitle("Waterfall Layout", displayMode: .inline)
@@ -102,7 +104,7 @@ struct WaterfallScreen: View
 						{
 							Button(action: {
 								withAnimation {
-									self.selectedItems.forEach { sectionIndex, selected in
+									self.selectedIndexes.forEach { sectionIndex, selected in
 										self.data[sectionIndex].remove(atOffsets: IndexSet(selected))
 									}
 								}

Sources/ASCollectionView/ASCollectionView+Modifiers.swift

@@ -140,6 +140,22 @@ public extension ASCollectionView
 		this.maintainScrollPositionOnOrientationChange = true
 		return this
 	}
+
+	/// Set whether the ASCollectionView should allow selection, default is true
+	func allowsSelection(_ allowsSelection: Bool) -> Self
+	{
+		var this = self
+		this.allowsSelection = allowsSelection
+		return this
+	}
+
+	/// Set whether the ASCollectionView should allow multiple selection, default is false
+	func allowsMultipleSelection(_ allowsMultipleSelection: Bool) -> Self
+	{
+		var this = self
+		this.allowsMultipleSelection = allowsMultipleSelection
+		return this
+	}
 }
 
 // MARK: PUBLIC layout modifier functions

Sources/ASCollectionView/ASSection+Initialisers.swift

@@ -24,7 +24,7 @@ public extension ASSection
 		data: DataCollection,
 		dataID dataIDKeyPath: KeyPath<DataCollection.Element, DataID>,
 		container: @escaping ((Content) -> Container),
-		selectedItems: Binding<Set<Int>>? = nil,
+		selectedIndexes: Binding<Set<Int>>? = nil,
 		shouldAllowSelection: ((_ index: Int) -> Bool)? = nil,
 		shouldAllowDeselection: ((_ index: Int) -> Bool)? = nil,
 		onCellEvent: OnCellEvent<DataCollection.Element>? = nil,
@@ -41,7 +41,7 @@ public extension ASSection
 			dataIDKeyPath: dataIDKeyPath,
 			container: container,
 			content: contentBuilder,
-			selectedItems: selectedItems,
+			selectedIndexes: selectedIndexes,
 			shouldAllowSelection: shouldAllowSelection,
 			shouldAllowDeselection: shouldAllowDeselection,
 			onCellEvent: onCellEvent,
@@ -55,7 +55,7 @@ public extension ASSection
 		id: SectionID,
 		data: DataCollection,
 		dataID dataIDKeyPath: KeyPath<DataCollection.Element, DataID>,
-		selectedItems: Binding<Set<Int>>? = nil,
+		selectedIndexes: Binding<Set<Int>>? = nil,
 		shouldAllowSelection: ((_ index: Int) -> Bool)? = nil,
 		shouldAllowDeselection: ((_ index: Int) -> Bool)? = nil,
 		onCellEvent: OnCellEvent<DataCollection.Element>? = nil,
@@ -66,7 +66,7 @@ public extension ASSection
 		@ViewBuilder contentBuilder: @escaping ((DataCollection.Element, ASCellContext) -> Content))
 		where DataCollection.Index == Int
 	{
-		self.init(id: id, data: data, dataID: dataIDKeyPath, container: { $0 }, selectedItems: selectedItems, shouldAllowSelection: shouldAllowSelection, shouldAllowDeselection: shouldAllowDeselection, onCellEvent: onCellEvent, dragDropConfig: dragDropConfig, shouldAllowSwipeToDelete: shouldAllowSwipeToDelete, onSwipeToDelete: onSwipeToDelete, contextMenuProvider: contextMenuProvider, contentBuilder: contentBuilder)
+		self.init(id: id, data: data, dataID: dataIDKeyPath, container: { $0 }, selectedIndexes: selectedIndexes, shouldAllowSelection: shouldAllowSelection, shouldAllowDeselection: shouldAllowDeselection, onCellEvent: onCellEvent, dragDropConfig: dragDropConfig, shouldAllowSwipeToDelete: shouldAllowSwipeToDelete, onSwipeToDelete: onSwipeToDelete, contextMenuProvider: contextMenuProvider, contentBuilder: contentBuilder)
 	}
 }
 
@@ -88,7 +88,7 @@ public extension ASCollectionViewSection
 		id: SectionID,
 		data: DataCollection,
 		container: @escaping ((Content) -> Container),
-		selectedItems: Binding<Set<Int>>? = nil,
+		selectedIndexes: Binding<Set<Int>>? = nil,
 		shouldAllowSelection: ((_ index: Int) -> Bool)? = nil,
 		shouldAllowDeselection: ((_ index: Int) -> Bool)? = nil,
 		onCellEvent: OnCellEvent<DataCollection.Element>? = nil,
@@ -99,13 +99,13 @@ public extension ASCollectionViewSection
 		@ViewBuilder contentBuilder: @escaping ((DataCollection.Element, ASCellContext) -> Content))
 		where DataCollection.Index == Int, DataCollection.Element: Identifiable
 	{
-		self.init(id: id, data: data, dataID: \.id, container: container, selectedItems: selectedItems, shouldAllowSelection: shouldAllowSelection, shouldAllowDeselection: shouldAllowDeselection, onCellEvent: onCellEvent, dragDropConfig: dragDropConfig, shouldAllowSwipeToDelete: shouldAllowSwipeToDelete, onSwipeToDelete: onSwipeToDelete, contextMenuProvider: contextMenuProvider, contentBuilder: contentBuilder)
+		self.init(id: id, data: data, dataID: \.id, container: container, selectedIndexes: selectedIndexes, shouldAllowSelection: shouldAllowSelection, shouldAllowDeselection: shouldAllowDeselection, onCellEvent: onCellEvent, dragDropConfig: dragDropConfig, shouldAllowSwipeToDelete: shouldAllowSwipeToDelete, onSwipeToDelete: onSwipeToDelete, contextMenuProvider: contextMenuProvider, contentBuilder: contentBuilder)
 	}
 
 	init<Content: View, DataCollection: RandomAccessCollection>(
 		id: SectionID,
 		data: DataCollection,
-		selectedItems: Binding<Set<Int>>? = nil,
+		selectedIndexes: Binding<Set<Int>>? = nil,
 		shouldAllowSelection: ((_ index: Int) -> Bool)? = nil,
 		shouldAllowDeselection: ((_ index: Int) -> Bool)? = nil,
 		onCellEvent: OnCellEvent<DataCollection.Element>? = nil,
@@ -116,7 +116,7 @@ public extension ASCollectionViewSection
 		@ViewBuilder contentBuilder: @escaping ((DataCollection.Element, ASCellContext) -> Content))
 		where DataCollection.Index == Int, DataCollection.Element: Identifiable
 	{
-		self.init(id: id, data: data, container: { $0 }, selectedItems: selectedItems, shouldAllowSelection: shouldAllowSelection, shouldAllowDeselection: shouldAllowDeselection, onCellEvent: onCellEvent, dragDropConfig: dragDropConfig, shouldAllowSwipeToDelete: shouldAllowSwipeToDelete, onSwipeToDelete: onSwipeToDelete, contextMenuProvider: contextMenuProvider, contentBuilder: contentBuilder)
+		self.init(id: id, data: data, container: { $0 }, selectedIndexes: selectedIndexes, shouldAllowSelection: shouldAllowSelection, shouldAllowDeselection: shouldAllowDeselection, onCellEvent: onCellEvent, dragDropConfig: dragDropConfig, shouldAllowSwipeToDelete: shouldAllowSwipeToDelete, onSwipeToDelete: onSwipeToDelete, contextMenuProvider: contextMenuProvider, contentBuilder: contentBuilder)
 	}
 }
 

Sources/ASCollectionView/Implementation/ASCollectionView.swift

@@ -39,6 +39,9 @@ public struct ASCollectionView<SectionID: Hashable>: UIViewControllerRepresentab
 	internal var alwaysBounceVertical: Bool = false
 	internal var alwaysBounceHorizontal: Bool = false
 
+	internal var allowsSelection: Bool = true
+	internal var allowsMultipleSelection: Bool = false
+
 	internal var scrollPositionSetter: Binding<ASCollectionViewScrollPosition?>?
 
 	internal var animateOnDataRefresh: Bool = true
@@ -139,6 +142,8 @@ public struct ASCollectionView<SectionID: Hashable>: UIViewControllerRepresentab
 		private var hasFiredBoundaryNotificationForBoundary: Set<Boundary> = []
 		private var haveRegisteredForSupplementaryOfKind: Set<String> = []
 
+		private var selectedIndexPaths: Set<IndexPath> = []
+
 		// MARK: Caching
 
 		private var autoCachingHostingControllers = ASPriorityCache<ASCollectionViewItemUniqueID, ASHostingControllerProtocol>()
@@ -190,14 +195,12 @@ public struct ASCollectionView<SectionID: Hashable>: UIViewControllerRepresentab
 			assignIfChanged(collectionView, \.dragInteractionEnabled, newValue: true)
 			assignIfChanged(collectionView, \.alwaysBounceVertical, newValue: parent.alwaysBounceVertical)
 			assignIfChanged(collectionView, \.alwaysBounceHorizontal, newValue: parent.alwaysBounceHorizontal)
+			assignIfChanged(collectionView, \.allowsSelection, newValue: parent.allowsSelection)
+			assignIfChanged(collectionView, \.allowsMultipleSelection, newValue: parent.allowsMultipleSelection)
 			assignIfChanged(collectionView, \.showsVerticalScrollIndicator, newValue: parent.verticalScrollIndicatorEnabled)
 			assignIfChanged(collectionView, \.showsHorizontalScrollIndicator, newValue: parent.horizontalScrollIndicatorEnabled)
 			assignIfChanged(collectionView, \.keyboardDismissMode, newValue: .onDrag)
 			updateCollectionViewContentInsets(collectionView)
-
-			let isEditing = parent.editMode?.wrappedValue.isEditing ?? false
-			assignIfChanged(collectionView, \.allowsSelection, newValue: isEditing)
-			assignIfChanged(collectionView, \.allowsMultipleSelection, newValue: isEditing)
 		}
 
 		func updateCollectionViewContentInsets(_ collectionView: UICollectionView)
@@ -331,7 +334,7 @@ public struct ASCollectionView<SectionID: Hashable>: UIViewControllerRepresentab
 				animated: parent.animateOnDataRefresh && transactionAnimationEnabled,
 				transaction: transaction)
 
-			updateSelectionBindings(cv)
+			updateSelection(cv, transaction: transaction)
 		}
 
 		func refreshVisibleCells()
@@ -678,18 +681,53 @@ public struct ASCollectionView<SectionID: Hashable>: UIViewControllerRepresentab
 
 		public func collectionView(_ collectionView: UICollectionView, didSelectItemAt indexPath: IndexPath)
 		{
-			updateSelectionBindings(collectionView)
+			updateSelection(collectionView)
 		}
 
 		public func collectionView(_ collectionView: UICollectionView, didDeselectItemAt indexPath: IndexPath)
 		{
-			updateSelectionBindings(collectionView)
+			updateSelection(collectionView)
 		}
 
-		func updateSelectionBindings(_ collectionView: UICollectionView)
+		func updateSelection(_ collectionView: UICollectionView, transaction: Transaction? = nil)
 		{
-			let selected = collectionView.indexPathsForSelectedItems ?? []
-			let selectionBySection = Dictionary(grouping: selected) { $0.section }
+			let selectedInDataSource = selectedIndexPathsInDataSource
+			let selectedInCollectionView = Set(collectionView.indexPathsForSelectedItems ?? [])
+			guard selectedInDataSource != selectedInCollectionView else { return }
+
+			let newSelection = threeWayMerge(base: selectedIndexPaths, dataSource: selectedInDataSource, collectionView: selectedInCollectionView)
+			let (toDeselect, toSelect) = selectionDifferences(oldSelectedIndexPaths: selectedInCollectionView, newSelectedIndexPaths: newSelection)
+
+			selectedIndexPaths = newSelection
+			updateSelectionBindings(newSelection)
+			updateSelectionInCollectionView(collectionView, indexPathsToDeselect: toDeselect, indexPathsToSelect: toSelect, transaction: transaction)
+		}
+
+		private var selectedIndexPathsInDataSource: Set<IndexPath>
+		{
+			parent.sections.enumerated().reduce(Set<IndexPath>())
+			{ (selectedIndexPaths, section) -> Set<IndexPath> in
+				guard let indexes = section.element.dataSource.getSelectedIndexes() else { return selectedIndexPaths }
+				let indexPaths = indexes.map { IndexPath(item: $0, section: section.offset) }
+				return selectedIndexPaths.union(indexPaths)
+			}
+		}
+
+		private func threeWayMerge(base: Set<IndexPath>, dataSource: Set<IndexPath>, collectionView: Set<IndexPath>) -> Set<IndexPath>
+		{
+			base == dataSource ? collectionView : dataSource
+		}
+
+		private func selectionDifferences(oldSelectedIndexPaths: Set<IndexPath>, newSelectedIndexPaths: Set<IndexPath>) -> (toDeselect: Set<IndexPath>, toSelect: Set<IndexPath>)
+		{
+			let toDeselect = oldSelectedIndexPaths.subtracting(newSelectedIndexPaths)
+			let toSelect = newSelectedIndexPaths.subtracting(oldSelectedIndexPaths)
+			return (toDeselect: toDeselect, toSelect: toSelect)
+		}
+
+		private func updateSelectionBindings(_ selectedIndexPaths: Set<IndexPath>)
+		{
+			let selectionBySection = Dictionary(grouping: selectedIndexPaths) { $0.section }
 				.mapValues
 			{
 				Set($0.map { $0.item })
@@ -699,6 +737,13 @@ public struct ASCollectionView<SectionID: Hashable>: UIViewControllerRepresentab
 			}
 		}
 
+		private func updateSelectionInCollectionView(_ collectionView: UICollectionView, indexPathsToDeselect: Set<IndexPath>, indexPathsToSelect: Set<IndexPath>, transaction: Transaction? = nil)
+		{
+			let isAnimated = (transaction?.animation != nil) && !(transaction?.disablesAnimations ?? false)
+			indexPathsToDeselect.forEach { collectionView.deselectItem(at: $0, animated: isAnimated) }
+			indexPathsToSelect.forEach { collectionView.selectItem(at: $0, animated: isAnimated, scrollPosition: []) }
+		}
+
 		func canDrop(at indexPath: IndexPath) -> Bool
 		{
 			guard !indexPath.isEmpty else { return false }

Sources/ASCollectionView/Implementation/ASSectionDataSource.swift

@@ -28,6 +28,7 @@ internal protocol ASSectionDataSourceProtocol
 	func getContextMenu(for indexPath: IndexPath) -> UIContextMenuConfiguration?
 	func getSelfSizingSettings(context: ASSelfSizingContext) -> ASSelfSizingConfig?
 
+	func getSelectedIndexes() -> Set<Int>?
 	func isSelected(index: Int) -> Bool
 	func updateSelection(_ indices: Set<Int>)
 	func shouldSelect(_ indexPath: IndexPath) -> Bool
@@ -55,7 +56,7 @@ internal struct ASSectionDataSource<DataCollection: RandomAccessCollection, Data
 	var container: (Content) -> Container
 	var content: (DataCollection.Element, ASCellContext) -> Content
 
-	var selectedItems: Binding<Set<Int>>?
+	var selectedIndexes: Binding<Set<Int>>?
 	var shouldAllowSelection: ((_ index: Int) -> Bool)?
 	var shouldAllowDeselection: ((_ index: Int) -> Bool)?
 
@@ -281,28 +282,37 @@ internal struct ASSectionDataSource<DataCollection: RandomAccessCollection, Data
 		selfSizingConfig?(context)
 	}
 
+	func getSelectedIndexes() -> Set<Int>?
+	{
+		selectedIndexes?.wrappedValue
+	}
+
 	func isSelected(index: Int) -> Bool
 	{
-		selectedItems?.wrappedValue.contains(index) ?? false
+		selectedIndexes?.wrappedValue.contains(index) ?? false
 	}
 
 	func updateSelection(_ indices: Set<Int>)
 	{
 		DispatchQueue.main.async {
-			self.selectedItems?.wrappedValue = Set(indices)
+			self.selectedIndexes?.wrappedValue = Set(indices)
 		}
 	}
 
 	func shouldSelect(_ indexPath: IndexPath) -> Bool
 	{
-		guard data.containsIndex(indexPath.item) else { return (selectedItems != nil) }
-		return shouldAllowSelection?(indexPath.item) ?? (selectedItems != nil)
+		guard data.containsIndex(indexPath.item) else { return isSelectable }
+		return shouldAllowSelection?(indexPath.item) ?? isSelectable
 	}
 
 	func shouldDeselect(_ indexPath: IndexPath) -> Bool
 	{
-		guard data.containsIndex(indexPath.item) else { return (selectedItems != nil) }
-		return shouldAllowDeselection?(indexPath.item) ?? (selectedItems != nil)
+		guard data.containsIndex(indexPath.item) else { return isSelectable }
+		return shouldAllowDeselection?(indexPath.item) ?? isSelectable
+	}
+
+	private var isSelectable: Bool {
+		selectedIndexes != nil
 	}
 }