Fix for two sticky headers positions after small scrolls

There is another corner case which is hard to handle without re-running the main sticky header logic. In order to not duplicate the logic lets just disable the relayout scroll when the first item is a sticky header.

Fixes: 313294874
Test: new test in LazyListHeadersTest
Change-Id: I61dd124a651f9820d9bd2c711d4c6628930bbca4

diff --git a/compose/foundation/foundation/src/androidInstrumentedTest/kotlin/androidx/compose/foundation/lazy/list/LazyListHeadersTest.kt b/compose/foundation/foundation/src/androidInstrumentedTest/kotlin/androidx/compose/foundation/lazy/list/LazyListHeadersTest.kt
index 881dd61..587e98b 100644
--- a/compose/foundation/foundation/src/androidInstrumentedTest/kotlin/androidx/compose/foundation/lazy/list/LazyListHeadersTest.kt
+++ b/compose/foundation/foundation/src/androidInstrumentedTest/kotlin/androidx/compose/foundation/lazy/list/LazyListHeadersTest.kt

@@ -19,6 +19,7 @@
 import androidx.compose.foundation.ExperimentalFoundationApi
 import androidx.compose.foundation.gestures.FlingBehavior
 import androidx.compose.foundation.gestures.ScrollableDefaults
+import androidx.compose.foundation.gestures.scrollBy
 import androidx.compose.foundation.layout.Arrangement
 import androidx.compose.foundation.layout.PaddingValues
 import androidx.compose.foundation.layout.Spacer
@@ -456,6 +457,48 @@
         rule.onNodeWithTag("0")
             .assertTopPositionInRootIsEqualTo(itemIndexDp * 3 / 2)
     }
+
+    @Test
+    fun lazyColumnShowsHeader_withoutBeyondBoundsItemCount2() {
+        val firstHeaderTag = "firstHeaderTag"
+        val secondHeaderTag = "secondHeaderTag"
+        val itemSizeDp = with(rule.density) { 100.toDp() }
+        val scrollDistance = 20
+        val scrollDistanceDp = with(rule.density) { scrollDistance.toDp() }
+        val state = LazyListState()
+
+        rule.setContent {
+            LazyColumn(Modifier.height(itemSizeDp * 3.5f), state) {
+                stickyHeader {
+                    Spacer(
+                        Modifier.height(itemSizeDp).fillParentMaxWidth()
+                            .testTag(firstHeaderTag)
+                    )
+                }
+                stickyHeader {
+                    Spacer(
+                        Modifier.height(itemSizeDp).fillParentMaxWidth()
+                            .testTag(secondHeaderTag)
+                    )
+                }
+
+                items(100) {
+                    Spacer(Modifier.height(itemSizeDp).fillParentMaxWidth().testTag(it.toString()))
+                }
+            }
+        }
+
+        rule.runOnIdle {
+            runBlocking { state.scrollBy(scrollDistance.toFloat()) }
+        }
+
+        rule.onNodeWithTag(firstHeaderTag)
+            .assertTopPositionInRootIsEqualTo(-scrollDistanceDp)
+        rule.onNodeWithTag(secondHeaderTag)
+            .assertTopPositionInRootIsEqualTo(itemSizeDp - scrollDistanceDp)
+        rule.onNodeWithTag("0")
+            .assertTopPositionInRootIsEqualTo(itemSizeDp * 2 - scrollDistanceDp)
+    }
 }
 
 @Composable

diff --git a/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/LazyListMeasureResult.kt b/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/LazyListMeasureResult.kt
index f484426..c01b0c9 100644
--- a/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/LazyListMeasureResult.kt
+++ b/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/LazyListMeasureResult.kt

@@ -19,9 +19,7 @@
 import androidx.compose.foundation.gestures.Orientation
 import androidx.compose.ui.layout.MeasureResult
 import androidx.compose.ui.unit.IntSize
-import androidx.compose.ui.util.fastFirst
 import androidx.compose.ui.util.fastForEach
-import androidx.compose.ui.util.fastLastOrNull
 
 /**
  * The result of the measure pass for lazy list layout.
@@ -86,8 +84,12 @@
         ) {
             return false
         }
-        val first = visibleItemsInfo.fastFirst { !it.nonScrollableItem }
-        val last = visibleItemsInfo.fastLastOrNull { !it.nonScrollableItem }!!
+        val first = visibleItemsInfo.first()
+        val last = visibleItemsInfo.last()
+        if (first.nonScrollableItem || last.nonScrollableItem) {
+            // non scrollable items like headers require special handling in the measurement.
+            return false
+        }
         val canApply = if (delta < 0) {
             // scrolling forward
             val deltaToFirstItemChange =