SimpleCollectionViewDataSource.swift

import Foundation
import UIKit

class SimpleCollectionViewDataSource<Cell: AnyObject, Object>: NSObject, UICollectionViewDataSource {
    
    var objects: [[Object]]
    
    func objectAtIndexPath(indexPath: NSIndexPath) -> Object {
        return objects[indexPath.section][indexPath.row]
    }
    
    typealias ConfigureCellBlock = (Cell, Object) -> Void

    init(objects: [[Object]], cellType: CellType, configureCell: ConfigureCellBlock) {
        self.objects = objects
        self.configureCell = configureCell
        cellRegistration = CollectionRegistrationState(resiterBlock: cellType.registerBlock())
    }
    
    private let configureCell: ConfigureCellBlock
    private var cellRegistration: CollectionRegistrationState
    
    // MARK: - UICollectionViewDataSource
    
    func numberOfSectionsInCollectionView(collectionView: UICollectionView) -> Int {
        return objects.count
    }
    
    func collectionView(collectionView: UICollectionView, numberOfItemsInSection section: Int) -> Int {
        return objects[section].count
    }
    
    func collectionView(collectionView: UICollectionView, cellForItemAtIndexPath indexPath: NSIndexPath) -> UICollectionViewCell {
        let cell = collectionView.dequeueReusableCellWithReuseIdentifier(
            cellRegistration.reuseIdentifierForClass(Cell.self, inCollectionView: collectionView),
            forIndexPath: indexPath)
        configureCell(cell as! Cell, objectAtIndexPath(indexPath))
        return cell
    }
}


enum CellType {
    case Nib
    case Class
    
    func registerBlock() -> RegisterBlock {
        return { collectionView, cellClass -> String in
            let identifier = String(cellClass)
            switch self {
            case .Nib:
                collectionView.registerNib(UINib(nibName: String(cellClass), bundle: nil), forCellWithReuseIdentifier: identifier)
            case .Class:
                collectionView.registerClass(cellClass, forCellWithReuseIdentifier: identifier)
            }
            return identifier
        }
    }
}

typealias RegisterBlock = (collectionView: UICollectionView, aClass: AnyClass) -> String // returns registered identifier

/**
 Helper enum to be used to register cell just once and than use it's identifier to dequeue from collection view.
 */
enum CollectionRegistrationState {
    case NotRegistered(RegisterBlock)
    case Registered(identifier: String)
    
    init(resiterBlock: RegisterBlock) {
        self = .NotRegistered(resiterBlock)
    }
    
    mutating func reuseIdentifierForClass(cellClass: AnyClass, inCollectionView collectionView: UICollectionView) -> String {
        switch self {
        case .NotRegistered(let register):
            let identifier = register(collectionView: collectionView, aClass: cellClass)
            self = .Registered(identifier: identifier)
            return identifier
        case .Registered(let identifier):
            return identifier
        }
    }
}