Merge "CacheWindow refill update #2" into androidx-main

Author: Treehugger Robot

compose/foundation/foundation/integration-tests/lazy-tests/src/androidTest/kotlin/androidx/compose/foundation/lazy/list/LazyListCacheWindowTest.kt

@@ -185,7 +185,7 @@ class LazyListCacheWindowTest(orientation: Orientation) :
     }
 
     @Test
-    fun datasetChanged_shouldMakeSureNestedItemsChanged() {
+    fun datasetChanged_shouldMakeSureNestedItemsChanged_afterScroll() {
         val items = mutableStateOf(listOf("a", "b", "c", "d", "e"))
 
         rule.setContent {
@@ -264,6 +264,80 @@ class LazyListCacheWindowTest(orientation: Orientation) :
         rule.onNodeWithTag("second-nested-e").assertExists() // nested prefetched
     }
 
+    @Test
+    fun datasetChanged_shouldMakeSureNestedItemsChanged_noScroll() {
+        val items = mutableStateOf(listOf("a", "b", "c", "d"))
+
+        rule.setContent {
+            @OptIn(ExperimentalFoundationApi::class)
+            state = rememberLazyListState(cacheWindow = viewportWindow)
+            LazyColumnOrRow(
+                Modifier.mainAxisSize(itemsSizeDp * 2f)
+                    .then(
+                        object : RemeasurementModifier {
+                            override fun onRemeasurementAvailable(remeasurement: Remeasurement) {
+                                remeasure = remeasurement
+                            }
+                        }
+                    ),
+                state,
+            ) {
+                items(items.value, key = { it }) {
+                    if (it == "e" || it == "f") {
+                        val state = rememberLazyListState(cacheWindow = viewportWindow)
+                        LazyRow(
+                            Modifier.mainAxisSize(itemsSizeDp).fillMaxCrossAxis().testTag(it),
+                            state = state,
+                        ) {
+                            item {
+                                Spacer(
+                                    Modifier.mainAxisSize(itemsSizeDp)
+                                        .fillMaxCrossAxis()
+                                        .testTag("first-nested-$it")
+                                )
+                            }
+
+                            item {
+                                Spacer(
+                                    Modifier.mainAxisSize(itemsSizeDp)
+                                        .fillMaxCrossAxis()
+                                        .testTag("second-nested-$it")
+                                )
+                            }
+                        }
+                    } else {
+                        Spacer(
+                            Modifier.mainAxisSize(itemsSizeDp)
+                                .fillMaxCrossAxis()
+                                .testTag(it)
+                                .layout { measurable, constraints ->
+                                    val placeable = measurable.measure(constraints)
+                                    layout(placeable.width, placeable.height) {
+                                        placeable.place(0, 0)
+                                    }
+                                }
+                        )
+                    }
+                }
+            }
+        }
+
+        rule.onNodeWithTag("a").assertIsDisplayed() // fully visible
+        rule.onNodeWithTag("b").assertIsDisplayed() // fully visible
+        rule.onNodeWithTag("c").assertExists() // part of the window
+        rule.onNodeWithTag("d").assertExists() // part of the window
+        println("Changing Dataset")
+        rule.runOnIdle { items.value = listOf("a", "b", "e", "f", "g", "h") }
+        rule.waitForIdle()
+
+        rule.onNodeWithTag("e").assertExists() // item e will take place of item c
+        rule.onNodeWithTag("first-nested-e").assertExists() // nested prefetched
+        rule.onNodeWithTag("second-nested-e").assertExists() // nested prefetched
+        rule.onNodeWithTag("f").assertExists() // item f will take place of item d
+        rule.onNodeWithTag("first-nested-f").assertExists() // nested prefetched
+        rule.onNodeWithTag("second-nested-f").assertExists() // nested prefetched
+    }
+
     @Test
     fun datasetChanged_noScrollHappened_shouldKeepAroundWithinBounds_notCrash() {
         val numItems = mutableStateOf(100)

compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/LazyListCacheWindowStrategy.kt

@@ -135,6 +135,10 @@ internal class LazyListCacheWindowScope() : CacheWindowScope {
 
     override fun getLastIndexInLine(lineIndex: Int): Int = lineIndex
 
+    override fun getVisibleLineKey(indexInVisibleLines: Int): Any {
+        return layoutInfo.visibleItemsInfo[indexInVisibleLines].key
+    }
+
     override fun getLastLineIndex(): Int {
         if (totalItemsCount == 0) return InvalidIndex
         return totalItemsCount - 1

compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/grid/LazyGridCacheWindowStrategy.kt

@@ -21,6 +21,7 @@ import androidx.compose.foundation.gestures.snapping.offsetOnMainAxis
 import androidx.compose.foundation.gestures.snapping.sizeOnMainAxis
 import androidx.compose.foundation.lazy.layout.CacheWindowLogic
 import androidx.compose.foundation.lazy.layout.CacheWindowScope
+import androidx.compose.foundation.lazy.layout.CachedItem
 import androidx.compose.foundation.lazy.layout.InvalidIndex
 import androidx.compose.foundation.lazy.layout.LazyLayoutCacheWindow
 import androidx.compose.foundation.lazy.layout.LazyLayoutPrefetchState.PrefetchHandle
@@ -137,6 +138,15 @@ private class LazyGridCacheWindowScope() : CacheWindowScope {
         return tallestItemSize
     }
 
+    override fun getVisibleLineKey(indexInVisibleLines: Int): Any {
+        // using the first item key to represent this line.
+        val laneIndex = indexInVisibleLines + firstVisibleLineIndex
+        return layoutInfo.visibleItemsInfo
+            .fastFilter { it.lineIndex == laneIndex }
+            .firstOrNull()
+            ?.key ?: CachedItem.NoKey
+    }
+
     override fun getVisibleItemLine(indexInVisibleLines: Int): Int =
         firstVisibleLineIndex + indexInVisibleLines
 

compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/layout/CacheWindowLogic.kt

@@ -16,7 +16,6 @@
 
 package androidx.compose.foundation.lazy.layout
 
-import androidx.collection.mutableIntIntMapOf
 import androidx.collection.mutableIntObjectMapOf
 import androidx.collection.mutableIntSetOf
 import androidx.compose.foundation.ExperimentalFoundationApi
@@ -43,7 +42,7 @@ internal abstract class CacheWindowLogic(
      * Cache for items sizes in the current window. Holds sizes for both visible and non-visible
      * items
      */
-    private val windowCache = mutableIntIntMapOf()
+    private val windowCache = mutableIntObjectMapOf<CachedItem>()
     private var previousPassDelta = 0f
     private var previousPassItemCount = UnsetItemCount
     private var hasUpdatedVisibleItemsOnce = false
@@ -120,15 +119,7 @@ internal abstract class CacheWindowLogic(
          * changed.
          */
         if (previousPassItemCount != UnsetItemCount && previousPassItemCount != totalItemsCount) {
-            debugLog { "Total Items Changed" }
-            shouldRefillWindow = true
-            prefetchWindowStartLine = prefetchWindowStartLine.coerceAtLeast(0)
-            val lastLineIndex = getLastLineIndex()
-            if (lastLineIndex != InvalidIndex) {
-                prefetchWindowEndLine = prefetchWindowEndLine.coerceAtMost(lastLineIndex)
-            }
-            /** Free up the space so the fill will happen and not re-use old data. */
-            removeOutOfBoundsItems(prefetchWindowEndLine, itemsCount - 1)
+            onDatasetChangedSize()
         }
 
         itemsCount = totalItemsCount
@@ -137,8 +128,8 @@ internal abstract class CacheWindowLogic(
         // by [cancelOutOfBounds]. If any items changed sizes we re-trigger the window filling
         // update.
         if (hasVisibleItems) {
-            forEachVisibleItem { index, mainAxisSize ->
-                if (index != InvalidIndex) cacheVisibleItemsInfo(index, mainAxisSize)
+            forEachVisibleItem { index, key, mainAxisSize ->
+                if (index != InvalidIndex) cacheVisibleItemsInfo(index, key, mainAxisSize)
             }
             if (shouldRefillWindow) {
                 // refill window in accordance with last pass delta
@@ -234,6 +225,25 @@ internal abstract class CacheWindowLogic(
         }
     }
 
+    private fun CacheWindowScope.onDatasetChangedSize() {
+        debugLog { "Total Items Changed" }
+        shouldRefillWindow = true
+        prefetchWindowStartLine = prefetchWindowStartLine.coerceAtLeast(0)
+        val lastLineIndex = getLastLineIndex()
+        if (lastLineIndex != InvalidIndex) {
+            prefetchWindowEndLine = prefetchWindowEndLine.coerceAtMost(lastLineIndex)
+        }
+
+        /**
+         * Resets the window state. We will refill the window on the direction of the last scroll.
+         */
+        if (previousPassDelta <= 0f) {
+            removeOutOfBoundsItems(lastVisibleLineIndex, itemsCount - 1)
+        } else {
+            removeOutOfBoundsItems(0, firstVisibleLineIndex)
+        }
+    }
+
     fun resetStrategy() {
         prefetchWindowStartLine = Int.MAX_VALUE
         prefetchWindowEndLine = Int.MIN_VALUE
@@ -283,7 +293,7 @@ internal abstract class CacheWindowLogic(
                 // If we get the same delta in the next frame, would we cover the extra space needed
                 // to actually need this item? If so, mark it as urgent
                 val isUrgent: Boolean =
-                    if (prefetchWindowEndLine + 1 == visibleWindowEnd + 1) {
+                    if (prefetchWindowEndLine + 1 == visibleWindowEnd + 1 && scrollDelta != 0.0f) {
                         scrollDelta.absoluteValue >= mainAxisExtraSpaceEnd
                     } else {
                         false
@@ -315,7 +325,9 @@ internal abstract class CacheWindowLogic(
                 // If we get the same delta in the next frame, would we cover the extra space needed
                 // to actually need this item? If so, mark it as urgent
                 val isUrgent: Boolean =
-                    if (prefetchWindowStartLine - 1 == visibleWindowStart - 1) {
+                    if (
+                        prefetchWindowStartLine - 1 == visibleWindowStart - 1 && scrollDelta != 0.0f
+                    ) {
                         scrollDelta.absoluteValue >= mainAxisExtraSpaceStart
                     } else {
                         false
@@ -358,7 +370,7 @@ internal abstract class CacheWindowLogic(
             while (prefetchWindowStartExtraSpace > 0 && prefetchWindowStartLine > 0) {
                 val item =
                     if (windowCache.containsKey(prefetchWindowStartLine - 1)) {
-                        windowCache[prefetchWindowStartLine - 1]
+                        windowCache[prefetchWindowStartLine - 1]!!.mainAxisSize
                     } else {
                         break
                     }
@@ -373,7 +385,7 @@ internal abstract class CacheWindowLogic(
             while (prefetchWindowEndExtraSpace > 0 && prefetchWindowEndLine < itemsCount - 1) {
                 val item =
                     if (windowCache.containsKey(prefetchWindowEndLine + 1)) {
-                        windowCache[prefetchWindowEndLine + 1]
+                        windowCache[prefetchWindowEndLine + 1]!!.mainAxisSize
                     } else {
                         break
                     }
@@ -386,13 +398,16 @@ internal abstract class CacheWindowLogic(
 
     private fun CacheWindowScope.getItemSizeOrPrefetch(index: Int, isUrgent: Boolean): Int {
         return if (windowCache.containsKey(index)) {
-            windowCache[index]
+            debugLog { "Item $index is Cached!" }
+            windowCache[index]!!.mainAxisSize
         } else if (prefetchWindowHandles.containsKey(index)) {
             // item is scheduled but didn't finish yet
+            debugLog { "Item=$index is already scheduled. isUrgent=$isUrgent" }
             if (isUrgent) prefetchWindowHandles[index]?.fastForEach { it.markAsUrgent() }
             InvalidItemSize
         } else {
             // item is not scheduled
+            debugLog { "Scheduling Prefetching for Item=$index. isUrgent=$isUrgent" }
             prefetchWindowHandles[index] =
                 schedulePrefetch(index) { prefetchedIndex, size ->
                     onItemPrefetched(prefetchedIndex, size)
@@ -404,7 +419,7 @@ internal abstract class CacheWindowLogic(
 
     /** Grows the window with measured items and prefetched items. */
     private fun cachePrefetchedItem(index: Int, size: Int) {
-        windowCache[index] = size
+        windowCache[index] = updateOrCreateCachedItem(index, size, CachedItem.NoKey)
         if (index > prefetchWindowEndLine) {
             prefetchWindowEndLine = index
             prefetchWindowEndExtraSpace -= size
@@ -414,18 +429,34 @@ internal abstract class CacheWindowLogic(
         }
     }
 
+    private fun updateOrCreateCachedItem(index: Int, size: Int, key: Any): CachedItem {
+        val cachedItem = windowCache[index]
+        return if (cachedItem != null) {
+            cachedItem.mainAxisSize = size
+            cachedItem.key = key
+            cachedItem
+        } else {
+            CachedItem(CachedItem.NoKey, size)
+        }
+    }
+
     /**
      * When caching visible items we need to check if the existing item changed sizes. If so, we
      * will set [shouldRefillWindow] which will trigger a complete window filling and cancel any out
-     * of bounds requests.
+     * of bounds requests. The same is valid if items are replaced (have the same size by key
+     * changed).
      */
-    private fun cacheVisibleItemsInfo(index: Int, size: Int) {
-        debugLog { "cacheVisibleItemsInfo item=$index size=$size" }
-        if (windowCache.containsKey(index) && windowCache[index] != size) {
-            shouldRefillWindow = true
+    private fun cacheVisibleItemsInfo(index: Int, key: Any, size: Int) {
+        debugLog { "cacheVisibleItemsInfo item=$index size=$size key=$key" }
+        if (windowCache.containsKey(index)) {
+            val cachedSize = windowCache[index]!!.mainAxisSize
+            val cachedKey = windowCache[index]!!.key
+            if (cachedSize != size || cachedKey != key) {
+                shouldRefillWindow = true
+            }
         }
 
-        windowCache[index] = size
+        windowCache[index] = updateOrCreateCachedItem(index, size, key)
         // We're caching a visible item, remove its handle since we won't need it anymore.
         prefetchWindowStartLine = minOf(prefetchWindowStartLine, index)
         prefetchWindowEndLine = maxOf(prefetchWindowEndLine, index)
@@ -439,6 +470,8 @@ internal abstract class CacheWindowLogic(
 
         windowCache.forEachKey { if (it in startLine..endLine) indicesToRemove.add(it) }
 
+        debugLog { "Indices to remove=$indicesToRemove" }
+
         indicesToRemove.forEach {
             prefetchWindowHandles.remove(it)?.fastForEach { it.cancel() }
             windowCache.remove(it)
@@ -501,15 +534,19 @@ internal interface CacheWindowScope {
 
     fun getVisibleItemLine(indexInVisibleLines: Int): Int
 
+    fun getVisibleLineKey(indexInVisibleLines: Int): Any
+
     fun getLastIndexInLine(lineIndex: Int): Int
 
     fun getLastLineIndex(): Int
 }
 
 internal inline fun CacheWindowScope.forEachVisibleItem(
-    action: (itemIndex: Int, mainAxisSize: Int) -> Unit
+    action: (itemIndex: Int, itemKey: Any, mainAxisSize: Int) -> Unit
 ) {
-    repeat(visibleLineCount) { action(getVisibleItemLine(it), getVisibleItemSize(it)) }
+    repeat(visibleLineCount) {
+        action(getVisibleItemLine(it), getVisibleLineKey(it), getVisibleItemSize(it))
+    }
 }
 
 private const val InvalidItemSize = -1
@@ -523,3 +560,7 @@ private inline fun debugLog(generateMsg: () -> String) {
         println("CacheWindowLogic: ${generateMsg()}")
     }
 }
+
+internal class CachedItem(var key: Any, var mainAxisSize: Int) {
+    companion object NoKey
+}

compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/pager/PagerCacheWindowLogic.kt

@@ -19,6 +19,7 @@ package androidx.compose.foundation.pager
 import androidx.compose.foundation.ExperimentalFoundationApi
 import androidx.compose.foundation.lazy.layout.CacheWindowLogic
 import androidx.compose.foundation.lazy.layout.CacheWindowScope
+import androidx.compose.foundation.lazy.layout.CachedItem
 import androidx.compose.foundation.lazy.layout.InvalidIndex
 import androidx.compose.foundation.lazy.layout.LazyLayoutCacheWindow
 import androidx.compose.foundation.lazy.layout.LazyLayoutPrefetchState
@@ -156,6 +157,30 @@ private class PagerCacheWindowScope(val itemCount: () -> Int) : CacheWindowScope
         return InvalidIndex
     }
 
+    override fun getVisibleLineKey(indexInVisibleLines: Int): Any {
+        val extraPagesBeforeCount = layoutInfo.extraPagesBefore.size
+
+        val visiblePagesCount = layoutInfo.visiblePagesInfo.size
+
+        if (indexInVisibleLines < extraPagesBeforeCount) {
+            return layoutInfo.extraPagesBefore[indexInVisibleLines].key
+        }
+
+        if (
+            indexInVisibleLines >= extraPagesBeforeCount &&
+                indexInVisibleLines < extraPagesBeforeCount + visiblePagesCount
+        ) {
+            return layoutInfo.visiblePagesInfo[indexInVisibleLines - extraPagesBeforeCount].key
+        }
+
+        if (indexInVisibleLines >= extraPagesBeforeCount + visiblePagesCount) {
+            return layoutInfo.extraPagesAfter[
+                    indexInVisibleLines - extraPagesBeforeCount - visiblePagesCount]
+                .key
+        }
+        return CachedItem.NoKey
+    }
+
     override fun getLastIndexInLine(lineIndex: Int): Int = lineIndex
 
     override fun getLastLineIndex(): Int {