Compose Table/Grid

Table.kt

package net.inferno.compose.view

import androidx.compose.foundation.layout.LayoutScopeMarker
import androidx.compose.runtime.*
import androidx.compose.ui.Alignment
import androidx.compose.ui.BiasAlignment
import androidx.compose.ui.Modifier
import androidx.compose.ui.geometry.Offset
import androidx.compose.ui.layout.*
import androidx.compose.ui.platform.InspectorInfo
import androidx.compose.ui.platform.InspectorValueInfo
import androidx.compose.ui.platform.debugInspectorInfo
import androidx.compose.ui.unit.Density
import androidx.compose.ui.unit.IntOffset
import androidx.compose.ui.unit.round

@Composable
fun VerticalTable(
    modifier: Modifier = Modifier,
    cells: TableCells,
    content: @Composable TableScope.() -> Unit,
) {
    val columnWidthsState = remember(cells) { mutableStateOf(emptyList<Int>()) }
    val measurePolicy = verticalTableMeasurePolicy(cells, columnWidthsState)

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

@Composable
fun HorizontalTable(
    modifier: Modifier = Modifier,
    cells: TableCells,
    content: @Composable TableScope.() -> Unit,
) {
    val rowHeightsState = remember(cells) { mutableStateOf(emptyList<Int>()) }
    val measurePolicy = horizontalTableMeasurePolicy(cells, rowHeightsState)

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

@LayoutScopeMarker
@Immutable
interface TableScope {
    @Stable
    fun Modifier.align(alignment: Alignment): Modifier

    @Stable
    fun Modifier.span(span: Int): Modifier
}

internal object TableScopeInstance : TableScope {
    @Stable
    override fun Modifier.align(alignment: Alignment) = this.then(
        TableAlignModifier(
            alignment = alignment,
            inspectorInfo = debugInspectorInfo {
                name = "align"
                value = alignment
            },
        )
    )

    @Stable
    override fun Modifier.span(span: Int): Modifier {
        require(span > 0) { "invalid span $span; must be greater than zero" }
        return this.then(
            TableSpanModifier(
                span = span,
                inspectorInfo = debugInspectorInfo {
                    name = "span"
                    value = span
                },
            )
        )
    }
}

internal data class TableParentData(
    var alignment: Alignment? = null,
    var span: Int = 1,
)

internal class TableAlignModifier(
    private val alignment: Alignment,
    inspectorInfo: InspectorInfo.() -> Unit,
) : ParentDataModifier, InspectorValueInfo(inspectorInfo) {
    override fun Density.modifyParentData(parentData: Any?): TableParentData {
        return ((parentData as? TableParentData) ?: TableParentData()).also {
            it.alignment = alignment
        }
    }

    override fun equals(other: Any?): Boolean {
        if (this === other) return true
        val otherModifier = other as? TableAlignModifier ?: return false
        return alignment == otherModifier.alignment
    }

    override fun hashCode(): Int = alignment.hashCode()

    override fun toString(): String = "TableAlignModifier(alignment=$alignment)"
}

internal class TableSpanModifier(
    private val span: Int,
    inspectorInfo: InspectorInfo.() -> Unit,
) : ParentDataModifier, InspectorValueInfo(inspectorInfo) {
    override fun Density.modifyParentData(parentData: Any?): TableParentData {
        return ((parentData as? TableParentData) ?: TableParentData()).also {
            it.span = span
        }
    }

    override fun equals(other: Any?): Boolean {
        if (this === other) return true
        val otherModifier = other as? TableSpanModifier ?: return false
        return span == otherModifier.span
    }

    override fun hashCode(): Int = span.hashCode()

    override fun toString(): String = "TableSpanModifier(span=$span)"
}

sealed class TableCells {
    abstract fun getSizes(maxSize: Int): List<Int>

    class Fixed(val count: Int) : TableCells() {
        override fun getSizes(maxSize: Int): List<Int> {
            return Array(count) { maxSize / count }.toList()
        }
    }

    class Weighted(vararg val weights: Int) : TableCells() {
        override fun getSizes(maxSize: Int): List<Int> {
            val weightSum = weights.sum()

            val weightUnit = maxSize / weightSum

            return weights.map { weight -> weightUnit * weight }
        }

        fun getSizes(maxWidth: Int, currentSizes: List<Int>): List<Int> {
            val weightSum = weights.sum()
            val wrapWidth =
                if (currentSizes.isNotEmpty()) weights.mapIndexed { index, weight ->
                    currentSizes[index].takeIf { weight == 0 } ?: 0
                }.sum()
                else 0

            val weightUnit =
                if (weightSum != 0) ((maxWidth - wrapWidth) / weightSum).coerceAtLeast(0)
                else 0

            return weights.mapIndexed { index, weight ->
                if (weight == 0) currentSizes.getOrNull(index) ?: maxWidth
                else {
                    weightUnit * weight
                }
            }
        }
    }
}

internal fun mapTable(
    maxAxis: Int,
    measurables: List<Measurable>,
    direction: TableDirection,
): MutableList<List<Measurable>> {
    val table = mutableListOf<List<Measurable>>()

    if (direction == TableDirection.VERTICAL) {
        var row = mutableListOf<Measurable>()
        var columnIndex = 0

        measurables.forEach {
            val span = it.data?.span?.coerceAtMost(maxAxis) ?: 1

            val remainingSpan = maxAxis - columnIndex

            if (remainingSpan < span && row.isNotEmpty()) {
                table += row

                row = mutableListOf()

                columnIndex = 0
            }

            row += it

            columnIndex += span

            if (columnIndex >= maxAxis) {
                columnIndex = 0

                table += row
                row = mutableListOf()
            }
        }

        table += row
    } else {
        var column = mutableListOf<Measurable>()
        var rowIndex = 0

        measurables.forEach {
            val span = it.data?.span?.coerceAtMost(maxAxis) ?: 1

            val remainingSpan = maxAxis - rowIndex

            if (remainingSpan < span && column.isNotEmpty()) {
                table += column

                column = mutableListOf()

                rowIndex = 0
            }

            column += it

            rowIndex += span

            if (rowIndex >= maxAxis) {
                rowIndex = 0

                table += column
                column = mutableListOf()
            }
        }

        if (column.isNotEmpty()) table += column
    }

    return table
}

internal val IntrinsicMeasurable.data: TableParentData?
    get() = parentData as? TableParentData

internal fun calculateOffset(
    alignment: BiasAlignment,
    columnX: Int,
    rowY: Int,
    columnWidth: Int,
    rowHeight: Int,
    width: Int,
    height: Int,
): IntOffset {
    return Offset(
        x = (columnX + (columnWidth - width) * (alignment.horizontalBias + 1) / 2),
        y = (rowY + (rowHeight - height) * (alignment.verticalBias + 1) / 2),
    ).round()
}

enum class TableDirection {
    HORIZONTAL,
    VERTICAL,
    ;
}

internal fun horizontalTableMeasurePolicy(
    cells: TableCells,
    rowHeightsState: MutableState<List<Int>>,
): MeasurePolicy {
    return MeasurePolicy { measurables, constraints ->
        //region Parameter Check
        check(constraints.hasBoundedHeight) {
            "Unbounded height not supported"
        }

        val rowCount = when (cells) {
            is TableCells.Fixed -> cells.count
            is TableCells.Weighted -> cells.weights.size
        }

        check(rowCount > 0) {
            "Rows count must be greater than Zero"
        }

        var hasWrapWeights = false

        if (cells is TableCells.Weighted) {
            cells.weights.forEach {
                if (it == 0) hasWrapWeights = true

                check(it >= 0) {
                    "Column Weight must not be below Zero"
                }
            }
        }
        //endregion

        //region Rows & Columns
        val table = mapTable(rowCount, measurables, TableDirection.HORIZONTAL)

        val currentRowHeights = rowHeightsState.value.toMutableList()

        val columns = mutableListOf<Int>()
        val rows = when (cells) {
            is TableCells.Fixed -> cells.getSizes(constraints.maxHeight)
            is TableCells.Weighted ->
                if (hasWrapWeights) cells.getSizes(
                    constraints.maxHeight,
                    currentRowHeights,
                )
                else cells.getSizes(constraints.maxHeight)
        }
        //endregion

        //region Weighted Calculation
        if (hasWrapWeights && currentRowHeights.isEmpty()) {
            val rowHeights = Array(rowCount) { 0 }

            table.forEachIndexed { columnIndex, list ->
                var rowIndex = 0

                list.forEachIndexed { index, measurable ->
                    val span = table[columnIndex][index].data?.span?.coerceAtMost(rowCount) ?: 1
                    val maxWidth = columns[rowIndex] * span

                    val itemConstraints = constraints.copy(
                        minWidth = 0,
                        maxWidth = maxWidth,
                    )

                    val placeable = measurable.measure(itemConstraints)

                    if (span == 1) {
                        rowHeights[rowIndex] = rowHeights[rowIndex]
                            .coerceAtLeast(
                                placeable.width
                            )
                            .coerceAtMost(
                                maxWidth - rowHeights.copyOfRange(0, rowIndex).sum()
                            )
                    }

                    rowIndex += span
                }
            }

            rowHeightsState.value = rowHeights.toList()

            return@MeasurePolicy layout(
                width = 0,
                height = constraints.maxHeight,
            ) {

            }
        }
        //endregion

        val placeables = table.mapIndexed { columnIndex, list ->
            var rowIndex = 0
            var columnWidth = 0

            val column = list.mapIndexed { index, measurable ->
                val span = table[columnIndex][index].data?.span?.coerceAtMost(rowCount) ?: 1

                val maxHeight =
                    if (span == 0) currentRowHeights[rowIndex]
                    else rows.subList(
                        rowIndex,
                        (rowIndex + span).coerceAtMost(rows.size),
                    ).sum()

                val itemConstraints = constraints.copy(
                    minHeight = 0,
                    maxHeight = maxHeight,
                )

                rowIndex += span

                measurable.measure(itemConstraints).also { placeable ->
                    columnWidth = columnWidth.coerceAtLeast(placeable.width)
                }
            }

            columns += columnWidth

            column
        }

        val width = columns.sum()

        layout(
            width = width,
            height = constraints.maxHeight,
        ) {
            var rowY: Int
            var columnX = 0

            placeables.forEachIndexed { columnIndex, column ->
                var rowIndex = 0
                rowY = 0
                val columnWidth = columns[columnIndex]

                column.forEachIndexed { index, placeable ->
                    val data = table[columnIndex][index].data

                    val rowHeight = rows.subList(
                        rowIndex,
                        (rowIndex + (data?.span ?: 1)).coerceAtMost(rows.size),
                    ).sum()

                    val offset = calculateOffset(
                        (data?.alignment ?: Alignment.TopStart) as BiasAlignment,
                        columnX,
                        rowY,
                        columnWidth,
                        rowHeight,
                        placeable.width,
                        placeable.height,
                    )

                    placeable.placeRelative(offset)

                    rowY += rowHeight

                    rowIndex += (data?.span ?: 1)
                }

                columnX += columnWidth
            }
        }
    }
}


internal fun verticalTableMeasurePolicy(
    cells: TableCells,
    columnWidthsState: MutableState<List<Int>>,
): MeasurePolicy {
    return MeasurePolicy { measurables, constraints ->
        //region Parameter Check
        check(constraints.hasBoundedWidth) {
            "Unbounded width not supported"
        }

        val columnCount = when (cells) {
            is TableCells.Fixed -> cells.count
            is TableCells.Weighted -> cells.weights.size
        }

        check(columnCount > 0) {
            "Columns count must be greater than Zero"
        }

        var hasWrapWeights = false

        if (cells is TableCells.Weighted) {
            cells.weights.forEach {
                if (it == 0) hasWrapWeights = true

                check(it >= 0) {
                    "Column Weight must not be below Zero"
                }
            }
        }
        //endregion

        //region Rows & Columns
        val table = mapTable(columnCount, measurables, TableDirection.VERTICAL)

        val currentColumnWidths = columnWidthsState.value.toMutableList()

        val rows = mutableListOf<Int>()
        val columns = when (cells) {
            is TableCells.Fixed -> cells.getSizes(constraints.maxWidth)
            is TableCells.Weighted ->
                if (hasWrapWeights) cells.getSizes(
                    constraints.maxWidth,
                    currentColumnWidths,
                )
                else cells.getSizes(constraints.maxWidth)
        }
        //endregion

        //region Weighted Calculation
        if (hasWrapWeights && currentColumnWidths.isEmpty()) {
            val columnWidths = Array(columnCount) { 0 }

            table.forEachIndexed { rowIndex, list ->
                var columnIndex = 0

                list.forEachIndexed { index, measurable ->
                    val span = table[rowIndex][index].data?.span?.coerceAtMost(columnCount) ?: 1
                    val maxWidth = columns[columnIndex] * span

                    val itemConstraints = constraints.copy(
                        minWidth = 0,
                        maxWidth = maxWidth,
                    )

                    val placeable = measurable.measure(itemConstraints)

                    if (span == 1) {
                        columnWidths[columnIndex] = columnWidths[columnIndex]
                            .coerceAtLeast(
                                placeable.width
                            )
                            .coerceAtMost(
                                maxWidth - columnWidths.copyOfRange(0, columnIndex).sum()
                            )
                    }

                    columnIndex += span
                }
            }

            columnWidthsState.value = columnWidths.toList()

            return@MeasurePolicy layout(
                width = constraints.maxWidth,
                height = 0,
            ) {

            }
        }
        //endregion

        val placeables = table.mapIndexed { rowIndex, list ->
            var columnIndex = 0
            var rowHeight = 0

            val row = list.mapIndexed { index, measurable ->
                val span = table[rowIndex][index].data?.span?.coerceAtMost(columnCount) ?: 1

                val maxWidth =
                    if (span == 0) currentColumnWidths[columnIndex]
                    else columns.subList(
                        columnIndex,
                        (columnIndex + span).coerceAtMost(columns.size),
                    ).sum()

                val itemConstraints = constraints.copy(
                    minWidth = 0,
                    maxWidth = maxWidth,
                )

                columnIndex += span

                measurable.measure(itemConstraints).also { placeable ->
                    rowHeight = rowHeight.coerceAtLeast(placeable.height)
                }
            }

            rows += rowHeight

            row
        }

        val height = rows.sum()

        layout(
            width = constraints.maxWidth,
            height = height,
        ) {
            var columnX: Int
            var rowY = 0

            placeables.forEachIndexed { rowIndex, row ->
                var columnIndex = 0
                columnX = 0
                val rowHeight = rows[rowIndex]

                row.forEachIndexed { index, placeable ->
                    val data = table[rowIndex][index].data

                    val columnWidth = columns.subList(
                        columnIndex,
                        (columnIndex + (data?.span ?: 1)).coerceAtMost(columns.size),
                    ).sum()

                    val offset = calculateOffset(
                        (data?.alignment ?: Alignment.TopStart) as BiasAlignment,
                        columnX,
                        rowY,
                        columnWidth,
                        rowHeight,
                        placeable.width,
                        placeable.height,
                    )

                    placeable.placeRelative(offset)

                    columnX += columnWidth

                    columnIndex += (data?.span ?: 1)
                }

                rowY += rowHeight
            }
        }
    }
}