Grid Picker Selection Animation Fix (1/2)

Fix grid picker selection animation issue, due to comparison issues in
diff util.

Bug: 273336388
Test: Manually verified that grid picker, color picker, and quick
affordance picker selection animation function correctly. Tested with
print statements to verify diff util functions as intended. Updated
field types in grid view model test.

Change-Id: If48da6985ca99e254fbe511f2062a32e946883c9
diff --git a/src/com/android/customization/model/grid/ui/binder/GridIconViewBinder.kt b/src/com/android/customization/model/grid/ui/binder/GridIconViewBinder.kt
new file mode 100644
index 0000000..fba89a7
--- /dev/null
+++ b/src/com/android/customization/model/grid/ui/binder/GridIconViewBinder.kt
@@ -0,0 +1,17 @@
+package com.android.customization.model.grid.ui.binder
+
+import android.widget.ImageView
+import com.android.customization.model.grid.ui.viewmodel.GridIconViewModel
+import com.android.customization.widget.GridTileDrawable
+
+object GridIconViewBinder {
+    fun bind(view: ImageView, viewModel: GridIconViewModel) {
+        view.setImageDrawable(
+            GridTileDrawable(
+                viewModel.columns,
+                viewModel.rows,
+                viewModel.path,
+            )
+        )
+    }
+}
diff --git a/src/com/android/customization/model/grid/ui/binder/GridScreenBinder.kt b/src/com/android/customization/model/grid/ui/binder/GridScreenBinder.kt
index 78e0376..78536ca 100644
--- a/src/com/android/customization/model/grid/ui/binder/GridScreenBinder.kt
+++ b/src/com/android/customization/model/grid/ui/binder/GridScreenBinder.kt
@@ -25,11 +25,10 @@
 import androidx.lifecycle.repeatOnLifecycle
 import androidx.recyclerview.widget.LinearLayoutManager
 import androidx.recyclerview.widget.RecyclerView
+import com.android.customization.model.grid.ui.viewmodel.GridIconViewModel
 import com.android.customization.model.grid.ui.viewmodel.GridScreenViewModel
 import com.android.customization.picker.common.ui.view.ItemSpacing
 import com.android.wallpaper.R
-import com.android.wallpaper.picker.common.icon.ui.viewbinder.IconViewBinder
-import com.android.wallpaper.picker.common.icon.ui.viewmodel.Icon
 import com.android.wallpaper.picker.option.ui.adapter.OptionItemAdapter
 import com.android.wallpaper.picker.option.ui.binder.OptionItemBinder
 import kotlinx.coroutines.CoroutineDispatcher
@@ -61,9 +60,9 @@
                         selectedColor = view.context.getColor(R.color.text_color_primary),
                         unselectedColor = view.context.getColor(R.color.text_color_secondary),
                     ),
-                bindIcon = { foregroundView: View, gridIcon: Icon ->
+                bindIcon = { foregroundView: View, gridIcon: GridIconViewModel ->
                     val imageView = foregroundView as? ImageView
-                    imageView?.let { IconViewBinder.bind(imageView, gridIcon) }
+                    imageView?.let { GridIconViewBinder.bind(imageView, gridIcon) }
                 }
             )
         optionView.adapter = adapter
diff --git a/src/com/android/customization/model/grid/ui/viewmodel/GridIconViewModel.kt b/src/com/android/customization/model/grid/ui/viewmodel/GridIconViewModel.kt
new file mode 100644
index 0000000..3942d7c
--- /dev/null
+++ b/src/com/android/customization/model/grid/ui/viewmodel/GridIconViewModel.kt
@@ -0,0 +1,24 @@
+/*
+ * Copyright (C) 2023 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.customization.model.grid.ui.viewmodel
+
+data class GridIconViewModel(
+    val columns: Int,
+    val rows: Int,
+    val path: String,
+)
diff --git a/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModel.kt b/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModel.kt
index 69d938c..c11a594 100644
--- a/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModel.kt
+++ b/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModel.kt
@@ -26,12 +26,11 @@
 import com.android.customization.model.ResourceConstants
 import com.android.customization.model.grid.domain.interactor.GridInteractor
 import com.android.customization.model.grid.shared.model.GridOptionItemsModel
-import com.android.customization.widget.GridTileDrawable
-import com.android.wallpaper.picker.common.icon.ui.viewmodel.Icon
 import com.android.wallpaper.picker.common.text.ui.viewmodel.Text
 import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
 import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.flowOf
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.launch
 
@@ -43,12 +42,12 @@
     @SuppressLint("StaticFieldLeak") // We're not leaking this context as it is the app context.
     private val applicationContext = context.applicationContext
 
-    val optionItems: Flow<List<OptionItemViewModel<Icon>>> =
+    val optionItems: Flow<List<OptionItemViewModel<GridIconViewModel>>> =
         interactor.options.map { model -> toViewModel(model) }
 
     private fun toViewModel(
         model: GridOptionItemsModel,
-    ): List<OptionItemViewModel<Icon>> {
+    ): List<OptionItemViewModel<GridIconViewModel>> {
         val iconShapePath =
             applicationContext.resources.getString(
                 Resources.getSystem()
@@ -63,17 +62,14 @@
             is GridOptionItemsModel.Loaded ->
                 model.options.map { option ->
                     val text = Text.Loaded(option.name)
-                    OptionItemViewModel<Icon>(
-                        key = flowOf("${option.cols}x${option.rows}"),
+                    OptionItemViewModel<GridIconViewModel>(
+                        key =
+                            MutableStateFlow("${option.cols}x${option.rows}") as StateFlow<String>,
                         payload =
-                            Icon.Loaded(
-                                drawable =
-                                    GridTileDrawable(
-                                        option.cols,
-                                        option.rows,
-                                        iconShapePath,
-                                    ),
-                                contentDescription = text
+                            GridIconViewModel(
+                                columns = option.cols,
+                                rows = option.rows,
+                                path = iconShapePath,
                             ),
                         text = text,
                         isSelected = option.isSelected,
diff --git a/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt b/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt
index 803663d..58bc858 100644
--- a/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt
+++ b/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt
@@ -31,9 +31,10 @@
 import kotlin.math.min
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.combine
-import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.launch
 
 /** Models UI state for a color picker experience. */
@@ -91,15 +92,19 @@
                                         colorOptionModel.colorOption as ColorSeedOption
                                     val colors =
                                         colorSeedOption.previewInfo.resolveColors(context.resources)
-                                    val isSelectedFlow: Flow<Boolean> =
-                                        interactor.activeColorOption.map {
-                                            it?.colorOption?.isEquivalent(
-                                                colorOptionModel.colorOption
-                                            )
-                                                ?: colorOptionModel.isSelected
-                                        }
+                                    val isSelectedFlow: StateFlow<Boolean> =
+                                        interactor.activeColorOption
+                                            .map {
+                                                it?.colorOption?.isEquivalent(
+                                                    colorOptionModel.colorOption
+                                                )
+                                                    ?: colorOptionModel.isSelected
+                                            }
+                                            .stateIn(viewModelScope)
                                     OptionItemViewModel<ColorOptionIconViewModel>(
-                                        key = flowOf(colorOptionModel.key),
+                                        key =
+                                            MutableStateFlow(colorOptionModel.key)
+                                                as StateFlow<String>,
                                         payload =
                                             ColorOptionIconViewModel(
                                                 colors[0],
@@ -141,15 +146,19 @@
                                         colorBundle.previewInfo.resolveSecondaryColor(
                                             context.resources
                                         )
-                                    val isSelectedFlow: Flow<Boolean> =
-                                        interactor.activeColorOption.map {
-                                            it?.colorOption?.isEquivalent(
-                                                colorOptionModel.colorOption
-                                            )
-                                                ?: colorOptionModel.isSelected
-                                        }
+                                    val isSelectedFlow: StateFlow<Boolean> =
+                                        interactor.activeColorOption
+                                            .map {
+                                                it?.colorOption?.isEquivalent(
+                                                    colorOptionModel.colorOption
+                                                )
+                                                    ?: colorOptionModel.isSelected
+                                            }
+                                            .stateIn(viewModelScope)
                                     OptionItemViewModel<ColorOptionIconViewModel>(
-                                        key = flowOf(colorOptionModel.key),
+                                        key =
+                                            MutableStateFlow(colorOptionModel.key)
+                                                as StateFlow<String>,
                                         payload =
                                             ColorOptionIconViewModel(
                                                 primaryColor,
diff --git a/src/com/android/customization/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModel.kt b/src/com/android/customization/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModel.kt
index cbc140e..14b6acc 100644
--- a/src/com/android/customization/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModel.kt
+++ b/src/com/android/customization/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModel.kt
@@ -148,7 +148,9 @@
                         selectedQuickAffordances =
                             selectedAffordances.map { affordanceModel ->
                                 OptionItemViewModel<Icon>(
-                                    key = flowOf("${slot.id}::${affordanceModel.id}"),
+                                    key =
+                                        MutableStateFlow("${slot.id}::${affordanceModel.id}")
+                                            as StateFlow<String>,
                                     payload =
                                         Icon.Loaded(
                                             drawable =
@@ -156,7 +158,7 @@
                                             contentDescription = null,
                                         ),
                                     text = Text.Loaded(affordanceModel.name),
-                                    isSelected = flowOf(true),
+                                    isSelected = MutableStateFlow(true) as StateFlow<Boolean>,
                                     onClicked = flowOf(null),
                                     onLongClicked = null,
                                     isEnabled = true,
@@ -196,7 +198,7 @@
     /** The list of all available quick affordances for the selected slot. */
     val quickAffordances: Flow<List<OptionItemViewModel<Icon>>> =
         quickAffordanceInteractor.affordances.map { affordances ->
-            val isNoneSelected = selectedAffordanceIds.map { it.isEmpty() }
+            val isNoneSelected = selectedAffordanceIds.map { it.isEmpty() }.stateIn(viewModelScope)
             listOf(
                 none(
                     slotId = selectedSlotId,
@@ -220,10 +222,15 @@
             ) +
                 affordances.map { affordance ->
                     val affordanceIcon = getAffordanceIcon(affordance.iconResourceId)
-                    val isSelectedFlow: Flow<Boolean> =
-                        selectedAffordanceIds.map { it.contains(affordance.id) }
+                    val isSelectedFlow: StateFlow<Boolean> =
+                        selectedAffordanceIds
+                            .map { it.contains(affordance.id) }
+                            .stateIn(viewModelScope)
                     OptionItemViewModel<Icon>(
-                        key = selectedSlotId.map { slotId -> "$slotId::${affordance.id}" },
+                        key =
+                            selectedSlotId
+                                .map { slotId -> "$slotId::${affordance.id}" }
+                                .stateIn(viewModelScope),
                         payload = Icon.Loaded(drawable = affordanceIcon, contentDescription = null),
                         text = Text.Loaded(affordance.name),
                         isSelected = isSelectedFlow,
@@ -359,13 +366,13 @@
 
     /** Returns a view-model for the special "None" option. */
     @SuppressLint("UseCompatLoadingForDrawables")
-    private fun none(
-        slotId: Flow<String>,
-        isSelected: Flow<Boolean>,
+    private suspend fun none(
+        slotId: StateFlow<String>,
+        isSelected: StateFlow<Boolean>,
         onSelected: Flow<(() -> Unit)?>,
     ): OptionItemViewModel<Icon> {
         return OptionItemViewModel<Icon>(
-            key = slotId.map { "$it::none" },
+            key = slotId.map { "$it::none" }.stateIn(viewModelScope),
             payload = Icon.Resource(res = R.drawable.link_off, contentDescription = null),
             text = Text.Resource(res = R.string.keyguard_affordance_none),
             isSelected = isSelected,
diff --git a/tests/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModelTest.kt b/tests/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModelTest.kt
index 301dbe8..58c5d99 100644
--- a/tests/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModelTest.kt
+++ b/tests/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModelTest.kt
@@ -22,7 +22,6 @@
 import com.android.customization.model.grid.data.repository.FakeGridRepository
 import com.android.customization.model.grid.domain.interactor.GridInteractor
 import com.android.customization.model.grid.domain.interactor.GridSnapshotRestorer
-import com.android.wallpaper.picker.common.icon.ui.viewmodel.Icon
 import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
 import com.android.wallpaper.testing.FakeSnapshotStore
 import com.android.wallpaper.testing.collectLastValue
@@ -95,13 +94,17 @@
             assertThat(getOnClick(optionItemsValue[1])).isNull()
         }
 
-    private fun TestScope.getSelectedIndex(optionItems: List<OptionItemViewModel<Icon>>): Int {
+    private fun TestScope.getSelectedIndex(
+        optionItems: List<OptionItemViewModel<GridIconViewModel>>
+    ): Int {
         return optionItems.indexOfFirst { optionItem ->
             collectLastValue(optionItem.isSelected).invoke() == true
         }
     }
 
-    private fun TestScope.getOnClick(optionItem: OptionItemViewModel<Icon>): (() -> Unit)? {
+    private fun TestScope.getOnClick(
+        optionItem: OptionItemViewModel<GridIconViewModel>
+    ): (() -> Unit)? {
         return collectLastValue(optionItem.onClicked).invoke()
     }
 }