Merge "fix performance regression on qs_new_tiles" into main
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/viewmodel/QSTileViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/viewmodel/QSTileViewModelTest.kt
index 9ce2e0f..c33e2a4 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/viewmodel/QSTileViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/viewmodel/QSTileViewModelTest.kt
@@ -92,7 +92,8 @@
             runCurrent()
 
             assertThat(states()).isNotEmpty()
-            assertThat(states().first().label).isEqualTo(testTileData)
+            assertThat(states().last()).isNotNull()
+            assertThat(states().last()!!.label).isEqualTo(testTileData)
             verify(qsTileLogger).logInitialRequest(eq(tileConfig.tileSpec))
         }
 
@@ -196,6 +197,7 @@
             qsTileLogger,
             FakeSystemClock(),
             testCoroutineDispatcher,
+            testCoroutineDispatcher,
             scope.backgroundScope,
         )
 
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/viewmodel/QSTileViewModelUserInputTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/viewmodel/QSTileViewModelUserInputTest.kt
index 6066d24..7955f2f 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/viewmodel/QSTileViewModelUserInputTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/viewmodel/QSTileViewModelUserInputTest.kt
@@ -256,6 +256,7 @@
             qsTileLogger,
             FakeSystemClock(),
             testCoroutineDispatcher,
+            testCoroutineDispatcher,
             scope.backgroundScope,
         )
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/base/viewmodel/QSTileViewModelFactory.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/base/viewmodel/QSTileViewModelFactory.kt
index 4c21080..8c75cf0 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/base/viewmodel/QSTileViewModelFactory.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/base/viewmodel/QSTileViewModelFactory.kt
@@ -17,6 +17,7 @@
 package com.android.systemui.qs.tiles.base.viewmodel
 
 import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.dagger.qualifiers.UiBackground
 import com.android.systemui.plugins.FalsingManager
 import com.android.systemui.qs.pipeline.shared.TileSpec
 import com.android.systemui.qs.tiles.base.analytics.QSTileAnalytics
@@ -61,6 +62,7 @@
         private val qsTileConfigProvider: QSTileConfigProvider,
         private val systemClock: SystemClock,
         @Background private val backgroundDispatcher: CoroutineDispatcher,
+        @UiBackground private val uiBackgroundDispatcher: CoroutineDispatcher,
         private val customTileComponentBuilder: CustomTileComponent.Builder,
     ) : QSTileViewModelFactory<CustomTileDataModel> {
 
@@ -86,6 +88,7 @@
                 qsTileLogger,
                 systemClock,
                 backgroundDispatcher,
+                uiBackgroundDispatcher,
                 component.coroutineScope(),
             )
         }
@@ -106,6 +109,7 @@
         private val qsTileConfigProvider: QSTileConfigProvider,
         private val systemClock: SystemClock,
         @Background private val backgroundDispatcher: CoroutineDispatcher,
+        @UiBackground private val uiBackgroundDispatcher: CoroutineDispatcher,
         private val coroutineScopeFactory: QSTileCoroutineScopeFactory,
     ) : QSTileViewModelFactory<T> {
 
@@ -136,6 +140,7 @@
                 qsTileLogger,
                 systemClock,
                 backgroundDispatcher,
+                uiBackgroundDispatcher,
                 coroutineScopeFactory.create(),
             )
     }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/base/viewmodel/QSTileViewModelImpl.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/base/viewmodel/QSTileViewModelImpl.kt
index 8782524..9e84f01 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/base/viewmodel/QSTileViewModelImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/base/viewmodel/QSTileViewModelImpl.kt
@@ -40,6 +40,7 @@
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.cancel
+import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -59,7 +60,9 @@
 import kotlinx.coroutines.flow.onStart
 import kotlinx.coroutines.flow.shareIn
 import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.flow.transformLatest
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
 
 /**
  * Provides a hassle-free way to implement new tiles according to current System UI architecture
@@ -81,6 +84,7 @@
     private val qsTileLogger: QSTileLogger,
     private val systemClock: SystemClock,
     private val backgroundDispatcher: CoroutineDispatcher,
+    uiBackgroundDispatcher: CoroutineDispatcher,
     private val tileScope: CoroutineScope,
 ) : QSTileViewModel, Dumpable {
 
@@ -93,18 +97,16 @@
 
     private val tileData: SharedFlow<DATA_TYPE> = createTileDataFlow()
 
-    override val state: SharedFlow<QSTileState> =
+    override val state: StateFlow<QSTileState?> =
         tileData
             .map { data ->
-                mapper().map(config, data).also { state ->
-                    qsTileLogger.logStateUpdate(spec, state, data)
-                }
+                withContext(uiBackgroundDispatcher) { mapper().map(config, data) }
+                    .also { state -> qsTileLogger.logStateUpdate(spec, state, data) }
             }
-            .flowOn(backgroundDispatcher)
-            .shareIn(
+            .stateIn(
                 tileScope,
                 SharingStarted.WhileSubscribed(),
-                replay = 1,
+                null,
             )
     override val isAvailable: StateFlow<Boolean> =
         users
@@ -147,26 +149,26 @@
 
     private fun createTileDataFlow(): SharedFlow<DATA_TYPE> =
         users
-            .flatMapLatest { user ->
-                val updateTriggers =
-                    merge(
-                            userInputFlow(user),
-                            forceUpdates
-                                .map { DataUpdateTrigger.ForceUpdate }
-                                .onEach { qsTileLogger.logForceUpdate(spec) },
-                        )
-                        .onStart {
-                            emit(DataUpdateTrigger.InitialRequest)
-                            qsTileLogger.logInitialRequest(spec)
-                        }
-                        .shareIn(tileScope, SharingStarted.WhileSubscribed())
-                tileDataInteractor()
-                    .tileData(user, updateTriggers)
-                    // combine makes sure updateTriggers is always listened even if
-                    // tileDataInteractor#tileData doesn't flatMapLatest on it
-                    .combine(updateTriggers) { data, _ -> data }
-                    .cancellable()
-                    .flowOn(backgroundDispatcher)
+            .transformLatest { user ->
+                coroutineScope {
+                    val updateTriggers: Flow<DataUpdateTrigger> =
+                        merge(
+                                userInputFlow(user),
+                                forceUpdates
+                                    .map { DataUpdateTrigger.ForceUpdate }
+                                    .onEach { qsTileLogger.logForceUpdate(spec) },
+                            )
+                            .onStart { qsTileLogger.logInitialRequest(spec) }
+                            .stateIn(this, SharingStarted.Eagerly, DataUpdateTrigger.InitialRequest)
+                    tileDataInteractor()
+                        .tileData(user, updateTriggers)
+                        // combine makes sure updateTriggers is always listened even if
+                        // tileDataInteractor#tileData doesn't transformLatest on it
+                        .combine(updateTriggers) { data, _ -> data }
+                        .cancellable()
+                        .flowOn(backgroundDispatcher)
+                        .collect { emit(it) }
+                }
             }
             .distinctUntilChanged()
             .shareIn(
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/qr/domain/interactor/QRCodeScannerTileDataInteractor.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/qr/domain/interactor/QRCodeScannerTileDataInteractor.kt
index 1e8ce58..233e913 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/qr/domain/interactor/QRCodeScannerTileDataInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/qr/domain/interactor/QRCodeScannerTileDataInteractor.kt
@@ -30,11 +30,9 @@
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.onStart
-import kotlinx.coroutines.flow.stateIn
 
 /** Observes one qr scanner state changes providing the [QRCodeScannerTileModel]. */
 class QRCodeScannerTileDataInteractor
@@ -66,11 +64,6 @@
             }
             .onStart { emit(generateModel()) }
             .flowOn(bgCoroutineContext)
-            .stateIn(
-                scope,
-                SharingStarted.WhileSubscribed(),
-                QRCodeScannerTileModel.TemporarilyUnavailable
-            )
 
     override fun availability(user: UserHandle): Flow<Boolean> =
         flowOf(qrController.isCameraAvailable)
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/viewmodel/QSTileState.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/viewmodel/QSTileState.kt
index ae6c014..30247c4 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/viewmodel/QSTileState.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/viewmodel/QSTileState.kt
@@ -54,18 +54,21 @@
             resources: Resources,
             theme: Theme,
             config: QSTileUIConfig,
-            build: Builder.() -> Unit
+            builder: Builder.() -> Unit
         ): QSTileState {
             val iconDrawable = resources.getDrawable(config.iconRes, theme)
             return build(
                 { Icon.Loaded(iconDrawable, null) },
                 resources.getString(config.labelRes),
-                build,
+                builder,
             )
         }
 
-        fun build(icon: () -> Icon?, label: CharSequence, build: Builder.() -> Unit): QSTileState =
-            Builder(icon, label).apply(build).build()
+        fun build(
+            icon: () -> Icon?,
+            label: CharSequence,
+            builder: Builder.() -> Unit
+        ): QSTileState = Builder(icon, label).apply { builder() }.build()
     }
 
     enum class ActivationState(val legacyState: Int) {
@@ -107,7 +110,9 @@
 
     sealed interface SideViewIcon {
         data class Custom(val icon: Icon) : SideViewIcon
+
         data object Chevron : SideViewIcon
+
         data object None : SideViewIcon
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/viewmodel/QSTileViewModel.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/viewmodel/QSTileViewModel.kt
index 226e2fa0..b1b0001 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/viewmodel/QSTileViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/viewmodel/QSTileViewModel.kt
@@ -17,7 +17,6 @@
 package com.android.systemui.qs.tiles.viewmodel
 
 import android.os.UserHandle
-import kotlinx.coroutines.flow.SharedFlow
 import kotlinx.coroutines.flow.StateFlow
 
 /**
@@ -31,7 +30,7 @@
 interface QSTileViewModel {
 
     /** State of the tile to be shown by the view. */
-    val state: SharedFlow<QSTileState>
+    val state: StateFlow<QSTileState?>
 
     val config: QSTileConfig
 
@@ -62,9 +61,7 @@
 }
 
 /**
- * Returns the immediate state of the tile or null if the state haven't been collected yet. Favor
- * reactive consumption over the [currentState], because there is no guarantee that current value
- * would be available at any time.
+ * Returns the immediate state of the tile or null if the state haven't been collected yet.
  */
 val QSTileViewModel.currentState: QSTileState?
-    get() = state.replayCache.lastOrNull()
+    get() = state.value
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/viewmodel/QSTileViewModelAdapter.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/viewmodel/QSTileViewModelAdapter.kt
index 7be13e0..ba0a8d6 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/viewmodel/QSTileViewModelAdapter.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/viewmodel/QSTileViewModelAdapter.kt
@@ -38,6 +38,7 @@
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.flow.collectIndexed
+import kotlinx.coroutines.flow.filterNotNull
 import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.map
@@ -168,6 +169,7 @@
                 if (listeningClients.size == 1) {
                     stateJob =
                         qsTileViewModel.state
+                            .filterNotNull()
                             .map { mapState(context, it, qsTileViewModel.config) }
                             .onEach { legacyState ->
                                 synchronized(callbacks) {
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/viewmodel/StubQSTileViewModel.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/viewmodel/StubQSTileViewModel.kt
index 64f1fd3..00b7e61 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/viewmodel/StubQSTileViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/viewmodel/StubQSTileViewModel.kt
@@ -17,12 +17,11 @@
 package com.android.systemui.qs.tiles.viewmodel
 
 import android.os.UserHandle
-import kotlinx.coroutines.flow.SharedFlow
 import kotlinx.coroutines.flow.StateFlow
 
 object StubQSTileViewModel : QSTileViewModel {
 
-    override val state: SharedFlow<QSTileState>
+    override val state: StateFlow<QSTileState?>
         get() = error("Don't call stubs")
 
     override val config: QSTileConfig
diff --git a/packages/SystemUI/src/com/android/systemui/util/kotlin/SysUICoroutinesModule.kt b/packages/SystemUI/src/com/android/systemui/util/kotlin/SysUICoroutinesModule.kt
index d10554f..b836016 100644
--- a/packages/SystemUI/src/com/android/systemui/util/kotlin/SysUICoroutinesModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/util/kotlin/SysUICoroutinesModule.kt
@@ -20,13 +20,16 @@
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.dagger.qualifiers.Tracing
+import com.android.systemui.dagger.qualifiers.UiBackground
 import dagger.Module
 import dagger.Provides
+import java.util.concurrent.Executor
 import kotlin.coroutines.CoroutineContext
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.DelicateCoroutinesApi
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.asCoroutineDispatcher
 import kotlinx.coroutines.newFixedThreadPoolContext
 import kotlinx.coroutines.plus
 
@@ -83,4 +86,25 @@
     ): CoroutineContext {
         return bgCoroutineDispatcher + tracingCoroutineContext
     }
+
+    /** Coroutine dispatcher for background operations on for UI. */
+    @Provides
+    @SysUISingleton
+    @UiBackground
+    @Deprecated(
+        "Use @UiBackground CoroutineContext instead",
+        ReplaceWith("uiBgCoroutineContext()", "kotlin.coroutines.CoroutineContext")
+    )
+    fun uiBgDispatcher(@UiBackground uiBgExecutor: Executor): CoroutineDispatcher =
+        uiBgExecutor.asCoroutineDispatcher()
+
+    @Provides
+    @UiBackground
+    @SysUISingleton
+    fun uiBgCoroutineContext(
+        @Tracing tracingCoroutineContext: CoroutineContext,
+        @UiBackground uiBgCoroutineDispatcher: CoroutineDispatcher,
+    ): CoroutineContext {
+        return uiBgCoroutineDispatcher + tracingCoroutineContext
+    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/base/viewmodel/QSTileViewModelImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/base/viewmodel/QSTileViewModelImplTest.kt
index 42b81de..c918ed8 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/base/viewmodel/QSTileViewModelImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/base/viewmodel/QSTileViewModelImplTest.kt
@@ -97,6 +97,7 @@
                 qsTileLogger,
                 FakeSystemClock(),
                 testCoroutineDispatcher,
+                testCoroutineDispatcher,
                 testScope.backgroundScope,
             )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/tiles/di/NewQSTileFactoryKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/tiles/di/NewQSTileFactoryKosmos.kt
index 419e781..dceb8bf 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/tiles/di/NewQSTileFactoryKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/tiles/di/NewQSTileFactoryKosmos.kt
@@ -33,7 +33,6 @@
 import com.android.systemui.util.mockito.whenever
 import javax.inject.Provider
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.SharedFlow
 import kotlinx.coroutines.flow.StateFlow
 
 var Kosmos.newFactoryTileMap by Kosmos.Fixture { emptyMap<String, Provider<QSTileViewModel>>() }
@@ -50,7 +49,7 @@
                         instanceIdSequenceFake.newInstanceId(),
                     )
                 object : QSTileViewModel {
-                    override val state: SharedFlow<QSTileState> =
+                    override val state: StateFlow<QSTileState?> =
                         MutableStateFlow(QSTileState.build({ null }, tileSpec.spec) {})
                     override val config: QSTileConfig = config
                     override val isAvailable: StateFlow<Boolean> = MutableStateFlow(true)
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/tiles/impl/qr/QRCodeScannerTileKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/tiles/impl/qr/QRCodeScannerTileKosmos.kt
index dcfcce7..537be4f 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/tiles/impl/qr/QRCodeScannerTileKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/tiles/impl/qr/QRCodeScannerTileKosmos.kt
@@ -69,6 +69,7 @@
             qsTileLogger,
             systemClock,
             testDispatcher,
+            testDispatcher,
             testScope.backgroundScope,
         )
     }