MVVMCollectionViewDataSource.swift

class MVVMCollectionViewDataSource: NSObject, UICollectionViewDataSource {
    
    var model: [SectionViewModel] = []
    
    func numberOfSections(in collectionView: UICollectionView) -> Int {
        return self.model.count
    }
    
    func collectionView(_ collectionView: UICollectionView, numberOfItemsInSection section: Int) -> Int {
        return self.model[section].items.count
    }
    
    func collectionView(_ collectionView: UICollectionView, cellForItemAt indexPath: IndexPath) -> UICollectionViewCell {
        
        let model = self.model[indexPath.section].items[indexPath.row]
        let cell = collectionView.dequeueReusableCell(withReuseIdentifier: model.reuseIdentifier, for: indexPath)
        model.setup(cell, in: collectionView, at: indexPath)
        return cell
    }
    
    func collectionView(_ collectionView: UICollectionView, viewForSupplementaryElementOfKind kind: String, at indexPath: IndexPath) -> UICollectionReusableView {
        
        let section = self.model[indexPath.section]
        guard let model = section.model(for: kind) else {
            return UICollectionReusableView()
        }

        let reusableView = collectionView.dequeueReusableSupplementaryView(ofKind: kind, withReuseIdentifier: model.reuseIdentifier, for: indexPath)
        model.setup(reusableView, in: collectionView, at: indexPath)
        return reusableView
    }
}

extension SectionViewModel {
    
    func model(for elementOfKind: String) -> ItemViewModel? {
        
        switch elementOfKind {
        case UICollectionElementKindSectionHeader:
            return self.headerItem
            
        case UICollectionElementKindSectionFooter:
            return self.footer
            
        default:
            return nil
        }
    }
}