Repository: android/compose-samples
File: JetNews/app/src/main/java/com/example/jetnews/ui/interests/InterestsScreen.kt (Kotlin)

Code before:
            val itemConstraints = if (itemWidth == Int.MAX_VALUE) {
                outerConstraints
            } else {
                outerConstraints.copy(maxWidth = itemWidth)
            }
            val placeable = measureable.measure(itemConstraints)
            // Update the height for each row
            val row = index.floorDiv(columns)
            rowHeights[row] = max(rowHeights[row], placeable.height)
            placeable
        }

        // Calculate maxHeight of the Interests layout. Heights of the row + top padding
        val layoutHeight = topPaddingPx + rowHeights.sumOf { it }
        // Calculate maxWidth of the Interests layout
        val layoutWidth = itemWidth * columns + (itemSpacingPx * (columns - 1))
        // Lay out given the max width and height
        layout(
            outerConstraints.maxWidth.coerceIn(0, layoutWidth),
            outerConstraints.maxHeight.coerceIn(0, layoutHeight)
        ) {
            // Track the y co-ord we have placed children up to
            var yPosition = topPaddingPx
            placeables.forEachIndexed { index, placeable ->
                // If there's only one column, then place items one after the other
                if (columns == 1) {
                    // Position item on the screen
                    placeable.placeRelative(x = 0, y = yPosition)
                    // Increment the y co-ord with the item's height
                    yPosition += placeable.height
                } else {
                    // If there are multiple columns, then place items relative to each other
                    if (index % columns == 0) {
                        // Position item on a new row
                        placeable.placeRelative(x = 0, y = yPosition)
                    } else {
                        // Place element in the same row, the previous item exists at this point
                        val previousPlaceable = placeables[index - 1]
                        placeable.placeRelative(
                            x = previousPlaceable.width + itemSpacingPx,
                            y = yPosition
                        )

                        // Increment the y co-ord with the item's height
                        yPosition += rowHeights[index.floorDiv(columns)]
                    }
                }
            }
        }
    }
}

Review comment (suggestion):
If I'm reading this correctly, I think you don't have to have the special case for one column, since `index % 1` will always equal 0.

Code after:
        val topPaddingPx = topPadding.roundToPx()
        val itemSpacingPx = itemSpacing.roundToPx()
        val itemMaxWidthPx = itemMaxWidth.roundToPx()

        // Number of columns to display on the screen. This is harcoded to 2 due to
        // the design mocks, but this logic could change in the future.
        val columns = if (outerConstraints.maxWidth < multipleColumnsBreakPointPx) 1 else 2
        // Max width for each item taking into account available space, spacing and `itemMaxWidth`
        val itemWidth = if (columns == 1) {
            outerConstraints.maxWidth
        } else {
            val maxWidthWithSpaces = outerConstraints.maxWidth - (columns - 1) * itemSpacingPx
            (maxWidthWithSpaces / columns).coerceIn(0, itemMaxWidthPx)
        }
        val itemConstraints = outerConstraints.copy(maxWidth = itemWidth)

        // Keep track of the height of each row to calculate the layout's final size
        val rowHeights = IntArray(measurables.size / columns + 1)
        // Measure elements with their maximum width and keep track of the height
        val placeables = measurables.mapIndexed { index, measureable ->
            val placeable = measureable.measure(itemConstraints)
            // Update the height for each row
            val row = index.floorDiv(columns)
            rowHeights[row] = max(rowHeights[row], placeable.height)
            placeable
        }

        // Calculate maxHeight of the Interests layout. Heights of the row + top padding
        val layoutHeight = topPaddingPx + rowHeights.sum()
        // Calculate maxWidth of the Interests layout
        val layoutWidth = itemWidth * columns + (itemSpacingPx * (columns - 1))

        // Lay out given the max width and height
        layout(
            width = outerConstraints.constrainWidth(layoutWidth),
            height = outerConstraints.constrainHeight(layoutHeight)
        ) {
            // Track the y co-ord we have placed children up to
            var yPosition = topPaddingPx
            // Split placeables in lists that don't exceed the number of columns
            // and place them taking into account their width and spacing
            placeables.chunked(columns).forEachIndexed { rowIndex, row ->
                var xPosition = 0
                row.forEach { placeable ->
                    placeable.placeRelative(x = xPosition, y = yPosition)
                    xPosition += placeable.width + itemSpacingPx
                }
                yPosition += rowHeights[rowIndex]
            }
        }
    }