Generic, boilerplate-reducing base classes for UIKit Diffable Data Sources

DiffableDatasource.swift

import UIKit

public protocol DiffableDatasource: AnyObject {
    typealias Snapshot = NSDiffableDataSourceSnapshot<Section, Item>
    typealias SnapshotProvider = () -> Snapshot
    
    associatedtype View
    associatedtype Datasource
    associatedtype Section: Hashable & Sendable
    associatedtype Item: Hashable & Sendable
    
    var datasource: Datasource? { get set }
    func snapshotData() -> SnapshotData<Section, Item>
}

fileprivate extension DiffableDatasource {
    func makeSnapshotProvider() -> SnapshotProvider {
        let data = snapshotData()
        let snapshotProvider: SnapshotProvider = {
            var snapshot = NSDiffableDataSourceSnapshot<Section, Item>()
            snapshot.appendSections(data.sections)
            for section in data.sections {
                if let sectionItems = data.items[section] {
                    snapshot.appendItems(sectionItems, toSection: section)
                }
            }
            return snapshot
        }
        return snapshotProvider
    }
}

public struct SnapshotData<Section: Hashable & Sendable, Item: Hashable & Sendable> {
    public let sections: [Section]
    public let items: [Section: [Item]]
    
    public init(sections: [Section], items: [Section : [Item]]) {
        self.sections = sections
        self.items = items
    }
}

public protocol TableDataSource: DiffableDatasource where View == UITableView, Datasource == UITableViewDiffableDataSource<Section, Item> {
    var tableView: View { get }
}

public extension TableDataSource {
    @MainActor
    func prepareDatasource(cellProvider: @escaping Datasource.CellProvider) {
        datasource = Datasource(tableView: tableView, cellProvider: cellProvider)
    }
    
    @MainActor
    func updateTableView(animated: Bool = false, completion: (() -> Void)? = nil) {
        applyChanges(snapshotProvider: makeSnapshotProvider(), animated: animated, completion: completion)
    }

    @MainActor
    func applyChanges(snapshotProvider: SnapshotProvider, animated: Bool = false, completion: (() -> Void)? = nil) {
        datasource?.apply(
            snapshotProvider(),
            animatingDifferences: animated,
            completion: completion
        )
    }
}

public protocol CollectionDataSource: DiffableDatasource where View == UICollectionView, Datasource == UICollectionViewDiffableDataSource<Section, Item> {
    var collectionView: View { get }
}

public extension CollectionDataSource {
    @MainActor
    func prepareDatasource(cellProvider: @escaping Datasource.CellProvider) {
        datasource = Datasource(collectionView: collectionView, cellProvider: cellProvider)
    }
    
    @MainActor
    func updateCollectionView(animated: Bool = false, completion: (() -> Void)? = nil) {
        applyChanges(snapshotProvider: makeSnapshotProvider(), animated: animated, completion: completion)
    }
    
    @MainActor
    func applyChanges(snapshotProvider: SnapshotProvider, animated: Bool = false, completion: (() -> Void)? = nil) {
        datasource?.apply(
            snapshotProvider(),
            animatingDifferences: animated,
            completion: completion
        )
    }
}

open class DiffableTableViewController<Section: Hashable & Sendable, Item: Hashable & Sendable>: UIViewController, @MainActor TableDataSource {
    public var datasource: UITableViewDiffableDataSource<Section, Item>?
    
    public lazy var tableView: UITableView = {
        let table = UITableView()
        table.translatesAutoresizingMaskIntoConstraints = false
        return table
    }()
    
    open override func viewDidLoad() {
        super.viewDidLoad()
        prepareDatasource { [weak self] tableView, indexPath, itemIdentifier in
            self?.configureCell(tableView: tableView, indexPath: indexPath, item: itemIdentifier)
        }
    }
    
    open override func viewDidAppear(_ animated: Bool) {
        super.viewDidAppear(animated)
        if tableView.superview == nil {
            fatalError("DiffableTableViewController requires its tableView to be added to the view hierarchy.")
        }
    }
    
    open func configureCell(tableView: UITableView, indexPath: IndexPath, item: Item) -> UITableViewCell? {
        fatalError("Subclasses must implement configureCell(tableView:indexPath:item:)")
    }
    
    open func snapshotData() -> SnapshotData<Section, Item> {
        fatalError("Subclasses must implement snapshotData()")
    }
}

open class DiffableCollectionViewController<Section: Hashable & Sendable, Item: Hashable & Sendable>: UIViewController, @MainActor CollectionDataSource {
    public var datasource: UICollectionViewDiffableDataSource<Section, Item>?
    
    private var _collectionView: UICollectionView?
    public var collectionView: UICollectionView {
        if let collection = _collectionView { return collection }
        let collection = prepareCollectionView()
        _collectionView = collection
        return collection
    }
    
    open override func viewDidLoad() {
        super.viewDidLoad()
        prepareDatasource { [weak self] collectionView, indexPath, itemIdentifier in
            self?.configureCell(collectionView: collectionView, indexPath: indexPath, item: itemIdentifier)
        }
    }
    
    open override func viewDidAppear(_ animated: Bool) {
        super.viewDidAppear(animated)
        if collectionView.superview == nil {
            fatalError("DiffableCollectionViewController requires its collectionView to be added to the view hierarchy.")
        }
    }
    
    open func prepareCollectionView() -> UICollectionView {
        fatalError("Subclasses must implement prepareCollectionView()")
    }
    
    open func configureCell(collectionView: UICollectionView, indexPath: IndexPath, item: Item) -> UICollectionViewCell? {
        fatalError("Subclasses must implement configureCell(collectionView:indexPath:item:)")
    }
    
    open func snapshotData() -> SnapshotData<Section, Item> {
        fatalError("Subclasses must implement snapshotData()")
    }
}