Eschew delegation. Embrace composition.

ComposableTableView.swift

//  ComposableTableView.swift
//  Created by Jake Hawken

import UIKit

class ComposableTableView<Item, SectionTypeID: Hashable>: UIView {
    
    typealias Section = SectionViewModel<Item, SectionTypeID>
    
    private let tableView: UITableView = UITableView()
    private let delegate = TableViewDelegate()
    private lazy var dataSource = TableViewDataSource()
        .numberOfSections { [weak self] _ -> Int in
            let count = self?.sections.count
            return count ?? 0
        }
        .numberOfRowsInSection { [weak self] _, sectionIndex -> Int in
            self?.sections[sectionIndex].items.count ?? 0
        }
        .titleForHeaderInSection { [weak self] _, sectionIndex -> String? in
            self?.sections[sectionIndex].title
        }
    
    private var sections = [Section]()
    
    /// - parameter cellTypes: An array of the cell types that will be used in the tableView.
    required init(cellTypes: [UITableViewCell.Type]) {
        super.init(frame: .zero)
        
        cellTypes.forEach { self.tableView.registerCellType($0) }

        tableView.delegate = delegate
        tableView.dataSource = dataSource
        
        tableView.translatesAutoresizingMaskIntoConstraints = false
        addSubview(subview)
        NSLayoutConstraint.activate([
            tableView.topAnchor.constraint(equalTo: topAnchor),
            tableView.leadingAnchor.constraint(equalTo: leadingAnchor),
            tableView.trailingAnchor.constraint(equalTo: trailingAnchor),
            tableView.bottomAnchor.constraint(equalTo: bottomAnchor),
        ])
    }
    
    override init(frame: CGRect) {
        Self.failInit()
    }
    
    required init?(coder: NSCoder) {
        print("Do not use ComposableTableView from interface builder.")
        Self.failInit()
    }
    
    private static func failInit(initName: String = #function) -> Never {
        return fatalError(
            "\(initName) has not been implemented. Please use .init(cellTypes:headerViewTypes:) instead."
        )
    }
    
    var separatorColor: UIColor? {
        get {
            tableView.separatorColor
        }
        set {
            tableView.separatorColor = newValue
        }
    }
    
    var allowsSelection: Bool {
        get {
            tableView.allowsSelection
        }
        set {
            tableView.allowsSelection = newValue
        }
    }
    
}

// MARK: - main interface

extension ComposableTableView {
    
    var sectionCount: Int {
        sections.count
    }
    
    var itemCount: Int {
        sections.reduce(0) { $0 + $1.items.count }
    }
    
    func updateSections(_ newSections: [Section]) {
        sections = newSections
        tableView.reloadData()
    }
    
    func updateSection(atIndex sectionIndex: Int, newSection: Section) {
        if sections.indices.contains(sectionIndex) {
            sections[sectionIndex] = newSection
        }
        else if sectionIndex == sections.count {
            sections.append(newSection)
        }
        else {
            assertionFailure("Invalid index passed to \(#function).")
        }
        tableView.reloadData()
        tableView.refreshControl?.endRefreshing()
    }
    
}

// MARK: - Composition / Declaration interface

extension ComposableTableView {
    
    @discardableResult func cellForIndexPathWithItem(
        _ implementation: @escaping (Item, SectionTypeID, IndexPath, UITableView) -> UITableViewCell
    ) -> Self {
        dataSource.cellForRowAtIndexPath { [weak self] tv, indexPath in
            guard let self else {
                let selfString = String(describing: ComposableTableView<Item, SectionTypeID>.self)
                fatalError("Instance of \(selfString) accessed by underlying tableView after being released from memory.")
            }
            guard let section = sections[safe: indexPath.section],
                  let item = section.items[safe: indexPath.row] else {
                fatalError("Bad data. No \(Item.self) found at IndexPath: \(indexPath)")
            }
            return implementation(item, section.sectionType, indexPath, tv)
        }
        return self
    }
    
    @discardableResult func canEditRowAtIndexPathWithItem(_ implementation: @escaping (IndexPath, Item, UITableView) -> Bool) -> Self {
        dataSource.canEditRowAtIndexPath { [weak self] tableView, indexPath -> Bool in
            guard let item = self?.sections[indexPath.section].items[indexPath.row] else {
                return false
            }
            return implementation(indexPath, item, tableView)
        }
        return self
    }
    
    @discardableResult func commitEditingStyleAtIndexPathWithItem(
        _ implementation: @escaping (UITableViewCell.EditingStyle, IndexPath, Item, UITableView) -> Void
    ) -> Self {
        dataSource.commitEditingStyleAtIndexPath { [weak self] tableView, editingStyle, indexPath in
            guard let item = self?.sections[indexPath.section].items[indexPath.row] else {
                return
            }
            implementation(editingStyle, indexPath, item, tableView)
        }
        return self
    }
    
    @discardableResult func didSelectRowWithItem(_ implementation: @escaping (Item, IndexPath, UITableView) -> Void) -> Self {
        delegate.didSelectRowAtIndexPath { [weak self] tableView, indexPath in
            guard let item = self?.sections[indexPath.section].items[indexPath.row] else {
                return
            }
            implementation(item, indexPath, tableView)
        }
        return self
    }
    
    @discardableResult func heightForRowWithItem(_ implementation: @escaping (Item, IndexPath, UITableView) -> CGFloat) -> Self {
        delegate.heightForRowAtIndexPath { [weak self] tableView, indexPath -> CGFloat in
            guard let item = self?.sections[indexPath.section].items[indexPath.row] else {
                return 0
            }
            return implementation(item, indexPath, tableView)
        }
        return self
    }

    typealias SectionIndex = Int
    @discardableResult func heightForHeaderInSection(_ implementation: @escaping (Section, SectionIndex, UITableView) -> CGFloat) -> Self {
        delegate.heightForHeaderInSection { [weak self] tableView, sectionIndex -> CGFloat in
            guard let section = self?.sections[sectionIndex] else {
                return 0
            }
            return implementation(section, sectionIndex, tableView)
        }
        return self
    }

    @discardableResult func heightForFooterInSection(_ implementation: @escaping ((Section, SectionIndex, UITableView) -> CGFloat)) -> Self {
        delegate.heightForFooterInSection { [weak self] tableView, sectionIndex -> CGFloat in
            guard let section = self?.sections[sectionIndex] else {
                return 0
            }
            return implementation(section, sectionIndex, tableView)
        }
        return self
    }
    
    /// Adds a refresh control and sets the action to be triggered when the user pulls to refresh.
    /// - parameter refreshAction: The closure that will be executed on pull to refresh. A nil value
    /// will remove any existing refresh control. A non-nil value will add a refresh control if one
    /// does not exist, or will replace the action of an existing refresh control.
    /// - returns: `@discardableResult self` to allow for operator chaining, to aid in composition.
    @discardableResult func onPullToRefresh(_ refreshAction: (() -> Void)?) -> Self {
        guard let refreshAction else {
            tableView.refreshControl?.removeFromSuperview()
            tableView.refreshControl = nil
            return self
        }
        if let control = tableView.refreshControl as? ClosureBasedRefreshControl {
            control.refreshAction = refreshAction
        } else {
            tableView.addRefreshControlWithBlock(refreshAction)
        }
        return self
    }
    
}

// MARK: - tableview helpers

extension UITableView {
    
    var registeredClasses: [String: Any] {
        let dict = value(forKey: "_cellClassDict") as? [String: Any]
        return dict ?? [:]
    }
    
    func reuseIdForCellType(_ cellType: UITableViewCell.Type) -> String {
        String(describing: cellType)
    }
    
    func registerCellType(_ cellType: UITableViewCell.Type) {
        if hasRegisteredClass(cellType) {
            return
        }
        let reuseID = reuseIdForCellType(cellType)
        register(cellType, forCellReuseIdentifier: reuseID)
    }
    
    func hasRegisteredClass(_ cellClass: UITableViewCell.Type) -> Bool {
        registeredClasses.values.contains { value -> Bool in
            guard let cellType = value as? UITableView.Type else {
                return false
            }
            return cellType == cellClass
        }
    }
    
    func dequeueCell<T: UITableViewCell>(_ cellType: T.Type = T.self, indexPath: IndexPath) -> T {
        let reuseID = reuseIdForCellType(cellType)
        guard let cell = dequeueReusableCell(withIdentifier: reuseID, for: indexPath) as? T else {
            fatalError("A cell of type \(cellType) could not be dequeued for this TableView.")
        }
        return cell
    }
    
    func dequeueCell<T: UITableViewCell>(_ cellType: T.Type = T.self) -> T {
        let reuseID = reuseIdForCellType(cellType)
        guard let cell = dequeueReusableCell(withIdentifier: reuseID) as? T else {
            fatalError("A cell of type \(cellType) could not be dequeued for this TableView.")
        }
        return cell
    }
    
}


// MARK: - TableView data source -

class TableViewDataSource: NSObject, UITableViewDataSource {
    
    typealias NumberOfSectionsClosure = ((UITableView) -> Int)
    private var numberOfSectionsCallback: NumberOfSectionsClosure?
    
    typealias NumberOfRowsInSectionClosure = ((UITableView, Int) -> Int)
    private var numberOfRowsInSectionCallback: NumberOfRowsInSectionClosure?
    
    typealias CellForRowAtIndexPathClosure = ((UITableView, IndexPath) -> UITableViewCell)
    private var cellForRowAtIndexPathCallback: CellForRowAtIndexPathClosure?
    
    typealias TitleForHeaderInSectionClosure = ((UITableView, Int) -> String?)
    private var titleForHeaderInSectionCallback: TitleForHeaderInSectionClosure?
    
    typealias CanEditRowAtIndexPathClosure = ((UITableView, IndexPath) -> Bool)
    private var canEditRowAtIndexPathCallback: CanEditRowAtIndexPathClosure?
    
    typealias CommitEditingStyleAtIndexPathClosure = ((UITableView, UITableViewCell.EditingStyle, IndexPath) -> Void)
    private var commitEditingStyleAtIndexPathCallback: CommitEditingStyleAtIndexPathClosure?
    
    // MARK: - number of sections
    
    // Set
    @discardableResult func numberOfSections(_ closure: @escaping NumberOfSectionsClosure) -> Self {
        numberOfSectionsCallback = closure
        return self
    }
    
    // Get
    func numberOfSections(in tableView: UITableView) -> Int {
        numberOfSectionsCallback?(tableView) ?? 0
    }
    
    // MARK: - number of rows in section
    
    // Set
    @discardableResult func numberOfRowsInSection(_ closure: @escaping NumberOfRowsInSectionClosure) -> Self {
        numberOfRowsInSectionCallback = closure
        return self
    }
    
    // Get
    func tableView(_ tableView: UITableView, numberOfRowsInSection section: Int) -> Int {
        numberOfRowsInSectionCallback?(tableView, section) ?? 0
    }
    
    // MARK: - cell for row at index path
    
    // Set
    @discardableResult func cellForRowAtIndexPath(_ closure: @escaping CellForRowAtIndexPathClosure) -> Self {
        cellForRowAtIndexPathCallback = closure
        return self
    }
    
    // Get
    func tableView(_ tableView: UITableView, cellForRowAt indexPath: IndexPath) -> UITableViewCell {
        guard let cellForRow = cellForRowAtIndexPathCallback else {
            fatalError(
                "Implementation for \(#function) not provided to instance of \(CollectionViewDataSouce.self)."
            )
        }
        return cellForRow(tableView, indexPath)
    }
    
    // MARK: - title for header in section
    
    // Set
    @discardableResult func titleForHeaderInSection(_ closure: @escaping TitleForHeaderInSectionClosure) -> Self {
        titleForHeaderInSectionCallback = closure
        return self
    }
    
    // Get
    func tableView(_ tableView: UITableView, titleForHeaderInSection section: Int) -> String? {
        titleForHeaderInSectionCallback?(tableView, section)
    }
    
    // MARK: - can edit row at index path
    
    // Set
    @discardableResult func canEditRowAtIndexPath(_ closure: @escaping CanEditRowAtIndexPathClosure) -> Self {
        canEditRowAtIndexPathCallback = closure
        return self
    }
    
    // Get
    func tableView(_ tableView: UITableView, canEditRowAt indexPath: IndexPath) -> Bool {
        canEditRowAtIndexPathCallback?(tableView, indexPath) ?? false
    }
    
    // MARK: - commit editing style at index path
    
    // Set
    @discardableResult func commitEditingStyleAtIndexPath(_ closure: @escaping CommitEditingStyleAtIndexPathClosure) -> Self {
        commitEditingStyleAtIndexPathCallback = closure
        return self
    }
    
    // Get
    func tableView(_ tableView: UITableView, commit editingStyle: UITableViewCell.EditingStyle, forRowAt indexPath: IndexPath) {
        commitEditingStyleAtIndexPathCallback?(tableView, editingStyle, indexPath)
    }
    
}

// MARK: - TableView delegate -

class TableViewDelegate: NSObject, UITableViewDelegate {
    
    typealias DidSelectRowAtIndexPathClosure = ((UITableView, IndexPath) -> Void)
    private var didSelectRowAtIndexPathCallback: DidSelectRowAtIndexPathClosure?
    
    typealias HeightForRowAtIndexPathClosure = ((UITableView, IndexPath) -> CGFloat)
    private var heightForRowAtIndexPathCallback: HeightForRowAtIndexPathClosure?
    
    typealias HeightForHeaderInSectionClosure = ((UITableView, Int) -> CGFloat)
    private var heightForHeaderInSectionCallback: HeightForHeaderInSectionClosure?
    
    typealias HeightForFooterInSectionClosure = ((UITableView, Int) -> CGFloat)
    private var heightForFooterInSectionCallback: HeightForFooterInSectionClosure?
    
    typealias DidScrollClosure = (UIScrollView) -> Void
    private var didScrollCallback: DidScrollClosure?
    
    // MARK: - did select row at index path
    
    // Set
    @discardableResult func didSelectRowAtIndexPath(_ closure: @escaping DidSelectRowAtIndexPathClosure) -> Self {
        didSelectRowAtIndexPathCallback = closure
        return self
    }
    
    // Get
    func tableView(_ tableView: UITableView, didSelectRowAt indexPath: IndexPath) {
        didSelectRowAtIndexPathCallback?(tableView, indexPath)
    }
    
    // MARK: - height for row at index path
    
    // Set
    @discardableResult func heightForRowAtIndexPath(_ closure: @escaping HeightForRowAtIndexPathClosure) -> Self {
        heightForRowAtIndexPathCallback = closure
        return self
    }
    
    // Get
    func tableView(_ tableView: UITableView, heightForRowAt indexPath: IndexPath) -> CGFloat {
        heightForRowAtIndexPathCallback?(tableView, indexPath) ?? UITableView.automaticDimension
    }
    
    // MARK: - height for header in section
    
    // Set
    @discardableResult func heightForHeaderInSection(_ closure: @escaping HeightForHeaderInSectionClosure) -> Self {
        heightForHeaderInSectionCallback = closure
        return self
    }
    
    // Get
    func tableView(_ tableView: UITableView, heightForHeaderInSection section: Int) -> CGFloat {
        heightForHeaderInSectionCallback?(tableView, section) ?? UITableView.automaticDimension
    }
    
    // MARK: - height for footer in section
    
    // Set
    @discardableResult func heightForFooterInSection(_ closure: @escaping HeightForFooterInSectionClosure) -> Self {
        heightForFooterInSectionCallback = closure
        return self
    }
    
    // Get
    func tableView(_ tableView: UITableView, heightForFooterInSection section: Int) -> CGFloat {
        heightForFooterInSectionCallback?(tableView, section) ?? UITableView.automaticDimension
    }
    
    // MARK: did scroll
    
    // Set
    @discardableResult func didScroll(_ closure: @escaping DidScrollClosure) -> Self {
        didScrollCallback = closure
        return self
    }
    
    // Get
    func scrollViewDidScroll(_ scrollView: UIScrollView) {
        didScrollCallback?(scrollView)
    }
    
}