Fix staggered grid measure when scrolled over limit

The staggered grid wasn't positioning items correctly when scroll delta was larger max content height and items on both lanes were positioned at the same offset. Staggered grid also chose the first min offset column, and didn't consider logical order of items with the same offset.

Fixes: 322734849
Test: Added a test in LazyStaggeredGridTest
(cherry picked from https://android-review.googlesource.com/q/commit:0d644fa2e6c5fff27ab3c02e692a8d3dc00ff9a0)
Merged-In: Ib6ac2bd43a87baba0feb541bfa125e9b6b9b38a8
Change-Id: Ib6ac2bd43a87baba0feb541bfa125e9b6b9b38a8
diff --git a/compose/foundation/foundation/src/androidInstrumentedTest/kotlin/androidx/compose/foundation/lazy/staggeredgrid/LazyStaggeredGridTest.kt b/compose/foundation/foundation/src/androidInstrumentedTest/kotlin/androidx/compose/foundation/lazy/staggeredgrid/LazyStaggeredGridTest.kt
index 3dd5ec6..a376481 100644
--- a/compose/foundation/foundation/src/androidInstrumentedTest/kotlin/androidx/compose/foundation/lazy/staggeredgrid/LazyStaggeredGridTest.kt
+++ b/compose/foundation/foundation/src/androidInstrumentedTest/kotlin/androidx/compose/foundation/lazy/staggeredgrid/LazyStaggeredGridTest.kt
@@ -2221,4 +2221,53 @@
             .assertMainAxisStartPositionInRootIsEqualTo(0.dp)
             .assertMainAxisSizeIsEqualTo(0.dp)
     }
+
+    @Test
+    fun itemsAreDistributedCorrectlyOnOverscrollPassWithSameOffset() {
+        val gridHeight = itemSizeDp * 11 // two big items + two small items
+        state = LazyStaggeredGridState()
+        rule.setContent {
+            LazyStaggeredGrid(
+                modifier = Modifier
+                    .mainAxisSize(gridHeight)
+                    .crossAxisSize(itemSizeDp * 2),
+                state = state,
+                lanes = 2,
+            ) {
+                items(20) {
+                    Spacer(
+                        Modifier
+                            .mainAxisSize(if (it % 2 == 0) itemSizeDp * 5 else itemSizeDp * 0.5f)
+                            .border(1.dp, Color.Red)
+                            .testTag("$it")
+                    )
+                }
+            }
+        }
+
+        // scroll to bottom
+        state.scrollBy(gridHeight * 2)
+
+        rule.onNodeWithTag("12")
+            .assertCrossAxisStartPositionInRootIsEqualTo(0.dp)
+            .assertMainAxisStartPositionInRootIsEqualTo(0.dp)
+
+        rule.onNodeWithTag("13")
+            .assertCrossAxisStartPositionInRootIsEqualTo(itemSizeDp)
+            .assertMainAxisStartPositionInRootIsEqualTo(0.dp)
+
+        // scroll a back a bit
+        state.scrollBy(-itemSizeDp * 5)
+
+        // scroll by a grid height
+        state.scrollBy(gridHeight)
+
+        rule.onNodeWithTag("12")
+            .assertCrossAxisStartPositionInRootIsEqualTo(0.dp)
+            .assertMainAxisStartPositionInRootIsEqualTo(0.dp)
+
+        rule.onNodeWithTag("13")
+            .assertCrossAxisStartPositionInRootIsEqualTo(itemSizeDp)
+            .assertMainAxisStartPositionInRootIsEqualTo(0.dp)
+    }
 }
diff --git a/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/staggeredgrid/LazyStaggeredGridMeasure.kt b/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/staggeredgrid/LazyStaggeredGridMeasure.kt
index e810b3b..c0c90ae 100644
--- a/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/staggeredgrid/LazyStaggeredGridMeasure.kt
+++ b/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/staggeredgrid/LazyStaggeredGridMeasure.kt
@@ -574,6 +574,7 @@
             firstItemIndices[laneIndex] = laneItems.firstOrNull()?.index ?: Unset
         }
 
+        // ensure no spacing for the last item
         if (currentItemIndices.any { it == itemCount - 1 }) {
             currentItemOffsets.offsetBy(-mainAxisSpacing)
         }
@@ -602,13 +603,21 @@
                 // Note that it is different from initial pass up where we selected largest index
                 // instead. The reason is that we already distributed items on downward pass and
                 // gap would be incorrect if those are moved.
-                val laneIndex = firstItemOffsets.indexOfMinValue()
+                var laneIndex = firstItemOffsets.indexOfMinValue()
+                val nextLaneIndex = firstItemIndices.indexOfMaxValue()
 
-                if (laneIndex != firstItemIndices.indexOfMaxValue()) {
-                    // If min offset lane doesn't have largest value, it means items are misaligned.
-                    // The correct thing here is to restart measure. We will measure up to the end
-                    // and restart measure from there after this pass.
-                    gapDetected = true
+                if (laneIndex != nextLaneIndex) {
+                    if (firstItemOffsets[laneIndex] == firstItemOffsets[nextLaneIndex]) {
+                        // If the offsets are the same, it means that there's no gap here.
+                        // In this case, we should choose the lane where item would go normally.
+                        laneIndex = nextLaneIndex
+                    } else {
+                        // If min offset lane doesn't have largest value, it means items are
+                        // misaligned.
+                        // The correct thing here is to restart measure. We will measure up to the
+                        // end and restart measure from there after this pass.
+                        gapDetected = true
+                    }
                 }
 
                 val currentIndex =