a simple implementation of waterfall layout using compositional layout

compositional_waterfall_layout.swift

import UIKit

extension UICollectionViewCompositionalLayout {
  static func waterfall(
    columnCount: Int,
    interItemSpacing: CGFloat = 0,
    lineSpacing: CGFloat = 0,
    itemCountProvider: @escaping (Int) -> Int,
    itemSizeProvider: @escaping (IndexPath) -> CGSize,
    sectionConfigurator: @escaping (inout NSCollectionLayoutSection) -> Void = { _ in },
    configuration: UICollectionViewCompositionalLayoutConfiguration = .init()
  ) -> UICollectionViewCompositionalLayout {
    .init(
      sectionProvider: { sectionIndex, environment in
        var section = NSCollectionLayoutSection.waterfall(
          sectionIndex,
          environment,
          columnCount: columnCount,
          interItemSpacing: interItemSpacing,
          lineSpacing: lineSpacing,
          itemCountProvider: itemCountProvider,
          itemSizeProvider: itemSizeProvider
        )
        sectionConfigurator(&section)
        return section
      },
      configuration: configuration
    )
  }
}

extension NSCollectionLayoutSection {
  static func waterfall(
    _ sectionIndex: Int,
    _ environment: NSCollectionLayoutEnvironment,
    columnCount: Int,
    interItemSpacing: CGFloat,
    lineSpacing: CGFloat,
    // contentInset
    itemCountProvider: @escaping (Int) -> Int,
    itemSizeProvider: @escaping (IndexPath) -> CGSize
  ) -> NSCollectionLayoutSection {
    let groupSize = NSCollectionLayoutSize(
      widthDimension: .fractionalWidth(1),
      heightDimension: .estimated(environment.container.effectiveContentSize.height)
    )
    
    let group = NSCollectionLayoutGroup.custom(layoutSize: groupSize) { environment in
      let columnCount = max(1, columnCount)
      let columnCountFloat = CGFloat(columnCount)
      let itemWidth = (environment.container.effectiveContentSize.width - (columnCountFloat - 1) * interItemSpacing) / columnCountFloat
      
      return (0..<itemCountProvider(sectionIndex))
        .reduce(into: (items: [NSCollectionLayoutGroupCustomItem](), columnHeights: [CGFloat](repeating: 0, count: columnCount))) {
          let indexPath = IndexPath(item: $1, section: sectionIndex)
          let estimatedItemSize = itemSizeProvider(indexPath)
          let itemHeight = itemWidth * estimatedItemSize.height / estimatedItemSize.width
          let minColumnHeight = $0.columnHeights.min() ?? 0
          let columnIndex = $0.columnHeights.firstIndex(of: minColumnHeight) ?? 0
          let columnIndexFloat = CGFloat(columnIndex)
          
          let frame = CGRect(
            x: itemWidth * columnIndexFloat + interItemSpacing * columnIndexFloat,
            y: minColumnHeight + lineSpacing,
            width: itemWidth,
            height: itemHeight
          )
          $0.items.append(.init(frame: frame))
          $0.columnHeights[columnIndex] = frame.maxY
        }
        .items
    }
    return NSCollectionLayoutSection(group: group)
  }
}