Waterfall Layout with SwiftUI Layout protocol

WaterfallLayout.swift

import SwiftUI

struct ContentView: View {
    typealias Item = (Color, CGFloat)

    @State private var columnCount = 3

    @State private var isWaterfall = true

    @State var items: [Item] = [
        (.red, 80),
        (.pink, 115),
        (.purple, 60),
        (.orange, 75),
        (.mint, 100),
        (.red, 120),
        (.gray, 200),
        (.green, 40),
        (.blue, 130),
        (.brown, 80),
    ]

    private var layout: AnyLayout {
        if isWaterfall {
            return AnyLayout(WaterfallLayout(columnCount: columnCount, spacing: 8))
        } else {
            return AnyLayout(VStack())
        }
    }

    var body: some View {
        VStack(spacing: .zero) {
            ScrollView {
                layout {
                    ForEach(Array(zip(items.indices, items)), id: \.0) { index, item in
                        item.0
                            .frame(height: item.1)
                    }
                    .clipShape(RoundedRectangle(cornerRadius: 8, style: .continuous))
                }
                .padding(.horizontal)
            }
            .frame(maxWidth: .infinity, maxHeight: .infinity)

            Divider()

            VStack {
                Toggle("Waterfall layout", isOn: $isWaterfall.transaction(.init(animation: .default)))

                if isWaterfall {
                    Stepper(value: $columnCount.transaction(.init(animation: .default)), in: (1...10)) {
                        Text("Column count")
                    }
                }

                Button("Add more") {
                    withAnimation {
                        items.append(contentsOf: (0..<20).map { _ in randomItem() })
                    }
                }
            }
            .padding()
        }
    }

    private func randomItem() -> Item {
        ([.red, .pink, .purple, .orange, .mint, .red, .gray, .green, .blue, .brown].randomElement()!, .random(in: 50...200))
    }
}

struct ContentView_Previews: PreviewProvider {
    static var previews: some View {
        ContentView()
    }
}

struct WaterfallLayout: Layout {
    let columnCount: Int

    let spacing: CGFloat

    func sizeThatFits(proposal: ProposedViewSize, subviews: Subviews, cache: inout ()) -> CGSize {
        let result = Waterfall(
            columnCount: columnCount,
            origin: .zero,
            width: proposal.replacingUnspecifiedDimensions().width,
            spacing: spacing,
            subviews: subviews
        )
        return result.bounds.size
    }

    func placeSubviews(in bounds: CGRect, proposal: ProposedViewSize, subviews: Subviews, cache: inout ()) {
        let result = Waterfall(
            columnCount: columnCount,
            origin: bounds.origin,
            width: proposal.replacingUnspecifiedDimensions().width,
            spacing: spacing,
            subviews: subviews
        )
        result.columns
            .flatMap { $0 }
            .forEach { subview, frame in
                subview.place(at: frame.origin, proposal: .init(width: frame.width, height: frame.height))
            }
    }

    struct Waterfall {
        typealias Column = [(LayoutSubview, CGRect)]
        let columns: [Column]

        let bounds: CGRect

        init(columnCount: Int, origin: CGPoint, width: CGFloat, spacing: CGFloat, subviews: Subviews) {
            let gapCount = max(0, columnCount - 1)
            let columnWidth = (width - CGFloat(gapCount) * spacing) / CGFloat(columnCount)
            func column(index: Int) -> Column {
                let sizes: [(LayoutSubview, CGSize)] = stride(from: index, to: subviews.count, by: columnCount).compactMap { index in
                    guard subviews.indices.contains(index) else { return nil }
                    let subview = subviews[index]
                    return (subview, subview.sizeThatFits(.init(width: width, height: nil)))
                }
                return sizes.reduce(into: []) { partialResult, subviewSize in
                    let (subview, size) = subviewSize
                    let lastFrame = partialResult.last?.1 ?? CGRect(
                        x: origin.x + columnWidth * CGFloat(index) + spacing * CGFloat(index),
                        y: origin.y - spacing,
                        width: columnWidth,
                        height: .zero
                    )
                    let frame = CGRect(
                        x: lastFrame.minX,
                        y: lastFrame.maxY + spacing,
                        width: lastFrame.width,
                        height: size.height
                    )
                    partialResult.append((subview, frame))
                }
            }
            let columns = (0..<columnCount).map(column(index:))
            self.columns = columns
            self.bounds = columns.compactMap(\.last?.1).reduce(into: .zero) { $0 = $0.union($1) }
        }
    }
}