UICollectionViewListCell实现展开/折叠功能时行为不可预测的问题求助

我用UICollectionViewListCell和Diffable Data Source实现了展开/折叠功能，但实际运行时行为很不稳定。预期效果是点击某一行时，该行展开，同时其他所有行自动折叠，但这个逻辑并没有每次都正常生效。可以通过下面的动图看到具体的异常表现：

以下是我的源码：

ViewController.swift #

import UIKit

enum Section : Codable {
    case note
}

extension UIView {
    static func getUINib() -> UINib {
        return UINib(nibName: String(describing: self), bundle: nil)
    }
}

class ViewController: UIViewController {
    private let strings = [
        "Helo world",
        "Good bye world",
        "Hotdog",
        "Egg",
        "The word collapse has multiple meanings, including:",
        "A very loooooong text. A very loooooong text. A very loooooong text. A very loooooong text. A very loooooong text. A very loooooong text. ",
        "short text"
        
    ]
    
    private var selectedIndexPath: IndexPath?
    
    lazy var cellRegistration = UICollectionView.CellRegistration<DemoListCell, String>(cellNib: DemoListCell.getUINib()) { [weak self] (demoListCell, indexPath, string) in
        
        guard let self = self else { return }
        
        if indexPath == self.selectedIndexPath {
            print(">>>>> expand - \(string)")
            demoListCell.expand()
        } else {
            print(">>>>> collapse - \(string)")
            demoListCell.collapse()
        }
        demoListCell.label.text = string
    }
    
    private typealias DataSource = UICollectionViewDiffableDataSource<Section, String>
    private typealias Snapshot = NSDiffableDataSourceSnapshot<Section, String>

    private var dataSource: DataSource!
    
    @IBOutlet weak var collectionView: UICollectionView!
    
    override func viewDidLoad() {
        super.viewDidLoad()
        
        collectionView.collectionViewLayout = layoutConfig()
        
        collectionView.delegate = self
        
        dataSource = makeDataSource()
        
        var snapshot = Snapshot()
        snapshot.appendSections([.note])
        snapshot.appendItems(strings, toSection: .note)
        dataSource.apply(snapshot, animatingDifferences: true)
    }

    private func layoutConfig() -> UICollectionViewCompositionalLayout {
        let layout = UICollectionViewCompositionalLayout { section, layoutEnvironment in
            var config = UICollectionLayoutListConfiguration(appearance: .insetGrouped)

            config.headerMode = .none
            config.footerMode = .none
            config.showsSeparators = true
            config.headerTopPadding = 0
            
            config.backgroundColor = nil
            
            let layoutSection = NSCollectionLayoutSection.list(using: config, layoutEnvironment: layoutEnvironment)
            
            return layoutSection
        }
        
        return layout
    }
    
    private func makeDataSource() -> DataSource {
        let cellRegistration = self.cellRegistration
        
        let dataSource = DataSource(
            collectionView: collectionView,
            cellProvider: { [weak self] (collectionView, indexPath, item) -> UICollectionViewCell? in
                guard let self = self else { return nil }
                
                return collectionView.dequeueConfiguredReusableCell(
                    using: cellRegistration,
                    for: indexPath,
                    item: item
                )
            }
        )
        
        return dataSource
    }
}

extension ViewController: UICollectionViewDelegate {
    func collectionView(_ collectionView: UICollectionView, didSelectItemAt indexPath: IndexPath) {
        
        self.selectedIndexPath = indexPath
        
        var snapshot = dataSource.snapshot()
        
        snapshot.reconfigureItems(snapshot.itemIdentifiers)

        print(">>>>> === apply ===")
        dataSource.apply(snapshot, animatingDifferences: true) {
            collectionView.collectionViewLayout.invalidateLayout()
        }
        
        collectionView.deselectItem(at: indexPath, animated: true)
    }
}

DemoListCell.swift #

class DemoListCell: UICollectionViewListCell {

    @IBOutlet weak var label: UILabel!
    @IBOutlet weak var extraView: UIView!
    
    override func awakeFromNib() {
        super.awakeFromNib()
        // Initialization code
    }

    func expand() {
        extraView.isHidden = false
    }
    
    func collapse() {
        extraView.isHidden = true
    }
}

我准备了完整可运行的源码，请问有没有大佬知道该怎么修复这个问题？感谢！

备注：内容来源于stack exchange，提问作者Cheok Yan Cheng