WaterfallLayout.swift

import UIKit
import PlaygroundSupport

protocol WaterfallLayoutDelegate: AnyObject {
        func numberOfColumns() -> Int
        func columnsSize(at indexPath: IndexPath) -> CGSize
        func columnSpace() -> CGFloat
}

final class WaterfallLayoutViewController: UIViewController, UICollectionViewDataSource {
    override func viewDidLoad() {
        super.viewDidLoad()
        waterfallLayout = self
        collectionView.translatesAutoresizingMaskIntoConstraints = false
        view.addSubview(collectionView)
        NSLayoutConstraint.activate([
            collectionView.leadingAnchor.constraint(equalTo: view.leadingAnchor),
            collectionView.topAnchor.constraint(equalTo: view.topAnchor),
            view.trailingAnchor.constraint(equalTo: collectionView.trailingAnchor),
            view.bottomAnchor.constraint(equalTo: collectionView.bottomAnchor),
        ])
    }

    // MARK: - UI

    weak var waterfallLayout: WaterfallLayoutDelegate?

    lazy var collectionViewLayout: UICollectionViewCompositionalLayout = {
        return UICollectionViewCompositionalLayout { [unowned self] (section, environment) -> NSCollectionLayoutSection? in
            let groupSize = NSCollectionLayoutSize(widthDimension: .fractionalWidth(1.0), heightDimension: .estimated(environment.container.effectiveContentSize.height))
            let group = NSCollectionLayoutGroup.custom(layoutSize: groupSize) { [unowned self] (environment) -> [NSCollectionLayoutGroupCustomItem] in
                var items: [NSCollectionLayoutGroupCustomItem] = []
                var layouts: [Int: CGFloat] = [:]
                let space: CGFloat = self.waterfallLayout.flatMap({ CGFloat($0.columnSpace()) }) ?? 1.0
                let numberOfColumn: CGFloat = self.waterfallLayout.flatMap({ CGFloat($0.numberOfColumns()) }) ?? 2.0
                let defaultSize = CGSize(width: 100, height: 100)

                (0 ..< self.collectionView.numberOfItems(inSection: section)).forEach {
                    let indexPath = IndexPath(item: $0, section: section)

                    let size = self.waterfallLayout?.columnsSize(at: indexPath) ?? defaultSize
                    let aspect = CGFloat(size.height) / CGFloat(size.width)

                    let width = (environment.container.effectiveContentSize.width - (numberOfColumn - 1) * space) / numberOfColumn
                    let height = width * aspect

                    let currentColumn = $0 % Int(numberOfColumn)
                    let y = layouts[currentColumn] ?? 0.0 + space
                    let x = width * CGFloat(currentColumn) + space * (CGFloat(currentColumn) - 1.0)

                    let frame = CGRect(x: x, y: y + space, width: width, height: height)
                    let item = NSCollectionLayoutGroupCustomItem(frame: frame)
                    items.append(item)

                    layouts[currentColumn] = frame.maxY
                }
                return items
            }
            return NSCollectionLayoutSection(group: group)
        }
    }()

    lazy var collectionView: UICollectionView = {
        let collectionView = UICollectionView(frame: view.bounds, collectionViewLayout: collectionViewLayout)
        collectionView.backgroundColor = .systemBackground
        collectionView.accessibilityIdentifier = #function
        collectionView.dataSource = self
        collectionView.register(UICollectionViewCell.self, forCellWithReuseIdentifier: "Cell")
        return collectionView
    }()

    func numberOfSections(in collectionView: UICollectionView) -> Int {
        return 1
    }

    func collectionView(_ collectionView: UICollectionView, numberOfItemsInSection section: Int) -> Int {
        return 1000
    }

    func collectionView(_ collectionView: UICollectionView, cellForItemAt indexPath: IndexPath) -> UICollectionViewCell {
        let cell = collectionView.dequeueReusableCell(withReuseIdentifier: "Cell", for: indexPath)
        cell.backgroundColor = .red
        return cell
    }
}

extension WaterfallLayoutViewController: WaterfallLayoutDelegate {
    func numberOfColumns() -> Int {
        3
    }

    func columnsSize(at indexPath: IndexPath) -> CGSize {
        let width = CGFloat.random(in: 1..<1000)
        let height = CGFloat.random(in: 1..<1000)
        return CGSize(width: width, height: height)
    }

    func columnSpace() -> CGFloat {
        3.0
    }
}

let viewController = WaterfallLayoutViewController()
PlaygroundPage.current.liveView = viewController

great work,
I have one problem though, when you give heights based on actual content, sometimes the last cell get aligned on the wrong column making a big empty space (if the number of columns is 2)

nice job bro! Thank you