Shareousel: Maintain cursor order for shated items

Add a position property to PreviewModel class to track relative order of
items. For each item, the initial value is artificial and derived from
the order of the initially shared items and is updated upon reading the
additional items cursor.
Upon sharing, If the selection has not change, the items will be shared
in their original order; If the selection has changed, the order of the
items will be affected by the observed items order in the cursor.

Fix: 329683774
Test: manual testing
Test: atest IntentResolver-tests-unit
Test: atest IntentResolver-tests-activity
Flag: android.service.chooser.chooser_payload_toggling
Change-Id: Ie552887702cde75cb1a05ed3ec5415f4f4a5c8dc

18 files changed, 282 insertions, 85 deletions

diff --git a/java/src/com/android/intentresolver/contentpreview/payloadtoggle/data/repository/PreviewSelectionsRepository.kt b/java/src/com/android/intentresolver/contentpreview/payloadtoggle/data/repository/PreviewSelectionsRepository.kt
index 48c06192..81c56d1e 100644
--- a/java/src/com/android/intentresolver/contentpreview/payloadtoggle/data/repository/PreviewSelectionsRepository.kt
+++ b/java/src/com/android/intentresolver/contentpreview/payloadtoggle/data/repository/PreviewSelectionsRepository.kt
@@ -16,6 +16,7 @@
 
 package com.android.intentresolver.contentpreview.payloadtoggle.data.repository
 
+import android.net.Uri
 import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewModel
 import dagger.hilt.android.scopes.ViewModelScoped
 import javax.inject.Inject
@@ -24,5 +25,5 @@ import kotlinx.coroutines.flow.MutableStateFlow
 /** Stores set of selected previews. */
 @ViewModelScoped
 class PreviewSelectionsRepository @Inject constructor() {
-    val selections = MutableStateFlow(emptyList<PreviewModel>())
+    val selections = MutableStateFlow(emptyMap<Uri, PreviewModel>())
 }
diff --git a/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/cursor/PayloadToggleCursorResolver.kt b/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/cursor/PayloadToggleCursorResolver.kt
index d9612696..148310e6 100644
--- a/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/cursor/PayloadToggleCursorResolver.kt
+++ b/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/cursor/PayloadToggleCursorResolver.kt
@@ -55,7 +55,7 @@ constructor(
                 )
             }
             .getOrNull()
-            ?.viewBy { readUri()?.let { uri -> CursorRow(uri, readSize()) } }
+            ?.viewBy { readUri()?.let { uri -> CursorRow(uri, readSize(), position) } }
     }
 
     private fun Cursor.readUri(): Uri? {
diff --git a/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/CursorPreviewsInteractor.kt b/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/CursorPreviewsInteractor.kt
index fa600c86..a475263c 100644
--- a/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/CursorPreviewsInteractor.kt
+++ b/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/CursorPreviewsInteractor.kt
@@ -56,6 +56,7 @@ class CursorPreviewsInteractor
 @Inject
 constructor(
     private val interactor: SetCursorPreviewsInteractor,
+    private val selectionInteractor: SelectionInteractor,
     @FocusedItemIndex private val focusedItemIdx: Int,
     private val uriMetadataReader: UriMetadataReader,
     @PageSize private val pageSize: Int,
@@ -287,19 +288,26 @@ constructor(
     private fun createPreviewModel(
         row: CursorRow,
         unclaimedRecords: MutableUnclaimedMap,
-    ): PreviewModel =
-        unclaimedRecords.remove(row.uri)?.second
-            ?: uriMetadataReader.getMetadata(row.uri).let { metadata ->
-                val size =
-                    row.previewSize
-                        ?: metadata.previewUri?.let { uriMetadataReader.readPreviewSize(it) }
-                PreviewModel(
-                    uri = row.uri,
-                    previewUri = metadata.previewUri,
-                    mimeType = metadata.mimeType,
-                    aspectRatio = size.aspectRatioOrDefault(1f),
-                )
+    ): PreviewModel = uriMetadataReader.getMetadata(row.uri).let { metadata ->
+            val size =
+                row.previewSize
+                    ?: metadata.previewUri?.let { uriMetadataReader.readPreviewSize(it) }
+            PreviewModel(
+                uri = row.uri,
+                previewUri = metadata.previewUri,
+                mimeType = metadata.mimeType,
+                aspectRatio = size.aspectRatioOrDefault(1f),
+                order = row.position,
+            )
+        }.also { updated ->
+            if (unclaimedRecords.remove(row.uri) != null) {
+                // unclaimedRecords contains initially shared (and thus selected) items with unknown
+                // cursor position. Update selection records when any of those items is encountered
+                // in the cursor to maintain proper selection order should other items also be
+                // selected.
+                selectionInteractor.updateSelection(updated)
             }
+        }
 
     private fun <M : MutablePreviewMap> M.putAllUnclaimedRight(unclaimed: UnclaimedMap): M =
         putAllUnclaimedWhere(unclaimed) { it >= focusedItemIdx }
diff --git a/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/FetchPreviewsInteractor.kt b/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/FetchPreviewsInteractor.kt
index c9c9a9b3..50086a23 100644
--- a/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/FetchPreviewsInteractor.kt
+++ b/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/FetchPreviewsInteractor.kt
@@ -25,7 +25,7 @@ import com.android.intentresolver.contentpreview.payloadtoggle.domain.model.Curs
 import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewModel
 import com.android.intentresolver.inject.ContentUris
 import com.android.intentresolver.inject.FocusedItemIndex
-import com.android.intentresolver.util.mapParallel
+import com.android.intentresolver.util.mapParallelIndexed
 import javax.inject.Inject
 import kotlinx.coroutines.async
 import kotlinx.coroutines.coroutineScope
@@ -45,7 +45,7 @@ constructor(
     suspend fun activate() = coroutineScope {
         val cursor = async { cursorResolver.getCursor() }
         val initialPreviewMap = getInitialPreviews()
-        selectionRepository.selections.value = initialPreviewMap
+        selectionRepository.selections.value = initialPreviewMap.associateBy { it.uri }
         setCursorPreviews.setPreviews(
             previews = initialPreviewMap,
             startIndex = focusedItemIdx,
@@ -61,7 +61,7 @@ constructor(
         selectedItems
             // Restrict parallelism so as to not overload the metadata reader; anecdotally, too
             // many parallel queries causes failures.
-            .mapParallel(parallelism = 4) { uri ->
+            .mapParallelIndexed(parallelism = 4) { index, uri ->
                 val metadata = uriMetadataReader.getMetadata(uri)
                 PreviewModel(
                     uri = uri,
@@ -70,8 +70,12 @@ constructor(
                     aspectRatio =
                         metadata.previewUri?.let {
                             uriMetadataReader.readPreviewSize(it).aspectRatioOrDefault(1f)
-                        }
-                            ?: 1f,
+                        } ?: 1f,
+                    order = when {
+                        index < focusedItemIdx -> Int.MIN_VALUE + index
+                        index == focusedItemIdx -> 0
+                        else -> Int.MAX_VALUE - selectedItems.size + index + 1
+                    }
                 )
             }
 }
diff --git a/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectablePreviewInteractor.kt b/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectablePreviewInteractor.kt
index 55a995f5..d52a71a1 100644
--- a/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectablePreviewInteractor.kt
+++ b/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectablePreviewInteractor.kt
@@ -29,7 +29,7 @@ class SelectablePreviewInteractor(
     val uri: Uri = key.uri
 
     /** Whether or not this preview is selected by the user. */
-    val isSelected: Flow<Boolean> = selectionInteractor.selections.map { key in it }
+    val isSelected: Flow<Boolean> = selectionInteractor.selections.map { key.uri in it }
 
     /** Sets whether this preview is selected by the user. */
     fun setSelected(isSelected: Boolean) {
diff --git a/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectionInteractor.kt b/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectionInteractor.kt
index 13af92cb..97d9fa66 100644
--- a/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectionInteractor.kt
+++ b/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectionInteractor.kt
@@ -16,6 +16,7 @@
 
 package com.android.intentresolver.contentpreview.payloadtoggle.domain.interactor
 
+import android.net.Uri
 import com.android.intentresolver.contentpreview.MimeTypeClassifier
 import com.android.intentresolver.contentpreview.payloadtoggle.data.repository.PreviewSelectionsRepository
 import com.android.intentresolver.contentpreview.payloadtoggle.domain.intent.TargetIntentModifier
@@ -23,8 +24,9 @@ import com.android.intentresolver.contentpreview.payloadtoggle.shared.ContentTyp
 import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewModel
 import javax.inject.Inject
 import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.flow.updateAndGet
 
 class SelectionInteractor
@@ -36,31 +38,41 @@ constructor(
     private val mimeTypeClassifier: MimeTypeClassifier,
 ) {
     /** List of selected previews. */
-    val selections: StateFlow<List<PreviewModel>>
-        get() = selectionsRepo.selections
+    val selections: Flow<Set<Uri>> =
+        selectionsRepo.selections.map { it.keys }.distinctUntilChanged()
 
     /** Amount of selected previews. */
     val amountSelected: Flow<Int> = selectionsRepo.selections.map { it.size }
 
-    val aggregateContentType: Flow<ContentType> = selections.map { aggregateContentType(it) }
+    val aggregateContentType: Flow<ContentType> =
+        selectionsRepo.selections.map { aggregateContentType(it.values) }
+
+    fun updateSelection(model: PreviewModel) {
+        selectionsRepo.selections.update {
+            if (it.containsKey(model.uri)) it + (model.uri to model) else it
+        }
+    }
 
     fun select(model: PreviewModel) {
-        updateChooserRequest(selectionsRepo.selections.updateAndGet { it + model })
+        updateChooserRequest(
+            selectionsRepo.selections.updateAndGet { it + (model.uri to model) }.values
+        )
     }
 
     fun unselect(model: PreviewModel) {
         if (selectionsRepo.selections.value.size > 1) {
-            updateChooserRequest(selectionsRepo.selections.updateAndGet { it - model })
+            updateChooserRequest(selectionsRepo.selections.updateAndGet { it - model.uri }.values)
         }
     }
 
-    private fun updateChooserRequest(selections: List<PreviewModel>) {
-        val intent = targetIntentModifier.intentFromSelection(selections)
+    private fun updateChooserRequest(selections: Collection<PreviewModel>) {
+        val sorted = selections.sortedBy { it.order }
+        val intent = targetIntentModifier.intentFromSelection(sorted)
         updateTargetIntentInteractor.updateTargetIntent(intent)
     }
 
     private fun aggregateContentType(
-        items: List<PreviewModel>,
+        items: Collection<PreviewModel>,
     ): ContentType {
         if (items.isEmpty()) {
             return ContentType.Other
diff --git a/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/model/CursorRow.kt b/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/model/CursorRow.kt
index f1d856ac..aae29102 100644
--- a/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/model/CursorRow.kt
+++ b/java/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/model/CursorRow.kt
@@ -20,4 +20,4 @@ import android.net.Uri
 import android.util.Size
 
 /** Represents additional content cursor row */
-data class CursorRow(val uri: Uri, val previewSize: Size?)
+data class CursorRow(val uri: Uri, val previewSize: Size?, val position: Int)
diff --git a/java/src/com/android/intentresolver/contentpreview/payloadtoggle/shared/model/PreviewModel.kt b/java/src/com/android/intentresolver/contentpreview/payloadtoggle/shared/model/PreviewModel.kt
index 85c70004..8a479156 100644
--- a/java/src/com/android/intentresolver/contentpreview/payloadtoggle/shared/model/PreviewModel.kt
+++ b/java/src/com/android/intentresolver/contentpreview/payloadtoggle/shared/model/PreviewModel.kt
@@ -27,4 +27,8 @@ data class PreviewModel(
     /** Mimetype for the data [uri] points to. */
     val mimeType: String?,
     val aspectRatio: Float = 1f,
+    /**
+     * Relative item position in the list that is used to determine items order in the target intent
+     */
+    val order: Int,
 )
diff --git a/java/src/com/android/intentresolver/util/ParallelIteration.kt b/java/src/com/android/intentresolver/util/ParallelIteration.kt
index 70c46c47..745bcdbf 100644
--- a/java/src/com/android/intentresolver/util/ParallelIteration.kt
+++ b/java/src/com/android/intentresolver/util/ParallelIteration.kt
@@ -48,3 +48,24 @@ private suspend fun <A, B> Iterable<A>.mapParallel(block: suspend (A) -> B): Lis
             }
             .awaitAll()
     }
+
+suspend fun <A, B> Iterable<A>.mapParallelIndexed(
+    parallelism: Int? = null,
+    block: suspend (Int, A) -> B,
+): List<B> =
+    parallelism?.let { permits ->
+        withSemaphore(permits = permits) {
+            mapParallelIndexed { idx, item -> withPermit { block(idx, item) } }
+        }
+    } ?: mapParallelIndexed(block)
+
+private suspend fun <A, B> Iterable<A>.mapParallelIndexed(block: suspend (Int, A) -> B): List<B> =
+    coroutineScope {
+        mapIndexed { index, item ->
+                async {
+                    yield()
+                    block(index, item)
+                }
+            }
+            .awaitAll()
+    }
diff --git a/tests/shared/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/PayloadToggleInteractorKosmos.kt b/tests/shared/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/PayloadToggleInteractorKosmos.kt
index 8f7c59de..cb88cd9e 100644
--- a/tests/shared/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/PayloadToggleInteractorKosmos.kt
+++ b/tests/shared/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/PayloadToggleInteractorKosmos.kt
@@ -48,6 +48,7 @@ val Kosmos.cursorPreviewsInteractor
     get() =
         CursorPreviewsInteractor(
             interactor = setCursorPreviewsInteractor,
+            selectionInteractor = selectionInteractor,
             focusedItemIdx = focusedItemIndex,
             uriMetadataReader = uriMetadataReader,
             pageSize = pageSize,
diff --git a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/cursor/PayloadToggleCursorResolverTest.kt b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/cursor/PayloadToggleCursorResolverTest.kt
index 9eaee233..5d81ec2a 100644
--- a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/cursor/PayloadToggleCursorResolverTest.kt
+++ b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/cursor/PayloadToggleCursorResolverTest.kt
@@ -26,6 +26,7 @@ import android.service.chooser.AdditionalContentContract.Columns.URI
 import android.util.Size
 import com.android.intentresolver.util.cursor.get
 import com.google.common.truth.Truth.assertThat
+import com.google.common.truth.Truth.assertWithMessage
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
 import org.mockito.kotlin.any
@@ -101,6 +102,37 @@ class PayloadToggleCursorResolverTest {
             assertThat(row.previewSize).isEqualTo(Size(100, 50))
         }
     }
+
+    @Test
+    fun testRowPositionValues() = runTest {
+        val rowCount = 10
+        val sourceCursor =
+            MatrixCursor(arrayOf(URI)).apply {
+                for (i in 1..rowCount) {
+                    addRow(arrayOf(createUri(i).toString()))
+                }
+            }
+        val fakeContentProvider =
+            mock<ContentInterface> {
+                on { query(eq(cursorUri), any(), any(), any()) } doReturn sourceCursor
+            }
+        val testSubject =
+            PayloadToggleCursorResolver(
+                fakeContentProvider,
+                cursorUri,
+                chooserIntent,
+            )
+
+        val cursor = testSubject.getCursor()
+        assertThat(cursor).isNotNull()
+        assertThat(cursor!!.count).isEqualTo(rowCount)
+        for (i in 0 until rowCount) {
+            cursor[i].let { row ->
+                assertWithMessage("Row $i").that(row).isNotNull()
+                assertWithMessage("Row $i").that(row!!.position).isEqualTo(i)
+            }
+        }
+    }
 }
 
 private fun createUri(id: Int) = Uri.parse("content://org.pkg/app/img-$id.png")
diff --git a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/CursorPreviewsInteractorTest.kt b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/CursorPreviewsInteractorTest.kt
index 0036e803..48e43190 100644
--- a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/CursorPreviewsInteractorTest.kt
+++ b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/CursorPreviewsInteractorTest.kt
@@ -27,6 +27,9 @@ import androidx.core.os.bundleOf
 import com.android.intentresolver.contentpreview.FileInfo
 import com.android.intentresolver.contentpreview.UriMetadataReader
 import com.android.intentresolver.contentpreview.payloadtoggle.data.repository.cursorPreviewsRepository
+import com.android.intentresolver.contentpreview.payloadtoggle.data.repository.previewSelectionsRepository
+import com.android.intentresolver.contentpreview.payloadtoggle.domain.intent.TargetIntentModifier
+import com.android.intentresolver.contentpreview.payloadtoggle.domain.intent.targetIntentModifier
 import com.android.intentresolver.contentpreview.payloadtoggle.domain.model.CursorRow
 import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewModel
 import com.android.intentresolver.contentpreview.readSize
@@ -59,6 +62,7 @@ class CursorPreviewsInteractorTest {
             this.focusedItemIndex = focusedItemIndex
             this.pageSize = pageSize
             this.maxLoadedPages = maxLoadedPages
+            this.targetIntentModifier = TargetIntentModifier { error("unexpected invocation") }
             uriMetadataReader =
                 object : UriMetadataReader {
                     override fun getMetadata(uri: Uri): FileInfo =
@@ -103,9 +107,15 @@ class CursorPreviewsInteractorTest {
                         )
                     }
                 }
-                .viewBy { getString(0)?.let { uriStr -> CursorRow(Uri.parse(uriStr), readSize()) } }
+                .viewBy {
+                    getString(0)?.let { uriStr ->
+                        CursorRow(Uri.parse(uriStr), readSize(), position)
+                    }
+                }
         val initialPreviews: List<PreviewModel> =
-            initialSelectionRange.map { i -> PreviewModel(uri = uri(i), mimeType = "image/bitmap") }
+            initialSelectionRange.map { i ->
+                PreviewModel(uri = uri(i), mimeType = "image/bitmap", order = i)
+            }
     }
 
     @Test
@@ -136,7 +146,8 @@ class CursorPreviewsInteractorTest {
                                         0 -> 2f
                                         3 -> 4f
                                         else -> 1f
-                                    }
+                                    },
+                                order = it,
                             )
                         }
                     )
@@ -257,6 +268,34 @@ class CursorPreviewsInteractorTest {
                 .isEqualTo(Uri.fromParts("scheme24", "ssp24", "fragment24"))
             assertThat(cursorPreviewsRepository.previewsModel.value!!.loadMoreLeft).isNull()
         }
+
+    @Test
+    fun unclaimedRecordsGotUpdatedInSelectionInteractor() =
+        runTestWithDeps(
+            initialSelection = listOf(1),
+            focusedItemIndex = 0,
+            cursor = listOf(0, 1),
+            cursorStartPosition = 1,
+        ) { deps ->
+            previewSelectionsRepository.selections.value =
+                PreviewModel(
+                    uri = uri(1),
+                    mimeType = "image/png",
+                    order = 0,
+                ).let { mapOf(it.uri to it) }
+            backgroundScope.launch {
+                cursorPreviewsInteractor.launch(deps.cursor, deps.initialPreviews)
+            }
+            runCurrent()
+
+            assertThat(previewSelectionsRepository.selections.value.values).containsExactly(
+                PreviewModel(
+                    uri = uri(1),
+                    mimeType = "image/bitmap",
+                    order = 1,
+                )
+            )
+        }
 }
 
 private fun uri(index: Int) = Uri.fromParts("scheme$index", "ssp$index", "fragment$index")
diff --git a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/FetchPreviewsInteractorTest.kt b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/FetchPreviewsInteractorTest.kt
index d04c984f..27c98dc0 100644
--- a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/FetchPreviewsInteractorTest.kt
+++ b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/FetchPreviewsInteractorTest.kt
@@ -27,6 +27,8 @@ import com.android.intentresolver.contentpreview.UriMetadataReader
 import com.android.intentresolver.contentpreview.payloadtoggle.data.repository.cursorPreviewsRepository
 import com.android.intentresolver.contentpreview.payloadtoggle.domain.cursor.CursorResolver
 import com.android.intentresolver.contentpreview.payloadtoggle.domain.cursor.payloadToggleCursorResolver
+import com.android.intentresolver.contentpreview.payloadtoggle.domain.intent.TargetIntentModifier
+import com.android.intentresolver.contentpreview.payloadtoggle.domain.intent.targetIntentModifier
 import com.android.intentresolver.contentpreview.payloadtoggle.domain.model.CursorRow
 import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewModel
 import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewsModel
@@ -76,6 +78,7 @@ class FetchPreviewsInteractorTest {
                 }
             this.pageSize = pageSize
             this.maxLoadedPages = maxLoadedPages
+            this.targetIntentModifier = TargetIntentModifier { error("unexpected invocation") }
             runKosmosTest { block() }
         }
     }
@@ -99,13 +102,15 @@ class FetchPreviewsInteractorTest {
                             newRow().add("uri", uri(i).toString())
                         }
                     }
-                    .viewBy { getString(0)?.let(Uri::parse)?.let { CursorRow(it, null) } }
+                    .viewBy { getString(0)?.let(Uri::parse)?.let { CursorRow(it, null, position) } }
             }
     }
 
     @Test
     fun setsInitialPreviews() =
-        runTest(previewSizes = mapOf(1 to Size(100, 50))) {
+        runTest(
+            initialSelection = (1..3),
+            previewSizes = mapOf(1 to Size(100, 50))) {
             backgroundScope.launch { fetchPreviewsInteractor.activate() }
             runCurrent()
 
@@ -117,18 +122,25 @@ class FetchPreviewsInteractorTest {
                                 PreviewModel(
                                     uri = Uri.fromParts("scheme1", "ssp1", "fragment1"),
                                     mimeType = "image/bitmap",
-                                    aspectRatio = 2f
+                                    aspectRatio = 2f,
+                                    order = Int.MIN_VALUE,
                                 ),
                                 PreviewModel(
                                     uri = Uri.fromParts("scheme2", "ssp2", "fragment2"),
                                     mimeType = "image/bitmap",
+                                    order = 0,
+                                ),
+                                PreviewModel(
+                                    uri = Uri.fromParts("scheme3", "ssp3", "fragment3"),
+                                    mimeType = "image/bitmap",
+                                    order = Int.MAX_VALUE,
                                 ),
                             ),
                         startIdx = 1,
                         loadMoreLeft = null,
                         loadMoreRight = null,
                         leftTriggerIndex = 0,
-                        rightTriggerIndex = 1,
+                        rightTriggerIndex = 2,
                     )
                 )
         }
@@ -148,19 +160,23 @@ class FetchPreviewsInteractorTest {
                 .containsExactly(
                     PreviewModel(
                         uri = Uri.fromParts("scheme0", "ssp0", "fragment0"),
-                        mimeType = "image/bitmap"
+                        mimeType = "image/bitmap",
+                        order = 0,
                     ),
                     PreviewModel(
                         uri = Uri.fromParts("scheme1", "ssp1", "fragment1"),
-                        mimeType = "image/bitmap"
+                        mimeType = "image/bitmap",
+                        order = 1,
                     ),
                     PreviewModel(
                         uri = Uri.fromParts("scheme2", "ssp2", "fragment2"),
-                        mimeType = "image/bitmap"
+                        mimeType = "image/bitmap",
+                        order = 2,
                     ),
                     PreviewModel(
                         uri = Uri.fromParts("scheme3", "ssp3", "fragment3"),
-                        mimeType = "image/bitmap"
+                        mimeType = "image/bitmap",
+                        order = 3,
                     ),
                 )
                 .inOrder()
diff --git a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectablePreviewInteractorTest.kt b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectablePreviewInteractorTest.kt
index 0275a9c3..f329b8a7 100644
--- a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectablePreviewInteractorTest.kt
+++ b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectablePreviewInteractorTest.kt
@@ -40,7 +40,11 @@ class SelectablePreviewInteractorTest {
         val underTest =
             SelectablePreviewInteractor(
                 key =
-                    PreviewModel(uri = Uri.fromParts("scheme", "ssp", "fragment"), mimeType = null),
+                    PreviewModel(
+                        uri = Uri.fromParts("scheme", "ssp", "fragment"),
+                        mimeType = null,
+                        order = 0,
+                    ),
                 selectionInteractor = selectionInteractor,
             )
         runCurrent()
@@ -56,7 +60,8 @@ class SelectablePreviewInteractorTest {
                 key =
                     PreviewModel(
                         uri = Uri.fromParts("scheme", "ssp", "fragment"),
-                        mimeType = "image/bitmap"
+                        mimeType = "image/bitmap",
+                        order = 0,
                     ),
                 selectionInteractor = selectionInteractor,
             )
@@ -64,12 +69,12 @@ class SelectablePreviewInteractorTest {
         assertThat(underTest.isSelected.first()).isFalse()
 
         previewSelectionsRepository.selections.value =
-            listOf(
-                PreviewModel(
+            PreviewModel(
                     uri = Uri.fromParts("scheme", "ssp", "fragment"),
-                    mimeType = "image/bitmap"
+                    mimeType = "image/bitmap",
+                    order = 0,
                 )
-            )
+                .let { mapOf(it.uri to it) }
         runCurrent()
 
         assertThat(underTest.isSelected.first()).isTrue()
@@ -84,7 +89,8 @@ class SelectablePreviewInteractorTest {
                 key =
                     PreviewModel(
                         uri = Uri.fromParts("scheme", "ssp", "fragment"),
-                        mimeType = "image/bitmap"
+                        mimeType = "image/bitmap",
+                        order = 0,
                     ),
                 selectionInteractor = selectionInteractor,
             )
@@ -92,13 +98,8 @@ class SelectablePreviewInteractorTest {
         underTest.setSelected(true)
         runCurrent()
 
-        assertThat(previewSelectionsRepository.selections.value)
-            .containsExactly(
-                PreviewModel(
-                    uri = Uri.fromParts("scheme", "ssp", "fragment"),
-                    mimeType = "image/bitmap"
-                )
-            )
+        assertThat(previewSelectionsRepository.selections.value.keys)
+            .containsExactly(Uri.fromParts("scheme", "ssp", "fragment"))
 
         assertThat(chooserRequestRepository.chooserRequest.value.targetIntent)
             .isSameInstanceAs(modifiedIntent)
diff --git a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectablePreviewsInteractorTest.kt b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectablePreviewsInteractorTest.kt
index 14b9c49c..c50d2d3f 100644
--- a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectablePreviewsInteractorTest.kt
+++ b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectablePreviewsInteractorTest.kt
@@ -43,10 +43,12 @@ class SelectablePreviewsInteractorTest {
                         PreviewModel(
                             uri = Uri.fromParts("scheme", "ssp", "fragment"),
                             mimeType = "image/bitmap",
+                            order = 0,
                         ),
                         PreviewModel(
                             uri = Uri.fromParts("scheme2", "ssp2", "fragment2"),
                             mimeType = "image/bitmap",
+                            order = 1,
                         ),
                     ),
                 startIdx = 0,
@@ -56,9 +58,12 @@ class SelectablePreviewsInteractorTest {
                 rightTriggerIndex = 1,
             )
         previewSelectionsRepository.selections.value =
-            listOf(
-                PreviewModel(uri = Uri.fromParts("scheme", "ssp", "fragment"), mimeType = null),
-            )
+            PreviewModel(
+                    uri = Uri.fromParts("scheme", "ssp", "fragment"),
+                    mimeType = null,
+                    order = 0,
+                )
+                .let { mapOf(it.uri to it) }
         targetIntentModifier = TargetIntentModifier { error("unexpected invocation") }
         val underTest = selectablePreviewsInteractor
         val keySet = underTest.previews.stateIn(backgroundScope)
@@ -68,11 +73,13 @@ class SelectablePreviewsInteractorTest {
             .containsExactly(
                 PreviewModel(
                     uri = Uri.fromParts("scheme", "ssp", "fragment"),
-                    mimeType = "image/bitmap"
+                    mimeType = "image/bitmap",
+                    order = 0,
                 ),
                 PreviewModel(
                     uri = Uri.fromParts("scheme2", "ssp2", "fragment2"),
-                    mimeType = "image/bitmap"
+                    mimeType = "image/bitmap",
+                    order = 1,
                 ),
             )
             .inOrder()
@@ -82,13 +89,21 @@ class SelectablePreviewsInteractorTest {
 
         val firstModel =
             underTest.preview(
-                PreviewModel(uri = Uri.fromParts("scheme", "ssp", "fragment"), mimeType = null)
+                PreviewModel(
+                    uri = Uri.fromParts("scheme", "ssp", "fragment"),
+                    mimeType = null,
+                    order = 0,
+                )
             )
         assertThat(firstModel.isSelected.first()).isTrue()
 
         val secondModel =
             underTest.preview(
-                PreviewModel(uri = Uri.fromParts("scheme2", "ssp2", "fragment2"), mimeType = null)
+                PreviewModel(
+                    uri = Uri.fromParts("scheme2", "ssp2", "fragment2"),
+                    mimeType = null,
+                    order = 1,
+                )
             )
         assertThat(secondModel.isSelected.first()).isFalse()
     }
@@ -96,16 +111,23 @@ class SelectablePreviewsInteractorTest {
     @Test
     fun keySet_reflectsRepositoryUpdate() = runKosmosTest {
         previewSelectionsRepository.selections.value =
-            listOf(
-                PreviewModel(uri = Uri.fromParts("scheme", "ssp", "fragment"), mimeType = null),
-            )
+            PreviewModel(
+                    uri = Uri.fromParts("scheme", "ssp", "fragment"),
+                    mimeType = null,
+                    order = 0,
+                )
+                .let { mapOf(it.uri to it) }
         targetIntentModifier = TargetIntentModifier { error("unexpected invocation") }
         val underTest = selectablePreviewsInteractor
 
         val previews = underTest.previews.stateIn(backgroundScope)
         val firstModel =
             underTest.preview(
-                PreviewModel(uri = Uri.fromParts("scheme", "ssp", "fragment"), mimeType = null)
+                PreviewModel(
+                    uri = Uri.fromParts("scheme", "ssp", "fragment"),
+                    mimeType = null,
+                    order = 0,
+                )
             )
 
         assertThat(previews.value).isNull()
@@ -120,10 +142,12 @@ class SelectablePreviewsInteractorTest {
                         PreviewModel(
                             uri = Uri.fromParts("scheme", "ssp", "fragment"),
                             mimeType = "image/bitmap",
+                            order = 0,
                         ),
                         PreviewModel(
                             uri = Uri.fromParts("scheme2", "ssp2", "fragment2"),
                             mimeType = "image/bitmap",
+                            order = 1,
                         ),
                     ),
                 startIdx = 5,
@@ -132,7 +156,7 @@ class SelectablePreviewsInteractorTest {
                 leftTriggerIndex = 0,
                 rightTriggerIndex = 1,
             )
-        previewSelectionsRepository.selections.value = emptyList()
+        previewSelectionsRepository.selections.value = emptyMap()
         runCurrent()
 
         assertThat(previews.value).isNotNull()
@@ -140,11 +164,13 @@ class SelectablePreviewsInteractorTest {
             .containsExactly(
                 PreviewModel(
                     uri = Uri.fromParts("scheme", "ssp", "fragment"),
-                    mimeType = "image/bitmap"
+                    mimeType = "image/bitmap",
+                    order = 0,
                 ),
                 PreviewModel(
                     uri = Uri.fromParts("scheme2", "ssp2", "fragment2"),
-                    mimeType = "image/bitmap"
+                    mimeType = "image/bitmap",
+                    order = 1,
                 ),
             )
             .inOrder()
diff --git a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectionInteractorTest.kt b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectionInteractorTest.kt
index a50efebf..87db243d 100644
--- a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectionInteractorTest.kt
+++ b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SelectionInteractorTest.kt
@@ -23,14 +23,19 @@ import com.android.intentresolver.contentpreview.payloadtoggle.data.repository.p
 import com.android.intentresolver.contentpreview.payloadtoggle.shared.model.PreviewModel
 import com.android.intentresolver.util.runKosmosTest
 import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.flow.first
 import org.junit.Test
 
 class SelectionInteractorTest {
     @Test
     fun singleSelection_removalPrevented() = runKosmosTest {
         val initialPreview =
-            PreviewModel(uri = Uri.fromParts("scheme", "ssp", "fragment"), mimeType = null)
-        previewSelectionsRepository.selections.value = listOf(initialPreview)
+            PreviewModel(
+                uri = Uri.fromParts("scheme", "ssp", "fragment"),
+                mimeType = null,
+                order = 0
+            )
+        previewSelectionsRepository.selections.value = mapOf(initialPreview.uri to initialPreview)
 
         val underTest =
             SelectionInteractor(
@@ -40,20 +45,29 @@ class SelectionInteractorTest {
                 mimetypeClassifier,
             )
 
-        assertThat(underTest.selections.value).containsExactly(initialPreview)
+        assertThat(underTest.selections.first()).containsExactly(initialPreview.uri)
 
         // Shouldn't do anything!
         underTest.unselect(initialPreview)
 
-        assertThat(underTest.selections.value).containsExactly(initialPreview)
+        assertThat(underTest.selections.first()).containsExactly(initialPreview.uri)
     }
 
     @Test
     fun multipleSelections_removalAllowed() = runKosmosTest {
-        val first = PreviewModel(uri = Uri.fromParts("scheme", "ssp", "fragment"), mimeType = null)
+        val first =
+            PreviewModel(
+                uri = Uri.fromParts("scheme", "ssp", "fragment"),
+                mimeType = null,
+                order = 0
+            )
         val second =
-            PreviewModel(uri = Uri.fromParts("scheme2", "ssp2", "fragment2"), mimeType = null)
-        previewSelectionsRepository.selections.value = listOf(first, second)
+            PreviewModel(
+                uri = Uri.fromParts("scheme2", "ssp2", "fragment2"),
+                mimeType = null,
+                order = 1
+            )
+        previewSelectionsRepository.selections.value = listOf(first, second).associateBy { it.uri }
 
         val underTest =
             SelectionInteractor(
@@ -65,6 +79,6 @@ class SelectionInteractorTest {
 
         underTest.unselect(first)
 
-        assertThat(underTest.selections.value).containsExactly(second)
+        assertThat(underTest.selections.first()).containsExactly(second.uri)
     }
 }
diff --git a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SetCursorPreviewsInteractorTest.kt b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SetCursorPreviewsInteractorTest.kt
index a165b41e..748459cb 100644
--- a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SetCursorPreviewsInteractorTest.kt
+++ b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/domain/interactor/SetCursorPreviewsInteractorTest.kt
@@ -39,6 +39,7 @@ class SetCursorPreviewsInteractorTest {
                         PreviewModel(
                             uri = Uri.fromParts("scheme", "ssp", "fragment"),
                             mimeType = null,
+                            order = 0,
                         )
                     ),
                 startIndex = 100,
@@ -60,6 +61,7 @@ class SetCursorPreviewsInteractorTest {
                     PreviewModel(
                         uri = Uri.fromParts("scheme", "ssp", "fragment"),
                         mimeType = null,
+                        order = 0
                     )
                 )
                 .inOrder()
@@ -76,6 +78,7 @@ class SetCursorPreviewsInteractorTest {
                             PreviewModel(
                                 uri = Uri.fromParts("scheme", "ssp", "fragment"),
                                 mimeType = null,
+                                order = 0,
                             )
                         ),
                     startIndex = 100,
diff --git a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/ui/viewmodel/ShareouselViewModelTest.kt b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/ui/viewmodel/ShareouselViewModelTest.kt
index a26b4288..bb67e084 100644
--- a/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/ui/viewmodel/ShareouselViewModelTest.kt
+++ b/tests/unit/src/com/android/intentresolver/contentpreview/payloadtoggle/ui/viewmodel/ShareouselViewModelTest.kt
@@ -85,12 +85,14 @@ class ShareouselViewModelTest {
                 PreviewModel(
                     uri = Uri.fromParts("scheme", "ssp", "fragment"),
                     mimeType = "image/png",
+                    order = 0,
                 ),
                 PreviewModel(
                     uri = Uri.fromParts("scheme1", "ssp1", "fragment1"),
                     mimeType = "image/jpeg",
+                    order = 1,
                 )
-            )
+            ).associateBy { it.uri }
         runCurrent()
         assertThat(shareouselViewModel.headline.first()).isEqualTo("IMAGES: 2")
     }
@@ -102,12 +104,14 @@ class ShareouselViewModelTest {
                 PreviewModel(
                     uri = Uri.fromParts("scheme", "ssp", "fragment"),
                     mimeType = "video/mpeg",
+                    order = 0,
                 ),
                 PreviewModel(
                     uri = Uri.fromParts("scheme1", "ssp1", "fragment1"),
                     mimeType = "video/mpeg",
+                    order = 1,
                 )
-            )
+            ).associateBy { it.uri }
         runCurrent()
         assertThat(shareouselViewModel.headline.first()).isEqualTo("VIDEOS: 2")
     }
@@ -119,12 +123,14 @@ class ShareouselViewModelTest {
                 PreviewModel(
                     uri = Uri.fromParts("scheme", "ssp", "fragment"),
                     mimeType = "image/jpeg",
+                    order = 0,
                 ),
                 PreviewModel(
                     uri = Uri.fromParts("scheme1", "ssp1", "fragment1"),
                     mimeType = "video/mpeg",
+                    order = 1,
                 )
-            )
+            ).associateBy { it.uri }
         runCurrent()
         assertThat(shareouselViewModel.headline.first()).isEqualTo("FILES: 2")
     }
@@ -154,10 +160,12 @@ class ShareouselViewModelTest {
                             PreviewModel(
                                 uri = Uri.fromParts("scheme", "ssp", "fragment"),
                                 mimeType = "image/png",
+                                order = 0,
                             ),
                             PreviewModel(
                                 uri = Uri.fromParts("scheme1", "ssp1", "fragment1"),
                                 mimeType = "video/mpeg",
+                                order = 1,
                             )
                         ),
                     startIdx = 1,
@@ -183,7 +191,8 @@ class ShareouselViewModelTest {
                 shareouselViewModel.preview.invoke(
                     PreviewModel(
                         uri = Uri.fromParts("scheme1", "ssp1", "fragment1"),
-                        mimeType = "video/mpeg"
+                        mimeType = "video/mpeg",
+                        order = 0,
                     ),
                     /* index = */ 1,
                     viewModelScope,
@@ -199,8 +208,7 @@ class ShareouselViewModelTest {
 
             previewVm.setSelected(true)
 
-            assertThat(previewSelectionsRepository.selections.value)
-                .comparingElementsUsingTransform("has uri of") { model: PreviewModel -> model.uri }
+            assertThat(previewSelectionsRepository.selections.value.keys)
                 .contains(Uri.fromParts("scheme1", "ssp1", "fragment1"))
         }
 
@@ -214,10 +222,12 @@ class ShareouselViewModelTest {
                             PreviewModel(
                                 uri = Uri.fromParts("scheme", "ssp", "fragment"),
                                 mimeType = "image/png",
+                                order = 0,
                             ),
                             PreviewModel(
                                 uri = Uri.fromParts("scheme1", "ssp1", "fragment1"),
                                 mimeType = "video/mpeg",
+                                order = 1,
                             )
                         ),
                     startIdx = 1,
@@ -232,7 +242,8 @@ class ShareouselViewModelTest {
                 shareouselViewModel.preview.invoke(
                     PreviewModel(
                         uri = Uri.fromParts("scheme", "ssp", "fragment"),
-                        mimeType = "video/mpeg"
+                        mimeType = "video/mpeg",
+                        order = 1,
                     ),
                     /* index = */ 1,
                     viewModelScope,
@@ -296,7 +307,11 @@ class ShareouselViewModelTest {
         this.pendingIntentSender = pendingIntentSender
         this.targetIntentModifier = targetIntentModifier
         previewSelectionsRepository.selections.value =
-            listOf(PreviewModel(uri = Uri.fromParts("scheme", "ssp", "fragment"), mimeType = null))
+            PreviewModel(
+                uri = Uri.fromParts("scheme", "ssp", "fragment"),
+                mimeType = null,
+                order = 0,
+            ).let { mapOf(it.uri to it) }
         payloadToggleImageLoader =
             FakeImageLoader(
                 initialBitmaps =