From da860e3d46ec848ba47fccf13168b32e9ab85c14 Mon Sep 17 00:00:00 2001
From: Uwais Ashraf <uwaisashraf@google.com>
Date: Fri, 9 May 2025 18:21:30 +0000
Subject: [PATCH] Reduce combine() usage as it is expensive

Flag: com.android.launcher3.enable_coroutine_threading_improvements
Bug: 416206104
Test: Perfetto trace comparison shows fewer continuations
Test: TaskViewModelTest
Change-Id: Id5c645afff9850453139b85dbcda907cd12c5d26
---
 .../recents/ui/viewmodel/TaskViewModel.kt     | 49 ++++++++++++++++++-
 .../launcher3/util/coroutines/Combine.kt      | 41 ++++++++++++++++
 2 files changed, 89 insertions(+), 1 deletion(-)
 create mode 100644 src/com/android/launcher3/util/coroutines/Combine.kt

diff --git a/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModel.kt b/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModel.kt
index 15727822f7..049556af28 100644
--- a/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModel.kt
+++ b/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModel.kt
@@ -19,6 +19,7 @@ package com.android.quickstep.recents.ui.viewmodel
 import android.annotation.ColorInt
 import android.util.Log
 import androidx.core.graphics.ColorUtils
+import com.android.launcher3.Flags.enableCoroutineThreadingImprovements
 import com.android.launcher3.util.coroutines.DispatcherProvider
 import com.android.quickstep.recents.domain.model.TaskId
 import com.android.quickstep.recents.domain.model.TaskModel
@@ -82,6 +83,16 @@ class TaskViewModel(
             combine(combinedTaskFlows, isLiveTile, ::mapToTaskData)
         }
 
+    private val taskModels =
+        taskIds.flatMapLatest { ids ->
+            // Combine Tasks requests
+            val taskFlows =
+                ids.map { id ->
+                    getTaskUseCase(id).distinctUntilChanged().map { taskModel -> id to taskModel }
+                }
+            combine(taskFlows) { taskArray -> taskArray }
+        }
+
     private val overlayEnabled =
         if (taskViewType == TaskViewType.DESKTOP) {
             flowOf(false)
@@ -94,8 +105,44 @@ class TaskViewModel(
                 .distinctUntilChanged()
         }
 
-    val state: Flow<TaskTileUiState> =
+    private val preThreadingImprovedState: Flow<TaskTileUiState> =
         combine(taskData, overlayEnabled, isCentralTask, ::mapToTaskTile)
+
+    private val threadingImprovedState: Flow<TaskTileUiState> =
+        com.android.launcher3.util.coroutines.combine(
+            taskModels,
+            recentsViewData.runningTaskIds,
+            recentsViewData.runningTaskShowScreenshot,
+            recentsViewData.overlayEnabled,
+            recentsViewData.settledFullyVisibleTaskIds,
+            recentsViewData.centralTaskIds,
+        ) {
+            taskModels: Array<Pair<Int, TaskModel?>>,
+            runningTaskIds: Set<Int>,
+            runningTaskShowScreenshot: Boolean,
+            isOverlayEnabled: Boolean,
+            settledFullyVisibleTaskIds: Set<Int>,
+            centralTaskIds: Set<Int> ->
+            val taskIds = taskModels.map { it.first }.toSet()
+            val isCentralTask = taskIds == centralTaskIds
+            val overlayEnabled =
+                if (taskViewType == TaskViewType.DESKTOP) {
+                    false
+                } else {
+                    isOverlayEnabled && settledFullyVisibleTaskIds.any { it in taskIds }
+                }
+            val isLiveTile = runningTaskIds == taskIds && !runningTaskShowScreenshot
+            val taskData = mapToTaskData(taskModels, isLiveTile)
+
+            mapToTaskTile(taskData, overlayEnabled, isCentralTask)
+        }
+
+    private val taskTileUiStateFlow =
+        if (enableCoroutineThreadingImprovements()) threadingImprovedState
+        else preThreadingImprovedState
+
+    val state: Flow<TaskTileUiState> =
+        taskTileUiStateFlow
             .distinctUntilChanged()
             .debounce { state ->
                 // Debouncing only when thumbnails are not present gives the best results.
diff --git a/src/com/android/launcher3/util/coroutines/Combine.kt b/src/com/android/launcher3/util/coroutines/Combine.kt
new file mode 100644
index 0000000000..dcd4b9b256
--- /dev/null
+++ b/src/com/android/launcher3/util/coroutines/Combine.kt
@@ -0,0 +1,41 @@
+/*
+ * Copyright (C) 2025 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.launcher3.util.coroutines
+
+import kotlinx.coroutines.flow.Flow
+
+/** A copy of [kotlinx.coroutines.flow.combine] that supports 6 flows. */
+inline fun <T1, T2, T3, T4, T5, T6, R> combine(
+    flow: Flow<T1>,
+    flow2: Flow<T2>,
+    flow3: Flow<T3>,
+    flow4: Flow<T4>,
+    flow5: Flow<T5>,
+    flow6: Flow<T6>,
+    crossinline transform: suspend (T1, T2, T3, T4, T5, T6) -> R,
+): Flow<R> =
+    kotlinx.coroutines.flow.combine(flow, flow2, flow3, flow4, flow5, flow6) { args: Array<*> ->
+        @Suppress("UNCHECKED_CAST")
+        transform(
+            args[0] as T1,
+            args[1] as T2,
+            args[2] as T3,
+            args[3] as T4,
+            args[4] as T5,
+            args[5] as T6,
+        )
+    }