CollectionViewConcatDataSource.swift

import DiffableDataSources
import UIKit

class CollectionViewConcatDataSource<SectionIdentifierType: Hashable>: NSObject, UICollectionViewDataSource {
    typealias ItemIdentifierType = AnyHashable
    public typealias SupplementaryViewProvider = (UICollectionView, String, IndexPath) -> UICollectionReusableView?

    private var innerDataSource: CollectionViewDiffableDataSource<SectionIdentifierType, AnyHashable>!
    private var adapterMap: [SectionIdentifierType: DiffableDataSourceSectionAdapterType] = [:]
    var defaultSupplementaryViewProvider: SupplementaryViewProvider? = nil

    init(collectionView: UICollectionView) {
        super.init()
        innerDataSource = .init(collectionView: collectionView, cellProvider: provideCell(_:_:_:))
        collectionView.dataSource = self
    }

    func appendAdapter<Item: Hashable>(_ adapter: DiffableDataSourceSectionAdapter<Item>, toSection sectionIdentifier: SectionIdentifierType) {
        if adapterMap[sectionIdentifier] != nil {
            assertionFailure()
            return
        }

        adapterMap[sectionIdentifier] = adapter
    }

    private func provideCell(_ collectionView: UICollectionView, _ indexPath: IndexPath, _ item: AnyHashable) -> UICollectionViewCell? {
        let snapshot = self.snapshot()
        let sectionIdentifier = snapshot.sectionIdentifiers[indexPath.section]
        guard let adapter = adapterMap[sectionIdentifier] else {
            fatalError()
        }

        guard let cell = adapter.provideCell(collectionView, indexPath: indexPath, item: item) else {
            fatalError("UICollectionView dataSource returned a nil cell for item at index path: \(indexPath), collectionView: \(collectionView), itemIdentifier: \(item)")
        }
        return cell
    }

    public func apply(_ snapshot: DiffableDataSourceSnapshot<SectionIdentifierType, ItemIdentifierType>, animatingDifferences: Bool = true, completion: (() -> Void)? = nil) {
        innerDataSource.apply(snapshot, animatingDifferences: animatingDifferences, completion: completion)
    }

    public func snapshot() -> DiffableDataSourceSnapshot<SectionIdentifierType, ItemIdentifierType> {
        innerDataSource.snapshot()
    }

    public func itemIdentifier(for indexPath: IndexPath) -> ItemIdentifierType? {
        innerDataSource.itemIdentifier(for: indexPath)
    }

    public func indexPath(for itemIdentifier: ItemIdentifierType) -> IndexPath? {
        innerDataSource.indexPath(for: itemIdentifier)
    }

    public func numberOfSections(in collectionView: UICollectionView) -> Int {
        innerDataSource.numberOfSections(in: collectionView)
    }

    public func collectionView(_ collectionView: UICollectionView, numberOfItemsInSection section: Int) -> Int {
        innerDataSource.collectionView(collectionView, numberOfItemsInSection: section)
    }

    open func collectionView(_ collectionView: UICollectionView, cellForItemAt indexPath: IndexPath) -> UICollectionViewCell {
        innerDataSource.collectionView(collectionView, cellForItemAt: indexPath)
    }

    open func collectionView(_ collectionView: UICollectionView, viewForSupplementaryElementOfKind kind: String, at indexPath: IndexPath) -> UICollectionReusableView {
        let snapshot = self.snapshot()
        let sectionIdentifier = snapshot.sectionIdentifiers[indexPath.section]
        guard let adapter = adapterMap[sectionIdentifier] else {
            return defaultSupplementaryViewProvider?(collectionView, kind, indexPath) ?? UICollectionViewCell()
        }

        return adapter.provideSupplementaryView(collectionView, viewForSupplementaryElementOfKind: kind, at: indexPath) ?? defaultSupplementaryViewProvider?(collectionView, kind, indexPath) ?? UICollectionViewCell()
    }

    func collectionView(_ collectionView: UICollectionView, didSelectItemAt indexPath: IndexPath) {
        let snapshot = self.snapshot()
        let sectionIdentifier = snapshot.sectionIdentifiers[indexPath.section]
        guard let adapter = adapterMap[sectionIdentifier] else {
            return
        }

        let item = snapshot.itemIdentifiers(inSection: sectionIdentifier)[indexPath.item]
        adapter.collectionView(collectionView, didSelectItemAt: indexPath, item: item)
    }
}

protocol DiffableDataSourceSectionAdapterType {
    func provideCell(_ collectionView: UICollectionView, indexPath: IndexPath, item: AnyHashable) -> UICollectionViewCell?
    func provideSupplementaryView(_ collectionView: UICollectionView, viewForSupplementaryElementOfKind kind: String, at indexPath: IndexPath) -> UICollectionReusableView?
    func collectionView(_ collectionView: UICollectionView, didSelectItemAt indexPath: IndexPath, item: AnyHashable)
}

struct DiffableDataSourceSectionAdapter<ItemIdentifierType: Hashable>: DiffableDataSourceSectionAdapterType {
    public typealias CellProvider = (UICollectionView, IndexPath, ItemIdentifierType) -> UICollectionViewCell?
    public typealias SupplementaryViewProvider = (UICollectionView, String, IndexPath) -> UICollectionReusableView?

    var cellProvider: CellProvider
    var supplementaryViewProvider: SupplementaryViewProvider? = nil
    var didSelectItemAt: ((UICollectionView, IndexPath, ItemIdentifierType) -> Void)? = nil

    func provideCell(_ collectionView: UICollectionView, indexPath: IndexPath, item: AnyHashable) -> UICollectionViewCell? {
        cellProvider(collectionView, indexPath, item.base as! ItemIdentifierType)
    }

    func provideSupplementaryView(_ collectionView: UICollectionView, viewForSupplementaryElementOfKind kind: String, at indexPath: IndexPath) -> UICollectionReusableView? {
        supplementaryViewProvider?(collectionView, kind, indexPath)
    }

    func collectionView(_ collectionView: UICollectionView, didSelectItemAt indexPath: IndexPath, item: AnyHashable) {
        didSelectItemAt?(collectionView, indexPath, item.base as! ItemIdentifierType)
    }
}

        dataSource = CollectionViewConcatDataSource<Section>(collectionView: collectionView)
        dataSource.appendAdapter(DiffableDataSourceSectionAdapter<Favorite>(
            cellProvider: { collectionView, indexPath, data in
                let cell = collectionView.dequeue(FavoriteCell.self, for: indexPath)
                cell.update(data: data)
                return cell
            },
            didSelectItemAt: { [weak self] _, _, data in
                self?.favoriteTapped(data)
            }),
            toSection: .favorite
        )

                var snapshot = DiffableDataSourceSnapshot<Section, AnyHashable>()
                if !contents.favorites.isEmpty {
                    snapshot.appendSections([.favorite])
                    snapshot.appendItems(contents.favorites) // ここに型の制約がない。セクションに応じた型をappendすることを強制できない
                }
                self.dataSource.apply(snapshot)