Add support for nested prefetching to LazyLayout/LazyList

This patch adds support for granular prefetching (precomposing) of items in nested LazyLayouts instead of composing them in premeasure of that item in the parent LazyLayout.

Performance results:
```
Before
======
frameCount                     min 506.0,   median 515.0,   max 537.0
gfxFrameJankPercent            min  29.7,   median  33.4,   max  39.7
gfxFrameTime50thPercentileMs   min  27.0,   median  28.0,   max  28.0
gfxFrameTime90thPercentileMs   min  40.0,   median  42.0,   max  42.0
gfxFrameTime95thPercentileMs   min  42.0,   median  44.0,   max  44.0
gfxFrameTime99thPercentileMs   min  46.0,   median  48.0,   max  57.0
gfxFrameTotalCount             min 505.0,   median 508.0,   max 539.0
frameDurationCpuMs             P50   19.2,   P90   31.9,   P95   34.3,   P99   38.8
frameOverrunMs                 P50    8.3,   P90   24.4,   P95   25.7,   P99   30.4

After
=====
frameCount                     min 545.0,   median 565.0,   max 574.0
gfxFrameJankPercent            min  16.5,   median  31.5,   max  44.6
gfxFrameTime50thPercentileMs   min  27.0,   median  28.0,   max  30.0
gfxFrameTime90thPercentileMs   min  36.0,   median  38.0,   max  40.0
gfxFrameTime95thPercentileMs   min  40.0,   median  42.0,   max  44.0
gfxFrameTime99thPercentileMs   min  46.0,   median  48.0,   max  69.0
gfxFrameTotalCount             min 547.0,   median 567.0,   max 576.0
frameDurationCpuMs             P50   19.5,   P90   31.1,   P95   32.4,   P99   37.5
frameOverrunMs                 P50    8.2,   P90   24.5,   P95   25.9,   P99   35.0
```

Test: - Unit tests
- Traces for LazyRows in LazyColumn to verify granular nested prefetch behavior

Relnote: Added support for prefetching items in nested LazyLists (e.g. a LazyColumn that renders nested LazyRows). This change is expected to reduce frame drops during scrolling for these LazyLists. The implementation default is to prefetch the first 2 nested items, however this behavior can be controlled by the new `LazyLayoutPrefetchStrategy(nestedPrefetchItemCount)` and `LazyListPrefetchStrategy#onNestedPrefetch` APIs.

Bug: 333430486

Change-Id: I519526a694d8e9a89a1a040cd179d0416fa2d6d9
diff --git a/compose/foundation/foundation/api/current.txt b/compose/foundation/foundation/api/current.txt
index e82b0af..9a45188 100644
--- a/compose/foundation/foundation/api/current.txt
+++ b/compose/foundation/foundation/api/current.txt
@@ -842,11 +842,16 @@
 
   @SuppressCompatibility @androidx.compose.foundation.ExperimentalFoundationApi public interface LazyListPrefetchStrategy {
     method public default androidx.compose.foundation.lazy.layout.PrefetchExecutor? getPrefetchExecutor();
+    method public void onNestedPrefetch(androidx.compose.foundation.lazy.layout.NestedPrefetchScope, int firstVisibleItemIndex);
     method public void onScroll(androidx.compose.foundation.lazy.LazyListPrefetchScope, float delta, androidx.compose.foundation.lazy.LazyListLayoutInfo layoutInfo);
     method public void onVisibleItemsUpdated(androidx.compose.foundation.lazy.LazyListPrefetchScope, androidx.compose.foundation.lazy.LazyListLayoutInfo layoutInfo);
     property public default androidx.compose.foundation.lazy.layout.PrefetchExecutor? prefetchExecutor;
   }
 
+  public final class LazyListPrefetchStrategyKt {
+    method @SuppressCompatibility @androidx.compose.foundation.ExperimentalFoundationApi public static androidx.compose.foundation.lazy.LazyListPrefetchStrategy LazyListPrefetchStrategy(optional int nestedPrefetchItemCount);
+  }
+
   @androidx.compose.foundation.lazy.LazyScopeMarker @kotlin.jvm.JvmDefaultWithCompatibility public interface LazyListScope {
     method public default void item(optional Object? key, optional Object? contentType, kotlin.jvm.functions.Function1<? super androidx.compose.foundation.lazy.LazyItemScope,kotlin.Unit> content);
     method @Deprecated public void item(optional Object? key, kotlin.jvm.functions.Function1<? super androidx.compose.foundation.lazy.LazyItemScope,kotlin.Unit> content);
@@ -1113,8 +1118,8 @@
   }
 
   @SuppressCompatibility @androidx.compose.foundation.ExperimentalFoundationApi @androidx.compose.runtime.Stable public final class LazyLayoutPrefetchState {
-    ctor public LazyLayoutPrefetchState(optional androidx.compose.foundation.lazy.layout.PrefetchExecutor? prefetchExecutor);
-    method public androidx.compose.foundation.lazy.layout.LazyLayoutPrefetchState.PrefetchHandle schedulePrefetch(int index, long constraints);
+    ctor public LazyLayoutPrefetchState(optional androidx.compose.foundation.lazy.layout.PrefetchExecutor? prefetchExecutor, optional kotlin.jvm.functions.Function1<? super androidx.compose.foundation.lazy.layout.NestedPrefetchScope,kotlin.Unit>? onNestedPrefetch);
+    method public androidx.compose.foundation.lazy.layout.LazyLayoutPrefetchState.PrefetchHandle schedulePrefetch(int index, optional androidx.compose.ui.unit.Constraints? constraints);
   }
 
   public static sealed interface LazyLayoutPrefetchState.PrefetchHandle {
@@ -1134,6 +1139,10 @@
     property public int size;
   }
 
+  @SuppressCompatibility @androidx.compose.foundation.ExperimentalFoundationApi public sealed interface NestedPrefetchScope {
+    method public void schedulePrefetch(int index, optional androidx.compose.ui.unit.Constraints? constraints);
+  }
+
   @SuppressCompatibility @androidx.compose.foundation.ExperimentalFoundationApi public interface PrefetchExecutor {
     method public void requestPrefetch(androidx.compose.foundation.lazy.layout.PrefetchRequest prefetchRequest);
   }
diff --git a/compose/foundation/foundation/api/restricted_current.txt b/compose/foundation/foundation/api/restricted_current.txt
index 5841c91..d5f80a9 100644
--- a/compose/foundation/foundation/api/restricted_current.txt
+++ b/compose/foundation/foundation/api/restricted_current.txt
@@ -844,11 +844,16 @@
 
   @SuppressCompatibility @androidx.compose.foundation.ExperimentalFoundationApi public interface LazyListPrefetchStrategy {
     method public default androidx.compose.foundation.lazy.layout.PrefetchExecutor? getPrefetchExecutor();
+    method public void onNestedPrefetch(androidx.compose.foundation.lazy.layout.NestedPrefetchScope, int firstVisibleItemIndex);
     method public void onScroll(androidx.compose.foundation.lazy.LazyListPrefetchScope, float delta, androidx.compose.foundation.lazy.LazyListLayoutInfo layoutInfo);
     method public void onVisibleItemsUpdated(androidx.compose.foundation.lazy.LazyListPrefetchScope, androidx.compose.foundation.lazy.LazyListLayoutInfo layoutInfo);
     property public default androidx.compose.foundation.lazy.layout.PrefetchExecutor? prefetchExecutor;
   }
 
+  public final class LazyListPrefetchStrategyKt {
+    method @SuppressCompatibility @androidx.compose.foundation.ExperimentalFoundationApi public static androidx.compose.foundation.lazy.LazyListPrefetchStrategy LazyListPrefetchStrategy(optional int nestedPrefetchItemCount);
+  }
+
   @androidx.compose.foundation.lazy.LazyScopeMarker @kotlin.jvm.JvmDefaultWithCompatibility public interface LazyListScope {
     method public default void item(optional Object? key, optional Object? contentType, kotlin.jvm.functions.Function1<? super androidx.compose.foundation.lazy.LazyItemScope,kotlin.Unit> content);
     method @Deprecated public void item(optional Object? key, kotlin.jvm.functions.Function1<? super androidx.compose.foundation.lazy.LazyItemScope,kotlin.Unit> content);
@@ -1115,8 +1120,8 @@
   }
 
   @SuppressCompatibility @androidx.compose.foundation.ExperimentalFoundationApi @androidx.compose.runtime.Stable public final class LazyLayoutPrefetchState {
-    ctor public LazyLayoutPrefetchState(optional androidx.compose.foundation.lazy.layout.PrefetchExecutor? prefetchExecutor);
-    method public androidx.compose.foundation.lazy.layout.LazyLayoutPrefetchState.PrefetchHandle schedulePrefetch(int index, long constraints);
+    ctor public LazyLayoutPrefetchState(optional androidx.compose.foundation.lazy.layout.PrefetchExecutor? prefetchExecutor, optional kotlin.jvm.functions.Function1<? super androidx.compose.foundation.lazy.layout.NestedPrefetchScope,kotlin.Unit>? onNestedPrefetch);
+    method public androidx.compose.foundation.lazy.layout.LazyLayoutPrefetchState.PrefetchHandle schedulePrefetch(int index, optional androidx.compose.ui.unit.Constraints? constraints);
   }
 
   public static sealed interface LazyLayoutPrefetchState.PrefetchHandle {
@@ -1136,6 +1141,10 @@
     property public int size;
   }
 
+  @SuppressCompatibility @androidx.compose.foundation.ExperimentalFoundationApi public sealed interface NestedPrefetchScope {
+    method public void schedulePrefetch(int index, optional androidx.compose.ui.unit.Constraints? constraints);
+  }
+
   @SuppressCompatibility @androidx.compose.foundation.ExperimentalFoundationApi public interface PrefetchExecutor {
     method public void requestPrefetch(androidx.compose.foundation.lazy.layout.PrefetchRequest prefetchRequest);
   }
diff --git a/compose/foundation/foundation/integration-tests/lazy-tests/src/androidTest/kotlin/androidx/compose/foundation/lazy/list/BaseLazyListTestWithOrientation.kt b/compose/foundation/foundation/integration-tests/lazy-tests/src/androidTest/kotlin/androidx/compose/foundation/lazy/list/BaseLazyListTestWithOrientation.kt
index a1110337..3748124 100644
--- a/compose/foundation/foundation/integration-tests/lazy-tests/src/androidTest/kotlin/androidx/compose/foundation/lazy/list/BaseLazyListTestWithOrientation.kt
+++ b/compose/foundation/foundation/integration-tests/lazy-tests/src/androidTest/kotlin/androidx/compose/foundation/lazy/list/BaseLazyListTestWithOrientation.kt
@@ -123,9 +123,10 @@
         flingBehavior: FlingBehavior = ScrollableDefaults.flingBehavior(),
         userScrollEnabled: Boolean = true,
         spacedBy: Dp = 0.dp,
+        isCrossAxis: Boolean = false,
         content: LazyListScope.() -> Unit
     ) {
-        if (vertical) {
+        if (vertical xor isCrossAxis) {
             val verticalArrangement = when {
                 spacedBy != 0.dp -> Arrangement.spacedBy(spacedBy)
                 reverseLayout xor reverseArrangement -> Arrangement.Bottom
diff --git a/compose/foundation/foundation/integration-tests/lazy-tests/src/androidTest/kotlin/androidx/compose/foundation/lazy/list/LazyListNestedPrefetchingTest.kt b/compose/foundation/foundation/integration-tests/lazy-tests/src/androidTest/kotlin/androidx/compose/foundation/lazy/list/LazyListNestedPrefetchingTest.kt
new file mode 100644
index 0000000..ebd489b
--- /dev/null
+++ b/compose/foundation/foundation/integration-tests/lazy-tests/src/androidTest/kotlin/androidx/compose/foundation/lazy/list/LazyListNestedPrefetchingTest.kt
@@ -0,0 +1,369 @@
+/*
+ * Copyright 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package androidx.compose.foundation.lazy.list
+
+import androidx.compose.foundation.AutoTestFrameClock
+import androidx.compose.foundation.ExperimentalFoundationApi
+import androidx.compose.foundation.gestures.Orientation
+import androidx.compose.foundation.gestures.scrollBy
+import androidx.compose.foundation.layout.Spacer
+import androidx.compose.foundation.lazy.LazyListLayoutInfo
+import androidx.compose.foundation.lazy.LazyListPrefetchScope
+import androidx.compose.foundation.lazy.LazyListPrefetchStrategy
+import androidx.compose.foundation.lazy.LazyListState
+import androidx.compose.foundation.lazy.layout.NestedPrefetchScope
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.DisposableEffect
+import androidx.compose.runtime.remember
+import androidx.compose.ui.Modifier
+import androidx.compose.ui.layout.layout
+import androidx.compose.ui.platform.testTag
+import androidx.compose.ui.test.onNodeWithTag
+import androidx.compose.ui.unit.Constraints
+import androidx.test.filters.LargeTest
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.runBlocking
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.Parameterized
+
+@LargeTest
+@RunWith(Parameterized::class)
+class LazyListNestedPrefetchingTest(
+    val config: Config
+) : BaseLazyListTestWithOrientation(config.orientation) {
+
+    companion object {
+        @JvmStatic
+        @Parameterized.Parameters(name = "{0}")
+        fun initParameters(): Array<Any> = arrayOf(
+            Config(Orientation.Vertical),
+            Config(Orientation.Horizontal),
+        )
+
+        class Config(
+            val orientation: Orientation,
+        ) {
+            override fun toString() = "orientation=$orientation"
+        }
+    }
+
+    sealed interface Action {
+        data class Compose(val index: Int, val nestedIndex: Int? = null) : Action
+        data class Measure(val index: Int, val nestedIndex: Int? = null) : Action
+    }
+
+    private val itemsSizePx = 30
+    private val itemsSizeDp = with(rule.density) { itemsSizePx.toDp() }
+    private val activeNodes = mutableSetOf<String>()
+    private val activeMeasuredNodes = mutableSetOf<String>()
+
+    @Test
+    fun nestedPrefetchingForwardAfterSmallScroll() {
+        val state = LazyListState()
+        composeList(state)
+
+        val prefetchIndex = 2
+        val actions = trackingActions {
+            rule.runOnIdle {
+                runBlocking {
+                    state.scrollBy(5f)
+                }
+            }
+
+            waitForPrefetch(tagFor(prefetchIndex))
+        }
+
+        // We want to make sure nested children were precomposed before the parent was premeasured
+        // (which would force them all to compose in a single block of work in premeasure)
+        assertThat(actions).containsExactly(
+            Action.Compose(prefetchIndex),
+            Action.Compose(prefetchIndex, 0),
+            Action.Compose(prefetchIndex, 1),
+            Action.Measure(prefetchIndex),
+            Action.Measure(prefetchIndex, 0),
+            Action.Measure(prefetchIndex, 1),
+        ).inOrder()
+
+        rule.onNodeWithTag(tagFor(prefetchIndex))
+            .assertExists()
+        rule.onNodeWithTag(tagFor(2, 0))
+            .assertExists()
+        rule.onNodeWithTag(tagFor(2, 1))
+            .assertExists()
+        rule.onNodeWithTag(tagFor(2, 2))
+            .assertDoesNotExist()
+    }
+
+    @Test
+    fun cancelingPrefetchCancelsItsNestedPrefetches() {
+        val state = LazyListState()
+        composeList(state)
+
+        rule.runOnIdle {
+            runBlocking {
+                // this will move the viewport so items 1-2 are visible
+                // and schedule a prefetching for 3
+                state.scrollBy(itemsSizePx.toFloat())
+            }
+        }
+
+        waitForPrefetch(tagFor(3))
+
+        rule.runOnIdle {
+            assertThat(activeNodes).contains(tagFor(3))
+            assertThat(activeNodes).contains(tagFor(3, 0))
+            assertThat(activeNodes).contains(tagFor(3, 1))
+        }
+
+        rule.runOnIdle {
+            runBlocking(AutoTestFrameClock()) {
+                // move viewport by screen size to items 4-5, so item 3 is just behind
+                // the first visible item
+                state.scrollBy(itemsSizePx * 3f)
+
+                // move scroll further to items 5-6, so item 3 is reused
+                state.scrollBy(itemsSizePx.toFloat())
+            }
+        }
+
+        waitForPrefetch(tagFor(7))
+
+        rule.runOnIdle {
+            runBlocking(AutoTestFrameClock()) {
+                // scroll again to ensure item 3 was dropped
+                state.scrollBy(itemsSizePx * 100f)
+            }
+        }
+
+        rule.runOnIdle {
+            assertThat(activeNodes).doesNotContain(tagFor(3))
+            assertThat(activeNodes).doesNotContain(tagFor(3, 0))
+            assertThat(activeNodes).doesNotContain(tagFor(3, 1))
+        }
+    }
+
+    @OptIn(ExperimentalFoundationApi::class)
+    @Test
+    fun overridingNestedPrefetchCountIsRespected() {
+        val state = LazyListState()
+        composeList(
+            state,
+            createNestedLazyListState = {
+                LazyListState(
+                    prefetchStrategy = LazyListPrefetchStrategy(1)
+                )
+            })
+
+        val prefetchIndex = 2
+        val actions = trackingActions {
+            rule.runOnIdle {
+                runBlocking {
+                    state.scrollBy(5f)
+                }
+            }
+
+            waitForPrefetch(tagFor(prefetchIndex))
+        }
+
+        // Since the nested prefetch count on the strategy is 1, we only expect index 0 to be
+        // precomposed before measure
+        assertThat(actions).containsExactly(
+            Action.Compose(prefetchIndex),
+            Action.Compose(prefetchIndex, 0),
+            Action.Measure(prefetchIndex),
+            Action.Measure(prefetchIndex, 0),
+            Action.Compose(prefetchIndex, 1),
+            Action.Measure(prefetchIndex, 1),
+        ).inOrder()
+    }
+
+    @OptIn(ExperimentalFoundationApi::class)
+    @Test
+    fun nestedPrefetchIsMeasuredWithProvidedConstraints() {
+        val nestedConstraints =
+            Constraints(minWidth = 20, minHeight = 20, maxWidth = 20, maxHeight = 20)
+        val state = LazyListState()
+        composeList(
+            state,
+            createNestedLazyListState = {
+                LazyListState(
+                    prefetchStrategy = NestedPrefetchWithConstraintsStrategy(nestedConstraints)
+                )
+            })
+
+        val prefetchIndex = 2
+        val actions = trackingActions {
+            rule.runOnIdle {
+                runBlocking {
+                    state.scrollBy(5f)
+                }
+            }
+
+            waitForPrefetch(tagFor(prefetchIndex))
+        }
+
+        assertThat(actions).containsExactly(
+            Action.Compose(prefetchIndex),
+            Action.Compose(prefetchIndex, 0),
+            Action.Measure(prefetchIndex, 0),
+            Action.Compose(prefetchIndex, 1),
+            Action.Measure(prefetchIndex, 1),
+            Action.Measure(prefetchIndex),
+            // Extra measure calls here since we didn't actually provide the right Constraints
+            Action.Measure(prefetchIndex, 0),
+            Action.Measure(prefetchIndex, 1),
+        ).inOrder()
+    }
+
+    @Test
+    fun nestedPrefetchStartsFromFirstVisibleItemIndex() {
+        val state = LazyListState()
+        composeList(
+            state,
+            createNestedLazyListState = {
+                LazyListState(firstVisibleItemIndex = 5)
+            })
+
+        val prefetchIndex = 2
+        val actions = trackingActions {
+            rule.runOnIdle {
+                runBlocking {
+                    state.scrollBy(5f)
+                }
+            }
+
+            waitForPrefetch(tagFor(prefetchIndex))
+        }
+
+        assertThat(actions).containsExactly(
+            Action.Compose(prefetchIndex),
+            Action.Compose(prefetchIndex, 5),
+            Action.Compose(prefetchIndex, 6),
+            Action.Measure(prefetchIndex),
+            Action.Measure(prefetchIndex, 5),
+            Action.Measure(prefetchIndex, 6),
+        ).inOrder()
+    }
+
+    private var actions: MutableList<Action>? = null
+
+    /**
+     * Returns the list of Actions performed during block()
+     */
+    private fun trackingActions(block: () -> Unit): List<Action> {
+        return mutableListOf<Action>().apply {
+            actions = this
+            block()
+            actions = null
+        }
+    }
+
+    private fun waitForPrefetch(tag: String) {
+        rule.waitUntil {
+            activeNodes.contains(tag) && activeMeasuredNodes.contains(tag)
+        }
+    }
+
+    fun tagFor(index: Int, nestedIndex: Int? = null): String {
+        return if (nestedIndex == null) {
+            "$index"
+        } else {
+            "$index:$nestedIndex"
+        }
+    }
+
+    private fun composeList(
+        lazyListState: LazyListState,
+        createNestedLazyListState: (index: Int) -> LazyListState = { LazyListState() }
+    ) {
+        rule.setContent {
+            LazyColumnOrRow(
+                modifier = Modifier.mainAxisSize(itemsSizeDp * 1.5f),
+                state = lazyListState
+            ) {
+                items(100) { index ->
+                    TrackActiveNodesEffect(index)
+                    val nestedState = remember(index) { createNestedLazyListState(index) }
+                    LazyColumnOrRow(
+                        modifier = Modifier
+                            .crossAxisSize(itemsSizeDp * 1.5f)
+                            .mainAxisSize(itemsSizeDp)
+                            .testTag(tagFor(index))
+                            .trackWhenMeasured(index),
+                        state = nestedState,
+                        isCrossAxis = true,
+                    ) {
+                        items(100) { nestedIndex ->
+                            TrackActiveNodesEffect(index, nestedIndex)
+                            Spacer(
+                                Modifier
+                                    .mainAxisSize(itemsSizeDp)
+                                    .crossAxisSize(itemsSizeDp)
+                                    .testTag(tagFor(index, nestedIndex))
+                                    .trackWhenMeasured(index, nestedIndex)
+                            )
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    @Composable
+    private fun TrackActiveNodesEffect(index: Int, nestedIndex: Int? = null) {
+        val tag = tagFor(index, nestedIndex)
+        DisposableEffect(tag) {
+            activeNodes.add(tag)
+            actions?.add(Action.Compose(index, nestedIndex))
+            onDispose {
+                activeNodes.remove(tag)
+                activeMeasuredNodes.remove(tag)
+            }
+        }
+    }
+
+    private fun Modifier.trackWhenMeasured(index: Int, nestedIndex: Int? = null): Modifier {
+        val tag = tagFor(index, nestedIndex)
+        return this then Modifier.layout { measurable, constraints ->
+            actions?.add(Action.Measure(index, nestedIndex))
+            val placeable = measurable.measure(constraints)
+            activeMeasuredNodes.add(tag)
+            layout(placeable.width, placeable.height) {
+                placeable.place(0, 0)
+            }
+        }
+    }
+
+    @OptIn(ExperimentalFoundationApi::class)
+    private class NestedPrefetchWithConstraintsStrategy(
+        private val childConstraints: Constraints,
+        private val nestedPrefetchItemCount: Int = 2
+    ) : LazyListPrefetchStrategy {
+        override fun LazyListPrefetchScope.onScroll(delta: Float, layoutInfo: LazyListLayoutInfo) {
+        }
+
+        override fun LazyListPrefetchScope.onVisibleItemsUpdated(layoutInfo: LazyListLayoutInfo) {
+        }
+
+        override fun NestedPrefetchScope.onNestedPrefetch(firstVisibleItemIndex: Int) {
+            repeat(nestedPrefetchItemCount) { i ->
+                schedulePrefetch(firstVisibleItemIndex + i, childConstraints)
+            }
+        }
+    }
+}
diff --git a/compose/foundation/foundation/integration-tests/lazy-tests/src/androidTest/kotlin/androidx/compose/foundation/lazy/list/LazyListPrefetchStrategyTest.kt b/compose/foundation/foundation/integration-tests/lazy-tests/src/androidTest/kotlin/androidx/compose/foundation/lazy/list/LazyListPrefetchStrategyTest.kt
index f64b86a..31dd332 100644
--- a/compose/foundation/foundation/integration-tests/lazy-tests/src/androidTest/kotlin/androidx/compose/foundation/lazy/list/LazyListPrefetchStrategyTest.kt
+++ b/compose/foundation/foundation/integration-tests/lazy-tests/src/androidTest/kotlin/androidx/compose/foundation/lazy/list/LazyListPrefetchStrategyTest.kt
@@ -28,6 +28,7 @@
 import androidx.compose.foundation.lazy.LazyListPrefetchStrategy
 import androidx.compose.foundation.lazy.LazyListState
 import androidx.compose.foundation.lazy.layout.LazyLayoutPrefetchState
+import androidx.compose.foundation.lazy.layout.NestedPrefetchScope
 import androidx.compose.foundation.lazy.list.LazyListPrefetchStrategyTest.RecordingLazyListPrefetchStrategy.Callback
 import androidx.compose.foundation.lazy.rememberLazyListState
 import androidx.compose.runtime.DisposableEffect
@@ -272,6 +273,8 @@
             _callbacks.add(Callback.OnVisibleItemsUpdated(layoutInfo.visibleIndices))
         }
 
+        override fun NestedPrefetchScope.onNestedPrefetch(firstVisibleItemIndex: Int) = Unit
+
         fun reset() {
             _callbacks.clear()
         }
@@ -298,6 +301,8 @@
         override fun LazyListPrefetchScope.onVisibleItemsUpdated(layoutInfo: LazyListLayoutInfo) =
             Unit
 
+        override fun NestedPrefetchScope.onNestedPrefetch(firstVisibleItemIndex: Int) = Unit
+
         private fun cancelPrefetch() {
             handle?.cancel()
             prefetchIndex = -1
diff --git a/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/LazyListPrefetchStrategy.kt b/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/LazyListPrefetchStrategy.kt
index cc4614b..5227192 100644
--- a/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/LazyListPrefetchStrategy.kt
+++ b/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/LazyListPrefetchStrategy.kt
@@ -18,6 +18,7 @@
 
 import androidx.compose.foundation.ExperimentalFoundationApi
 import androidx.compose.foundation.lazy.layout.LazyLayoutPrefetchState
+import androidx.compose.foundation.lazy.layout.NestedPrefetchScope
 import androidx.compose.foundation.lazy.layout.PrefetchExecutor
 import androidx.compose.runtime.Stable
 
@@ -45,17 +46,38 @@
      * If the visible items have also changed, then this will be invoked in the same frame *after*
      * [onVisibleItemsUpdated].
      *
-     * [delta] can be used to understand scroll direction: delta < 0 indicates scrolling down while
+     * @param delta the change in scroll direction. Delta < 0 indicates scrolling down while
      * delta > 0 indicates scrolling up.
+     * @param layoutInfo the current [LazyListLayoutInfo]
      */
     fun LazyListPrefetchScope.onScroll(delta: Float, layoutInfo: LazyListLayoutInfo)
 
     /**
      * onVisibleItemsUpdated is invoked when the LazyList scrolls if the visible items have changed.
-     * Info about these visible items can be found in [layoutInfo]'s
-     * [LazyListLayoutInfo.visibleItemsInfo].
+     *
+     * @param layoutInfo the current [LazyListLayoutInfo]. Info about the updated visible items can
+     * be found in [LazyListLayoutInfo.visibleItemsInfo].
      */
     fun LazyListPrefetchScope.onVisibleItemsUpdated(layoutInfo: LazyListLayoutInfo)
+
+    /**
+     * onNestedPrefetch is invoked when a parent LazyLayout has prefetched content which contains
+     * this LazyList. It gives this LazyList a chance to request prefetch for some of its own
+     * children before coming onto screen.
+     *
+     * Implementations can use [NestedPrefetchScope.schedulePrefetch] to schedule child
+     * prefetches. For example, this is useful if this LazyList is a LazyRow that is a child of a
+     * LazyColumn: in that case, [onNestedPrefetch] can schedule the children it expects to be
+     * visible when it comes onto screen, giving the LazyLayout infra a chance to compose these
+     * children ahead of time and reduce jank.
+     *
+     * Generally speaking, [onNestedPrefetch] should only request prefetch for children that it
+     * expects to actually be visible when this list is scrolled into view.
+     *
+     * @param firstVisibleItemIndex the index of the first visible item. It should be used to start
+     * prefetching from the correct index in case the list has been created at a non-zero offset.
+     */
+    fun NestedPrefetchScope.onNestedPrefetch(firstVisibleItemIndex: Int)
 }
 
 /**
@@ -71,17 +93,34 @@
      * [LazyLayoutPrefetchState.PrefetchHandle.cancel].
      *
      * See [PrefetchExecutor].
+     *
+     * @param index the index of the child to prefetch
      */
     fun schedulePrefetch(index: Int): LazyLayoutPrefetchState.PrefetchHandle
 }
 
 /**
+ * Creates an instance of the default [LazyListPrefetchStrategy], allowing for customization of the
+ * nested prefetch count.
+ *
+ * @param nestedPrefetchItemCount specifies how many inner items should be prefetched when this
+ * LazyList is nested inside another LazyLayout. For example, if this is the state for a horizontal
+ * LazyList nested in a vertical LazyList, you might want to set this to the number of items that
+ * will be visible when this list is scrolled into view.
+ */
+@ExperimentalFoundationApi
+fun LazyListPrefetchStrategy(
+    nestedPrefetchItemCount: Int = 2
+): LazyListPrefetchStrategy = DefaultLazyListPrefetchStrategy(nestedPrefetchItemCount)
+
+/**
  * The default prefetching strategy for LazyLists - this will be used automatically if no other
  * strategy is provided.
  */
 @OptIn(ExperimentalFoundationApi::class)
 @Stable
-internal class DefaultLazyListPrefetchStrategy : LazyListPrefetchStrategy {
+private class DefaultLazyListPrefetchStrategy(private val nestedPrefetchItemCount: Int = 2) :
+    LazyListPrefetchStrategy {
 
     /**
      * The index scheduled to be prefetched (or the last prefetched index if the prefetch is done).
@@ -140,4 +179,10 @@
             }
         }
     }
+
+    override fun NestedPrefetchScope.onNestedPrefetch(firstVisibleItemIndex: Int) {
+        repeat(nestedPrefetchItemCount) { i ->
+            schedulePrefetch(firstVisibleItemIndex + i)
+        }
+    }
 }
diff --git a/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/LazyListState.kt b/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/LazyListState.kt
index f743c88..e6f2ee3 100644
--- a/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/LazyListState.kt
+++ b/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/LazyListState.kt
@@ -43,6 +43,7 @@
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.neverEqualPolicy
+import androidx.compose.runtime.remember
 import androidx.compose.runtime.saveable.Saver
 import androidx.compose.runtime.saveable.listSaver
 import androidx.compose.runtime.saveable.rememberSaveable
@@ -102,9 +103,9 @@
 fun rememberLazyListState(
     initialFirstVisibleItemIndex: Int = 0,
     initialFirstVisibleItemScrollOffset: Int = 0,
-    prefetchStrategy: LazyListPrefetchStrategy = DefaultLazyListPrefetchStrategy(),
+    prefetchStrategy: LazyListPrefetchStrategy = remember { LazyListPrefetchStrategy() },
 ): LazyListState {
-    return rememberSaveable(saver = LazyListState.saver(prefetchStrategy)) {
+    return rememberSaveable(prefetchStrategy, saver = LazyListState.saver(prefetchStrategy)) {
         LazyListState(
             initialFirstVisibleItemIndex,
             initialFirstVisibleItemScrollOffset,
@@ -129,7 +130,7 @@
 class LazyListState @ExperimentalFoundationApi constructor(
     firstVisibleItemIndex: Int = 0,
     firstVisibleItemScrollOffset: Int = 0,
-    private val prefetchStrategy: LazyListPrefetchStrategy = DefaultLazyListPrefetchStrategy(),
+    private val prefetchStrategy: LazyListPrefetchStrategy = LazyListPrefetchStrategy(),
 ) : ScrollableState {
 
     /**
@@ -140,7 +141,7 @@
     constructor(
         firstVisibleItemIndex: Int = 0,
         firstVisibleItemScrollOffset: Int = 0
-    ) : this(firstVisibleItemIndex, firstVisibleItemScrollOffset, DefaultLazyListPrefetchStrategy())
+    ) : this(firstVisibleItemIndex, firstVisibleItemScrollOffset, LazyListPrefetchStrategy())
 
     internal var hasLookaheadPassOccurred: Boolean = false
         private set
@@ -265,7 +266,11 @@
 
     internal val beyondBoundsInfo = LazyLayoutBeyondBoundsInfo()
 
-    internal val prefetchState = LazyLayoutPrefetchState(prefetchStrategy.prefetchExecutor)
+    internal val prefetchState = LazyLayoutPrefetchState(prefetchStrategy.prefetchExecutor) {
+        with(prefetchStrategy) {
+            onNestedPrefetch(Snapshot.withoutReadObservation { firstVisibleItemIndex })
+        }
+    }
 
     private val prefetchScope: LazyListPrefetchScope = object : LazyListPrefetchScope {
         override fun schedulePrefetch(index: Int): LazyLayoutPrefetchState.PrefetchHandle {
diff --git a/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/layout/LazyLayout.kt b/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/layout/LazyLayout.kt
index f1d49fd..b51f4ab 100644
--- a/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/layout/LazyLayout.kt
+++ b/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/layout/LazyLayout.kt
@@ -107,7 +107,7 @@
 
         SubcomposeLayout(
             subcomposeLayoutState,
-            modifier,
+            modifier.traversablePrefetchState(prefetchState),
             remember(itemContentFactory, measurePolicy) {
                 { constraints ->
                     with(
diff --git a/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/layout/LazyLayoutPrefetchState.kt b/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/layout/LazyLayoutPrefetchState.kt
index 619c8e4..e3b1586 100644
--- a/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/layout/LazyLayoutPrefetchState.kt
+++ b/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/layout/LazyLayoutPrefetchState.kt
@@ -19,7 +19,12 @@
 import androidx.compose.foundation.ExperimentalFoundationApi
 import androidx.compose.foundation.lazy.layout.LazyLayoutPrefetchState.PrefetchHandle
 import androidx.compose.runtime.Stable
+import androidx.compose.ui.Modifier
 import androidx.compose.ui.layout.SubcomposeLayoutState
+import androidx.compose.ui.node.ModifierNodeElement
+import androidx.compose.ui.node.TraversableNode
+import androidx.compose.ui.node.TraversableNode.Companion.TraverseDescendantsAction
+import androidx.compose.ui.platform.InspectorInfo
 import androidx.compose.ui.unit.Constraints
 import androidx.compose.ui.util.trace
 import kotlin.system.measureNanoTime
@@ -33,10 +38,16 @@
  *
  * @param prefetchExecutor the PrefetchExecutor implementation to use to execute prefetch requests.
  * If null is provided, the default PrefetchExecutor for the platform will be used.
+ * @param onNestedPrefetch a callback which will be invoked when this LazyLayout is prefetched in
+ * context of a parent LazyLayout, giving a chance to recursively prefetch its own children. See
+ * [NestedPrefetchScope].
  */
 @ExperimentalFoundationApi
 @Stable
-class LazyLayoutPrefetchState(internal val prefetchExecutor: PrefetchExecutor? = null) {
+class LazyLayoutPrefetchState(
+    internal val prefetchExecutor: PrefetchExecutor? = null,
+    private val onNestedPrefetch: (NestedPrefetchScope.() -> Unit)? = null
+) {
 
     private val prefetchMetrics: PrefetchMetrics = PrefetchMetrics()
     internal var prefetchHandleProvider: PrefetchHandleProvider? = null
@@ -45,13 +56,23 @@
      * Schedules precomposition and premeasure for the new item.
      *
      * @param index item index to prefetch.
-     * @param constraints [Constraints] to use for premeasuring.
+     * @param constraints [Constraints] to use for premeasuring. If null, the child will not
+     * be premeasured.
      */
-    fun schedulePrefetch(index: Int, constraints: Constraints): PrefetchHandle {
+    fun schedulePrefetch(index: Int, constraints: Constraints? = null): PrefetchHandle {
         return prefetchHandleProvider?.schedulePrefetch(index, constraints, prefetchMetrics)
             ?: DummyHandle
     }
 
+    internal fun collectNestedPrefetchRequests(): List<PrefetchRequest> {
+        val onNestedPrefetch = onNestedPrefetch ?: return listOf()
+
+        return NestedPrefetchScopeImpl().run {
+            onNestedPrefetch()
+            requests
+        }
+    }
+
     sealed interface PrefetchHandle {
         /**
          * Notifies the prefetcher that previously scheduled item is no longer needed. If the item
@@ -59,6 +80,39 @@
          */
         fun cancel()
     }
+
+    private inner class NestedPrefetchScopeImpl : NestedPrefetchScope {
+
+        val requests: List<PrefetchRequest>
+            get() = _requests
+        private val _requests: MutableList<PrefetchRequest> = mutableListOf()
+
+        override fun schedulePrefetch(index: Int, constraints: Constraints?) {
+            val prefetchHandleProvider = prefetchHandleProvider ?: return
+            _requests.add(
+                prefetchHandleProvider.createNestedPrefetchRequest(
+                    index,
+                    constraints,
+                    prefetchMetrics
+                )
+            )
+        }
+    }
+}
+
+/**
+ * A scope which allows nested prefetches to be requested for the precomposition of a LazyLayout.
+ */
+@ExperimentalFoundationApi
+sealed interface NestedPrefetchScope {
+
+    /**
+     * Requests a child index to be prefetched as part of the prefetch of a parent LazyLayout.
+     * @param index the index of the child to prefetch.
+     * @param constraints [Constraints] to use for premeasuring. If null, the child will not
+     * be premeasured.
+     */
+    fun schedulePrefetch(index: Int, constraints: Constraints? = null)
 }
 
 /**
@@ -131,17 +185,24 @@
 ) {
     fun schedulePrefetch(
         index: Int,
-        constraints: Constraints,
+        constraints: Constraints?,
         prefetchMetrics: PrefetchMetrics
     ): PrefetchHandle =
         HandleAndRequestImpl(index, constraints, prefetchMetrics).also {
             executor.requestPrefetch(it)
         }
 
+    fun createNestedPrefetchRequest(
+        index: Int,
+        constraints: Constraints?,
+        prefetchMetrics: PrefetchMetrics,
+    ): PrefetchRequest =
+        HandleAndRequestImpl(index, constraints = constraints, prefetchMetrics)
+
     @ExperimentalFoundationApi
     private inner class HandleAndRequestImpl(
         private val index: Int,
-        private val constraints: Constraints,
+        private val constraints: Constraints?,
         private val prefetchMetrics: PrefetchMetrics,
     ) : PrefetchHandle, PrefetchRequest {
 
@@ -149,6 +210,9 @@
         private var isMeasured = false
         private var isCanceled = false
         private val isComposed get() = precomposeHandle != null
+        private var hasResolvedNestedPrefetches = false
+        private var nestedPrefetchController: NestedPrefetchController? = null
+
         private val isValid
             get() = !isCanceled &&
                 index in 0 until itemContentFactory.itemProvider().itemCount
@@ -178,11 +242,30 @@
                 }
             }
 
-            if (!isMeasured) {
+            // Nested prefetch logic is best-effort: if nested LazyLayout children are
+            // added/removed/updated after we've resolved nested prefetch states here or resolved
+            // nestedPrefetchRequests below, those changes won't be taken into account.
+            if (!hasResolvedNestedPrefetches) {
+                if (availableTimeNanos > 0) {
+                    trace("compose:lazy:prefetch:resolve-nested") {
+                        nestedPrefetchController = resolveNestedPrefetchStates()
+                        hasResolvedNestedPrefetches = true
+                    }
+                } else {
+                    return true
+                }
+            }
+
+            val hasMoreWork = nestedPrefetchController?.run { executeNestedPrefetches() } ?: false
+            if (hasMoreWork) {
+                return true
+            }
+
+            if (!isMeasured && constraints != null) {
                 if (prefetchMetrics.averageMeasureTimeNanos < availableTimeNanos) {
                     prefetchMetrics.recordMeasureTiming {
                         trace("compose:lazy:prefetch:measure") {
-                            performMeasure()
+                            performMeasure(constraints)
                         }
                     }
                 } else {
@@ -207,7 +290,7 @@
             precomposeHandle = subcomposeLayoutState.precompose(key, content)
         }
 
-        private fun performMeasure() {
+        private fun performMeasure(constraints: Constraints) {
             require(!isCanceled) {
                 "Callers should check whether the request is still valid before calling " +
                     "performMeasure()"
@@ -222,8 +305,117 @@
             }
         }
 
+        private fun resolveNestedPrefetchStates(): NestedPrefetchController? {
+            val precomposedSlotHandle = requireNotNull(precomposeHandle) {
+                "Should precompose before resolving nested prefetch states"
+            }
+
+            var nestedStates: MutableList<LazyLayoutPrefetchState>? = null
+            precomposedSlotHandle.traverseDescendants(TraversablePrefetchStateNodeKey) {
+                val prefetchState = (it as TraversablePrefetchStateNode).prefetchState
+                nestedStates =
+                    nestedStates?.apply { add(prefetchState) } ?: mutableListOf(prefetchState)
+                TraverseDescendantsAction.SkipSubtreeAndContinueTraversal
+            }
+            return nestedStates?.let { NestedPrefetchController(it) }
+        }
+
         override fun toString(): String =
             "HandleAndRequestImpl { index = $index, constraints = $constraints, " +
                 "isComposed = $isComposed, isMeasured = $isMeasured, isCanceled = $isCanceled }"
+
+        private inner class NestedPrefetchController(
+            private val states: List<LazyLayoutPrefetchState>
+        ) {
+
+            // This array is parallel to nestedPrefetchStates, so index 0 in nestedPrefetchStates
+            // corresponds to index 0 in this array, etc.
+            private val requestsByState: Array<List<PrefetchRequest>?> = arrayOfNulls(states.size)
+            private var stateIndex: Int = 0
+            private var requestIndex: Int = 0
+
+            init {
+                require(states.isNotEmpty()) {
+                    "NestedPrefetchController shouldn't be created with no states"
+                }
+            }
+
+            fun PrefetchRequestScope.executeNestedPrefetches(): Boolean {
+                if (stateIndex >= states.size) {
+                    return false
+                }
+                check(!isCanceled) { "Should not execute nested prefetch on canceled request" }
+
+                trace("compose:lazy:prefetch:nested") {
+                    while (stateIndex < states.size) {
+                        if (requestsByState[stateIndex] == null) {
+                            if (availableTimeNanos <= 0) {
+                                // When we have time again, we'll resolve nested requests for this
+                                // state
+                                return true
+                            }
+
+                            requestsByState[stateIndex] =
+                                states[stateIndex].collectNestedPrefetchRequests()
+                        }
+
+                        val nestedRequests = requestsByState[stateIndex]!!
+                        while (requestIndex < nestedRequests.size) {
+                            val hasMoreWork = with(nestedRequests[requestIndex]) { execute() }
+                            if (hasMoreWork) {
+                                return true
+                            } else {
+                                requestIndex++
+                            }
+                        }
+
+                        requestIndex = 0
+                        stateIndex++
+                    }
+                }
+
+                return false
+            }
+        }
+    }
+}
+
+private const val TraversablePrefetchStateNodeKey =
+    "androidx.compose.foundation.lazy.layout.TraversablePrefetchStateNode"
+
+/**
+ * A modifier which lets the [LazyLayoutPrefetchState] for a [LazyLayout] to be discoverable via
+ * [TraversableNode] traversal.
+ */
+@ExperimentalFoundationApi
+internal fun Modifier.traversablePrefetchState(
+    lazyLayoutPrefetchState: LazyLayoutPrefetchState?
+): Modifier {
+    return lazyLayoutPrefetchState?.let {
+        this then TraversablePrefetchStateModifierElement(it)
+    } ?: this
+}
+
+@ExperimentalFoundationApi
+private class TraversablePrefetchStateNode(
+    var prefetchState: LazyLayoutPrefetchState,
+) : Modifier.Node(), TraversableNode {
+
+    override val traverseKey: String = TraversablePrefetchStateNodeKey
+}
+
+@ExperimentalFoundationApi
+private data class TraversablePrefetchStateModifierElement(
+    private val prefetchState: LazyLayoutPrefetchState,
+) : ModifierNodeElement<TraversablePrefetchStateNode>() {
+    override fun create() = TraversablePrefetchStateNode(prefetchState)
+
+    override fun update(node: TraversablePrefetchStateNode) {
+        node.prefetchState = prefetchState
+    }
+
+    override fun InspectorInfo.inspectableProperties() {
+        name = "traversablePrefetchState"
+        value = prefetchState
     }
 }
diff --git a/compose/ui/ui/api/current.txt b/compose/ui/ui/api/current.txt
index f0f1961..fe5c4f7 100644
--- a/compose/ui/ui/api/current.txt
+++ b/compose/ui/ui/api/current.txt
@@ -2485,6 +2485,7 @@
     method public void dispose();
     method public default int getPlaceablesCount();
     method public default void premeasure(int index, long constraints);
+    method public default void traverseDescendants(Object? key, kotlin.jvm.functions.Function1<? super androidx.compose.ui.node.TraversableNode,? extends androidx.compose.ui.node.TraversableNode.Companion.TraverseDescendantsAction> block);
     property public default int placeablesCount;
   }
 
diff --git a/compose/ui/ui/api/restricted_current.txt b/compose/ui/ui/api/restricted_current.txt
index 3a8128c..c07c0bf 100644
--- a/compose/ui/ui/api/restricted_current.txt
+++ b/compose/ui/ui/api/restricted_current.txt
@@ -2492,6 +2492,7 @@
     method public void dispose();
     method public default int getPlaceablesCount();
     method public default void premeasure(int index, long constraints);
+    method public default void traverseDescendants(Object? key, kotlin.jvm.functions.Function1<? super androidx.compose.ui.node.TraversableNode,? extends androidx.compose.ui.node.TraversableNode.Companion.TraverseDescendantsAction> block);
     property public default int placeablesCount;
   }
 
diff --git a/compose/ui/ui/src/commonMain/kotlin/androidx/compose/ui/layout/SubcomposeLayout.kt b/compose/ui/ui/src/commonMain/kotlin/androidx/compose/ui/layout/SubcomposeLayout.kt
index c34d127..b72e48b 100644
--- a/compose/ui/ui/src/commonMain/kotlin/androidx/compose/ui/layout/SubcomposeLayout.kt
+++ b/compose/ui/ui/src/commonMain/kotlin/androidx/compose/ui/layout/SubcomposeLayout.kt
@@ -43,8 +43,11 @@
 import androidx.compose.ui.node.LayoutNode
 import androidx.compose.ui.node.LayoutNode.LayoutState
 import androidx.compose.ui.node.LayoutNode.UsageByParent
+import androidx.compose.ui.node.TraversableNode
+import androidx.compose.ui.node.TraversableNode.Companion.TraverseDescendantsAction
 import androidx.compose.ui.node.checkMeasuredSize
 import androidx.compose.ui.node.requireOwner
+import androidx.compose.ui.node.traverseDescendants
 import androidx.compose.ui.platform.createSubcomposition
 import androidx.compose.ui.unit.Constraints
 import androidx.compose.ui.unit.LayoutDirection
@@ -256,6 +259,18 @@
          * @param constraints Constraints to measure this placeable with.
          */
         fun premeasure(index: Int, constraints: Constraints) {}
+
+        /**
+         * Conditionally executes [block] for each [Modifier.Node] of this Composition that is a
+         * [TraversableNode] with a matching [key].
+         *
+         * See [androidx.compose.ui.node.traverseDescendants] for the complete semantics of this
+         * function.
+         */
+        fun traverseDescendants(
+            key: Any?,
+            block: (TraversableNode) -> TraverseDescendantsAction
+        ) {}
     }
 }
 
@@ -814,6 +829,13 @@
                     }
                 }
             }
+
+            override fun traverseDescendants(
+                key: Any?,
+                block: (TraversableNode) -> TraverseDescendantsAction
+            ) {
+                precomposeMap[slotId]?.nodes?.head?.traverseDescendants(key, block)
+            }
         }
     }