Simple take on a compositional layout with 2 column variable height items waterfall

CompositionalTwoColumnWaterfall.swift

import UIKit

class ViewController: UIViewController {

    enum Section {
        case main
    }
    
    struct Item: Hashable {
        let height: CGFloat
        let color: UIColor
        private let identifier = UUID()
        func hash(into hasher: inout Hasher) {
            hasher.combine(identifier)
        }
        static func == (lhs:Item, rhs:Item) -> Bool {
            return lhs.identifier == rhs.identifier
        }
    }
    
    var currentSnapshot: NSDiffableDataSourceSnapshot<Section, Item>!
    var dataSource: UICollectionViewDiffableDataSource<Section, Item>!
    var collectionView: UICollectionView!
    
    override func viewDidLoad() {
        super.viewDidLoad()
        configureHierarchy()
        configureDataSource()
    }
}

extension ViewController {
    
    func configureHierarchy() {
        collectionView = UICollectionView(frame: view.bounds, collectionViewLayout: twoColumnWaterfallLayout())
        collectionView.autoresizingMask = [.flexibleWidth, .flexibleWidth]
        view.addSubview(collectionView)
    }
    
    func configureDataSource() {
        let reuseIdentifier = "cell-idententifier"
        collectionView.register(UICollectionViewCell.self, forCellWithReuseIdentifier: reuseIdentifier)
        dataSource = UICollectionViewDiffableDataSource<Section, Item>(collectionView: collectionView) { (collectionView: UICollectionView, indexPath: IndexPath, item: Item) -> UICollectionViewCell? in
            let cell = collectionView.dequeueReusableCell(withReuseIdentifier: reuseIdentifier, for: indexPath)
            cell.contentView.backgroundColor = item.color
            cell.contentView.layer.borderColor = UIColor.black.cgColor
            cell.contentView.layer.borderWidth = 1
            return cell
        }
        currentSnapshot = intialSnapshot()
        dataSource.apply(currentSnapshot, animatingDifferences: false)
    }
    
    func intialSnapshot() -> NSDiffableDataSourceSnapshot<Section, Item> {
        let itemCount = 200
        var items = [Item]()
        for _ in 0..<itemCount {
            let height = CGFloat.random(in: 88..<121)
            let color = UIColor(hue:CGFloat.random(in: 0.1..<0.9), saturation: 1.0, brightness: 1.0, alpha: 1.0)
            let item = Item(height: height, color: color)
            items.append(item)
        }
        
        let snapshot = NSDiffableDataSourceSnapshot<Section, Item>()
        snapshot.appendSections([.main])
        snapshot.appendItems(items)
        return snapshot
    }
    
    
    //    +----------------------+     +----------------------+
    //    |Leading Vertical Group|     |  Trailing Vertical   |
    //    +----------------------+     +----------------------+
    //                |                            |
    //                |                            |
    //                |                            |
    //                v                            v
    //   +--------------------------+--------------------------+
    //   |                          |                          |
    //   |                          |                          |
    //   |                          |                          |
    //   +--------------------------+--------------------------+
    //   |                          |                          |
    //   +--------------------------+                          |
    //   |                          |                          |
    //   |                          +--------------------------+
    //   |                          |                          |
    //   |                          |                          |
    //   +--------------------------+                          |
    //   |                          |                          |
    //   |                          |                          |
    //   |                          |                          |
    //   +--------------------------+--------------------------+
    //   |                          |                          |
    //   |                          |                          |            +-----------------------------+
    //   +--------------------------+--------------------------+ <----------| Horizontal Container Group  |
    //   |                          |                          |            +-----------------------------+
    //   |                          |                          |
    //   |                          +--------------------------+
    //   +--------------------------+                          |
    //   |                          |                          |
    //   |                          |                          |
    //   +--------------------------+--------------------------+
    //   |                          |                          |
    //   |                          |                          |
    //   |                          +--------------------------+
    //   +--------------------------+                          |
    //   |                          |                          |
    //   |                          |                          |
    //   |                          +--------------------------+
    //   +--------------------------+                          |
    //   |                          |                          |
    //   |                          |                          |
    //   +--------------------------+--------------------------+
    //
    //
    //
    //   +---------------------------------------------------------------------------------------------------------+
    //   |*Container group is horizontal with Leading + Trailing vertical groups                                   |
    //   |                                                                                                         |
    //   |* Alternate between the leading + trailing group adding items from metadata about each item's height     |
    //   |                                                                                                         |
    //   |* When updates occur, the sectionProvider is involved                                                    |
    //   |                                                                                                         |
    //   |* This is still very fast: the definitions are very cheap to create and the layout is optimized for even |
    //   |large groups                                                                                             |
    //   |                                                                                                         |
    //   +---------------------------------------------------------------------------------------------------------+
    
    func twoColumnWaterfallLayout() -> UICollectionViewLayout {
        let sectionProvider = { [weak self] (sectionIndex: Int, layoutEnvironment: NSCollectionLayoutEnvironment) -> NSCollectionLayoutSection? in
            guard let self = self else { return nil }
            
            var leadingGroupHeight = CGFloat(0.0)
            var trailingGroupHeight = CGFloat(0.0)
            var leadingGroupItems = [NSCollectionLayoutItem]()
            var trailingGroupItems = [NSCollectionLayoutItem]()
            let items = self.currentSnapshot.itemIdentifiers
            
            let totalHeight = items.reduce(0) { $0 + $1.height }
            let columnHeight = CGFloat(totalHeight / 2.0)
            
            // could get a bit fancier and balance the columns if they are too different height-wise -  here is just a simple take on this
            
            var runningHeight = CGFloat(0.0)
            for index in 0..<self.currentSnapshot.numberOfItems {
                let item = items[index]
                let isLeading = runningHeight < columnHeight
                let layoutSize = NSCollectionLayoutSize(widthDimension: .fractionalWidth(1.0), heightDimension: .absolute(item.height))
                let layoutItem = NSCollectionLayoutItem(layoutSize: layoutSize)
                
                runningHeight += item.height
                
                if isLeading {
                    leadingGroupItems.append(layoutItem)
                    leadingGroupHeight += item.height
                } else {
                    trailingGroupItems.append(layoutItem)
                    trailingGroupHeight += item.height
                }
            }
            
            let leadingGroupSize = NSCollectionLayoutSize(widthDimension: .fractionalWidth(0.5), heightDimension: .absolute(leadingGroupHeight))
            let leadingGroup = NSCollectionLayoutGroup.vertical(layoutSize: leadingGroupSize, subitems:leadingGroupItems)
            
            let trailingGroupSize = NSCollectionLayoutSize(widthDimension: .fractionalWidth(0.5), heightDimension: .absolute(trailingGroupHeight))
            let trailingGroup = NSCollectionLayoutGroup.vertical(layoutSize: trailingGroupSize, subitems: trailingGroupItems)
            
            let containerGroupSize = NSCollectionLayoutSize(widthDimension: .fractionalWidth(1.0), heightDimension: .absolute(max(leadingGroupHeight, trailingGroupHeight)))
            let containerGroup = NSCollectionLayoutGroup.horizontal(layoutSize: containerGroupSize, subitems: [leadingGroup, trailingGroup])
            
            let section = NSCollectionLayoutSection(group: containerGroup)
            return section
        }
        let layout = UICollectionViewCompositionalLayout(sectionProvider: sectionProvider)
        return layout
    }
}