DiffableInfiniteClippingCollectionView.swift

import UIKit

protocol InnerCollectionViewCell: UICollectionViewCell {
    associatedtype Item: Hashable
    func configure(with item: Item, index: Int)
}

extension UICollectionViewFlowLayout {
    func rect(at index: Int) -> CGRect {
        CGRect(
            x: sectionInset.left + CGFloat(index) * (itemSize.width + minimumLineSpacing),
            y: sectionInset.top,
            width: itemSize.width,
            height: itemSize.height
        )
    }

    func totalWidth(of items: Int) -> CGFloat {
        CGFloat(items) * (itemSize.width + minimumLineSpacing)
    }
}

class InfiniteCell<Cell: InnerCollectionViewCell>: UICollectionViewCell, UICollectionViewDelegate {
    /// The total number of items being shown.
    private var numberOfItems: Int { dataSource.snapshot().numberOfItems }
    /// The number of additional items over the original items.count.
    private var numberOfAdditionalItems: Int { numberOfItems - items.count }
    /// At what index to load more data, this number will be subtracted from total.
    private var triggerIndex: Int = 0
    /// The original items that we want to loop over.
    private var items: [Cell.Item] = []
    private lazy var flowLayout = UICollectionViewFlowLayout()
    private lazy var collectionView = UICollectionView(frame: .zero, collectionViewLayout: flowLayout)
    private lazy var collectionViewHeightConstraint = collectionView.heightAnchor.constraint(equalToConstant: 0)

    struct Config {
        let itemSize: CGSize
        let itemSpacing: CGFloat
        let sectionInset: UIEdgeInsets
        let triggerIndex: Int
    }

    private struct OnlySection: Identifiable, Hashable {
        let id = 0
    }

    private struct UniqueItem: Identifiable, Hashable {
        let id: Int
        let wrappedItem: Cell.Item
    }

    private lazy var cellRegistration = UICollectionView.CellRegistration<Cell, UniqueItem> { cell, indexPath, item in
        cell.configure(with: item.wrappedItem, index: indexPath.row)
    }

    private lazy var dataSource = UICollectionViewDiffableDataSource<OnlySection, UniqueItem>(
        collectionView: collectionView,
        cellProvider: cellRegistration.cellProvider
    )

    override func willMove(toSuperview newSuperview: UIView?) {
        super.willMove(toSuperview: newSuperview)
        if newSuperview != nil {
            collectionView.translatesAutoresizingMaskIntoConstraints = false
            contentView.addSubview(collectionView)
            NSLayoutConstraint.activate([
                collectionView.topAnchor.constraint(equalTo: contentView.topAnchor),
                collectionView.leadingAnchor.constraint(equalTo: contentView.leadingAnchor),
                collectionView.trailingAnchor.constraint(equalTo: contentView.trailingAnchor),
                collectionViewHeightConstraint,
                collectionView.bottomAnchor.constraint(equalTo: contentView.bottomAnchor),
            ])
            collectionView.dataSource = dataSource
            collectionView.delegate = self
            collectionView.showsHorizontalScrollIndicator = false
            collectionView.contentInset = .zero
            collectionView.contentInsetAdjustmentBehavior = .never
            collectionView.insetsLayoutMarginsFromSafeArea = false

            flowLayout.sectionInset = .zero
            flowLayout.scrollDirection = .horizontal
            flowLayout.minimumInteritemSpacing = 0
            flowLayout.minimumLineSpacing = 0
        }
    }

    func configure(
        with newItems: [Cell.Item],
        config: Config
    ) {
        triggerIndex = config.triggerIndex
        collectionViewHeightConstraint.constant = config.sectionInset.top + config.itemSize.height + config.sectionInset.bottom
        flowLayout.itemSize = config.itemSize
        flowLayout.sectionInset = config.sectionInset
        flowLayout.minimumLineSpacing = config.itemSpacing
        updateConstraintsIfNeeded()

        items = newItems
        if numberOfItems > 0 {
            apply(count: numberOfItems)
        } else {
            apply(count: newItems.count)
        }
    }

    func collectionView(_ collectionView: UICollectionView, willDisplay cell: UICollectionViewCell, forItemAt indexPath: IndexPath) {
        if indexPath.row == numberOfItems - triggerIndex {
            DispatchQueue.main.async {
                self.apply(count: self.numberOfItems + self.items.count)
            }
        }
    }

    func scrollViewDidEndDecelerating(_ scrollView: UIScrollView) {
        removeAdditionalItems()
    }

    func scrollViewDidEndDragging(_ scrollView: UIScrollView, willDecelerate decelerate: Bool) {
        if !decelerate {
            removeAdditionalItems()
        }
    }

    private func apply(count: Int) {
        var snapshot = NSDiffableDataSourceSnapshot<OnlySection, UniqueItem>()
        let section = OnlySection()
        snapshot.appendSections([section])
        snapshot.appendItems((0..<count).map { index in
            UniqueItem(id: index, wrappedItem: items[index % items.count])
        }, toSection: section)
        dataSource.apply(snapshot)
    }

    private func removeAdditionalItems() {
        guard numberOfAdditionalItems > 0 else { return }
        let triggerEnd = numberOfAdditionalItems
        let triggerStart = triggerEnd - triggerIndex
        let triggerRange = triggerStart..<triggerEnd
        let triggerRects = triggerRange.map(flowLayout.rect(at:))
        let visibleRect = CGRect(
            origin: CGPoint(x: collectionView.contentOffset.x, y: 0),
            size: collectionView.frame.size
        )
        guard triggerRects.contains(where: visibleRect.intersects) else {
            return resetAdditionalItems(to: 0)
        }
        guard numberOfItems > items.count * 2 else { return }
        resetAdditionalItems(to: items.count)
    }

    private func resetAdditionalItems(to count: Int) {
        var contentOffset = collectionView.contentOffset
        contentOffset.x = contentOffset.x.truncatingRemainder(dividingBy: flowLayout.totalWidth(of: items.count))
        DispatchQueue.main.async {
            self.collectionView.contentOffset = contentOffset
            self.apply(count: self.numberOfItems + count)
        }
    }
}

class ViewController: UIViewController {
    private var colors: [UIColor] = [
        .red,
        .blue,
        .green,
        .orange,
        .magenta,
        .purple,
        .yellow,
        .black,
    ]

    private lazy var collectionView = UICollectionView(frame: .zero, collectionViewLayout: UICollectionViewFlowLayout())

    struct Section: Hashable, Identifiable {
        let id: UUID = UUID()
        let name: String
    }

    struct Item: Hashable, Identifiable {
        let id: UUID = UUID()
        let colors: [UIColor]
    }

    lazy var dataSource = UICollectionViewDiffableDataSource<Section, Item>(
        collectionView: collectionView,
        cellProvider: cellRegistration.cellProvider
    )

    let wideConfig = InfiniteCell<InnerCell>.Config(
        itemSize: CGSize(width: 210, height: 140),
        itemSpacing: 10,
        sectionInset: UIEdgeInsets(top: 10, left: 10, bottom: 0, right: 10),
        triggerIndex: 3
    )

    let tallConfig = InfiniteCell<InnerCell>.Config(
        itemSize: CGSize(width: 140, height: 210),
        itemSpacing: 5,
        sectionInset: UIEdgeInsets(top: 5, left: 10, bottom: 0, right: 10),
        triggerIndex: 4
    )

    lazy var cellRegistration = UICollectionView.CellRegistration<InfiniteCell<InnerCell>, Item> { [unowned self] cell, indexPath, item in
        cell.configure(
            with: item.colors,
            config: indexPath.section % 2 == 0 ? self.wideConfig : self.tallConfig
        )
    }

    lazy var sectionRegistration = UICollectionView.SupplementaryRegistration<Header>(elementKind: UICollectionView.elementKindSectionHeader) { [unowned self] supplementaryView, elementKind, indexPath in
        supplementaryView.configure(with: self.dataSource.sectionIdentifier(for: indexPath.section)!.name)
    }

    lazy var layoutConfig = UICollectionLayoutListConfiguration(appearance: .plain)

    lazy var layout = UICollectionViewCompositionalLayout.list(using: layoutConfig)

    override func viewDidLoad() {
        super.viewDidLoad()
        collectionView.translatesAutoresizingMaskIntoConstraints = false
        view.addSubview(collectionView)
        NSLayoutConstraint.activate([
            collectionView.topAnchor.constraint(equalTo: view.safeAreaLayoutGuide.topAnchor),
            collectionView.leadingAnchor.constraint(equalTo: view.safeAreaLayoutGuide.leadingAnchor),
            collectionView.trailingAnchor.constraint(equalTo: view.safeAreaLayoutGuide.trailingAnchor),
            collectionView.bottomAnchor.constraint(equalTo: view.safeAreaLayoutGuide.bottomAnchor),
        ])

        layoutConfig.showsSeparators = false
        layoutConfig.headerMode = .supplementary
        layoutConfig.footerMode = .none

        collectionView.collectionViewLayout = layout

        dataSource.supplementaryViewProvider = sectionRegistration.supplementaryProvider
    }

    override func viewWillAppear(_ animated: Bool) {
        super.viewWillAppear(animated)

        var snapshot = NSDiffableDataSourceSnapshot<Section, Item>()
        let a = Section(name: "AAA")
        let b = Section(name: "BBB")
        let c = Section(name: "CCC")
        let d = Section(name: "DDD")
        let e = Section(name: "EEE")
        let f = Section(name: "FFF")
        snapshot.appendSections([a, b, c, d, e, f])
        snapshot.appendItems([Item(colors: colors)], toSection: a)
        snapshot.appendItems([Item(colors: colors)], toSection: b)
        snapshot.appendItems([Item(colors: colors)], toSection: c)
        snapshot.appendItems([Item(colors: colors)], toSection: d)
        snapshot.appendItems([Item(colors: colors)], toSection: e)
        snapshot.appendItems([Item(colors: colors)], toSection: f)

        dataSource.apply(snapshot)
    }
}

class Header: UICollectionReusableView {
    private lazy var textLabel = UILabel()

    override func willMove(toSuperview newSuperview: UIView?) {
        super.willMove(toSuperview: newSuperview)
        guard newSuperview != nil else { return }
        textLabel.translatesAutoresizingMaskIntoConstraints = false
        textLabel.font = .preferredFont(forTextStyle: .subheadline, compatibleWith: traitCollection)
        addSubview(textLabel)
        NSLayoutConstraint.activate([
            textLabel.leadingAnchor.constraint(equalTo: leadingAnchor, constant: 10),
            textLabel.trailingAnchor.constraint(equalTo: trailingAnchor, constant: -10),
            textLabel.topAnchor.constraint(equalTo: topAnchor),
            textLabel.bottomAnchor.constraint(equalTo: bottomAnchor),
        ])
    }

    func configure(with text: String) {
        textLabel.text = text
    }
}

class InnerCell: UICollectionViewCell, InnerCollectionViewCell {
    private lazy var label = UILabel()

    override func willMove(toSuperview newSuperview: UIView?) {
        super.willMove(toSuperview: newSuperview)
        label.translatesAutoresizingMaskIntoConstraints = false
        label.textColor = .white
        label.textAlignment = .center
        contentView.addSubview(label)
        NSLayoutConstraint.activate([
            label.leadingAnchor.constraint(equalTo: contentView.leadingAnchor),
            label.trailingAnchor.constraint(equalTo: contentView.trailingAnchor),
            label.topAnchor.constraint(equalTo: contentView.topAnchor),
            label.bottomAnchor.constraint(equalTo: contentView.bottomAnchor),
        ])
    }

    func configure(with color: UIColor, index: Int) {
        label.text = "\(index)"
        contentView.backgroundColor = color
        backgroundColor = color
        layer.masksToBounds = true
        layer.cornerRadius = 15
    }
}

extension UICollectionView.CellRegistration {
    var cellProvider: (UICollectionView, IndexPath, Item) -> Cell {
        return { collectionView, indexPath, item in
            collectionView.dequeueConfiguredReusableCell(
                using: self,
                for: indexPath,
                item: item
            )
        }
    }
}

extension UICollectionView.SupplementaryRegistration {
    var supplementaryProvider: (UICollectionView, String, IndexPath) -> Supplementary {
        return { collectionView, kind, indexPath in
            return collectionView.dequeueConfiguredReusableSupplementary(
                using: self,
                for: indexPath
            )
        }
    }
}