A UITableViewDiffableDataSource subclass with support for section titles

SectionTitledDiffableDataSource.swift

import UIKit

open class SectionTitledDiffableDataSource<SectionIdentifierType, ItemIdentifierType>: UITableViewDiffableDataSource<SectionIdentifierType, ItemIdentifierType>
    where SectionIdentifierType : Hashable, ItemIdentifierType : Hashable {
    
    public typealias SectionTitleProvider = (UITableView, SectionIdentifierType) -> String?
    
    open var sectionTitleProvider: SectionTitleProvider?
    open var useSectionIndex: Bool = false
    
    open override func sectionIndexTitles(for tableView: UITableView) -> [String]? {
        guard useSectionIndex, let sectionTitleProvider = sectionTitleProvider else { return nil }
        return snapshot().sectionIdentifiers.compactMap { sectionTitleProvider(tableView, $0) }
    }
    
    open override func tableView(_ tableView: UITableView, titleForHeaderInSection section: Int) -> String? {
        sectionTitleProvider?(tableView, snapshot().sectionIdentifiers[section])
    }

    open override func tableView(_ tableView: UITableView, sectionForSectionIndexTitle title: String, at index: Int) -> Int {
        guard useSectionIndex else { return 0 }
        return snapshot().sectionIdentifiers.firstIndex(where: { sectionTitleProvider?(tableView, $0) == title }) ?? 0
    }
}