UICollectionViewCellSizeManager.swift

import UIKit

/// Struct containing the arguments for the UIView.systemLayoutSizeFitting() method
public struct LayoutSizeFitting {
    let targetSize: CGSize
    let horizontalFittingPriority: UILayoutPriority
    let verticalFittingPriority: UILayoutPriority

    public static let `default` = LayoutSizeFitting(targetSize: UIView.layoutFittingCompressedSize,
                                                    horizontalFittingPriority: .fittingSizeLevel,
                                                    verticalFittingPriority: .fittingSizeLevel)
}

/// Cell size manager for collection views
///
/// NB: The cell must be defined in a xib file or
/// as a plain class that sets up its own constraints.
/// This will not work with cells defined in storyboards.
///
/// The manager calculates cell sizes independently of the collection view
/// either by instantiating the cells from their respective xib files or by
/// instantiating them directly using their constructor.
/// Once instantiated, the cells are then cached for reuse.
///
/// When the size method is called, it configures the cell
/// using the supplied configure() callback, then it
/// creates a hash key for that specific configuration.
/// The next time a size is requested for a cell that has
/// already been configured in the exact same way,
/// the size will be fetched from a cache instead of
/// the cell size being recalculated.
open class UICollectionViewCellSizeManager {
    private var cellCache: [String : UICollectionViewCell]
    private var sizeCache: [SizeCacheKey : CGSize]

    public init() {
        self.cellCache = [:]
        self.sizeCache = [:]
    }

    private func fetchCell<Cell: UICollectionViewCell>(_ type: Cell.Type, reuseIdentifier: String) -> Cell {
        let cell: Cell
        if let cachedCell = cellCache[reuseIdentifier] as? Cell {
            cell = cachedCell
        } else {
            cell = Cell()
            cellCache[reuseIdentifier] = cell
        }
        cell.prepareForReuse()
        return cell
    }

    private func fetchNibCell<Cell: UICollectionViewCell & NibLoadableView>(_ type: Cell.Type) -> Cell {
        let cell: Cell
        if let cachedCell = cellCache[Cell.defaultReuseIdentifier] as? Cell {
            cell = cachedCell
        } else {
            cell = Cell.fromNib()
            cellCache[Cell.defaultReuseIdentifier] = cell
        }
        cell.prepareForReuse()
        return cell
    }

    private func size(for cell: UICollectionViewCell, withKey key: SizeCacheKey,
                      layoutSizeFitting: LayoutSizeFitting = .default) -> CGSize {
        if let size = sizeCache[key] {
            return size
        }
        let size = cell.contentView.systemLayoutSizeFitting(
            layoutSizeFitting.targetSize,
            withHorizontalFittingPriority: layoutSizeFitting.horizontalFittingPriority,
            verticalFittingPriority: layoutSizeFitting.verticalFittingPriority)
        sizeCache[key] = size
        return size
    }

    public func size<Cell: UICollectionViewCell>(of cellType: Cell.Type,
                                                 layoutSizeFitting: LayoutSizeFitting = .default,
                                                 reuseIdentifier: String,
                                                 configure: (Cell) -> AnyHashable) -> CGSize {
        let cell = fetchCell(cellType, reuseIdentifier: reuseIdentifier)
        let configurable = configure(cell)
        let key = SizeCacheKey(reuseIdentifier: reuseIdentifier, configurable: configurable)
        return size(for: cell, withKey: key, layoutSizeFitting: layoutSizeFitting)
    }

    public func size<Cell: UICollectionViewCell & NibLoadableView>(of cellType: Cell.Type,
                                                                   layoutSizeFitting: LayoutSizeFitting = .default,
                                                                   configure: (Cell) -> AnyHashable) -> CGSize {
        let cell = fetchNibCell(cellType)
        let configurable = configure(cell)
        let key = SizeCacheKey(reuseIdentifier: Cell.defaultReuseIdentifier, configurable: configurable)
        return size(for: cell, withKey: key, layoutSizeFitting: layoutSizeFitting)
    }

    /// The sizes are cached based on the hash of two properties:
    ///  1. The reuse identifier of the cell
    ///     (assumed to be unique to one type of cell.)
    ///  2. The data the cell was configured with
    ///     e.g. a view model - or more generally,
    ///     anything that conforms to Hashable.
    private struct SizeCacheKey: Hashable {
        let reuseIdentifier: String
        let configurable: AnyHashable
    }
}

The nib view code depends on the following protocol:

public protocol NibLoadableView: class {
    static var nibName: String { get }

    static func fromNib() -> Self
}

extension NibLoadableView where Self: UIView {

    public static var nibName: String {
        return String(describing: self)
    }

    public static func fromNib() -> Self {
        guard let nib = Bundle(for: self).loadNibNamed(nibName, owner: nil, options: nil) else {
            fatalError("Failed loading the nib named \(nibName) for 'NibLoadableView' view of type '\(self)'.")
        }
        guard let view = (nib.first { $0 is Self }) as? Self else {
            fatalError("Did not find 'NibLoadableView' view of type '\(self)' inside '\(nibName).xib'.")
        }
        return view
    }
}