Отображение иерархических данных в виде списка с возможностью поиска

Примеры исходного когда содержат только ключевые моменты описания работы алгоритма, финальная рабочая версия с примерами находится по ссылке выше.

class UserModel {
    let name: String
    let children: [UserModel]
    
    init(name: String, children: [UserModel]) {
        self.name = name
        self.children = children
    }
}

struct DisplayItem {
    var name: String
    var isChild: Bool = false
}

func flattenList(_ list: [UserModel]) -> [DisplayItem] {
    var output = [DisplayItem]()
    for item in list {
        output.append( DisplayItem(name: item.name) )
        for child in item.children {
            output.append( DisplayItem(name: child.name, isChild: true) )
        }
    }
    return output
}

var displayItems = [DisplayItem]()

override func viewDidLoad() {
    super.viewDidLoad()
    displayItems = flattenList(hierarchicalObjects)
}

func tableView(_: UITableView, numberOfRowsInSection _: Int) -> Int {
    return displayItems.count
}

func tableView(_ tableView: UITableView, cellForRowAt indexPath: IndexPath) -> UITableViewCell {
    let cell = tableView.dequeueReusableCell(
        withIdentifier: Cell.identifier,
        for: indexPath
    ) as! Cell
    let item = displayItems[indexPath.row]
    cell.titleLabel.text = item.name
    cell.titleLabelLeadingConstraint.constant = item.isChild ? 20 : 0
    return cell
}

class TreeItem<Item: Hashable>: Hashable {
    let id = UUID()
    var value: Item
    var subitems: [TreeItem]
    weak var parent: TreeItem?
    fileprivate(set) var level: Int = 0
    
    init(name: String, parent: TreeItem?, items: [TreeItem] = []) {
        self.name = name
        self.parent = parent
        self.subitems = items
        self.level = level(of: self)
    }
    
   func level(of item: TreeItem) -> Int {
        // Traverse up to next parent to find level. root element has `0` level.
        var counter: Int = 0
        var currentItem: TreeItem = item
        while let parent = currentItem.parent {
            counter += 1
            currentItem = parent
        }
        return counter
    }

    public static func == (lhs: TreeItem<Item>, rhs: TreeItem<Item>) -> Bool {
        return lhs.id == rhs.id && lhs.value == rhs.value
    }
    
    public func hash(into hasher: inout Hasher) {
        hasher.combine(id)
        hasher.combine(value)
    }
}

class ListTreeDataSource<Item> where Item: Hashable {
    private(set) var backingStore: [TreeItem<Item>] = []
    private(set) var items: [TreeItemType] = [] // TODO: Display items for UI framework
    private var lookupTable: [Item: TreeItem<Item>] = [:]

    public func append(_ items: [Item], to parent: Item? = nil) {
        func append(items: [Item], into insertionArray: inout [TreeItem], parentBackingItem: TreeItem?) {
            let treeItems = items.map { item in TreeItem(value: item, parent: parentBackingItem) }
            cacheTreeItems(treeItems)
            insertionArray.append(contentsOf: treeItems)
        }

        if let parent = parent, let parentBackingItem = self.lookupTable[parent] {
            append(items: items, into: &parentBackingItem.subitems, parentBackingItem: parentBackingItem)
        } else {
            append(items: items, into: &backingStore, parentBackingItem: nil)
        }
    }
    
    private func cacheTreeItems(_ treeItems: [TreeItem<Item>]) {
        treeItems.forEach { lookupTable[$0.value] = $0 }
    }
}

func createDataSource() -> ListTreeDataSource<UserModel> {
    var dataSource = ListTreeDataSource<UserModel>()
    
    func addItems(_ items: [UserModel], to parent: UserModel?) {
        dataSource.append(items, to: parent)
        for item in items {
            addItems(item.children, to: item)
        }
    }
    
    addItems(items, to: nil)
    return dataSource
}

func depthFirstFlattened(_ list: [Item]) -> [Item] {
    var output = [Item]()
    for item in list {
        output.append( item )
        output.append(contentsOf: depthFirstFlattened(item.children) )
    }
    return output
}

extension TreeDataSource {
    func reload() {
self.items = depthFirstFlattened(items: self.backingStore, itemChildren: \.subitems)
    }
}

    func reload() {
        self.items = depthFirstFlattened(items: self.backingStore, itemChildren: { $0.isExpanded ? $0.subitems : [] })
    }

override func tableView(_ tableView: UITableView, cellForRowAt indexPath: IndexPath) -> UITableViewCell {
    let cell = tableView.dequeueReusableCell(withIdentifier: "Cell", for: indexPath) as! Cell
    
    let item = treeDataSource.items[indexPath.row]
    
    let left = 11 + 10 * item.level
    cell.lbl?.text = item.value.title
    cell.lblLeadingConstraint.constant = CGFloat(left)
    cell.disclosureImageView.isHidden = item.subitems.isEmpty
    
    let transform = CGAffineTransform.init(rotationAngle: item.isExpanded ? CGFloat.pi/2.0 : 0)
    cell.disclosureImageView.transform = transform
    
    return cell
}

override func tableView(_ tableView: UITableView, didSelectRowAt indexPath: IndexPath) {
    let item = treeDataSource.items[indexPath.row]
    
    // update the state
    item.isExpanded = !item.isExpanded
    // rebuild shown items
    treeDataSource.reload()
    
    if let cell = tableView.cellForRow(at: indexPath) as? Cell {
        UIView.animate(withDuration: animationDuration) {
            let transform = CGAffineTransform.init(rotationAngle: item.isExpanded ? CGFloat.pi/2.0 : 0)
            cell.disclosureImageView.transform = transform
        }
    }
    
    tableView.reloadData()
}

public func addItems<Item>(_ items: [Item], itemChildren: (Item) -> [Item], to dataSource: TreeDataSource<Item>) {
    dataSource.append(items, to: nil)
    
    var queue: Queue<Item> = .init()
    queue.enqueue(items)
    
    while let current = queue.dequeue() {
        dataSource.append(itemChildren(current), to: current)
        queue.enqueue(itemChildren(current))
    }
    
    dataSource.reload()
}

func depthFirstFlattened<Item>(items: [Item], itemChildren: (Item) -> [Item]) -> [Item] {
    var outFlatStore: Array< Item > = []
    
    var stack: Array< Item > = items.reversed()
    while let current = stack.last {
        stack.removeLast()
        
        outFlatStore.append(current)
        stack.append(contentsOf: itemChildren(current).reversed())
    }
    return outFlatStore
}

MacBook Pro 2019, 2,6 GHz 6-Core Intel Core i7, 16 GB 2667 MHz DDR4

class FilterableTreeDataSource<Item>: ListTreeDataSource<Item> where Item: Hashable {
    public override init() { super.init() }
    private(set) var allFlattenedItemStore: [TreeItem<Item>] = []
    private(set) var filterPredicate: ((ItemIdentifierType) -> Bool)?
    
    private var dispatchQueue = DispatchQueue(label: "FilterableTreeDataSource")
    private var dispatchWorkItem: DispatchWorkItem?

   public private(set) var filteredTargetItems: [TreeItem<Item>] = []
    public private(set) var targetItemsTraversedParentSet: Set<TreeItem<Item>> = .init()

    public func filterItemsKeepingParents(by predicate: @escaping (ItemIdentifierType) -> Bool, completion: @escaping (() -> Void)) {
        self.filterPredicate = predicate
        self.dispatchWorkItem?.cancel()
        
        var workItem: DispatchWorkItem?
        workItem = DispatchWorkItem { [weak self] in
            guard let self = self else { return }
            if workItem?.isCancelled ?? true { workItem = nil; return }

            // collapse all elements
            self.allFlattenedItemStore.forEach { $0.isExpanded = false }
            
            self.filteredTargetItems = self.allFlattenedItemStore.filter { predicate($0.value) }
            self.targetItemsTraversedParentSet = Set( self.filteredTargetItems.flatMap { $0.allParents } )
            
            // expand all traversed parents for the target items
            self.targetItemsTraversedParentSet.forEach { $0.isExpanded = true }
                                    
            DispatchQueue.main.async {
                if workItem?.isCancelled ?? true { workItem = nil; return }

                self.rebuildShownItemsStore()
                workItem = nil
                completion()
            }
        }
        self.dispatchWorkItem = workItem
        self.dispatchQueue.async(execute: workItem!)
    }
    
    override func rebuildShownItemsStore() {
        // Depth first search + filtering. Tries to open expanded nested items matching `isIncludedInExpand` predicate.
        let outFlatStore = depthFirstFlattened(items: self.backingStore, itemChildren: { $0.isExpanded ? $0.subitems : [] }, filter: isIncludedInExpand)
        self.setShownFlatItems(outFlatStore)
    }

    private func isIncludedInExpand(_ item: TreeItemType) -> Bool {
        guard let filterPredicate = filterPredicate else { return true } // allow for all elements when no filtering.
        
        //  Root
        //      RootChild_1
        //          ItemMatching[SearchText]
        //          Item2Matching[SearchText]
        //      RootChild_2
        //          ItemMatching4[SearchText]
        //          ItemMatching5[SearchText]
        //              ItemNoMatchingButMightBeSearchTarget1
        //              ItemNoMatchingButMightBeSearchTarget2
        //  Root2

        switch (item.parent, item) {
        // inсluded root elements when matches predicate or contained in traversed parent set (e.g. Root, Root2).
        case let (.none, rootItem):
            return filterPredicate(rootItem.value) || self.targetItemsTraversedParentSet.contains(rootItem)

        // inсluded all items contained in traversed parent set, to display full path (e.g. RootChild_1, RootChild_2).
        case let (.some(_), item) where targetItemsTraversedParentSet.contains(item):
            return true
            
        case let (.some(parent), item):
            // inсluded items with parent from traversed parent set & matching to `filterPredicate`
            // (e.g. item: ItemMatching4[SearchText], parent: RootChild_2).
            if self.targetItemsTraversedParentSet.contains(parent) {
                return filterPredicate(item.value)
            } else {
                // user might inсlude any additional descendant items (e.g. ItemNoMatchingButMightBeSearchTarget1, ItemNoMatchingButMightBeSearchTarget2)
                return true
            }
        }
    }
    
    private func resetFilteringState() {
        filterPredicate = nil
        dispatchWorkItem?.cancel()
        dispatchWorkItem = nil
    }
    
    private func rebuildAllFlattenedItemStore() {
        self.allFlattenedItemStore = depthFirstFlattened(items: self.backingStore, itemChildren: { $0.subitems })
    }
}