focus_lost.swift

//
//  ViewController.swift
//  DiffableDataSourcePlayground
//
//  Created by Majd Taby on 5/7/23.
//

import UIKit
import IdentifiedCollections
import SwiftUI

struct Item: Identifiable {
    let id: UUID
    let title: String
}

/*
 This sample app demonstrates a bug on iOS, where having a text field in the section header of a diffable data source
 causes the focus to be lost whenever the snapshot is updated. This doesn't happen in regular cells though. To test,
 try to type a T or two in one of the cells to demonstrate focus is maintained after filtering. Then try to focus
 on the section header and do the same, and observe that the focus is lost after every key press and snapshot update.
 */
class ViewController: UIViewController {

    let model: IdentifiedArrayOf<Item> = IdentifiedArrayOf<Item>(uniqueElements: [
        Item(id: UUID(), title: "T"), Item(id: UUID(), title: "TT"), Item(id: UUID(), title: "TTTT"), Item(id: UUID(), title: "TTTTT")
    ])

    enum Section { case main }

    var collectionView: UICollectionView!
    var collectionViewLayout: UICollectionViewCompositionalLayout!
    var cellRegistration: UICollectionView.CellRegistration<UICollectionViewListCell, Item.ID>!
    var headerRegistration: UICollectionView.SupplementaryRegistration<UICollectionViewListCell>!

    var dataSource: UICollectionViewDiffableDataSource<Section, Item.ID>!

    var count = 0

    override func viewDidLoad() {
        super.viewDidLoad()

        // Set up compositional layout
        let groupSize = NSCollectionLayoutSize(
            widthDimension: .fractionalWidth(1.0),
            heightDimension: .estimated(50)
        )

        let item = NSCollectionLayoutItem(
            layoutSize: NSCollectionLayoutSize(
                widthDimension: .fractionalWidth(1.0),
                heightDimension: .fractionalHeight(1.0)
            )
        )

        let group = NSCollectionLayoutGroup.horizontal(layoutSize: groupSize, subitems: [item])

        let section = NSCollectionLayoutSection(group: group)

        let headerSize = NSCollectionLayoutSize(widthDimension: .fractionalWidth(1.0), heightDimension: .estimated(40))
        let header = NSCollectionLayoutBoundarySupplementaryItem(layoutSize: headerSize, elementKind: UICollectionView.elementKindSectionHeader, alignment: .top)
        section.boundarySupplementaryItems = [header]

        collectionViewLayout = UICollectionViewCompositionalLayout(section: section)

        // Initialize collection view
        collectionView = UICollectionView(frame: .zero, collectionViewLayout: collectionViewLayout)
        collectionView.translatesAutoresizingMaskIntoConstraints = false
        collectionView.backgroundColor = .black
        view.addSubview(collectionView)

        cellRegistration = UICollectionView.CellRegistration(handler: { cell, indexPath, itemIdentifier in
            cell.contentConfiguration = UIHostingConfiguration(content: {
                // Typing in the text field of this cell works
                TestCell(item: self.model[id: itemIdentifier]!, delegate: self)
            })
        })

        headerRegistration = UICollectionView.SupplementaryRegistration<UICollectionViewListCell>(elementKind: UICollectionView.elementKindSectionHeader) { headerView, elementKind, indexPath in
            headerView.contentConfiguration = UIHostingConfiguration(content: {
                // Typing in the text field of this cell doesn't work
                VStack {
                    Text("All the cells have a text field that filters the cells. This top one is the header, it loses focus when you type, but the cells don't")
                    TestCell(item: Item(id: UUID(), title: "Type here to filter."), delegate: self)
                }
            })
            .margins(.all, 20)
        }

        dataSource.supplementaryViewProvider = { collectionView, elementKind, indexPath in
            if elementKind == UICollectionView.elementKindSectionHeader {
                let view = collectionView.dequeueConfiguredReusableSupplementary(using: self.headerRegistration, for: indexPath)
                return view
            }

            return nil
        }

        dataSource = UICollectionViewDiffableDataSource<Section, Item.ID>(
            collectionView: collectionView,
            cellProvider: { collectionView, indexPath, itemID in
                return collectionView.dequeueConfiguredReusableCell(
                    using: self.cellRegistration,
                    for: indexPath,
                    item: itemID
                )
            }
        )

        collectionView.dataSource = dataSource

        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),
            ]
        )
    }

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

        var snapshot = NSDiffableDataSourceSnapshot<Section, Item.ID>()
        snapshot.appendSections([.main])
        snapshot.appendItems(model.elements.map { $0.id })
        dataSource.apply(snapshot, animatingDifferences: false)
    }
}

extension ViewController: Delegate {
    func doWork(query: String) {
        var snapshot = dataSource.snapshot(for: .main)
        snapshot.deleteAll()
        snapshot.append(model.elements.map { $0.id }.filter { model[id: $0]!.title.contains(query)})
        dataSource.apply(snapshot, to: .main, animatingDifferences: true)
    }
}

protocol Delegate {
    func doWork(query: String)
}

struct TestCell: View {
    let item: Item
    @State var text: String = ""
    let delegate: Delegate
    var body: some View {
        TextField(item.title, text: $text)
            .onChange(of: text) { newValue in
                delegate.doWork(query: newValue)
            }
    }
}

Demo of bug:

CleanShot.2023-05-17.at.12.06.31.mp4

Backtrace demonstrating flow to lose focus