Merge "Add ability for SpannedGrids to provide row and column for each item" into main

4 files changed, 106 insertions, 78 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/grid/ui/compose/SpannedGrids.kt b/packages/SystemUI/src/com/android/systemui/grid/ui/compose/SpannedGrids.kt
index 96ef03c996a9..e06c228e24b2 100644
--- a/packages/SystemUI/src/com/android/systemui/grid/ui/compose/SpannedGrids.kt
+++ b/packages/SystemUI/src/com/android/systemui/grid/ui/compose/SpannedGrids.kt
@@ -13,9 +13,9 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-
 package com.android.systemui.grid.ui.compose
 
+import androidx.collection.IntIntPair
 import androidx.compose.foundation.layout.Arrangement
 import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.BoxScope
@@ -24,7 +24,6 @@ import androidx.compose.runtime.key
 import androidx.compose.runtime.remember
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.layout.Layout
-import androidx.compose.ui.layout.Placeable
 import androidx.compose.ui.semantics.CollectionInfo
 import androidx.compose.ui.semantics.CollectionItemInfo
 import androidx.compose.ui.semantics.collectionInfo
@@ -34,6 +33,8 @@ import androidx.compose.ui.unit.Constraints
 import androidx.compose.ui.unit.Dp
 import androidx.compose.ui.unit.LayoutDirection
 import androidx.compose.ui.unit.dp
+import androidx.compose.ui.util.fastForEachIndexed
+import androidx.compose.ui.util.fastMapIndexed
 import kotlin.math.max
 
 /**
@@ -65,7 +66,11 @@ fun HorizontalSpannedGrid(
     spans: List<Int>,
     modifier: Modifier = Modifier,
     keys: (spanIndex: Int) -> Any = { it },
-    composables: @Composable BoxScope.(spanIndex: Int) -> Unit,
+    composables:
+        @Composable
+        BoxScope.(
+            spanIndex: Int, row: Int, isFirstInColumn: Boolean, isLastInColumn: Boolean,
+        ) -> Unit,
 ) {
     SpannedGrid(
         primarySpaces = rows,
@@ -80,7 +85,7 @@ fun HorizontalSpannedGrid(
 }
 
 /**
- * Horizontal (non lazy) grid that supports [spans] for its elements.
+ * Vertical (non lazy) grid that supports [spans] for its elements.
  *
  * The elements will be laid down horizontally first, and then by rows. So assuming LTR layout, it
  * will be (for a span list `[2, 1, 2, 1, 1, 1, 1, 1]` and 4 columns):
@@ -107,7 +112,9 @@ fun VerticalSpannedGrid(
     spans: List<Int>,
     modifier: Modifier = Modifier,
     keys: (spanIndex: Int) -> Any = { it },
-    composables: @Composable BoxScope.(spanIndex: Int) -> Unit,
+    composables:
+        @Composable
+        BoxScope.(spanIndex: Int, column: Int, isFirstInRow: Boolean, isLastInRow: Boolean) -> Unit,
 ) {
     SpannedGrid(
         primarySpaces = columns,
@@ -130,7 +137,9 @@ private fun SpannedGrid(
     isVertical: Boolean,
     modifier: Modifier = Modifier,
     keys: (spanIndex: Int) -> Any = { it },
-    composables: @Composable BoxScope.(spanIndex: Int) -> Unit,
+    composables:
+        @Composable
+        BoxScope.(spanIndex: Int, secondaryAxis: Int, isFirst: Boolean, isLast: Boolean) -> Unit,
 ) {
     val crossAxisArrangement = Arrangement.spacedBy(crossAxisSpacing)
     spans.forEachIndexed { index, span ->
@@ -139,7 +148,6 @@ private fun SpannedGrid(
                 "expected rance of [1, $primarySpaces]"
         }
     }
-
     if (isVertical) {
         check(crossAxisSpacing >= 0.dp) { "Negative columnSpacing $crossAxisSpacing" }
         check(mainAxisSpacing >= 0.dp) { "Negative rowSpacing $mainAxisSpacing" }
@@ -147,29 +155,30 @@ private fun SpannedGrid(
         check(mainAxisSpacing >= 0.dp) { "Negative columnSpacing $mainAxisSpacing" }
         check(crossAxisSpacing >= 0.dp) { "Negative rowSpacing $crossAxisSpacing" }
     }
-
-    val totalMainAxisGroups: Int =
+    // List of primary axis index to secondary axis index
+    // This is keyed to the size of the spans list for performance reasons as we don't expect the
+    // spans value to change outside of edit mode.
+    val positions = remember(spans.size) { Array(spans.size) { IntIntPair(0, 0) } }
+    val totalMainAxisGroups =
         remember(primarySpaces, spans) {
-            var currentAccumulated = 0
-            var groups = 1
-            spans.forEach { span ->
-                if (currentAccumulated + span <= primarySpaces) {
-                    currentAccumulated += span
-                } else {
-                    groups += 1
-                    currentAccumulated = span
+            var mainAxisGroup = 0
+            var currentSlot = 0
+            spans.fastForEachIndexed { index, span ->
+                if (currentSlot + span > primarySpaces) {
+                    currentSlot = 0
+                    mainAxisGroup += 1
                 }
+                positions[index] = IntIntPair(mainAxisGroup, currentSlot)
+                currentSlot += span
             }
-            groups
+            mainAxisGroup + 1
         }
-
     val slotPositionsAndSizesCache = remember {
         object {
             var sizes = IntArray(0)
             var positions = IntArray(0)
         }
     }
-
     Layout(
         {
             (0 until spans.size).map { spanIndex ->
@@ -184,7 +193,13 @@ private fun SpannedGrid(
                                 }
                         }
                     ) {
-                        composables(spanIndex)
+                        val position = positions[spanIndex]
+                        composables(
+                            spanIndex,
+                            position.second,
+                            position.second == 0,
+                            positions.getOrNull(spanIndex + 1)?.first != position.first,
+                        )
                     }
                 }
             }
@@ -205,7 +220,6 @@ private fun SpannedGrid(
             slotPositionsAndSizesCache.sizes,
         )
         val cellSizesInCrossAxis = slotPositionsAndSizesCache.sizes
-
         // with is needed because of the double receiver (Density, Arrangement).
         with(crossAxisArrangement) {
             arrange(
@@ -216,68 +230,73 @@ private fun SpannedGrid(
             )
         }
         val startPositions = slotPositionsAndSizesCache.positions
-
         val mainAxisSpacingPx = mainAxisSpacing.roundToPx()
         val mainAxisTotalGaps = (totalMainAxisGroups - 1) * mainAxisSpacingPx
-        val mainAxisSize = if (isVertical) constraints.maxHeight else constraints.maxWidth
+        val mainAxisMaxSize = if (isVertical) constraints.maxHeight else constraints.maxWidth
         val mainAxisElementConstraint =
-            if (mainAxisSize == Constraints.Infinity) {
+            if (mainAxisMaxSize == Constraints.Infinity) {
                 Constraints.Infinity
             } else {
-                max(0, (mainAxisSize - mainAxisTotalGaps) / totalMainAxisGroups)
+                max(0, (mainAxisMaxSize - mainAxisTotalGaps) / totalMainAxisGroups)
             }
 
-        val mainAxisSizes = IntArray(totalMainAxisGroups) { 0 }
-
-        var currentSlot = 0
-        var mainAxisGroup = 0
+        var mainAxisTotalSize = mainAxisTotalGaps
+        var currentMainAxis = 0
+        var currentMainAxisMax = 0
         val placeables =
-            measurables.mapIndexed { index, measurable ->
+            measurables.fastMapIndexed { index, measurable ->
                 val span = spans[index]
-                if (currentSlot + span > primarySpaces) {
-                    currentSlot = 0
-                    mainAxisGroup += 1
-                }
+                val position = positions[index]
                 val crossAxisConstraint =
-                    calculateWidth(cellSizesInCrossAxis, startPositions, currentSlot, span)
-                PlaceResult(
-                        measurable.measure(
-                            makeConstraint(
-                                isVertical,
-                                mainAxisElementConstraint,
-                                crossAxisConstraint,
-                            )
-                        ),
-                        currentSlot,
-                        mainAxisGroup,
+                    calculateWidth(cellSizesInCrossAxis, startPositions, position.second, span)
+
+                measurable
+                    .measure(
+                        makeConstraint(isVertical, mainAxisElementConstraint, crossAxisConstraint)
                     )
                     .also {
-                        currentSlot += span
-                        mainAxisSizes[mainAxisGroup] =
-                            max(
-                                mainAxisSizes[mainAxisGroup],
-                                if (isVertical) it.placeable.height else it.placeable.width,
-                            )
+                        val placeableSize = if (isVertical) it.height else it.width
+                        if (position.first != currentMainAxis) {
+                            // New row -- Add the max size to the total and reset the max
+                            mainAxisTotalSize += currentMainAxisMax
+                            currentMainAxisMax = placeableSize
+                            currentMainAxis = position.first
+                        } else {
+                            currentMainAxisMax = max(currentMainAxisMax, placeableSize)
+                        }
                     }
             }
+        mainAxisTotalSize += currentMainAxisMax
 
-        val mainAxisTotalSize = mainAxisTotalGaps + mainAxisSizes.sum()
-        val mainAxisStartingPoints =
-            mainAxisSizes.runningFold(0) { acc, value -> acc + value + mainAxisSpacingPx }
         val height = if (isVertical) mainAxisTotalSize else crossAxisSize
         val width = if (isVertical) crossAxisSize else mainAxisTotalSize
 
         layout(width, height) {
-            placeables.forEach { (placeable, slot, mainAxisGroup) ->
+            var previousMainAxis = 0
+            var currentMainAxisPosition = 0
+            var currentMainAxisMax = 0
+            placeables.forEachIndexed { index, placeable ->
+                val slot = positions[index].second
+                val mainAxisSize = if (isVertical) placeable.height else placeable.width
+
+                if (positions[index].first != previousMainAxis) {
+                    // Move up a row + padding
+                    currentMainAxisPosition += currentMainAxisMax + mainAxisSpacingPx
+                    currentMainAxisMax = mainAxisSize
+                    previousMainAxis = positions[index].first
+                } else {
+                    currentMainAxisMax = max(currentMainAxisMax, mainAxisSize)
+                }
+
                 val x =
                     if (isVertical) {
                         startPositions[slot]
                     } else {
-                        mainAxisStartingPoints[mainAxisGroup]
+                        currentMainAxisPosition
                     }
                 val y =
                     if (isVertical) {
-                        mainAxisStartingPoints[mainAxisGroup]
+                        currentMainAxisPosition
                     } else {
                         startPositions[slot]
                     }
@@ -321,9 +340,3 @@ private fun calculateCellsCrossAxisSize(
         outArray[index] = slotSize + if (index < remainingPixels) 1 else 0
     }
 }
-
-private data class PlaceResult(
-    val placeable: Placeable,
-    val slotIndex: Int,
-    val mainAxisGroup: Int,
-)
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/BounceableInfo.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/BounceableInfo.kt
index eb6f97942ec0..b80e6b403cf3 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/BounceableInfo.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/BounceableInfo.kt
@@ -35,11 +35,16 @@ fun List<BounceableTileViewModel>.bounceableInfo(
     index: Int,
     column: Int,
     columns: Int,
+    isFirstInRow: Boolean,
+    isLastInRow: Boolean,
 ): BounceableInfo {
-    // Only look for neighbor bounceables if they are on the same row
+    // A tile may be the last in the row without being on the last column
     val onLastColumn = sizedTile.onLastColumn(column, columns)
-    val previousTile = getOrNull(index - 1)?.takeIf { column != 0 }
-    val nextTile = getOrNull(index + 1)?.takeIf { !onLastColumn }
+
+    // Only look for neighbor bounceables if they are on the same row
+    val previousTile = getOrNull(index - 1)?.takeIf { !isFirstInRow }
+    val nextTile = getOrNull(index + 1)?.takeIf { !isLastInRow }
+
     return BounceableInfo(this[index], previousTile, nextTile, !onLastColumn)
 }
 
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/QuickQuickSettings.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/QuickQuickSettings.kt
index 495870f0a978..cdc03bb9be35 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/QuickQuickSettings.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/QuickQuickSettings.kt
@@ -57,7 +57,6 @@ fun ContentScope.QuickQuickSettings(
         onDispose { tiles.forEach { it.stopListening(token) } }
     }
     val columns = viewModel.columns
-    var cellIndex = 0
     Box(modifier = modifier) {
         GridAnchor()
         VerticalSpannedGrid(
@@ -67,17 +66,23 @@ fun ContentScope.QuickQuickSettings(
             spans = spans,
             modifier = Modifier.sysuiResTag("qqs_tile_layout"),
             keys = { sizedTiles[it].tile.spec },
-        ) { spanIndex ->
+        ) { spanIndex, column, isFirstInColumn, isLastInColumn ->
             val it = sizedTiles[spanIndex]
-            val column = cellIndex % columns
-            cellIndex += it.width
             Element(it.tile.spec.toElementKey(spanIndex), Modifier) {
                 Tile(
                     tile = it.tile,
                     iconOnly = it.isIcon,
                     squishiness = { squishiness },
                     coroutineScope = scope,
-                    bounceableInfo = bounceables.bounceableInfo(it, spanIndex, column, columns),
+                    bounceableInfo =
+                        bounceables.bounceableInfo(
+                            it,
+                            index = spanIndex,
+                            column = column,
+                            columns = columns,
+                            isFirstInRow = isFirstInColumn,
+                            isLastInRow = isLastInColumn,
+                        ),
                     tileHapticsViewModelFactoryProvider =
                         viewModel.tileHapticsViewModelFactoryProvider,
                     // There should be no QuickQuickSettings when the details view is enabled.
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/infinitegrid/InfiniteGridLayout.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/infinitegrid/InfiniteGridLayout.kt
index dfee497655d1..0503049382a8 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/infinitegrid/InfiniteGridLayout.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/infinitegrid/InfiniteGridLayout.kt
@@ -85,8 +85,6 @@ constructor(
             remember(sizedTiles) { List(sizedTiles.size) { BounceableTileViewModel() } }
         val squishiness by viewModel.squishinessViewModel.squishiness.collectAsStateWithLifecycle()
         val scope = rememberCoroutineScope()
-        var cellIndex = 0
-
         val spans by remember(sizedTiles) { derivedStateOf { sizedTiles.fastMap { it.width } } }
 
         VerticalSpannedGrid(
@@ -95,10 +93,9 @@ constructor(
             rowSpacing = dimensionResource(R.dimen.qs_tile_margin_vertical),
             spans = spans,
             keys = { sizedTiles[it].tile.spec },
-        ) { spanIndex ->
+        ) { spanIndex, column, isFirstInColumn, isLastInColumn ->
             val it = sizedTiles[spanIndex]
-            val column = cellIndex % columns
-            cellIndex += it.width
+
             Element(it.tile.spec.toElementKey(spanIndex), Modifier) {
                 Tile(
                     tile = it.tile,
@@ -106,7 +103,15 @@ constructor(
                     squishiness = { squishiness },
                     tileHapticsViewModelFactoryProvider = tileHapticsViewModelFactoryProvider,
                     coroutineScope = scope,
-                    bounceableInfo = bounceables.bounceableInfo(it, spanIndex, column, columns),
+                    bounceableInfo =
+                        bounceables.bounceableInfo(
+                            it,
+                            index = spanIndex,
+                            column = column,
+                            columns = columns,
+                            isFirstInRow = isFirstInColumn,
+                            isLastInRow = isLastInColumn,
+                        ),
                     detailsViewModel = detailsViewModel,
                 )
             }