Preserve order of shared items in preview

Fix: 280480016
Test: manual testing, unit tests
Change-Id: I39ed4a6db6bd68f2d89d53b03698d76c966a8001

2 files changed, 302 insertions, 35 deletions

diff --git a/java/src/com/android/intentresolver/widget/ScrollableImagePreviewView.kt b/java/src/com/android/intentresolver/widget/ScrollableImagePreviewView.kt
index e760e6d0..1f5be601 100644
--- a/java/src/com/android/intentresolver/widget/ScrollableImagePreviewView.kt
+++ b/java/src/com/android/intentresolver/widget/ScrollableImagePreviewView.kt
@@ -28,6 +28,7 @@ import android.view.View
 import android.view.ViewGroup
 import android.widget.ImageView
 import android.widget.TextView
+import androidx.annotation.VisibleForTesting
 import androidx.constraintlayout.widget.ConstraintLayout
 import androidx.recyclerview.widget.LinearLayoutManager
 import androidx.recyclerview.widget.RecyclerView
@@ -45,8 +46,6 @@ import kotlinx.coroutines.flow.takeWhile
 import kotlinx.coroutines.joinAll
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.plus
-import java.util.ArrayDeque
-import kotlin.math.roundToInt
 
 private const val TRANSITION_NAME = "screenshot_preview_image"
 private const val PLURALS_COUNT = "count"
@@ -149,14 +148,17 @@ class ScrollableImagePreviewView : RecyclerView, ImagePreviewView {
         previewAdapter.reset(0, imageLoader)
         batchLoader?.cancel()
         batchLoader = BatchPreviewLoader(
-            previewAdapter,
             imageLoader,
             previews,
             otherItemCount,
-        ) {
-            onNoPreviewCallback?.run()
-        }
-        .apply {
+            onReset = { totalItemCount -> previewAdapter.reset(totalItemCount, imageLoader) },
+            onUpdate = previewAdapter::addPreviews,
+            onCompletion = {
+                if (!previewAdapter.hasPreviews) {
+                    onNoPreviewCallback?.run()
+                }
+            }
+        ).apply {
             if (isMeasured) {
                 loadAspectRatios(getMaxWidth(), this@ScrollableImagePreviewView::updatePreviewSize)
             }
@@ -409,14 +411,17 @@ class ScrollableImagePreviewView : RecyclerView, ImagePreviewView {
         }
     }
 
-    private class BatchPreviewLoader(
-        private val adapter: Adapter,
+    @VisibleForTesting
+    class BatchPreviewLoader(
         private val imageLoader: CachingImageLoader,
         previews: List<Preview>,
         otherItemCount: Int,
-        private val onNoPreviewCallback: (() -> Unit)
+        private val onReset: (Int) -> Unit,
+        private val onUpdate: (List<Preview>) -> Unit,
+        private val onCompletion: () -> Unit,
     ) {
-        private val pendingPreviews = ArrayDeque<Preview>(previews)
+        private val previews: List<Preview> =
+            if (previews is RandomAccess) previews else ArrayList(previews)
         private val totalItemCount = previews.size + otherItemCount
         private var scope: CoroutineScope? = MainScope() + Dispatchers.Main.immediate
 
@@ -427,52 +432,75 @@ class ScrollableImagePreviewView : RecyclerView, ImagePreviewView {
 
         fun loadAspectRatios(maxWidth: Int, previewSizeUpdater: (Preview, Int, Int) -> Int) {
             val scope = this.scope ?: return
-            val updates = ArrayDeque<Preview>(pendingPreviews.size)
+            // -1 encodes that the preview has not been processed,
+            // 0 means failed, > 0 is a preview width
+            val previewWidths = IntArray(previews.size) { -1 }
+            var blockStart = 0 // inclusive
+            var blockEnd = 0 // exclusive
+
             // replay 2 items to guarantee that we'd get at least one update
             val reportFlow = MutableSharedFlow<Any>(replay = 2)
-            var isFirstUpdate = true
             val updateEvent = Any()
             val completedEvent = Any()
-            // throttle adapter updates by waiting on the channel, the channel first notified
-            // when enough preview elements is loaded and then periodically with a delay
+
+            // throttle adapter updates using flow; the flow first emits when enough preview
+            // elements is loaded to fill the viewport and then each time a subsequent block of
+            // previews is loaded
             scope.launch(Dispatchers.Main) {
                 reportFlow
                     .takeWhile { it !== completedEvent }
                     .throttle(ADAPTER_UPDATE_INTERVAL_MS)
                     .onCompletion { cause ->
-                        if (cause == null && !adapter.hasPreviews) {
-                            onNoPreviewCallback()
+                        if (cause == null) {
+                            onCompletion()
                         }
                     }
                     .collect {
-                        if (isFirstUpdate) {
-                            isFirstUpdate = false
-                            adapter.reset(totalItemCount, imageLoader)
+                        if (blockStart == 0) {
+                            onReset(totalItemCount)
+                        }
+                        val updates = ArrayList<Preview>(blockEnd - blockStart)
+                        while (blockStart < blockEnd) {
+                            if (previewWidths[blockStart] > 0) {
+                                updates.add(previews[blockStart])
+                            }
+                            blockStart++
                         }
                         if (updates.isNotEmpty()) {
-                            adapter.addPreviews(updates)
-                            updates.clear()
+                            onUpdate(updates)
                         }
                     }
             }
 
             scope.launch {
-                var loadedPreviewWidth = 0
+                var blockWidth = 0
+                var isFirstBlock = true
+                var nextIdx = 0
                 List<Job>(4) {
                     launch {
-                        while (pendingPreviews.isNotEmpty()) {
-                            val preview = pendingPreviews.poll() ?: continue
-                            val isVisible = loadedPreviewWidth < maxWidth
-                            val bitmap = runCatching {
+                        while (true) {
+                            val i = nextIdx++
+                            if (i >= previews.size) break
+                            val preview = previews[i]
+
+                            previewWidths[i] = runCatching {
                                 // TODO: decide on adding a timeout
-                                imageLoader(preview.uri, isVisible)
-                            }.getOrNull() ?: continue
-                            val previewWidth =
-                                previewSizeUpdater(preview, bitmap.width, bitmap.height)
-                            updates.add(preview)
-                            if (isVisible) {
-                                loadedPreviewWidth += previewWidth
-                                if (loadedPreviewWidth >= maxWidth) {
+                                imageLoader(preview.uri, isFirstBlock)?.let { bitmap ->
+                                    previewSizeUpdater(preview, bitmap.width, bitmap.height)
+                                } ?: 0
+                            }.getOrDefault(0)
+
+                            if (blockEnd != i) continue
+                            while (
+                                blockEnd < previewWidths.size
+                                    && previewWidths[blockEnd] >= 0
+                            ) {
+                                blockWidth += previewWidths[blockEnd]
+                                blockEnd++
+                            }
+                            if (isFirstBlock) {
+                                if (blockWidth >= maxWidth) {
+                                    isFirstBlock = false
                                     // notify that the preview now can be displayed
                                     reportFlow.emit(updateEvent)
                                 }
diff --git a/java/tests/src/com/android/intentresolver/widget/BatchPreviewLoaderTest.kt b/java/tests/src/com/android/intentresolver/widget/BatchPreviewLoaderTest.kt
new file mode 100644
index 00000000..c1d7451f
--- /dev/null
+++ b/java/tests/src/com/android/intentresolver/widget/BatchPreviewLoaderTest.kt
@@ -0,0 +1,239 @@
+/*
+ * Copyright (C) 2023 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 com.android.intentresolver.widget
+
+import android.graphics.Bitmap
+import android.net.Uri
+import com.android.intentresolver.widget.ScrollableImagePreviewView.BatchPreviewLoader
+import com.android.intentresolver.widget.ScrollableImagePreviewView.Preview
+import com.android.intentresolver.widget.ScrollableImagePreviewView.PreviewType
+import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.flow.MutableSharedFlow
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.test.UnconfinedTestDispatcher
+import org.junit.Test
+import com.android.intentresolver.mock
+import com.android.intentresolver.captureMany
+import com.android.intentresolver.withArgCaptor
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.resetMain
+import kotlinx.coroutines.test.setMain
+import org.junit.After
+import org.junit.Before
+import org.mockito.Mockito.atLeast
+import org.mockito.Mockito.times
+import org.mockito.Mockito.verify
+import com.google.common.truth.Truth.assertThat
+
+@OptIn(ExperimentalCoroutinesApi::class)
+class BatchPreviewLoaderTest {
+    private val dispatcher = UnconfinedTestDispatcher()
+    private val testScope = CoroutineScope(dispatcher)
+    private val onCompletion = mock<() -> Unit>()
+    private val onReset = mock<(Int) -> Unit>()
+    private val onUpdate = mock<(List<Preview>) -> Unit>()
+
+    @Before
+    fun setup() {
+        Dispatchers.setMain(dispatcher)
+    }
+
+    @After
+    fun cleanup() {
+        testScope.cancel()
+        Dispatchers.resetMain()
+    }
+
+    @Test
+    fun test_allImagesWithinViewPort_oneUpdate() {
+        val imageLoader = TestImageLoader(testScope)
+        val uriOne = createUri(1)
+        val uriTwo = createUri(2)
+        imageLoader.setUriLoadingOrder(succeed(uriTwo), succeed(uriOne))
+        val testSubject = BatchPreviewLoader(
+            imageLoader,
+            previews(uriOne, uriTwo),
+            0,
+            onReset,
+            onUpdate,
+            onCompletion
+        )
+        testSubject.loadAspectRatios(200) { _, _, _ -> 100 }
+        dispatcher.scheduler.advanceUntilIdle()
+
+        verify(onCompletion, times(1)).invoke()
+        verify(onReset, times(1)).invoke(2)
+        val list = withArgCaptor {
+            verify(onUpdate, times(1)).invoke(capture())
+        }.map { it.uri }
+        assertThat(list).containsExactly(uriOne, uriTwo).inOrder()
+    }
+
+    @Test
+    fun test_allImagesWithinViewPortOneFailed_failedPreviewIsNotUpdated() {
+        val imageLoader = TestImageLoader(testScope)
+        val uriOne = createUri(1)
+        val uriTwo = createUri(2)
+        val uriThree = createUri(3)
+        imageLoader.setUriLoadingOrder(succeed(uriThree), fail(uriTwo), succeed(uriOne))
+        val testSubject = BatchPreviewLoader(
+            imageLoader,
+            previews(uriOne, uriTwo, uriThree),
+            0,
+            onReset,
+            onUpdate,
+            onCompletion
+        )
+        testSubject.loadAspectRatios(200) { _, _, _ -> 100 }
+        dispatcher.scheduler.advanceUntilIdle()
+
+        verify(onCompletion, times(1)).invoke()
+        verify(onReset, times(1)).invoke(3)
+        val list = withArgCaptor {
+            verify(onUpdate, times(1)).invoke(capture())
+        }.map { it.uri }
+        assertThat(list).containsExactly(uriOne, uriThree).inOrder()
+    }
+
+    @Test
+    fun test_imagesLoadedNotInOrder_updatedInOrder() {
+        val imageLoader = TestImageLoader(testScope)
+        val uris = Array(10) { createUri(it) }
+        val loadingOrder = Array(uris.size) { i ->
+            val uriIdx = when {
+                i % 2 == 1 -> i - 1
+                i % 2 == 0 && i < uris.size - 1 -> i + 1
+                else -> i
+            }
+            succeed(uris[uriIdx])
+        }
+        imageLoader.setUriLoadingOrder(*loadingOrder)
+        val testSubject = BatchPreviewLoader(
+            imageLoader,
+            previews(*uris),
+            0,
+            onReset,
+            onUpdate,
+            onCompletion
+        )
+        testSubject.loadAspectRatios(200) { _, _, _ -> 100 }
+        dispatcher.scheduler.advanceUntilIdle()
+
+        verify(onCompletion, times(1)).invoke()
+        verify(onReset, times(1)).invoke(uris.size)
+        val list = captureMany {
+            verify(onUpdate, atLeast(1)).invoke(capture())
+        }.fold(ArrayList<Preview>()) { acc, update ->
+            acc.apply {
+                addAll(update)
+            }
+        }.map { it.uri }
+        assertThat(list).containsExactly(*uris).inOrder()
+    }
+
+    @Test
+    fun test_imagesLoadedNotInOrderSomeFailed_updatedInOrder() {
+        val imageLoader = TestImageLoader(testScope)
+        val uris = Array(10) { createUri(it) }
+        val loadingOrder = Array(uris.size) { i ->
+            val uriIdx = when {
+                i % 2 == 1 -> i - 1
+                i % 2 == 0 && i < uris.size - 1 -> i + 1
+                else -> i
+            }
+            if (uriIdx % 2 == 0) fail(uris[uriIdx]) else succeed(uris[uriIdx])
+        }
+        val expectedUris = Array(uris.size / 2) { createUri(it * 2 + 1) }
+        imageLoader.setUriLoadingOrder(*loadingOrder)
+        val testSubject = BatchPreviewLoader(
+            imageLoader,
+            previews(*uris),
+            0,
+            onReset,
+            onUpdate,
+            onCompletion
+        )
+        testSubject.loadAspectRatios(200) { _, _, _ -> 100 }
+        dispatcher.scheduler.advanceUntilIdle()
+
+        verify(onCompletion, times(1)).invoke()
+        verify(onReset, times(1)).invoke(uris.size)
+        val list = captureMany {
+            verify(onUpdate, atLeast(1)).invoke(capture())
+        }.fold(ArrayList<Preview>()) { acc, update ->
+            acc.apply {
+                addAll(update)
+            }
+        }.map { it.uri }
+        assertThat(list).containsExactly(*expectedUris).inOrder()
+    }
+
+    private fun createUri(idx: Int): Uri = Uri.parse("content://org.pkg.app/image-$idx.png")
+
+    private fun fail(uri: Uri) = uri to false
+    private fun succeed(uri: Uri) = uri to true
+    private fun previews(vararg uris: Uri) =
+        uris.fold(ArrayList<Preview>(uris.size)) { acc, uri ->
+            acc.apply {
+                add(Preview(PreviewType.Image, uri))
+            }
+        }
+}
+
+private class TestImageLoader(
+    scope: CoroutineScope
+) : suspend (Uri, Boolean) -> Bitmap? {
+    private val loadingOrder = ArrayDeque<Pair<Uri, Boolean>>()
+    private val pendingRequests = LinkedHashMap<Uri, CompletableDeferred<Bitmap?>>()
+    private val flow = MutableSharedFlow<Unit>(replay = 1)
+    private val bitmap by lazy {
+        Bitmap.createBitmap(100, 100, Bitmap.Config.ARGB_8888)
+    }
+
+    init {
+        scope.launch {
+            flow.collect {
+                while (true) {
+                    val (nextUri, isLoaded) = loadingOrder.firstOrNull() ?: break
+                    val deferred = pendingRequests.remove(nextUri) ?: break
+                    loadingOrder.removeFirst()
+                    deferred.complete(if (isLoaded) bitmap else null)
+                }
+                if (loadingOrder.isEmpty()) {
+                    pendingRequests.forEach { (uri, deferred) ->
+                        deferred.complete(bitmap)
+                    }
+                    pendingRequests.clear()
+                }
+            }
+        }
+    }
+
+    fun setUriLoadingOrder(vararg uris: Pair<Uri, Boolean>) {
+        loadingOrder.clear()
+        loadingOrder.addAll(uris)
+    }
+
+    override suspend fun invoke(uri: Uri, cache: Boolean): Bitmap? {
+        val deferred = pendingRequests.getOrPut(uri) { CompletableDeferred() }
+        flow.tryEmit(Unit)
+        return deferred.await()
+    }
+}