Merge "Move the Tile composable in InfiniteGridLayout." into main

12 files changed, 224 insertions, 295 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/dagger/PanelsModule.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/dagger/PanelsModule.kt
index ee4eeb8b9e32..e3ba36fe6e32 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/dagger/PanelsModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/dagger/PanelsModule.kt
@@ -18,20 +18,31 @@ package com.android.systemui.qs.panels.dagger
 
 import com.android.systemui.qs.panels.data.repository.IconTilesRepository
 import com.android.systemui.qs.panels.data.repository.IconTilesRepositoryImpl
-import com.android.systemui.qs.panels.shared.model.GridLayoutTypeKey
+import com.android.systemui.qs.panels.shared.model.GridLayoutType
 import com.android.systemui.qs.panels.shared.model.InfiniteGridLayoutType
 import com.android.systemui.qs.panels.ui.compose.GridLayout
 import com.android.systemui.qs.panels.ui.compose.InfiniteGridLayout
 import dagger.Binds
 import dagger.Module
-import dagger.multibindings.IntoMap
+import dagger.Provides
+import dagger.multibindings.IntoSet
 
 @Module
 interface PanelsModule {
     @Binds fun bindIconTilesRepository(impl: IconTilesRepositoryImpl): IconTilesRepository
 
-    @Binds
-    @IntoMap
-    @GridLayoutTypeKey(InfiniteGridLayoutType::class)
-    fun bindGridLayout(impl: InfiniteGridLayout): GridLayout
+    companion object {
+        @Provides
+        @IntoSet
+        fun provideGridLayout(gridLayout: InfiniteGridLayout): Pair<GridLayoutType, GridLayout> {
+            return Pair(InfiniteGridLayoutType, gridLayout)
+        }
+
+        @Provides
+        fun provideGridLayoutMap(
+            entries: Set<@JvmSuppressWildcards Pair<GridLayoutType, GridLayout>>
+        ): Map<GridLayoutType, GridLayout> {
+            return entries.toMap()
+        }
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/GridLayout.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/GridLayout.kt
index 920cbe77e42f..68ce5d830570 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/GridLayout.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/GridLayout.kt
@@ -25,6 +25,5 @@ interface GridLayout {
     fun TileGrid(
         tiles: List<TileViewModel>,
         modifier: Modifier,
-        tile: @Composable (TileViewModel) -> Unit,
     )
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/InfiniteGridLayout.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/InfiniteGridLayout.kt
index 4d0089e70e80..e2143e090b21 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/InfiniteGridLayout.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/InfiniteGridLayout.kt
@@ -16,32 +16,77 @@
 
 package com.android.systemui.qs.panels.ui.compose
 
+import android.graphics.drawable.Animatable
+import android.text.TextUtils
+import androidx.appcompat.content.res.AppCompatResources
+import androidx.compose.animation.graphics.ExperimentalAnimationGraphicsApi
+import androidx.compose.animation.graphics.res.animatedVectorResource
+import androidx.compose.animation.graphics.res.rememberAnimatedVectorPainter
+import androidx.compose.animation.graphics.vector.AnimatedImageVector
+import androidx.compose.foundation.Image
+import androidx.compose.foundation.background
+import androidx.compose.foundation.basicMarquee
+import androidx.compose.foundation.clickable
 import androidx.compose.foundation.layout.Arrangement
+import androidx.compose.foundation.layout.Column
+import androidx.compose.foundation.layout.Row
+import androidx.compose.foundation.layout.fillMaxHeight
+import androidx.compose.foundation.layout.fillMaxWidth
+import androidx.compose.foundation.layout.height
+import androidx.compose.foundation.layout.padding
+import androidx.compose.foundation.layout.size
 import androidx.compose.foundation.lazy.grid.GridCells
 import androidx.compose.foundation.lazy.grid.GridItemSpan
 import androidx.compose.foundation.lazy.grid.LazyVerticalGrid
+import androidx.compose.foundation.shape.RoundedCornerShape
+import androidx.compose.material3.Text
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.DisposableEffect
+import androidx.compose.runtime.LaunchedEffect
+import androidx.compose.runtime.collectAsState
+import androidx.compose.runtime.getValue
+import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.remember
+import androidx.compose.runtime.setValue
+import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
+import androidx.compose.ui.draw.clip
+import androidx.compose.ui.graphics.Color
+import androidx.compose.ui.graphics.ColorFilter
+import androidx.compose.ui.platform.LocalContext
 import androidx.compose.ui.res.dimensionResource
 import androidx.compose.ui.res.integerResource
+import com.android.compose.theme.colorAttr
+import com.android.systemui.common.shared.model.Icon
+import com.android.systemui.common.ui.compose.Icon
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.qs.panels.domain.interactor.IconTilesInteractor
+import com.android.systemui.qs.panels.ui.viewmodel.TileUiState
 import com.android.systemui.qs.panels.ui.viewmodel.TileViewModel
+import com.android.systemui.qs.panels.ui.viewmodel.toUiState
+import com.android.systemui.qs.tileimpl.QSTileImpl
 import com.android.systemui.res.R
 import javax.inject.Inject
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.flow.mapLatest
 
-class InfiniteGridLayout @Inject constructor() : GridLayout {
+@SysUISingleton
+class InfiniteGridLayout @Inject constructor(private val iconTilesInteractor: IconTilesInteractor) :
+    GridLayout {
 
     @Composable
     override fun TileGrid(
         tiles: List<TileViewModel>,
         modifier: Modifier,
-        tile: @Composable (TileViewModel) -> Unit
     ) {
         DisposableEffect(tiles) {
             val token = Any()
             tiles.forEach { it.startListening(token) }
             onDispose { tiles.forEach { it.stopListening(token) } }
         }
+        val iconTilesSpecs by
+            iconTilesInteractor.iconTilesSpecs.collectAsState(initial = emptySet())
 
         LazyVerticalGrid(
             columns =
@@ -54,14 +99,128 @@ class InfiniteGridLayout @Inject constructor() : GridLayout {
                 Arrangement.spacedBy(dimensionResource(R.dimen.qs_tile_margin_horizontal)),
             modifier = modifier
         ) {
-            tiles.forEach { item(span = { it.span() }) { tile(it) } }
+            items(
+                tiles.size,
+                span = { index ->
+                    val iconOnly = iconTilesSpecs.contains(tiles[index].spec)
+                    if (iconOnly) {
+                        GridItemSpan(1)
+                    } else {
+                        GridItemSpan(2)
+                    }
+                }
+            ) { index ->
+                Tile(
+                    tiles[index],
+                    iconTilesSpecs.contains(tiles[index].spec),
+                    Modifier.height(dimensionResource(id = R.dimen.qs_tile_height))
+                )
+            }
         }
     }
 
-    private fun TileViewModel.span(): GridItemSpan =
-        if (iconOnly) {
-            GridItemSpan(1)
-        } else {
-            GridItemSpan(2)
+    @OptIn(ExperimentalCoroutinesApi::class)
+    @Composable
+    private fun Tile(
+        tile: TileViewModel,
+        iconOnly: Boolean,
+        modifier: Modifier,
+    ) {
+        val state: TileUiState by
+            tile.state
+                .mapLatest { it.toUiState() }
+                .collectAsState(initial = tile.currentState.toUiState())
+        val context = LocalContext.current
+        val horizontalAlignment =
+            if (iconOnly) {
+                Alignment.CenterHorizontally
+            } else {
+                Alignment.Start
+            }
+
+        Row(
+            modifier =
+                modifier
+                    .fillMaxWidth()
+                    .clip(RoundedCornerShape(dimensionResource(R.dimen.qs_corner_radius)))
+                    .clickable { tile.onClick(null) }
+                    .background(colorAttr(state.colors.background))
+                    .padding(
+                        horizontal = dimensionResource(id = R.dimen.qs_label_container_margin)
+                    ),
+            verticalAlignment = Alignment.CenterVertically,
+            horizontalArrangement =
+                Arrangement.spacedBy(
+                    space = dimensionResource(id = R.dimen.qs_label_container_margin),
+                    alignment = horizontalAlignment
+                )
+        ) {
+            val icon =
+                remember(state.icon) {
+                    state.icon.get().let {
+                        if (it is QSTileImpl.ResourceIcon) {
+                            Icon.Resource(it.resId, null)
+                        } else {
+                            Icon.Loaded(it.getDrawable(context), null)
+                        }
+                    }
+                }
+            TileIcon(icon, colorAttr(state.colors.icon))
+
+            if (!iconOnly) {
+                Column(
+                    verticalArrangement = Arrangement.Center,
+                    modifier = Modifier.fillMaxHeight()
+                ) {
+                    Text(
+                        state.label.toString(),
+                        color = colorAttr(state.colors.label),
+                        modifier = Modifier.basicMarquee(),
+                    )
+                    if (!TextUtils.isEmpty(state.secondaryLabel)) {
+                        Text(
+                            state.secondaryLabel.toString(),
+                            color = colorAttr(state.colors.secondaryLabel),
+                            modifier = Modifier.basicMarquee(),
+                        )
+                    }
+                }
+            }
+        }
+    }
+
+    @OptIn(ExperimentalAnimationGraphicsApi::class)
+    @Composable
+    private fun TileIcon(icon: Icon, color: Color) {
+        val modifier = Modifier.size(dimensionResource(id = R.dimen.qs_icon_size))
+        val context = LocalContext.current
+        val loadedDrawable =
+            remember(icon, context) {
+                when (icon) {
+                    is Icon.Loaded -> icon.drawable
+                    is Icon.Resource -> AppCompatResources.getDrawable(context, icon.res)
+                }
+            }
+        if (loadedDrawable !is Animatable) {
+            Icon(
+                icon = icon,
+                tint = color,
+                modifier = modifier,
+            )
+        } else if (icon is Icon.Resource) {
+            val image = AnimatedImageVector.animatedVectorResource(id = icon.res)
+            var atEnd by remember(icon.res) { mutableStateOf(false) }
+            LaunchedEffect(key1 = icon.res) {
+                delay(350)
+                atEnd = true
+            }
+            val painter = rememberAnimatedVectorPainter(animatedImageVector = image, atEnd = atEnd)
+            Image(
+                painter = painter,
+                contentDescription = null,
+                colorFilter = ColorFilter.tint(color = color),
+                modifier = modifier
+            )
         }
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/Tile.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/Tile.kt
deleted file mode 100644
index 35f29705fab2..000000000000
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/Tile.kt
+++ /dev/null
@@ -1,155 +0,0 @@
-/*
- * Copyright (C) 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.systemui.qs.panels.ui.compose
-
-import android.graphics.drawable.Animatable
-import android.text.TextUtils
-import androidx.appcompat.content.res.AppCompatResources
-import androidx.compose.animation.graphics.ExperimentalAnimationGraphicsApi
-import androidx.compose.animation.graphics.res.animatedVectorResource
-import androidx.compose.animation.graphics.res.rememberAnimatedVectorPainter
-import androidx.compose.animation.graphics.vector.AnimatedImageVector
-import androidx.compose.foundation.Image
-import androidx.compose.foundation.background
-import androidx.compose.foundation.basicMarquee
-import androidx.compose.foundation.clickable
-import androidx.compose.foundation.layout.Arrangement
-import androidx.compose.foundation.layout.Column
-import androidx.compose.foundation.layout.Row
-import androidx.compose.foundation.layout.fillMaxHeight
-import androidx.compose.foundation.layout.fillMaxWidth
-import androidx.compose.foundation.layout.padding
-import androidx.compose.foundation.layout.size
-import androidx.compose.foundation.shape.RoundedCornerShape
-import androidx.compose.material3.Text
-import androidx.compose.runtime.Composable
-import androidx.compose.runtime.LaunchedEffect
-import androidx.compose.runtime.collectAsState
-import androidx.compose.runtime.getValue
-import androidx.compose.runtime.mutableStateOf
-import androidx.compose.runtime.remember
-import androidx.compose.runtime.setValue
-import androidx.compose.ui.Alignment
-import androidx.compose.ui.Modifier
-import androidx.compose.ui.draw.clip
-import androidx.compose.ui.graphics.Color
-import androidx.compose.ui.graphics.ColorFilter
-import androidx.compose.ui.platform.LocalContext
-import androidx.compose.ui.res.dimensionResource
-import com.android.compose.theme.colorAttr
-import com.android.systemui.common.shared.model.Icon as IconModel
-import com.android.systemui.common.ui.compose.Icon
-import com.android.systemui.qs.panels.ui.viewmodel.TileUiState
-import com.android.systemui.qs.panels.ui.viewmodel.TileViewModel
-import com.android.systemui.qs.tileimpl.QSTileImpl
-import com.android.systemui.res.R
-import kotlinx.coroutines.delay
-
-@Composable
-fun Tile(
-    tileViewModel: TileViewModel,
-    modifier: Modifier = Modifier,
-) {
-    val state: TileUiState by tileViewModel.state.collectAsState(tileViewModel.currentState)
-    val context = LocalContext.current
-    val horizontalAlignment =
-        if (state.iconOnly) {
-            Alignment.CenterHorizontally
-        } else {
-            Alignment.Start
-        }
-
-    Row(
-        modifier =
-            modifier
-                .fillMaxWidth()
-                .clip(RoundedCornerShape(dimensionResource(R.dimen.qs_corner_radius)))
-                .clickable { tileViewModel.onClick(null) }
-                .background(colorAttr(state.colors.background))
-                .padding(horizontal = dimensionResource(id = R.dimen.qs_label_container_margin)),
-        verticalAlignment = Alignment.CenterVertically,
-        horizontalArrangement =
-            Arrangement.spacedBy(
-                space = dimensionResource(id = R.dimen.qs_label_container_margin),
-                alignment = horizontalAlignment
-            )
-    ) {
-        val icon =
-            remember(state.icon) {
-                state.icon.get().let {
-                    if (it is QSTileImpl.ResourceIcon) {
-                        IconModel.Resource(it.resId, null)
-                    } else {
-                        IconModel.Loaded(it.getDrawable(context), null)
-                    }
-                }
-            }
-        TileIcon(icon, colorAttr(state.colors.icon))
-
-        if (!state.iconOnly) {
-            Column(verticalArrangement = Arrangement.Center, modifier = Modifier.fillMaxHeight()) {
-                Text(
-                    state.label.toString(),
-                    color = colorAttr(state.colors.label),
-                    modifier = Modifier.basicMarquee(),
-                )
-                if (!TextUtils.isEmpty(state.secondaryLabel)) {
-                    Text(
-                        state.secondaryLabel.toString(),
-                        color = colorAttr(state.colors.secondaryLabel),
-                        modifier = Modifier.basicMarquee(),
-                    )
-                }
-            }
-        }
-    }
-}
-
-@OptIn(ExperimentalAnimationGraphicsApi::class)
-@Composable
-private fun TileIcon(icon: IconModel, color: Color) {
-    val modifier = Modifier.size(dimensionResource(id = R.dimen.qs_icon_size))
-    val context = LocalContext.current
-    val loadedDrawable =
-        remember(icon, context) {
-            when (icon) {
-                is IconModel.Loaded -> icon.drawable
-                is IconModel.Resource -> AppCompatResources.getDrawable(context, icon.res)
-            }
-        }
-    if (loadedDrawable !is Animatable) {
-        Icon(
-            icon = icon,
-            tint = color,
-            modifier = modifier,
-        )
-    } else if (icon is IconModel.Resource) {
-        val image = AnimatedImageVector.animatedVectorResource(id = icon.res)
-        var atEnd by remember(icon.res) { mutableStateOf(false) }
-        LaunchedEffect(key1 = icon.res) {
-            delay(350)
-            atEnd = true
-        }
-        val painter = rememberAnimatedVectorPainter(animatedImageVector = image, atEnd = atEnd)
-        Image(
-            painter = painter,
-            contentDescription = null,
-            colorFilter = ColorFilter.tint(color = color),
-            modifier = modifier
-        )
-    }
-}
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/TileGrid.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/TileGrid.kt
index a528eed4a346..2f32d7264350 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/TileGrid.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/TileGrid.kt
@@ -16,21 +16,16 @@
 
 package com.android.systemui.qs.panels.ui.compose
 
-import androidx.compose.foundation.layout.height
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.collectAsState
 import androidx.compose.runtime.getValue
 import androidx.compose.ui.Modifier
-import androidx.compose.ui.res.dimensionResource
 import com.android.systemui.qs.panels.ui.viewmodel.TileGridViewModel
-import com.android.systemui.res.R
 
 @Composable
 fun TileGrid(viewModel: TileGridViewModel, modifier: Modifier = Modifier) {
-    val gridLayout by viewModel.gridLayout.collectAsState(InfiniteGridLayout())
+    val gridLayout by viewModel.gridLayout.collectAsState()
     val tiles by viewModel.tileViewModels.collectAsState(emptyList())
 
-    gridLayout.TileGrid(tiles, modifier) {
-        Tile(it, modifier = Modifier.height(dimensionResource(id = R.dimen.qs_tile_height)))
-    }
+    gridLayout.TileGrid(tiles, modifier)
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/TileGridViewModel.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/TileGridViewModel.kt
index fc134602e9cb..5eee691dcf79 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/TileGridViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/TileGridViewModel.kt
@@ -17,34 +17,39 @@
 package com.android.systemui.qs.panels.ui.viewmodel
 
 import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.qs.panels.domain.interactor.GridLayoutTypeInteractor
-import com.android.systemui.qs.panels.domain.interactor.IconTilesInteractor
 import com.android.systemui.qs.panels.shared.model.GridLayoutType
 import com.android.systemui.qs.panels.ui.compose.GridLayout
 import com.android.systemui.qs.panels.ui.compose.InfiniteGridLayout
 import com.android.systemui.qs.pipeline.domain.interactor.CurrentTilesInteractor
 import javax.inject.Inject
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.mapLatest
+import kotlinx.coroutines.flow.stateIn
 
+@OptIn(ExperimentalCoroutinesApi::class)
 @SysUISingleton
 class TileGridViewModel
 @Inject
 constructor(
     gridLayoutTypeInteractor: GridLayoutTypeInteractor,
-    gridLayoutMap: Map<Class<out GridLayoutType>, @JvmSuppressWildcards GridLayout>,
+    gridLayoutMap: Map<GridLayoutType, @JvmSuppressWildcards GridLayout>,
     tilesInteractor: CurrentTilesInteractor,
-    iconTilesInteractor: IconTilesInteractor,
+    defaultGridLayout: InfiniteGridLayout,
+    @Application private val applicationScope: CoroutineScope
 ) {
-    val gridLayout: Flow<GridLayout> =
-        gridLayoutTypeInteractor.layout.map {
-            gridLayoutMap[it::class.java] ?: InfiniteGridLayout()
-        }
+    val gridLayout: StateFlow<GridLayout> =
+        gridLayoutTypeInteractor.layout
+            .map { gridLayoutMap[it] ?: defaultGridLayout }
+            .stateIn(applicationScope, SharingStarted.Eagerly, defaultGridLayout)
     val tileViewModels: Flow<List<TileViewModel>> =
-        combine(tilesInteractor.currentTiles, iconTilesInteractor.iconTilesSpecs) {
-            tiles,
-            iconTilesSpecs ->
-            tiles.map { TileViewModel(it.tile, iconTilesSpecs.contains(it.spec)) }
+        tilesInteractor.currentTiles.mapLatest { tiles ->
+            tiles.map { TileViewModel(it.tile, it.spec) }
         }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/TileUiState.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/TileUiState.kt
index f4b7255693e2..58d07c368456 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/TileUiState.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/TileUiState.kt
@@ -24,15 +24,13 @@ data class TileUiState(
     val secondaryLabel: CharSequence,
     val colors: TileColorAttributes,
     val icon: Supplier<QSTile.Icon>,
-    val iconOnly: Boolean,
 )
 
-fun QSTile.State.toUiState(iconOnly: Boolean): TileUiState {
+fun QSTile.State.toUiState(): TileUiState {
     return TileUiState(
         label ?: "",
         secondaryLabel ?: "",
         colors(),
         icon?.let { Supplier { icon } } ?: iconSupplier,
-        iconOnly,
     )
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/TileViewModel.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/TileViewModel.kt
index 08e9119a6c76..8fc66d37ef86 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/TileViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/TileViewModel.kt
@@ -19,16 +19,16 @@ package com.android.systemui.qs.panels.ui.viewmodel
 import android.view.View
 import android.view.View.OnLongClickListener
 import com.android.systemui.plugins.qs.QSTile
+import com.android.systemui.qs.pipeline.shared.TileSpec
 import com.android.systemui.utils.coroutines.flow.conflatedCallbackFlow
 import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.distinctUntilChanged
-import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.onStart
 
-class TileViewModel(private val tile: QSTile, val iconOnly: Boolean = false) :
+class TileViewModel(private val tile: QSTile, val spec: TileSpec) :
     OnLongClickListener, View.OnClickListener {
-    val state: Flow<TileUiState> =
+    val state: Flow<QSTile.State> =
         conflatedCallbackFlow {
                 val callback = QSTile.Callback { trySend(it.copy()) }
 
@@ -37,11 +37,10 @@ class TileViewModel(private val tile: QSTile, val iconOnly: Boolean = false) :
                 awaitClose { tile.removeCallback(callback) }
             }
             .onStart { emit(tile.state) }
-            .map { it.toUiState(iconOnly) }
             .distinctUntilChanged()
 
-    val currentState: TileUiState
-        get() = tile.state.toUiState(iconOnly)
+    val currentState: QSTile.State
+        get() = tile.state
 
     override fun onClick(view: View?) {
         tile.click(view)
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/panels/ui/viewmodel/TileGridViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/qs/panels/ui/viewmodel/TileGridViewModelTest.kt
deleted file mode 100644
index e8c5fd978319..000000000000
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/panels/ui/viewmodel/TileGridViewModelTest.kt
+++ /dev/null
@@ -1,84 +0,0 @@
-/*
- * Copyright (C) 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.systemui.qs.panels.ui.viewmodel
-
-import android.testing.AndroidTestingRunner
-import androidx.test.filters.SmallTest
-import com.android.systemui.SysuiTestCase
-import com.android.systemui.coroutines.collectLastValue
-import com.android.systemui.kosmos.testScope
-import com.android.systemui.plugins.qs.QSTile
-import com.android.systemui.qs.FakeQSFactory
-import com.android.systemui.qs.pipeline.data.repository.tileSpecRepository
-import com.android.systemui.qs.pipeline.domain.interactor.FakeQSTile
-import com.android.systemui.qs.pipeline.shared.TileSpec
-import com.android.systemui.qs.qsTileFactory
-import com.android.systemui.testKosmos
-import kotlinx.coroutines.test.runTest
-import org.junit.Assert
-import org.junit.Test
-import org.junit.runner.RunWith
-
-@SmallTest
-@RunWith(AndroidTestingRunner::class)
-class TileGridViewModelTest : SysuiTestCase() {
-
-    private val kosmos = testKosmos().apply { qsTileFactory = FakeQSFactory(::tileCreator) }
-    private val underTest = with(kosmos) { tileGridViewModel }
-
-    @Test
-    fun noIconTiles() =
-        with(kosmos) {
-            testScope.runTest {
-                val latest by collectLastValue(underTest.tileViewModels)
-
-                tileSpecRepository.setTiles(
-                    0,
-                    listOf(
-                        TileSpec.create("bluetooth"),
-                        TileSpec.create("internet"),
-                        TileSpec.create("alarm")
-                    )
-                )
-
-                latest!!.forEach { Assert.assertFalse(it.iconOnly) }
-            }
-        }
-
-    @Test
-    fun withIconTiles() =
-        with(kosmos) {
-            testScope.runTest {
-                val latest by collectLastValue(underTest.tileViewModels)
-
-                tileSpecRepository.setTiles(
-                    0,
-                    listOf(
-                        TileSpec.create("airplane"),
-                        TileSpec.create("flashlight"),
-                        TileSpec.create("rotation")
-                    )
-                )
-
-                latest!!.forEach { Assert.assertTrue(it.iconOnly) }
-            }
-        }
-
-    private fun tileCreator(spec: String): QSTile {
-        return FakeQSTile(0).apply { tileSpec = spec }
-    }
-}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/GridLayoutTypeInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/GridLayoutTypeInteractorKosmos.kt
index e44fa0717602..c9516429553b 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/GridLayoutTypeInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/GridLayoutTypeInteractorKosmos.kt
@@ -18,6 +18,12 @@ package com.android.systemui.qs.panels.domain.interactor
 
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.qs.panels.data.repository.gridLayoutTypeRepository
+import com.android.systemui.qs.panels.shared.model.GridLayoutType
+import com.android.systemui.qs.panels.shared.model.InfiniteGridLayoutType
+import com.android.systemui.qs.panels.ui.compose.GridLayout
 
 val Kosmos.gridLayoutTypeInteractor by
     Kosmos.Fixture { GridLayoutTypeInteractor(gridLayoutTypeRepository) }
+
+val Kosmos.gridLayoutMap: Map<GridLayoutType, GridLayout> by
+    Kosmos.Fixture { mapOf(Pair(InfiniteGridLayoutType, infiniteGridLayout)) }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/shared/model/GridLayoutTypeKey.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/InfiniteGridLayoutKosmos.kt
index 0dbaaba9dfac..1893c302c332 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/shared/model/GridLayoutTypeKey.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/InfiniteGridLayoutKosmos.kt
@@ -14,15 +14,9 @@
  * limitations under the License.
  */
 
-package com.android.systemui.qs.panels.shared.model
+package com.android.systemui.qs.panels.domain.interactor
 
-import dagger.MapKey
-import kotlin.reflect.KClass
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.qs.panels.ui.compose.InfiniteGridLayout
 
-/**
- * Dagger map key to associate a [GridLayoutType] with its
- * [com.android.systemui.qs.panels.ui.compose.GridLayout].
- */
-@Retention(AnnotationRetention.RUNTIME)
-@MapKey
-annotation class GridLayoutTypeKey(val value: KClass<out GridLayoutType>)
+val Kosmos.infiniteGridLayout by Kosmos.Fixture { InfiniteGridLayout(iconTilesInteractor) }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/ui/viewmodel/TileGridViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/ui/viewmodel/TileGridViewModelKosmos.kt
index 9851f043e39a..5fd876222852 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/ui/viewmodel/TileGridViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/ui/viewmodel/TileGridViewModelKosmos.kt
@@ -17,9 +17,10 @@
 package com.android.systemui.qs.panels.ui.viewmodel
 
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.applicationCoroutineScope
+import com.android.systemui.qs.panels.domain.interactor.gridLayoutMap
 import com.android.systemui.qs.panels.domain.interactor.gridLayoutTypeInteractor
 import com.android.systemui.qs.panels.domain.interactor.iconTilesInteractor
-import com.android.systemui.qs.panels.shared.model.InfiniteGridLayoutType
 import com.android.systemui.qs.panels.ui.compose.InfiniteGridLayout
 import com.android.systemui.qs.pipeline.domain.interactor.currentTilesInteractor
 
@@ -27,8 +28,9 @@ val Kosmos.tileGridViewModel by
     Kosmos.Fixture {
         TileGridViewModel(
             gridLayoutTypeInteractor,
-            mapOf(Pair(InfiniteGridLayoutType::class.java, InfiniteGridLayout())),
+            gridLayoutMap,
             currentTilesInteractor,
-            iconTilesInteractor
+            InfiniteGridLayout(iconTilesInteractor),
+            applicationCoroutineScope,
         )
     }