Merge changes I0d568a6c,I092d9b67 into main

* changes:
  Add PayloadToggleInteractor implementation
  Shareousel selection tracker component

6 files changed, 522 insertions, 27 deletions

diff --git a/java/src/com/android/intentresolver/contentpreview/CursorUriReader.kt b/java/src/com/android/intentresolver/contentpreview/CursorUriReader.kt
index 30495b8b..dbf27a88 100644
--- a/java/src/com/android/intentresolver/contentpreview/CursorUriReader.kt
+++ b/java/src/com/android/intentresolver/contentpreview/CursorUriReader.kt
@@ -34,9 +34,11 @@ class CursorUriReader(
     private val predicate: (Uri) -> Boolean,
 ) : PayloadToggleInteractor.CursorReader {
     override val count = cursor.count
-    // the first position of the next unread page on the right
+    // Unread ranges are:
+    // - left: [0, leftPos);
+    // - right: [rightPos, count)
+    // i.e. read range is: [leftPos, rightPos)
     private var rightPos = startPos.coerceIn(0, count)
-    // the first position of the next from the leftmost unread page on the left
     private var leftPos = rightPos
 
     override val hasMoreBefore
@@ -74,7 +76,7 @@ class CursorUriReader(
             return SparseArray()
         }
         val result = SparseArray<Uri>(leftPos - startPos)
-        for (pos in startPos ..< leftPos) {
+        for (pos in startPos until leftPos) {
             cursor
                 .getString(0)
                 ?.let(Uri::parse)
diff --git a/java/src/com/android/intentresolver/contentpreview/PayloadToggleInteractor.kt b/java/src/com/android/intentresolver/contentpreview/PayloadToggleInteractor.kt
index ca868226..3393dcfc 100644
--- a/java/src/com/android/intentresolver/contentpreview/PayloadToggleInteractor.kt
+++ b/java/src/com/android/intentresolver/contentpreview/PayloadToggleInteractor.kt
@@ -16,44 +16,357 @@
 
 package com.android.intentresolver.contentpreview
 
+import android.content.Intent
 import android.net.Uri
 import android.service.chooser.ChooserAction
+import android.util.Log
 import android.util.SparseArray
 import java.io.Closeable
+import java.util.LinkedList
+import java.util.concurrent.atomic.AtomicBoolean
+import java.util.concurrent.atomic.AtomicReference
+import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.awaitCancellation
+import kotlinx.coroutines.channels.BufferOverflow.DROP_LATEST
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.flowOf
+import kotlinx.coroutines.flow.filter
+import kotlinx.coroutines.flow.flow
 import kotlinx.coroutines.flow.map
-import kotlinx.coroutines.flow.update
+import kotlinx.coroutines.launch
 
-class PayloadToggleInteractor {
+private const val TAG = "PayloadToggleInteractor"
 
-    private val storage = MutableStateFlow<Map<Any, Item>>(emptyMap()) // TODO: implement
-    private val selectedKeys = MutableStateFlow<Set<Any>>(emptySet())
+class PayloadToggleInteractor(
+    // must use single-thread dispatcher (or we should enforce it with a lock)
+    private val scope: CoroutineScope,
+    private val initiallySharedUris: List<Uri>,
+    private val focusedUriIdx: Int,
+    private val mimeTypeClassifier: MimeTypeClassifier,
+    private val cursorReaderProvider: suspend () -> CursorReader,
+    private val uriMetadataReader: (Uri) -> FileInfo,
+    private val targetIntentModifier: (List<Item>) -> Intent,
+    private val selectionCallback: (Intent) -> CallbackResult?,
+) {
+    private var cursorDataRef = CompletableDeferred<CursorData?>()
+    private val records = LinkedList<Record>()
+    private val prevPageLoadingGate = AtomicBoolean(true)
+    private val nextPageLoadingGate = AtomicBoolean(true)
+    private val notifySelectionJobRef = AtomicReference<Job?>()
+    private val emptyState =
+        State(
+            emptyList(),
+            hasMoreItemsBefore = false,
+            hasMoreItemsAfter = false,
+            allowSelectionChange = false
+        )
+
+    private val stateFlowSource = MutableStateFlow(emptyState)
+
+    val customActions =
+        MutableSharedFlow<List<ChooserAction>>(replay = 1, onBufferOverflow = DROP_LATEST)
+
+    val stateFlow: Flow<State>
+        get() = stateFlowSource.filter { it !== emptyState }
+
+    val targetPosition: Flow<Int> = stateFlow.map { it.targetPos }
+    val previewKeys: Flow<List<Item>> = stateFlow.map { it.items }
+
+    fun getKey(item: Any): Int = (item as Item).key
+
+    fun selected(key: Item): Flow<Boolean> = (key as Record).isSelected
 
-    val targetPosition: Flow<Int> = flowOf(0) // TODO: implement
-    val previewKeys: Flow<List<Any>> = flowOf(emptyList()) // TODO: implement
+    fun previewUri(key: Item): Flow<Uri?> = flow { emit(key.previewUri) }
 
-    fun setSelected(key: Any, isSelected: Boolean) {
-        if (isSelected) {
-            selectedKeys.update { it + key }
+    fun previewInteractor(key: Any): PayloadTogglePreviewInteractor {
+        val state = stateFlowSource.value
+        if (state === emptyState) {
+            Log.wtf(TAG, "Requesting item preview before any item has been published")
         } else {
-            selectedKeys.update { it - key }
+            if (state.hasMoreItemsBefore && key === state.items.firstOrNull()) {
+                loadMorePreviousItems()
+            }
+            if (state.hasMoreItemsAfter && key == state.items.lastOrNull()) {
+                loadMoreNextItems()
+            }
         }
+        return PayloadTogglePreviewInteractor(key as Item, this)
     }
 
-    fun selected(key: Any): Flow<Boolean> = previewKeys.map { key in it }
+    init {
+        scope
+            .launch { awaitCancellation() }
+            .invokeOnCompletion {
+                cursorDataRef.cancel()
+                runCatching {
+                        if (cursorDataRef.isCompleted && !cursorDataRef.isCancelled) {
+                            cursorDataRef.getCompleted()
+                        } else {
+                            null
+                        }
+                    }
+                    .getOrNull()
+                    ?.reader
+                    ?.close()
+            }
+    }
 
-    fun previewInteractor(key: Any) = PayloadTogglePreviewInteractor(key, this)
+    fun start() {
+        scope.launch {
+            publishInitialState()
+            val cursorReader = cursorReaderProvider()
+            val selectedItems =
+                initiallySharedUris.map { uri ->
+                    val fileInfo = uriMetadataReader(uri)
+                    Record(
+                        0, // artificial key for the pending record, it should not be used anywhere
+                        uri,
+                        fileInfo.previewUri,
+                        fileInfo.mimeType,
+                    )
+                }
+            val cursorData =
+                CursorData(
+                    cursorReader,
+                    SelectionTracker(selectedItems, focusedUriIdx, cursorReader.count) { uri },
+                )
+            if (cursorDataRef.complete(cursorData)) {
+                doLoadMorePreviousItems()
+                val startPos = records.size
+                doLoadMoreNextItems()
+                prevPageLoadingGate.set(false)
+                nextPageLoadingGate.set(false)
+                publishSnapshot(startPos)
+            } else {
+                cursorReader.close()
+            }
+        }
+    }
 
-    fun previewUri(key: Any): Flow<Uri?> = storage.map { it[key]?.previewUri }
+    private suspend fun publishInitialState() {
+        stateFlowSource.emit(
+            State(
+                if (0 <= focusedUriIdx && focusedUriIdx < initiallySharedUris.size) {
+                    val fileInfo = uriMetadataReader(initiallySharedUris[focusedUriIdx])
+                    listOf(
+                        Record(
+                            // a unique key that won't appear anywhere after more items are loaded
+                            -initiallySharedUris.size - 1,
+                            initiallySharedUris[focusedUriIdx],
+                            fileInfo.previewUri,
+                            fileInfo.mimeType,
+                            fileInfo.mimeType?.mimeTypeToItemType() ?: ItemType.File,
+                        ),
+                    )
+                } else {
+                    emptyList()
+                },
+                hasMoreItemsBefore = true,
+                hasMoreItemsAfter = true,
+                allowSelectionChange = false,
+            )
+        )
+    }
+
+    fun loadMorePreviousItems() {
+        invokeAsyncIfNotRunning(prevPageLoadingGate) {
+            doLoadMorePreviousItems()
+            publishSnapshot()
+        }
+    }
+
+    fun loadMoreNextItems() {
+        invokeAsyncIfNotRunning(nextPageLoadingGate) {
+            doLoadMoreNextItems()
+            publishSnapshot()
+        }
+    }
+
+    fun setSelected(item: Item, isSelected: Boolean) {
+        val record = item as Record
+        record.isSelected.value = isSelected
+        scope.launch {
+            val (_, selectionTracker) = waitForCursorData() ?: return@launch
+            selectionTracker.setItemSelection(record.key, record, isSelected)
+            val targetIntent = targetIntentModifier(selectionTracker.getSelection())
+
+            val newJob = scope.launch { notifySelectionChanged(targetIntent) }
+            notifySelectionJobRef.getAndSet(newJob)?.cancel()
+        }
+    }
+
+    private fun invokeAsyncIfNotRunning(guardingFlag: AtomicBoolean, block: suspend () -> Unit) {
+        if (guardingFlag.compareAndSet(false, true)) {
+            scope.launch { block() }.invokeOnCompletion { guardingFlag.set(false) }
+        }
+    }
+
+    private suspend fun doLoadMorePreviousItems() {
+        val (reader, selectionTracker) = waitForCursorData() ?: return
+        if (!reader.hasMoreBefore) return
+
+        val newItems = reader.readPageBefore().toRecords()
+        selectionTracker.onStartItemsAdded(newItems)
+        for (i in newItems.size() - 1 downTo 0) {
+            records.add(
+                0,
+                (newItems.valueAt(i) as Record).apply {
+                    isSelected.value = selectionTracker.isItemSelected(key)
+                }
+            )
+        }
+        if (!reader.hasMoreBefore && !reader.hasMoreAfter) {
+            val pendingItems = selectionTracker.getPendingItems()
+            val newRecords =
+                pendingItems.foldIndexed(SparseArray<Item>()) { idx, acc, item ->
+                    assert(item is Record) { "Unexpected pending item type: ${item.javaClass}" }
+                    val rec = item as Record
+                    val key = idx - pendingItems.size
+                    acc.append(
+                        key,
+                        Record(
+                            key,
+                            rec.uri,
+                            rec.previewUri,
+                            rec.mimeType,
+                            rec.mimeType?.mimeTypeToItemType() ?: ItemType.File
+                        )
+                    )
+                    acc
+                }
 
-    private data class Item(
-        val previewUri: Uri?,
+            selectionTracker.onStartItemsAdded(newRecords)
+            for (i in (newRecords.size() - 1) downTo 0) {
+                records.add(0, (newRecords.valueAt(i) as Record).apply { isSelected.value = true })
+            }
+        }
+    }
+
+    private suspend fun doLoadMoreNextItems() {
+        val (reader, selectionTracker) = waitForCursorData() ?: return
+        if (!reader.hasMoreAfter) return
+
+        val newItems = reader.readPageAfter().toRecords()
+        selectionTracker.onEndItemsAdded(newItems)
+        for (i in 0 until newItems.size()) {
+            val key = newItems.keyAt(i)
+            records.add(
+                (newItems.valueAt(i) as Record).apply {
+                    isSelected.value = selectionTracker.isItemSelected(key)
+                }
+            )
+        }
+        if (!reader.hasMoreBefore && !reader.hasMoreAfter) {
+            val items =
+                selectionTracker.getPendingItems().let { items ->
+                    items.foldIndexed(SparseArray<Item>(items.size)) { i, acc, item ->
+                        val key = reader.count + i
+                        val record = item as Record
+                        acc.append(
+                            key,
+                            Record(key, record.uri, record.previewUri, record.mimeType, record.type)
+                        )
+                        acc
+                    }
+                }
+            selectionTracker.onEndItemsAdded(items)
+            for (i in 0 until items.size()) {
+                records.add((items.valueAt(i) as Record).apply { isSelected.value = true })
+            }
+        }
+    }
+
+    private fun SparseArray<Uri>.toRecords(): SparseArray<Item> {
+        val items = SparseArray<Item>(size())
+        for (i in 0 until size()) {
+            val key = keyAt(i)
+            val uri = valueAt(i)
+            val fileInfo = uriMetadataReader(uri)
+            items.append(
+                key,
+                Record(
+                    key,
+                    uri,
+                    fileInfo.previewUri,
+                    fileInfo.mimeType,
+                    fileInfo.mimeType?.mimeTypeToItemType() ?: ItemType.File
+                )
+            )
+        }
+        return items
+    }
+
+    private suspend fun waitForCursorData() = cursorDataRef.await()
+
+    private fun notifySelectionChanged(targetIntent: Intent) {
+        selectionCallback(targetIntent)?.customActions?.let { customActions.tryEmit(it) }
+    }
+
+    private suspend fun publishSnapshot(startPos: Int = -1) {
+        val (reader, _) = waitForCursorData() ?: return
+        // TODO: publish a view into the list as it can only grow on each side thus a view won't be
+        // invalidated
+        val items = ArrayList<Item>(records)
+        stateFlowSource.emit(
+            State(
+                items,
+                reader.hasMoreBefore,
+                reader.hasMoreAfter,
+                allowSelectionChange = true,
+                targetPos = startPos,
+            )
+        )
+    }
+
+    private fun String.mimeTypeToItemType(): ItemType =
+        when {
+            mimeTypeClassifier.isImageType(this) -> ItemType.Image
+            mimeTypeClassifier.isVideoType(this) -> ItemType.Video
+            else -> ItemType.File
+        }
+
+    class State(
+        val items: List<Item>,
+        val hasMoreItemsBefore: Boolean,
+        val hasMoreItemsAfter: Boolean,
+        val allowSelectionChange: Boolean,
+        val targetPos: Int = -1,
     )
 
+    sealed interface Item {
+        val key: Int
+        val uri: Uri
+        val previewUri: Uri?
+        val mimeType: String?
+        val type: ItemType
+    }
+
+    enum class ItemType {
+        Image,
+        Video,
+        File,
+    }
+
+    private class Record(
+        override val key: Int,
+        override val uri: Uri,
+        override val previewUri: Uri? = uri,
+        override val mimeType: String?,
+        override val type: ItemType = ItemType.Image,
+    ) : Item {
+        val isSelected = MutableStateFlow(false)
+    }
+
     data class CallbackResult(val customActions: List<ChooserAction>?)
 
+    private data class CursorData(
+        val reader: CursorReader,
+        val selectionTracker: SelectionTracker<Item>,
+    )
+
     interface CursorReader : Closeable {
         val count: Int
         val hasMoreBefore: Boolean
@@ -66,13 +379,17 @@ class PayloadToggleInteractor {
 }
 
 class PayloadTogglePreviewInteractor(
-    private val key: Any,
+    private val item: PayloadToggleInteractor.Item,
     private val interactor: PayloadToggleInteractor,
 ) {
     fun setSelected(selected: Boolean) {
-        interactor.setSelected(key, selected)
+        interactor.setSelected(item, selected)
     }
 
-    val previewUri: Flow<Uri?> = interactor.previewUri(key)
-    val selected: Flow<Boolean> = interactor.selected(key)
+    val previewUri: Flow<Uri?>
+        get() = interactor.previewUri(item)
+    val selected: Flow<Boolean>
+        get() = interactor.selected(item)
+    val key
+        get() = item.key
 }
diff --git a/java/src/com/android/intentresolver/contentpreview/PreviewViewModel.kt b/java/src/com/android/intentresolver/contentpreview/PreviewViewModel.kt
index d855ea16..77cf0ac9 100644
--- a/java/src/com/android/intentresolver/contentpreview/PreviewViewModel.kt
+++ b/java/src/com/android/intentresolver/contentpreview/PreviewViewModel.kt
@@ -65,7 +65,7 @@ constructor(
     }
 
     override val payloadToggleInteractor: PayloadToggleInteractor? by lazy {
-        PayloadToggleInteractor()
+        null // TODO: initialize PayloadToggleInteractor()
     }
 
     companion object {
diff --git a/java/src/com/android/intentresolver/contentpreview/SelectionTracker.kt b/java/src/com/android/intentresolver/contentpreview/SelectionTracker.kt
new file mode 100644
index 00000000..4ce006ec
--- /dev/null
+++ b/java/src/com/android/intentresolver/contentpreview/SelectionTracker.kt
@@ -0,0 +1,175 @@
+/*
+ * Copyright 2024 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.contentpreview
+
+import android.net.Uri
+import android.util.SparseArray
+import android.util.SparseIntArray
+import androidx.core.util.containsKey
+import androidx.core.util.isNotEmpty
+
+/**
+ * Tracks selected items (including those that has not been read frm the cursor) and their relative
+ * order.
+ */
+class SelectionTracker<Item>(
+    selectedItems: List<Item>,
+    private val focusedItemIdx: Int,
+    private val cursorCount: Int,
+    private val getUri: Item.() -> Uri,
+) {
+    /** Contains selected items keys. */
+    private val selections = SparseArray<Item>(selectedItems.size)
+
+    /**
+     * A set of initially selected items that has not yet been observed by the lazy read of the
+     * cursor and thus has unknown key (cursor position). Initially, all [selectedItems] are put in
+     * this map with items at the index less than [focusedItemIdx] with negative keys (to the left
+     * of all cursor items) and items at the index more or equal to [focusedItemIdx] with keys more
+     * or equal to [cursorCount] (to the right of all cursor items) in their relative order. Upon
+     * reading the cursor, [onEndItemsAdded]/[onStartItemsAdded], all pending items from that
+     * collection in the corresponding direction get their key assigned and gets removed from the
+     * map. Items that were missing from the cursor get removed from the map by
+     * [getPendingItems] + [onStartItemsAdded]/[onEndItemsAdded] combination.
+     */
+    private val pendingKeys = HashMap<Uri, SparseIntArray>()
+
+    init {
+        selectedItems.forEachIndexed { i, item ->
+            // all items before focusedItemIdx gets "positioned" before all the cursor items
+            // and all the reset after all the cursor items in their relative order.
+            // Also see the comments to pendingKeys property.
+            val key =
+                if (i < focusedItemIdx) {
+                    i - focusedItemIdx
+                } else {
+                    i + cursorCount - focusedItemIdx
+                }
+            selections.append(key, item)
+            pendingKeys.getOrPut(item.getUri()) { SparseIntArray(1) }.append(key, key)
+        }
+    }
+
+    /** Update selections based on the set of items read from the end of the cursor */
+    fun onEndItemsAdded(items: SparseArray<Item>) {
+        for (i in 0 until items.size()) {
+            val item = items.valueAt(i)
+            pendingKeys[item.getUri()]
+                // if only one pending (unmatched) item with this URI is left, removed this URI
+                ?.also {
+                    if (it.size() <= 1) {
+                        pendingKeys.remove(item.getUri())
+                    }
+                }
+                // a safeguard, we should not observe empty arrays at this point
+                ?.takeIf { it.isNotEmpty() }
+                // pick a matching pending items from the right side
+                ?.let { pendingUriPositions ->
+                    val key = items.keyAt(i)
+                    val insertPos =
+                        pendingUriPositions
+                            .findBestKeyPosition(key)
+                            .coerceIn(0, pendingUriPositions.size() - 1)
+                    // select next pending item from the right, if not such item exists then
+                    // the data is inconsistent and we pick the closes one from the left
+                    val keyPlaceholder = pendingUriPositions.keyAt(insertPos)
+                    pendingUriPositions.removeAt(insertPos)
+                    selections.remove(keyPlaceholder)
+                    selections[key] = item
+                }
+        }
+    }
+
+    /** Update selections based on the set of items read from the head of the cursor */
+    fun onStartItemsAdded(items: SparseArray<Item>) {
+        for (i in (items.size() - 1) downTo 0) {
+            val item = items.valueAt(i)
+            pendingKeys[item.getUri()]
+                // if only one pending (unmatched) item with this URI is left, removed this URI
+                ?.also {
+                    if (it.size() <= 1) {
+                        pendingKeys.remove(item.getUri())
+                    }
+                }
+                // a safeguard, we should not observe empty arrays at this point
+                ?.takeIf { it.isNotEmpty() }
+                // pick a matching pending items from the left side
+                ?.let { pendingUriPositions ->
+                    val key = items.keyAt(i)
+                    val insertPos =
+                        pendingUriPositions
+                            .findBestKeyPosition(key)
+                            .coerceIn(1, pendingUriPositions.size())
+                    // select next pending item from the left, if not such item exists then
+                    // the data is inconsistent and we pick the closes one from the right
+                    val keyPlaceholder = pendingUriPositions.keyAt(insertPos - 1)
+                    pendingUriPositions.removeAt(insertPos - 1)
+                    selections.remove(keyPlaceholder)
+                    selections[key] = item
+                }
+        }
+    }
+
+    /** Updated selection status for the given item */
+    fun setItemSelection(key: Int, item: Item, isSelected: Boolean): Boolean {
+        val idx = selections.indexOfKey(key)
+        if (isSelected && idx < 0) {
+            selections[key] = item
+            return true
+        }
+        if (!isSelected && idx >= 0) {
+            selections.removeAt(idx)
+            return true
+        }
+        return false
+    }
+
+    /** Return selection status for the given item */
+    fun isItemSelected(key: Int): Boolean = selections.containsKey(key)
+
+    fun getSelection(): List<Item> =
+        buildList(selections.size()) {
+            for (i in 0 until selections.size()) {
+                add(selections.valueAt(i))
+            }
+        }
+
+    /** Return all selected items that has not yet been read from the cursor */
+    fun getPendingItems(): List<Item> =
+        if (pendingKeys.isEmpty()) {
+            emptyList()
+        } else {
+            buildList {
+                for (i in 0 until selections.size()) {
+                    val item = selections.valueAt(i) ?: continue
+                    if (isPending(item, selections.keyAt(i))) {
+                        add(item)
+                    }
+                }
+            }
+        }
+
+    private fun isPending(item: Item, key: Int): Boolean {
+        val keys = pendingKeys[item.getUri()] ?: return false
+        return keys.containsKey(key)
+    }
+
+    private fun SparseIntArray.findBestKeyPosition(key: Int): Int =
+        // undocumented, but indexOfKey behaves in the same was as
+        // java.util.Collections#binarySearch()
+        indexOfKey(key).let { if (it < 0) it.inv() else it }
+}
diff --git a/java/src/com/android/intentresolver/contentpreview/shareousel/ui/composable/ShareouselComposable.kt b/java/src/com/android/intentresolver/contentpreview/shareousel/ui/composable/ShareouselComposable.kt
index c83c10b0..f636966e 100644
--- a/java/src/com/android/intentresolver/contentpreview/shareousel/ui/composable/ShareouselComposable.kt
+++ b/java/src/com/android/intentresolver/contentpreview/shareousel/ui/composable/ShareouselComposable.kt
@@ -15,7 +15,6 @@
  */
 package com.android.intentresolver.contentpreview.shareousel.ui.composable
 
-import android.os.Parcelable
 import androidx.compose.foundation.Image
 import androidx.compose.foundation.border
 import androidx.compose.foundation.clickable
@@ -64,7 +63,7 @@ fun Shareousel(viewModel: ShareouselViewModel) {
                 Modifier.fillMaxWidth()
                     .height(dimensionResource(R.dimen.chooser_preview_image_height_tall))
         ) {
-            items(previewKeys, key = { (it as? Parcelable) ?: Unit }) { key ->
+            items(previewKeys, key = viewModel.previewRowKey) { key ->
                 ShareouselCard(viewModel.previewForKey(key))
             }
         }
diff --git a/java/src/com/android/intentresolver/contentpreview/shareousel/ui/viewmodel/ShareouselViewModel.kt b/java/src/com/android/intentresolver/contentpreview/shareousel/ui/viewmodel/ShareouselViewModel.kt
index 4592ea6d..ff22a6fd 100644
--- a/java/src/com/android/intentresolver/contentpreview/shareousel/ui/viewmodel/ShareouselViewModel.kt
+++ b/java/src/com/android/intentresolver/contentpreview/shareousel/ui/viewmodel/ShareouselViewModel.kt
@@ -29,6 +29,7 @@ data class ShareouselViewModel(
     val actions: Flow<List<ActionChipViewModel>>,
     val centerIndex: Flow<Int>,
     val previewForKey: (key: Any) -> ShareouselImageViewModel,
+    val previewRowKey: (Any) -> Any
 )
 
 data class ActionChipViewModel(val label: String, val icon: ComposeIcon?, val onClick: () -> Unit)
@@ -56,6 +57,7 @@ fun PayloadToggleInteractor.toShareouselViewModel(imageLoader: ImageLoader): Sha
                 setSelected = { isSelected -> previewInteractor.setSelected(isSelected) },
                 onActionClick = {},
             )
-        }
+        },
+        previewRowKey = { getKey(it) },
     )
 }