DoubleColumnVStackLayout.swift

import SwiftUI

public struct DoubleColumnVStackLayout: Layout {
    
    private let alignment: HorizontalAlignment
    private let spacing: CGFloat?
    
    public init(alignment: HorizontalAlignment = .center, spacing: CGFloat? = nil) {
        self.alignment = alignment
        self.spacing = spacing
    }
    
    public func sizeThatFits(proposal: ProposedViewSize, subviews: Subviews, cache: inout Cache) -> CGSize {
        let rows = arrangeRows(proposal: proposal, subviews: subviews, cache: &cache)
        
        guard !rows.isEmpty else {
            return .zero
        }
        
        let width = min(proposal.width, proposal.height) ?? rows.map(\.width).reduce(.zero, max)
        var height: CGFloat = .zero
        
        if let lastRow = rows.last {
            height = lastRow.yOffset + lastRow.height
        }
        
        return CGSize(width: width, height: height)
    }
    
    public func placeSubviews(in bounds: CGRect, proposal: ProposedViewSize, subviews: Subviews, cache: inout Cache) {
        let rows = arrangeRows(proposal: proposal, subviews: subviews, cache: &cache)
        let anchor = UnitPoint(alignment)
        
        for row in rows {
            for element in row.elements {
                let x = element.xOffset + anchor.x * (bounds.width - row.width)
                let y = row.yOffset + anchor.y * (row.height - element.size.height)
                let point = CGPoint(x: x + bounds.minX, y: y + bounds.minY)
                
                let sizeProposal = ProposedViewSize(element.size)
                subviews[element.index].place(at: point, anchor: .topLeading, proposal: sizeProposal)
            }
        }
    }
}

extension DoubleColumnVStackLayout {
    
    public struct Cache {
        fileprivate var rows: (Int, [Row])?
    }
    
    public static var layoutProperties: LayoutProperties {
        var properties = LayoutProperties()
        properties.stackOrientation = .horizontal
        return properties
    }
    
    public func makeCache(subviews: Subviews) -> Cache {
        Cache()
    }
}

private extension DoubleColumnVStackLayout {
    
    struct Row {
        var elements: [(index: Int, size: CGSize, xOffset: CGFloat)] = []
        var yOffset: CGFloat = .zero
        var width: CGFloat = .zero
        var height: CGFloat = .zero
    }
    
    func arrangeRows(
        proposal: ProposedViewSize,
        subviews: Subviews,
        cache: inout Cache) -> [Row]
    {
        if subviews.isEmpty {
            return []
        }
        
        let count = subviews.count
        if let (oldCount, oldRows) = cache.rows, oldCount == count {
            return oldRows
        }
        
        var sizes = [CGSize]()
        var currentX: CGFloat = .zero
        var currentRow = Row()
        var rows = [Row]()
        
        for (index, subview) in subviews.enumerated() {
            let spacingToPreviousElement: CGFloat?
            if let previousIndex = currentRow.elements.last?.index {
                spacingToPreviousElement = horizontalSpacing(subviews[previousIndex], subview)
            } else {
                spacingToPreviousElement = nil
            }
            
            var spacing = spacingToPreviousElement ?? .zero
            let proposedSpacing = spacingToPreviousElement ?? horizontalSpacing(subview, subview)
            let proposedLength = min(proposal.width, proposal.height).map { ($0 - proposedSpacing) / 2 }
            let sizeProposal = ProposedViewSize(width: proposedLength, height: proposedLength)
            let size = subview.sizeThatFits(sizeProposal)
            
            if currentX + size.width + spacing > proposal.width ?? .infinity, !currentRow.elements.isEmpty {
                currentRow.width = currentX
                rows.append(currentRow)
                currentRow = Row()
                spacing = .zero
                currentX = .zero
            }
            
            currentRow.elements.append((index, size, currentX + spacing))
            currentX += size.width + spacing
            sizes.append(size)
        }
        
        if !currentRow.elements.isEmpty {
            currentRow.width = currentX
            rows.append(currentRow)
        }
        
        var currentY: CGFloat = .zero
        var previousMaxHeightIndex: Int?
        
        for index in rows.indices {
            let maxHeightIndex = rows[index].elements
                .max { $0.size.height < $1.size.height }!
                .index
            
            let size = sizes[maxHeightIndex]
            
            var spacing: CGFloat = .zero
            if let previousMaxHeightIndex {
                spacing = verticalSpacing(subviews[previousMaxHeightIndex], subviews[maxHeightIndex])
            }
            
            rows[index].yOffset = currentY + spacing
            currentY += size.height + spacing
            rows[index].height = size.height
            previousMaxHeightIndex = maxHeightIndex
        }
        
        cache.rows = (count, rows)
        
        return rows
    }
    
    func horizontalSpacing(_ lhs: LayoutSubview, _ rhs: LayoutSubview) -> CGFloat {
        spacing ?? lhs.spacing.distance(to: rhs.spacing, along: .horizontal)
    }
    
    func verticalSpacing(_ lhs: LayoutSubview, _ rhs: LayoutSubview) -> CGFloat {
        spacing ?? lhs.spacing.distance(to: rhs.spacing, along: .vertical)
    }
}

private extension UnitPoint {
    
    init(_ alignment: HorizontalAlignment) {
        switch alignment {
        case .leading:
            self = .leading
        case .trailing:
            self = .trailing
        default:
            self = .center
        }
    }
}

private func min<T: Comparable>(_ x: T?, _ y: T?) -> T? {
    guard let x else { return y }
    guard let y else { return x }
    return min(x, y)
}

#Preview {
    ScrollView {
        DoubleColumnVStackLayout(alignment: .center, spacing: 8) {
            ForEach(0..<9) { _ in
                Text("Hello, World!")
                    .frame(maxWidth: .infinity, maxHeight: .infinity)
                    .background(Color.accentColor)
            }
        }
        .background(Color(white: 0.2))
        .padding()
        //.environment(\.layoutDirection, .rightToLeft)
    }
}