Don't use flow for one-off get of ThumbnailData

Bug: 381317629
Flag: com.android.launcher3.enable_refactor_task_thumbnail
Test: TasksRepositoryTest
Test: Performance tests
Change-Id: I1990ad1972beea8c81b44cdd50c7f674f225c69d
diff --git a/quickstep/src/com/android/quickstep/recents/data/RecentTasksRepository.kt b/quickstep/src/com/android/quickstep/recents/data/RecentTasksRepository.kt
index 3b59864..53969c5 100644
--- a/quickstep/src/com/android/quickstep/recents/data/RecentTasksRepository.kt
+++ b/quickstep/src/com/android/quickstep/recents/data/RecentTasksRepository.kt
@@ -37,6 +37,12 @@
     fun getThumbnailById(taskId: Int): Flow<ThumbnailData?>
 
     /**
+     * Gets the [ThumbnailData] associated with a task that has id [taskId]. Flow will settle on
+     * null if the task was not found or is invisible.
+     */
+    fun getCurrentThumbnailById(taskId: Int): ThumbnailData?
+
+    /**
      * Sets the tasks that are visible, indicating that properties relating to visuals need to be
      * populated e.g. icons/thumbnails etc.
      */
diff --git a/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt b/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
index 6c627ef..3aae760 100644
--- a/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
+++ b/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
@@ -72,6 +72,8 @@
     override fun getThumbnailById(taskId: Int) =
         getTaskDataById(taskId).map { it?.thumbnail }.distinctUntilChangedBy { it?.snapshotId }
 
+    override fun getCurrentThumbnailById(taskId: Int) = tasks.value[taskId]?.thumbnail
+
     override fun setVisibleTasks(visibleTaskIdList: Set<Int>) {
         val tasksNoLongerVisible = taskRequests.keys.subtract(visibleTaskIdList)
         val newlyVisibleTasks = visibleTaskIdList.subtract(taskRequests.keys)
diff --git a/quickstep/src/com/android/quickstep/recents/usecase/GetThumbnailPositionUseCase.kt b/quickstep/src/com/android/quickstep/recents/usecase/GetThumbnailPositionUseCase.kt
index f060d7d..bea1d07 100644
--- a/quickstep/src/com/android/quickstep/recents/usecase/GetThumbnailPositionUseCase.kt
+++ b/quickstep/src/com/android/quickstep/recents/usecase/GetThumbnailPositionUseCase.kt
@@ -24,18 +24,17 @@
 import com.android.quickstep.recents.usecase.ThumbnailPositionState.MatrixScaling
 import com.android.quickstep.recents.usecase.ThumbnailPositionState.MissingThumbnail
 import com.android.systemui.shared.recents.utilities.PreviewPositionHelper
-import kotlinx.coroutines.flow.firstOrNull
 
 /** Use case for retrieving [Matrix] for positioning Thumbnail in a View */
 class GetThumbnailPositionUseCase(
     private val deviceProfileRepository: RecentsDeviceProfileRepository,
     private val rotationStateRepository: RecentsRotationStateRepository,
     private val tasksRepository: RecentTasksRepository,
-    private val previewPositionHelper: PreviewPositionHelper = PreviewPositionHelper()
+    private val previewPositionHelper: PreviewPositionHelper = PreviewPositionHelper(),
 ) {
-    suspend fun run(taskId: Int, width: Int, height: Int, isRtl: Boolean): ThumbnailPositionState {
+    fun run(taskId: Int, width: Int, height: Int, isRtl: Boolean): ThumbnailPositionState {
         val thumbnailData =
-            tasksRepository.getThumbnailById(taskId).firstOrNull() ?: return MissingThumbnail
+            tasksRepository.getCurrentThumbnailById(taskId) ?: return MissingThumbnail
         val thumbnail = thumbnailData.thumbnail ?: return MissingThumbnail
         previewPositionHelper.updateThumbnailMatrix(
             Rect(0, 0, thumbnail.width, thumbnail.height),
@@ -44,11 +43,11 @@
             height,
             deviceProfileRepository.getRecentsDeviceProfile().isLargeScreen,
             rotationStateRepository.getRecentsRotationState().activityRotation,
-            isRtl
+            isRtl,
         )
         return MatrixScaling(
             previewPositionHelper.matrix,
-            previewPositionHelper.isOrientationChanged
+            previewPositionHelper.isOrientationChanged,
         )
     }
 }
diff --git a/quickstep/src/com/android/quickstep/recents/usecase/GetThumbnailUseCase.kt b/quickstep/src/com/android/quickstep/recents/usecase/GetThumbnailUseCase.kt
index 3aa808e..b9e9e02 100644
--- a/quickstep/src/com/android/quickstep/recents/usecase/GetThumbnailUseCase.kt
+++ b/quickstep/src/com/android/quickstep/recents/usecase/GetThumbnailUseCase.kt
@@ -18,13 +18,9 @@
 
 import android.graphics.Bitmap
 import com.android.quickstep.recents.data.RecentTasksRepository
-import kotlinx.coroutines.flow.firstOrNull
-import kotlinx.coroutines.runBlocking
 
 /** Use case for retrieving thumbnail. */
 class GetThumbnailUseCase(private val taskRepository: RecentTasksRepository) {
     /** Returns the latest thumbnail associated with [taskId] if loaded, or null otherwise */
-    fun run(taskId: Int): Bitmap? = runBlocking {
-        taskRepository.getThumbnailById(taskId).firstOrNull()?.thumbnail
-    }
+    fun run(taskId: Int): Bitmap? = taskRepository.getCurrentThumbnailById(taskId)?.thumbnail
 }
diff --git a/quickstep/src/com/android/quickstep/recents/usecase/SysUiStatusNavFlagsUseCase.kt b/quickstep/src/com/android/quickstep/recents/usecase/SysUiStatusNavFlagsUseCase.kt
index 1d19c7d..5be5f4a 100644
--- a/quickstep/src/com/android/quickstep/recents/usecase/SysUiStatusNavFlagsUseCase.kt
+++ b/quickstep/src/com/android/quickstep/recents/usecase/SysUiStatusNavFlagsUseCase.kt
@@ -22,14 +22,11 @@
 import com.android.launcher3.util.SystemUiController.FLAG_LIGHT_NAV
 import com.android.launcher3.util.SystemUiController.FLAG_LIGHT_STATUS
 import com.android.quickstep.recents.data.RecentTasksRepository
-import kotlinx.coroutines.flow.firstOrNull
-import kotlinx.coroutines.runBlocking
 
 /** UseCase to calculate flags for status bar and navigation bar */
 class SysUiStatusNavFlagsUseCase(private val taskRepository: RecentTasksRepository) {
     fun getSysUiStatusNavFlags(taskId: Int): Int {
-        val thumbnailData =
-            runBlocking { taskRepository.getThumbnailById(taskId).firstOrNull() } ?: return 0
+        val thumbnailData = taskRepository.getCurrentThumbnailById(taskId) ?: return 0
 
         val thumbnailAppearance = thumbnailData.appearance
         var flags = 0
diff --git a/quickstep/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModel.kt b/quickstep/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModel.kt
index 4e13d1c..14359db 100644
--- a/quickstep/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModel.kt
+++ b/quickstep/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModel.kt
@@ -28,7 +28,6 @@
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.map
-import kotlinx.coroutines.runBlocking
 
 /** View model for TaskOverlay */
 class TaskOverlayViewModel(
@@ -41,7 +40,7 @@
         combine(
                 recentsViewData.overlayEnabled,
                 recentsViewData.settledFullyVisibleTaskIds.map { it.contains(task.key.id) },
-                recentTasksRepository.getThumbnailById(task.key.id)
+                recentTasksRepository.getThumbnailById(task.key.id),
             ) { isOverlayEnabled, isFullyVisible, thumbnailData ->
                 if (isOverlayEnabled && isFullyVisible) {
                     Enabled(
@@ -55,24 +54,22 @@
             .distinctUntilChanged()
 
     fun getThumbnailPositionState(width: Int, height: Int, isRtl: Boolean): ThumbnailPositionState {
-        return runBlocking {
-            val matrix: Matrix
-            val isRotated: Boolean
-            when (
-                val thumbnailPositionState =
-                    getThumbnailPositionUseCase.run(task.key.id, width, height, isRtl)
-            ) {
-                is MatrixScaling -> {
-                    matrix = thumbnailPositionState.matrix
-                    isRotated = thumbnailPositionState.isRotated
-                }
-                is MissingThumbnail -> {
-                    matrix = Matrix.IDENTITY_MATRIX
-                    isRotated = false
-                }
+        val matrix: Matrix
+        val isRotated: Boolean
+        when (
+            val thumbnailPositionState =
+                getThumbnailPositionUseCase.run(task.key.id, width, height, isRtl)
+        ) {
+            is MatrixScaling -> {
+                matrix = thumbnailPositionState.matrix
+                isRotated = thumbnailPositionState.isRotated
             }
-            ThumbnailPositionState(matrix, isRotated)
+            is MissingThumbnail -> {
+                matrix = Matrix.IDENTITY_MATRIX
+                isRotated = false
+            }
         }
+        return ThumbnailPositionState(matrix, isRotated)
     }
 
     data class ThumbnailPositionState(val matrix: Matrix, val isRotated: Boolean)
diff --git a/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModelImpl.kt b/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModelImpl.kt
index e5bad67..b5b2fc9 100644
--- a/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModelImpl.kt
+++ b/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModelImpl.kt
@@ -44,7 +44,6 @@
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.map
-import kotlinx.coroutines.runBlocking
 
 @OptIn(ExperimentalCoroutinesApi::class)
 class TaskThumbnailViewModelImpl(
@@ -109,17 +108,14 @@
         splashProgress.value = splashAlphaUseCase.execute(taskId)
     }
 
-    override fun getThumbnailPositionState(width: Int, height: Int, isRtl: Boolean): Matrix {
-        return runBlocking {
-            when (
-                val thumbnailPositionState =
-                    getThumbnailPositionUseCase.run(taskId, width, height, isRtl)
-            ) {
-                is ThumbnailPositionState.MatrixScaling -> thumbnailPositionState.matrix
-                is ThumbnailPositionState.MissingThumbnail -> Matrix.IDENTITY_MATRIX
-            }
+    override fun getThumbnailPositionState(width: Int, height: Int, isRtl: Boolean): Matrix =
+        when (
+            val thumbnailPositionState =
+                getThumbnailPositionUseCase.run(taskId, width, height, isRtl)
+        ) {
+            is ThumbnailPositionState.MatrixScaling -> thumbnailPositionState.matrix
+            is ThumbnailPositionState.MissingThumbnail -> Matrix.IDENTITY_MATRIX
         }
-    }
 
     private fun isBackgroundOnly(task: Task): Boolean = task.isLocked || task.thumbnail == null
 
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTasksRepository.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTasksRepository.kt
index d6688d6..1c9ce0b 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTasksRepository.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTasksRepository.kt
@@ -49,6 +49,9 @@
     override fun getThumbnailById(taskId: Int): Flow<ThumbnailData?> =
         getTaskDataById(taskId).map { it?.thumbnail }
 
+    override fun getCurrentThumbnailById(taskId: Int): ThumbnailData? =
+        tasks.value.firstOrNull { it.key.id == taskId }?.thumbnail
+
     override fun setVisibleTasks(visibleTaskIdList: Set<Int>) {
         visibleTasks.value = visibleTaskIdList
         tasks.value =
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt
index 357df6e..e0e7f28 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt
@@ -87,6 +87,48 @@
         }
 
     @Test
+    fun getThumbnailByIdReturnsNullWithNoLoadedThumbnails() =
+        testScope.runTest {
+            recentsModel.seedTasks(defaultTaskList)
+            systemUnderTest.getAllTaskData(forceRefresh = true)
+
+            assertThat(systemUnderTest.getThumbnailById(1).first()).isNull()
+        }
+
+    @Test
+    fun getCurrentThumbnailByIdReturnsNullWithNoLoadedThumbnails() =
+        testScope.runTest {
+            recentsModel.seedTasks(defaultTaskList)
+            systemUnderTest.getAllTaskData(forceRefresh = true)
+
+            assertThat(systemUnderTest.getCurrentThumbnailById(1)).isNull()
+        }
+
+    @Test
+    fun getThumbnailByIdReturnsThumbnailWithLoadedThumbnails() =
+        testScope.runTest {
+            recentsModel.seedTasks(defaultTaskList)
+            systemUnderTest.getAllTaskData(forceRefresh = true)
+            val bitmap1 = taskThumbnailDataSource.taskIdToBitmap[1]
+
+            systemUnderTest.setVisibleTasks(setOf(1))
+
+            assertThat(systemUnderTest.getThumbnailById(1).first()!!.thumbnail).isEqualTo(bitmap1)
+        }
+
+    @Test
+    fun getCurrentThumbnailByIdReturnsThumbnailWithLoadedThumbnails() =
+        testScope.runTest {
+            recentsModel.seedTasks(defaultTaskList)
+            systemUnderTest.getAllTaskData(forceRefresh = true)
+            val bitmap1 = taskThumbnailDataSource.taskIdToBitmap[1]
+
+            systemUnderTest.setVisibleTasks(setOf(1))
+
+            assertThat(systemUnderTest.getCurrentThumbnailById(1)?.thumbnail).isEqualTo(bitmap1)
+        }
+
+    @Test
     fun setVisibleTasksPopulatesThumbnails() =
         testScope.runTest {
             recentsModel.seedTasks(defaultTaskList)