An example of a UICollectionView-like view made from UIStackViews for easy layout of small grids

StackCollectionView.swift

//
//  StackCollectionView.swift
//
//  Created by Albert Bori on 6/23/17.
//

import UIKit

@IBDesignable
class StackCollectionView: UIView {
    private var _containerStackView = UIStackView()
    
    var rowAlignment: UIStackViewAlignment = UIStackViewAlignment.leading
    @IBInspectable var sectionSpacing: CGFloat = 0 {
        didSet {
            _containerStackView.spacing = sectionSpacing
        }
    }
    @IBInspectable var rowSpacing: CGFloat = 0
    @IBInspectable var itemSpacing: CGFloat = 0
    
    weak var dataSource: StackCollectionViewDataSource?
    weak var delegate: StackCollectionViewDelegate?
    
    init() {
        super.init(frame: CGRect.zero)
        self.addAndConstrainSubview(_containerStackView)
        _containerStackView.axis = .vertical
        _containerStackView.spacing = sectionSpacing
    }
    
    required init?(coder: NSCoder) {
        super.init(coder: coder)
        self.addAndConstrainSubview(_containerStackView)
        _containerStackView.axis = .vertical
        _containerStackView.spacing = sectionSpacing
    }
    
    override func layoutSubviews() {
        super.layoutSubviews()
        reloadData()
    }
    
    func reloadData() {
        let subviews = _containerStackView.arrangedSubviews
        for subview in subviews {
            _containerStackView.removeArrangedSubview(subview)
            subview.removeFromSuperview()
        }
        
        guard let dataSource = dataSource else {  return }
        let sectionCount = dataSource.numberOfSections(in: self)
        for sectionIndex in 0..<sectionCount {
            
            let sectionStackView = UIStackView()
            sectionStackView.axis = .vertical
            sectionStackView.alignment = rowAlignment
            sectionStackView.spacing = rowSpacing
            
            _containerStackView.addArrangedSubview(sectionStackView)
            
            if let headerView = dataSource.stackCollectionView(self, viewForHeaderInSection: sectionIndex) {
                sectionStackView.addArrangedSubview(headerView)
            }
            
            let viewCount = dataSource.stackCollectionView(self, numberOfViewsInSection: sectionIndex)
            var currentRowWidth: CGFloat = 0
            
            var itemsStackView = getNewItemStackView()
            for itemIndex in 0..<viewCount {
                let itemView = dataSource.stackCollectionView(self, viewAtIndex: (section: sectionIndex, item: itemIndex))
                itemView.cachedSectionIndex = sectionIndex
                itemView.cachedItemIndex = itemIndex
                let itemViewSize = itemView.systemLayoutSizeFitting(UILayoutFittingCompressedSize)
                
                if itemsStackView.arrangedSubviews.count == 0 || currentRowWidth + itemSpacing + itemViewSize.width < self.frame.width {
                    itemsStackView.addArrangedSubview(itemView)
                    currentRowWidth += itemViewSize.width + (itemsStackView.arrangedSubviews.count > 0 ? itemSpacing : 0)
                } else {
                    sectionStackView.addArrangedSubview(itemsStackView)
                    itemsStackView = getNewItemStackView()
                    itemsStackView.addArrangedSubview(itemView)
                    currentRowWidth = itemViewSize.width
                }
                
                itemView.addGestureRecognizer(UITapGestureRecognizer(target: self, action: #selector(didSelectItem(tap:))))
            }
            
            //close up any open row
            if itemsStackView.arrangedSubviews.count > 0 {
                sectionStackView.addArrangedSubview(itemsStackView)
            }
            
            if let footerView = dataSource.stackCollectionView(self, viewForFooterInSection: sectionIndex) {
                sectionStackView.addArrangedSubview(footerView)
            }
        }
    }
    
    @objc private func didSelectItem(tap: UITapGestureRecognizer) {
        guard let view = tap.view else { return }
        delegate?.stackCollectionView(self, didSelectViewAtindex: (section: view.cachedSectionIndex, item: view.cachedItemIndex))
    }
    
    private func getNewItemStackView() -> UIStackView {
        let itemsStackView = UIStackView()
        itemsStackView.axis = .horizontal
        itemsStackView.spacing = itemSpacing
        return itemsStackView
    }
}

private extension UIView
{
    
    private static var cachedSectionIndexKey = "cachedSectionIndexKey"
    var cachedSectionIndex: Int {
        get {
            return objc_getAssociatedObject( self, &UIView.cachedSectionIndexKey ) as? Int ?? 0
        }
        set {
            objc_setAssociatedObject( self, &UIView.cachedSectionIndexKey,  newValue, .OBJC_ASSOCIATION_RETAIN)
        }
    }
    
    private static var cachedItemIndexKey = "cachedItemIndexKey"
    var cachedItemIndex: Int {
        get {
            return objc_getAssociatedObject( self, &UIView.cachedItemIndexKey ) as? Int ?? 0
        }
        set {
            objc_setAssociatedObject( self, &UIView.cachedItemIndexKey,  newValue, .OBJC_ASSOCIATION_RETAIN)
        }
    }
}

protocol StackCollectionViewDataSource: class {
    func numberOfSections(in collectionView: StackCollectionView) -> Int
    func stackCollectionView(_ collectionView: StackCollectionView, numberOfViewsInSection section: Int) -> Int
    func stackCollectionView(_ collectionView: StackCollectionView, viewAtIndex index: (section: Int, item: Int)) -> UIView
    func stackCollectionView(_ collectionView: StackCollectionView, viewForHeaderInSection section: Int) -> UIView?
    func stackCollectionView(_ collectionView: StackCollectionView, viewForFooterInSection section: Int) -> UIView?
}

protocol StackCollectionViewDelegate: class {
    func stackCollectionView(_ collectionView: StackCollectionView, didSelectViewAtindex index: (section: Int, item: Int))
}

extension StackCollectionViewDataSource { //define default behavior for optional functions
    func stackCollectionView(_ collectionView: StackCollectionView, viewForHeaderInSection section: Int) -> UIView? {
        return nil
    }
    func stackCollectionView(_ collectionView: StackCollectionView, viewForFooterInSection section: Int) -> UIView? {
        return nil
    }
}