CachedItemHeightDatasource.swift

import Foundation
import UIKit

class CachedItemHeightDatasource: NSObject, UITableViewDelegate, UITableViewDataSource {
    
    var heightAtIndexPath: [IndexPath: CGFloat] = [:]
    let tableView: UITableView
    
    func tableView(_ tableView: UITableView, estimatedHeightForRowAt indexPath: IndexPath) -> CGFloat {
        return heightAtIndexPath[indexPath] ?? UITableViewAutomaticDimension
    }
    
    func tableView(_ tableView: UITableView, willDisplay cell: UITableViewCell, forRowAt indexPath: IndexPath) {
        heightAtIndexPath[indexPath] = cell.frame.size.height
    }
    
    func update(rowAt indexPath: IndexPath) {
        tableView.beginUpdates()
        heightAtIndexPath[indexPath] = nil
        reloadRows(at: [indexPath], with: .automatic) // fetches new cell via tableView(_:cellForRowAt:)
        tableView.endUpdates()
    }
    
}