Move the Tile composable in InfiniteGridLayout.

The Tile implementation is specific to the GridLayout implementation, so moving there makes the most sense.
Similarly, the icon only tiles are specific to InfiniteGridLayout, so moving that logic out of TileViewModel as well.

Flag: ACONFIG com.android.systemui.qs_ui_refactor DEVELOPMENT
Test: launching QSActivity
Bug: 331598956
Change-Id: I83b3573f9a5bf3af137c57aeb43e0d190225cb75
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 ee4eeb8..e3ba36f 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 @@
 
 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 920cbe7..68ce5d8 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 @@
     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 4d0089e7..e2143e0 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 @@
                 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 35f2970..0000000
--- 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 a528eed..2f32d72 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 fc13460..5eee691 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 f4b7255..58d07c3 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 @@
     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 08e9119..8fc66d3 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 @@
 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 @@
                 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 e8c5fd9..0000000
--- 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 e44fa07..c951642 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 @@
 
 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
similarity index 63%
rename from packages/SystemUI/src/com/android/systemui/qs/panels/shared/model/GridLayoutTypeKey.kt
rename to packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/InfiniteGridLayoutKosmos.kt
index 0dbaaba..1893c30 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 9851f04..5fd8762 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 @@
     Kosmos.Fixture {
         TileGridViewModel(
             gridLayoutTypeInteractor,
-            mapOf(Pair(InfiniteGridLayoutType::class.java, InfiniteGridLayout())),
+            gridLayoutMap,
             currentTilesInteractor,
-            iconTilesInteractor
+            InfiniteGridLayout(iconTilesInteractor),
+            applicationCoroutineScope,
         )
     }