Fix other colors loading

Make sure colors are only returned after preset colors are loaded by
fixing usage of coroutine. Previously, the coroutine launch was wrapped
by a try catch block, but because the coroutine removes preset color
loading out of the try catch scope, the try catch finishes immediately,
and returns a list of colors without the presets. With fix, there is
slightly longer loading time for colors, but preset colors are returned
correctly.

Also simplify combine function syntax.

Flag: EXEMPT Bug fix
Bug: 384393408
Bug: 397250302
Bug: 392014907
Test: Manually verified with new picker flag on and off, see bugs
Change-Id: I3203cecc2708533e4233656094ddf3f98d1ddc0c
diff --git a/src/com/android/customization/model/color/ColorProvider.kt b/src/com/android/customization/model/color/ColorProvider.kt
index ef5e3d6..5c2b891 100644
--- a/src/com/android/customization/model/color/ColorProvider.kt
+++ b/src/com/android/customization/model/color/ColorProvider.kt
@@ -117,16 +117,19 @@
         }
 
         scope.launch {
+            // Wait for the previous preset color loading job to finish before evaluating whether to
+            // start a new one
             loaderJob?.join()
             if (presetColorBundles == null || reload) {
-                try {
-                    loaderJob = launch { loadPreset(isNewPickerUi) }
-                } catch (e: Throwable) {
-                    colorsAvailable = false
-                    callback?.onError(e)
-                    return@launch
+                loaderJob = launch {
+                    try {
+                        loadPreset(isNewPickerUi)
+                        callback?.onOptionsLoaded(buildFinalList(isNewPickerUi))
+                    } catch (e: Throwable) {
+                        colorsAvailable = false
+                        callback?.onError(e)
+                    }
                 }
-                callback?.onOptionsLoaded(buildFinalList(isNewPickerUi))
             } else {
                 callback?.onOptionsLoaded(buildFinalList(isNewPickerUi))
             }
diff --git a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
index 43d510c..7e63834 100644
--- a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
+++ b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
@@ -33,7 +33,6 @@
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.combine
-import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.suspendCancellableCoroutine
 
 // TODO (b/262924623): refactor to remove dependency on ColorCustomizationManager & ColorOption
@@ -58,68 +57,61 @@
 
     override val colorOptions: Flow<Map<ColorType, List<ColorOptionModel>>> =
         combine(homeWallpaperColors, lockWallpaperColors) { homeColors, lockColors ->
-                homeColors to lockColors
-            }
-            .map { (homeColors, lockColors) ->
-                suspendCancellableCoroutine { continuation ->
-                    if (
-                        homeColors is WallpaperColorsModel.Loading ||
-                            lockColors is WallpaperColorsModel.Loading
-                    ) {
-                        continuation.resumeWith(
-                            Result.success(
-                                mapOf(
-                                    ColorType.WALLPAPER_COLOR to listOf(),
-                                    ColorType.PRESET_COLOR to listOf(),
-                                )
+            suspendCancellableCoroutine { continuation ->
+                if (
+                    homeColors is WallpaperColorsModel.Loading ||
+                        lockColors is WallpaperColorsModel.Loading
+                ) {
+                    continuation.resumeWith(
+                        Result.success(
+                            mapOf(
+                                ColorType.WALLPAPER_COLOR to listOf(),
+                                ColorType.PRESET_COLOR to listOf(),
                             )
                         )
-                        return@suspendCancellableCoroutine
-                    }
-                    val homeColorsLoaded = homeColors as WallpaperColorsModel.Loaded
-                    val lockColorsLoaded = lockColors as WallpaperColorsModel.Loaded
-                    colorManager.setWallpaperColors(
-                        homeColorsLoaded.colors,
-                        lockColorsLoaded.colors,
                     )
-                    colorManager.fetchOptions(
-                        object : CustomizationManager.OptionsFetchedListener<ColorOption?> {
-                            override fun onOptionsLoaded(options: MutableList<ColorOption?>?) {
-                                val wallpaperColorOptions: MutableList<ColorOptionModel> =
-                                    mutableListOf()
-                                val presetColorOptions: MutableList<ColorOptionModel> =
-                                    mutableListOf()
-                                options?.forEach { option ->
-                                    when ((option as ColorOptionImpl).type) {
-                                        ColorType.WALLPAPER_COLOR ->
-                                            wallpaperColorOptions.add(option.toModel())
-                                        ColorType.PRESET_COLOR ->
-                                            presetColorOptions.add(option.toModel())
-                                    }
-                                }
-                                continuation.resumeWith(
-                                    Result.success(
-                                        mapOf(
-                                            ColorType.WALLPAPER_COLOR to wallpaperColorOptions,
-                                            ColorType.PRESET_COLOR to presetColorOptions,
-                                        )
-                                    )
-                                )
-                            }
-
-                            override fun onError(throwable: Throwable?) {
-                                Log.e(TAG, "Error loading theme bundles", throwable)
-                                continuation.resumeWith(
-                                    Result.failure(
-                                        throwable ?: Throwable("Error loading theme bundles")
-                                    )
-                                )
-                            }
-                        },
-                        /* reload= */ false,
-                    )
+                    return@suspendCancellableCoroutine
                 }
+                val homeColorsLoaded = homeColors as WallpaperColorsModel.Loaded
+                val lockColorsLoaded = lockColors as WallpaperColorsModel.Loaded
+                colorManager.setWallpaperColors(homeColorsLoaded.colors, lockColorsLoaded.colors)
+                colorManager.fetchOptions(
+                    object : CustomizationManager.OptionsFetchedListener<ColorOption?> {
+                        override fun onOptionsLoaded(options: MutableList<ColorOption?>?) {
+                            val wallpaperColorOptions: MutableList<ColorOptionModel> =
+                                mutableListOf()
+                            val presetColorOptions: MutableList<ColorOptionModel> = mutableListOf()
+                            options?.forEach { option ->
+                                when ((option as ColorOptionImpl).type) {
+                                    ColorType.WALLPAPER_COLOR ->
+                                        wallpaperColorOptions.add(option.toModel())
+                                    ColorType.PRESET_COLOR ->
+                                        presetColorOptions.add(option.toModel())
+                                }
+                            }
+                            continuation.resumeWith(
+                                Result.success(
+                                    mapOf(
+                                        ColorType.WALLPAPER_COLOR to wallpaperColorOptions,
+                                        ColorType.PRESET_COLOR to presetColorOptions,
+                                    )
+                                )
+                            )
+                        }
+
+                        override fun onError(throwable: Throwable?) {
+                            Log.e(TAG, "Error loading theme bundles", throwable)
+                            continuation.resumeWith(
+                                Result.failure(
+                                    throwable ?: Throwable("Error loading theme bundles")
+                                )
+                            )
+                        }
+                    },
+                    /* reload= */ false,
+                )
             }
+        }
 
     override suspend fun select(colorOptionModel: ColorOptionModel) {
         _isApplyingSystemColor.value = true
diff --git a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl2.kt b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl2.kt
index 96f9368..c925748 100644
--- a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl2.kt
+++ b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl2.kt
@@ -78,8 +78,7 @@
             .map { (_, colors) -> colors }
 
     override val colorOptions: Flow<Map<ColorType, List<ColorOption>>> =
-        combine(homeWallpaperColors, lockWallpaperColors, ::Pair)
-            .map { (homeColors, lockColors) ->
+        combine(homeWallpaperColors, lockWallpaperColors) { homeColors, lockColors ->
                 suspendCancellableCoroutine { continuation ->
                     colorManager.setWallpaperColors(homeColors, lockColors)
                     colorManager.fetchOptions(