Non-lazy Grid (Primitive)

GridScope.kt

package net.peanuuutz.compose.desktop

import androidx.compose.runtime.Immutable
import androidx.compose.ui.Modifier

@Immutable
interface GridScope {
    fun Modifier.span(
        span: GridSpanScope.() -> Int
    ): Modifier
}

interface GridSpanScope {
    val maxSpan: Int

    val remainingSpan: Int
}

// ======== Internal ========

// This scope object should not be @Stable as it's used everywhere
internal object ReusableGridSpanScope : GridSpanScope {
    override var maxSpan: Int = 0

    override var remainingSpan: Int = 0
}

HorizontalGrid.kt

@file:OptIn(ExperimentalComposeUiApi::class)
@file:Suppress("INVISIBLE_REFERENCE", "INVISIBLE_MEMBER")

package net.peanuuutz.compose.desktop

import androidx.compose.foundation.layout.Arrangement
import androidx.compose.foundation.lazy.grid.GridCells
import androidx.compose.runtime.Composable
import androidx.compose.runtime.Immutable
import androidx.compose.runtime.collection.MutableVector
import androidx.compose.runtime.remember
import androidx.compose.runtime.snapshots.fastForEach
import androidx.compose.ui.Alignment
import androidx.compose.ui.ExperimentalComposeUiApi
import androidx.compose.ui.Modifier
import androidx.compose.ui.layout.Layout
import androidx.compose.ui.layout.Measurable
import androidx.compose.ui.layout.MeasurePolicy
import androidx.compose.ui.layout.Placeable
import androidx.compose.ui.node.ModifierNodeElement
import androidx.compose.ui.node.ParentDataModifierNode
import androidx.compose.ui.unit.Constraints
import androidx.compose.ui.unit.Constraints.Companion.Infinity
import androidx.compose.ui.unit.Density
import kotlin.math.max
import kotlin.math.min

@Composable
inline fun HorizontalGrid(
    rows: GridCells,
    modifier: Modifier = Modifier,
    horizontalArrangement: Arrangement.Horizontal = Arrangement.Start,
    verticalArrangement: Arrangement.Vertical = Arrangement.Top,
    horizontalAlignment: Alignment.Horizontal = Alignment.Start,
    content: @Composable HorizontalGridScope.() -> Unit
) {
    val measurePolicy = rememberHorizontalGridMeasurePolicy(
        rows = rows,
        horizontalArrangement = horizontalArrangement,
        verticalArrangement = verticalArrangement,
        horizontalAlignment = horizontalAlignment
    )

    Layout(
        content = { HorizontalGridScopeImpl.content() },
        modifier = modifier,
        measurePolicy = measurePolicy
    )
}

@Immutable
interface HorizontalGridScope : GridScope {
    fun Modifier.align(
        alignment: GridSpanScope.() -> Alignment.Horizontal
    ): Modifier
}

// ======== Internal ========

@PublishedApi
internal object HorizontalGridScopeImpl : HorizontalGridScope {
    override fun Modifier.span(span: GridSpanScope.() -> Int): Modifier {
        return this then HorizontalGridChildElement(span = span)
    }

    override fun Modifier.align(alignment: GridSpanScope.() -> Alignment.Horizontal): Modifier {
        return this then HorizontalGridChildElement(alignment = alignment)
    }
}

@PublishedApi
@Composable
internal fun rememberHorizontalGridMeasurePolicy(
    rows: GridCells,
    horizontalArrangement: Arrangement.Horizontal,
    verticalArrangement: Arrangement.Vertical,
    horizontalAlignment: Alignment.Horizontal
): MeasurePolicy {
    return remember(rows, horizontalArrangement, verticalArrangement, horizontalAlignment) {
        horizontalGridMeasurePolicy(
            rows = rows,
            horizontalArrangement = horizontalArrangement,
            verticalArrangement = verticalArrangement,
            horizontalAlignment = horizontalAlignment
        )
    }
}

// Compose REALLY should add this pre-allocated lambda
private val EmptyPlacement: Placeable.PlacementScope.() -> Unit = {}

private fun horizontalGridMeasurePolicy(
    rows: GridCells,
    horizontalArrangement: Arrangement.Horizontal,
    verticalArrangement: Arrangement.Vertical,
    horizontalAlignment: Alignment.Horizontal
): MeasurePolicy {
    return MeasurePolicy { measurables, constraints ->
        val minWidth = constraints.minWidth
        val maxWidth = constraints.maxWidth
        val minHeight = constraints.minHeight
        val maxHeight = constraints.maxHeight
        require(maxHeight < Infinity) { "Cannot measure HorizontalGrid with unlimited height" }
        if (measurables.isEmpty()) {
            return@MeasurePolicy layout(
                width = minWidth,
                height = minHeight,
                placementBlock = EmptyPlacement
            )
        }
        val requiredVerticalSpacing = verticalArrangement.spacing.roundToPx()
        val rowHeights = with(rows) {
            calculateCrossAxisCellSizes(
                availableSize = maxHeight,
                spacing = requiredVerticalSpacing
            )
        }
        if (rowHeights.isEmpty()) {
            return@MeasurePolicy layout(
                width = minWidth,
                height = minHeight,
                placementBlock = EmptyPlacement
            )
        }
        val scope = ReusableGridSpanScope
        val requiredHorizontalSpacing = horizontalArrangement.spacing.roundToPx()
        val childCount = measurables.size
        val rowCount = rowHeights.size
        val gridHeight = rowHeights.sum().coerceIn(minHeight, maxHeight)
        val rowYs = IntArray(rowHeights.size)
        with(verticalArrangement) {
            arrange(
                totalSize = gridHeight,
                sizes = rowHeights.toIntArray(),
                outPositions = rowYs
            )
        }
        val placeables = MutableVector<Placeable>(childCount)
        // We cannot calculate column count by childCount / rowCount because some
        // children can specify custom span, thus breaking this formula. Instead,
        // we give it a minimum count and potentially increase on the go.
        val minimumColumnCount = childCount / rowCount + 1
        val columnWidths = MutableVector<Int>(minimumColumnCount)
        val columnCellSpans = MutableVector<IntArray>(minimumColumnCount)
        var occupiedWidth = 0
        var extraHorizontalSpace = 0
        var currentColumnWidth = 0
        // There is rare situation where user specify a large span, so we give it
        // an initial capacity of row count to ensure that most columns don't
        // need to expand backed array on the go
        val currentColumnCellSpans = MutableVector<Int>(rowCount)
        var remainingSpan = rowCount
        measurables.fastForEach { measurable ->
            scope.maxSpan = rowCount
            scope.remainingSpan = remainingSpan
            val spanProvider = measurable.span
            val span = if (spanProvider == null) {
                1
            } else {
                scope.spanProvider().coerceIn(1, remainingSpan)
            }
            val rowIndex = rowCount - remainingSpan
            val currentHeight = if (span == 1) {
                rowHeights[rowIndex]
            } else {
                val lastRowIndex = rowIndex + span - 1
                rowYs[lastRowIndex] - rowYs[rowIndex] + rowHeights[lastRowIndex]
            }
            val currentMaxWidth = if (maxWidth < Infinity) {
                maxWidth - occupiedWidth
            } else {
                Infinity
            }
            val currentConstraints = Constraints(
                minWidth = 0,
                maxWidth = currentMaxWidth,
                minHeight = currentHeight,
                maxHeight = currentHeight
            )
            val placeable = measurable.measure(currentConstraints)
            currentColumnWidth = max(currentColumnWidth, placeable.width)
            currentColumnCellSpans.add(span)
            placeables.add(placeable)
            if (remainingSpan != span) {
                remainingSpan -= span
            } else {
                extraHorizontalSpace = min(requiredHorizontalSpacing, maxWidth - occupiedWidth - currentColumnWidth)
                occupiedWidth += currentColumnWidth + extraHorizontalSpace
                columnWidths.add(currentColumnWidth)
                columnCellSpans.add(currentColumnCellSpans.toIntArray())
                currentColumnWidth = 0
                currentColumnCellSpans.clear()
                remainingSpan = rowCount
            }
        }
        if (remainingSpan != rowCount) {
            occupiedWidth += currentColumnWidth
            columnWidths.add(currentColumnWidth)
            columnCellSpans.add(currentColumnCellSpans.toIntArray())
        } else {
            occupiedWidth -= extraHorizontalSpace
        }
        val gridWidth = occupiedWidth.coerceIn(minWidth, maxWidth)
        val columnXs = IntArray(columnWidths.size)
        with(horizontalArrangement) {
            arrange(
                totalSize = gridWidth,
                sizes = columnWidths.toIntArray(),
                layoutDirection = layoutDirection,
                outPositions = columnXs
            )
        }
        layout(gridWidth, gridHeight) {
            var placementCurrentColumnIndex = 0
            var placementCurrentColumnCellSpans = columnCellSpans[0]
            var placementCurrentCellIndex = 0
            var placementRemainingSpan = rowCount
            placeables.forEachIndexed { index, placeable ->
                scope.maxSpan = rowCount
                scope.remainingSpan = placementRemainingSpan
                val span = placementCurrentColumnCellSpans[placementCurrentCellIndex]
                val alignmentProvider = measurables[index].alignment
                val alignment = if (alignmentProvider == null) {
                    horizontalAlignment
                } else {
                    scope.alignmentProvider()
                }
                val columnLocalCellX = alignment.align(
                    size = placeable.width,
                    space = columnWidths[placementCurrentColumnIndex],
                    layoutDirection = layoutDirection
                )
                placeable.place(
                    x = columnXs[placementCurrentColumnIndex] + columnLocalCellX,
                    y = rowYs[rowCount - placementRemainingSpan]
                )
                if (placementRemainingSpan != span) {
                    placementCurrentCellIndex++
                    placementRemainingSpan -= span
                } else {
                    placementCurrentColumnIndex++
                    placementCurrentColumnCellSpans = columnCellSpans[placementCurrentColumnIndex]
                    placementCurrentCellIndex = 0
                    placementRemainingSpan = rowCount
                }
            }
        }
    }
}

private data class HorizontalGridChildElement(
    val span: (GridSpanScope.() -> Int)? = null,
    val alignment: (GridSpanScope.() -> Alignment.Horizontal)? = null
) : ModifierNodeElement<HorizontalGridChildNode>() {
    override fun create(): HorizontalGridChildNode {
        return HorizontalGridChildNode(
            span = span,
            alignment = alignment
        )
    }

    override fun update(node: HorizontalGridChildNode): HorizontalGridChildNode {
        node.span = span
        node.alignment = alignment
        return node
    }
}

private val Measurable.span: (GridSpanScope.() -> Int)?
    get() = (parentData as? HorizontalGridChildData)?.span

private val Measurable.alignment: (GridSpanScope.() -> Alignment.Horizontal)?
    get() = (parentData as? HorizontalGridChildData)?.alignment

private class HorizontalGridChildData(
    var span: (GridSpanScope.() -> Int)?,
    var alignment: (GridSpanScope.() -> Alignment.Horizontal)?
)

private class HorizontalGridChildNode(
    var span: (GridSpanScope.() -> Int)?,
    var alignment: (GridSpanScope.() -> Alignment.Horizontal)?
) : Modifier.Node(), ParentDataModifierNode {
    override fun Density.modifyParentData(parentData: Any?): Any {
        return if (parentData is HorizontalGridChildData) {
            if (parentData.span == null) {
                parentData.span = span
            }
            if (parentData.alignment == null) {
                parentData.alignment = alignment
            }
            parentData
        } else {
            HorizontalGridChildData(
                span = span,
                alignment = alignment
            )
        }
    }
}

VerticalGrid.kt

@file:OptIn(ExperimentalComposeUiApi::class)
@file:Suppress("INVISIBLE_REFERENCE", "INVISIBLE_MEMBER")

package net.peanuuutz.compose.desktop

import androidx.compose.foundation.layout.Arrangement
import androidx.compose.foundation.lazy.grid.GridCells
import androidx.compose.runtime.Composable
import androidx.compose.runtime.Immutable
import androidx.compose.runtime.collection.MutableVector
import androidx.compose.runtime.remember
import androidx.compose.runtime.snapshots.fastForEach
import androidx.compose.ui.Alignment
import androidx.compose.ui.ExperimentalComposeUiApi
import androidx.compose.ui.Modifier
import androidx.compose.ui.layout.Layout
import androidx.compose.ui.layout.Measurable
import androidx.compose.ui.layout.MeasurePolicy
import androidx.compose.ui.layout.Placeable
import androidx.compose.ui.node.ModifierNodeElement
import androidx.compose.ui.node.ParentDataModifierNode
import androidx.compose.ui.unit.Constraints
import androidx.compose.ui.unit.Constraints.Companion.Infinity
import androidx.compose.ui.unit.Density
import kotlin.math.max
import kotlin.math.min

@Composable
inline fun VerticalGrid(
    columns: GridCells,
    modifier: Modifier = Modifier,
    verticalArrangement: Arrangement.Vertical = Arrangement.Top,
    horizontalArrangement: Arrangement.Horizontal = Arrangement.Start,
    verticalAlignment: Alignment.Vertical = Alignment.Top,
    content: @Composable VerticalGridScope.() -> Unit
) {
    val measurePolicy = rememberVerticalGridMeasurePolicy(
        columns = columns,
        verticalArrangement = verticalArrangement,
        horizontalArrangement = horizontalArrangement,
        verticalAlignment = verticalAlignment
    )

    Layout(
        content = { VerticalGridScopeImpl.content() },
        modifier = modifier,
        measurePolicy = measurePolicy
    )
}

@Immutable
interface VerticalGridScope : GridScope {
    fun Modifier.align(
        alignment: GridSpanScope.() -> Alignment.Vertical
    ): Modifier
}

// ======== Internal ========

@PublishedApi
internal object VerticalGridScopeImpl : VerticalGridScope {
    override fun Modifier.span(span: GridSpanScope.() -> Int): Modifier {
        return this then VerticalGridChildElement(span = span)
    }

    override fun Modifier.align(alignment: GridSpanScope.() -> Alignment.Vertical): Modifier {
        return this then VerticalGridChildElement(alignment = alignment)
    }
}

@PublishedApi
@Composable
internal fun rememberVerticalGridMeasurePolicy(
    columns: GridCells,
    verticalArrangement: Arrangement.Vertical,
    horizontalArrangement: Arrangement.Horizontal,
    verticalAlignment: Alignment.Vertical
): MeasurePolicy {
    return remember(columns, verticalArrangement, horizontalArrangement, verticalAlignment) {
        verticalGridMeasurePolicy(
            columns = columns,
            verticalArrangement = verticalArrangement,
            horizontalArrangement = horizontalArrangement,
            verticalAlignment = verticalAlignment
        )
    }
}

// Compose REALLY should add this pre-allocated lambda
private val EmptyPlacement: Placeable.PlacementScope.() -> Unit = {}

private fun verticalGridMeasurePolicy(
    columns: GridCells,
    verticalArrangement: Arrangement.Vertical,
    horizontalArrangement: Arrangement.Horizontal,
    verticalAlignment: Alignment.Vertical
): MeasurePolicy {
    return MeasurePolicy { measurables, constraints ->
        val minWidth = constraints.minWidth
        val maxWidth = constraints.maxWidth
        val minHeight = constraints.minHeight
        val maxHeight = constraints.maxHeight
        require(maxWidth < Infinity) { "Cannot measure VerticalGrid with unlimited width" }
        if (measurables.isEmpty()) {
            return@MeasurePolicy layout(
                width = minWidth,
                height = minHeight,
                placementBlock = EmptyPlacement
            )
        }
        val requiredHorizontalSpacing = horizontalArrangement.spacing.roundToPx()
        val columnWidths = with(columns) {
            calculateCrossAxisCellSizes(
                availableSize = maxWidth,
                spacing = requiredHorizontalSpacing
            )
        }
        if (columnWidths.isEmpty()) {
            return@MeasurePolicy layout(
                width = minWidth,
                height = minHeight,
                placementBlock = EmptyPlacement
            )
        }
        val scope = ReusableGridSpanScope
        val requiredVerticalSpacing = verticalArrangement.spacing.roundToPx()
        val childCount = measurables.size
        val columnCount = columnWidths.size
        val gridWidth = columnWidths.sum().coerceIn(minWidth, maxWidth)
        val columnXs = IntArray(columnWidths.size)
        with(horizontalArrangement) {
            arrange(
                totalSize = gridWidth,
                sizes = columnWidths.toIntArray(),
                layoutDirection = layoutDirection,
                outPositions = columnXs
            )
        }
        val placeables = MutableVector<Placeable>(childCount)
        // We cannot calculate row count by childCount / columnCount because some
        // children can specify custom span, thus breaking this formula. Instead,
        // we give it a minimum count and potentially increase on the go.
        val minimumRowCount = childCount / columnCount + 1
        val rowHeights = MutableVector<Int>(minimumRowCount)
        val rowCellSpans = MutableVector<IntArray>(minimumRowCount)
        var occupiedHeight = 0
        var extraVerticalSpace = 0
        var currentRowHeight = 0
        // There is rare situation where user specify a large span, so we give it
        // an initial capacity of column count to ensure that most rows don't
        // need to expand backed array on the go
        val currentRowCellSpans = MutableVector<Int>(columnCount)
        var remainingSpan = columnCount
        measurables.fastForEach { measurable ->
            scope.maxSpan = columnCount
            scope.remainingSpan = remainingSpan
            val spanProvider = measurable.span
            val span = if (spanProvider == null) {
                1
            } else {
                scope.spanProvider().coerceIn(1, remainingSpan)
            }
            val columnIndex = columnCount - remainingSpan
            val currentWidth = if (span == 1) {
                columnWidths[columnIndex]
            } else {
                val lastColumnIndex = columnIndex + span - 1
                columnXs[lastColumnIndex] - columnXs[columnIndex] + columnWidths[lastColumnIndex]
            }
            val currentMaxHeight = if (maxHeight < Infinity) {
                maxHeight - occupiedHeight
            } else {
                Infinity
            }
            val currentConstraints = Constraints(
                minWidth = currentWidth,
                maxWidth = currentWidth,
                minHeight = 0,
                maxHeight = currentMaxHeight
            )
            val placeable = measurable.measure(currentConstraints)
            currentRowHeight = max(currentRowHeight, placeable.height)
            currentRowCellSpans.add(span)
            placeables.add(placeable)
            if (remainingSpan != span) {
                remainingSpan -= span
            } else {
                extraVerticalSpace = min(requiredVerticalSpacing, maxHeight - occupiedHeight - currentRowHeight)
                occupiedHeight += currentRowHeight + extraVerticalSpace
                rowHeights.add(currentRowHeight)
                rowCellSpans.add(currentRowCellSpans.toIntArray())
                currentRowHeight = 0
                currentRowCellSpans.clear()
                remainingSpan = columnCount
            }
        }
        if (remainingSpan != columnCount) {
            occupiedHeight += currentRowHeight
            rowHeights.add(currentRowHeight)
            rowCellSpans.add(currentRowCellSpans.toIntArray())
        } else {
            occupiedHeight -= extraVerticalSpace
        }
        val gridHeight = occupiedHeight.coerceIn(minHeight, maxHeight)
        val rowYs = IntArray(rowHeights.size)
        with(verticalArrangement) {
            arrange(
                totalSize = gridHeight,
                sizes = rowHeights.toIntArray(),
                outPositions = rowYs
            )
        }
        layout(gridWidth, gridHeight) {
            var placementCurrentRowIndex = 0
            var placementCurrentRowCellSpans = rowCellSpans[0]
            var placementCurrentCellIndex = 0
            var placementRemainingSpan = columnCount
            placeables.forEachIndexed { index, placeable ->
                scope.maxSpan = columnCount
                scope.remainingSpan = placementRemainingSpan
                val span = placementCurrentRowCellSpans[placementCurrentCellIndex]
                val alignmentProvider = measurables[index].alignment
                val alignment = if (alignmentProvider == null) {
                    verticalAlignment
                } else {
                    scope.alignmentProvider()
                }
                val rowLocalCellY = alignment.align(
                    size = placeable.height,
                    space = rowHeights[placementCurrentRowIndex]
                )
                placeable.place(
                    x = columnXs[columnCount - placementRemainingSpan],
                    y = rowYs[placementCurrentRowIndex] + rowLocalCellY
                )
                if (placementRemainingSpan != span) {
                    placementCurrentCellIndex++
                    placementRemainingSpan -= span
                } else {
                    placementCurrentRowIndex++
                    placementCurrentRowCellSpans = rowCellSpans[placementCurrentRowIndex]
                    placementCurrentCellIndex = 0
                    placementRemainingSpan = columnCount
                }
            }
        }
    }
}

private data class VerticalGridChildElement(
    val span: (GridSpanScope.() -> Int)? = null,
    val alignment: (GridSpanScope.() -> Alignment.Vertical)? = null
) : ModifierNodeElement<VerticalGridChildNode>() {
    override fun create(): VerticalGridChildNode {
        return VerticalGridChildNode(
            span = span,
            alignment = alignment
        )
    }

    override fun update(node: VerticalGridChildNode): VerticalGridChildNode {
        node.span = span
        node.alignment = alignment
        return node
    }
}

private val Measurable.span: (GridSpanScope.() -> Int)?
    get() = (parentData as? VerticalGridChildData)?.span

private val Measurable.alignment: (GridSpanScope.() -> Alignment.Vertical)?
    get() = (parentData as? VerticalGridChildData)?.alignment

private class VerticalGridChildData(
    var span: (GridSpanScope.() -> Int)?,
    var alignment: (GridSpanScope.() -> Alignment.Vertical)?
)

private class VerticalGridChildNode(
    var span: (GridSpanScope.() -> Int)?,
    var alignment: (GridSpanScope.() -> Alignment.Vertical)?
) : Modifier.Node(), ParentDataModifierNode {
    override fun Density.modifyParentData(parentData: Any?): Any {
        return if (parentData is VerticalGridChildData) {
            if (parentData.span == null) {
                parentData.span = span
            }
            if (parentData.alignment == null) {
                parentData.alignment = alignment
            }
            parentData
        } else {
            VerticalGridChildData(
                span = span,
                alignment = alignment
            )
        }
    }
}