Custom UICollectionViewLayout getting stale IndexPaths in conjunction with DiffableDataSource

I've been chasing a few similar-looking crashes for a while, and finally found a way to reproduce them. Then, I found what I thought was a fix (on iOS 16.4.1), but that turned out to throw an exception on iOS 16.3, and now I'm no longer sure whether this fix is appropriate.

First, the cause of the crash. Since the reproducing code below is rather long, here's a quick summary:

I have a custom layout subclass of UICollectionViewFlowLayout that has its own delegate protocol to query some extra information required for layout. It uses that information to adjust layout attributes returned by the superclass.

I use this layout in conjunction with a UICollectionViewDiffableDataSource. After the data changes (toggling state in the example, which adds/removes a section), I create and apply a new snapshot.

And now we come to the crash: While applying the updated snapshot, the layout's layoutAttributesForElements(in:) is called, calls through to super, which in turn returns layout attributes for index paths that are no longer present in the new snapshot

When the layout then calls out to its delegate for more information, it does so with index paths that the delegate thinks should not exist according to its state, which leads to out of bounds array accesses and similar crashes (represented by precondition in the example code).

The potential solution I found is this: Wrapping the state update and the call to the diffable data source's apply in a collectionView.performBatchUpdates call prevented any out-of-date index paths being returned from the layout superclass.

I tested this on my iOS 16.4.1 device, and pushed out a beta, which revealed that on iOS 16.3, the code would lead to an exception being thrown with this message:

UICollectionView must be updated via the UICollectionViewDiffableDataSource APIs when acting as the UICollectionView's dataSource: please do not call mutation APIs directly on UICollectionView

So, finally, my actual question:

Does the fact that now works what was so decisively rejected by earlier iOS versions mean that my solution is appropriate, but I need to limit it to iOS 16.4.1 and above (I also tested 16.5 beta, where the solution also works)?

Or is the fact that this no longer throws an exception a bug and I should investigate other ways to address this issue (e.g. caching delegate responses in my layout so that by the time the indices are no longer valid the layout doesn't need to consult the delegate)?

Full example app code below. To reproduce, run the app on an iOS sim/device, tap one of the cells at the bottom of the screen, scroll all the way down, tap one of the cells again, observe crash. Then, uncomment the performBatchUpdates in didSelectItem and observe how the crash disappears.

import SwiftUI

protocol MyBuggyLayoutDelegate: UICollectionViewDelegateFlowLayout {
    func moreInfoAboutItem(at indexPath: IndexPath) -> String
}

class MyBuggyLayout: UICollectionViewFlowLayout {
    override func layoutAttributesForElements(in rect: CGRect) -> [UICollectionViewLayoutAttributes]? {
        guard let attrs = super.layoutAttributesForElements(in: rect) else { return nil }
        for a in attrs {
            if let del = collectionView?.delegate as? MyBuggyLayoutDelegate {
                let info = del.moreInfoAboutItem(at: a.indexPath)
                // do something with info...
            }
        }
        return attrs
    }
}

class VC: UICollectionViewController, MyBuggyLayoutDelegate {
    var state: Bool
    var source: UICollectionViewDiffableDataSource<Int, IndexPath>!

    init() {
        state = false
        super.init(collectionViewLayout: MyBuggyLayout())
    }

    required init?(coder: NSCoder) { fatalError("init(coder:) has not been implemented") }

    override func viewDidLoad() {
        super.viewDidLoad()

        collectionView.register(UICollectionViewCell.self, forCellWithReuseIdentifier: "cell")

        source = UICollectionViewDiffableDataSource(collectionView: collectionView) { (collectionView, indexPath, itemIdentifier) in
            let c = collectionView.dequeueReusableCell(withReuseIdentifier: "cell", for: indexPath)
            c.backgroundColor = itemIdentifier.item == 0 ? .blue : .red
            return c
        }

        source.apply(makeSnapshot(), animatingDifferences: false)
    }

    func makeSnapshot() -> NSDiffableDataSourceSnapshot<Int, IndexPath> {
        var snap = NSDiffableDataSourceSnapshot<Int, IndexPath>()

        // flipping state inserts/removes the first section
        for section in (state ? 0 : 1)...1 {
            snap.appendSections([section])
            for item in 0..<2 { snap.appendItems([[section, item]], toSection: section) }
        }

        return snap
    }

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

        collectionView.contentInset = UIEdgeInsets(
            top: collectionView.frame.height - 100, left: 0, bottom: 0, right: 0
        )
    }

    override func collectionView(_ collectionView: UICollectionView, didSelectItemAt indexPath: IndexPath) {
        // uncomment this to fix the error on iOS 16.4.1/16.5 beta, but throw an exception on iOS 16.3
//        collectionView.performBatchUpdates {
        state.toggle()
        source.apply(makeSnapshot(), animatingDifferences: true)
//        }
    }

    func collectionView(_ collectionView: UICollectionView, layout collectionViewLayout: UICollectionViewLayout, sizeForItemAt indexPath: IndexPath) -> CGSize {
        CGSize(width: collectionView.frame.width, height: 50)
    }

    func moreInfoAboutItem(at indexPath: IndexPath) -> String {
        // if state is false, makeSnapshot only produces a single section, so it
        // would stand to reason that this precondition should never be violated
        precondition(indexPath.section <= (state ? 1 : 0))
        return indexPath.item == 0 ? "lol" : "rofl"
    }
}

struct VCWrap: UIViewControllerRepresentable {
    typealias UIViewControllerType = VC
    func makeUIViewController(context: Context) -> VC { VC() }
    func updateUIViewController(_ uiViewController: VC, context: Context) {}
}

@main struct CVBugTestApp: App {
    var body: some Scene {  WindowGroup { VCWrap() } }
}