Merge "Fix editing current wallpaper crop on foldable" into main
diff --git a/src/com/android/wallpaper/picker/customization/data/content/WallpaperClient.kt b/src/com/android/wallpaper/picker/customization/data/content/WallpaperClient.kt
index 7ca9791..649a140 100644
--- a/src/com/android/wallpaper/picker/customization/data/content/WallpaperClient.kt
+++ b/src/com/android/wallpaper/picker/customization/data/content/WallpaperClient.kt
@@ -17,6 +17,7 @@
 
 package com.android.wallpaper.picker.customization.data.content
 
+import android.app.WallpaperColors
 import android.app.WallpaperManager
 import android.graphics.Bitmap
 import android.graphics.Point
@@ -99,4 +100,7 @@
         displaySizes: List<Point>,
         @WallpaperManager.SetWallpaperFlags which: Int
     ): Map<Point, Rect>?
+
+    /** Returns the wallpaper colors for preview a bitmap with a set of crop hints */
+    suspend fun getWallpaperColors(bitmap: Bitmap, cropHints: Map<Point, Rect>?): WallpaperColors?
 }
diff --git a/src/com/android/wallpaper/picker/customization/data/content/WallpaperClientImpl.kt b/src/com/android/wallpaper/picker/customization/data/content/WallpaperClientImpl.kt
index 955e82f..e9f97b0 100644
--- a/src/com/android/wallpaper/picker/customization/data/content/WallpaperClientImpl.kt
+++ b/src/com/android/wallpaper/picker/customization/data/content/WallpaperClientImpl.kt
@@ -17,6 +17,7 @@
 
 package com.android.wallpaper.picker.customization.data.content
 
+import android.app.WallpaperColors
 import android.app.WallpaperManager
 import android.app.WallpaperManager.FLAG_LOCK
 import android.app.WallpaperManager.FLAG_SYSTEM
@@ -571,6 +572,13 @@
         return cropHintsMap
     }
 
+    override suspend fun getWallpaperColors(
+        bitmap: Bitmap,
+        cropHints: Map<Point, Rect>?
+    ): WallpaperColors? {
+        return wallpaperManager.getWallpaperColors(bitmap, cropHints)
+    }
+
     fun WallpaperDestination.asString(): String {
         return when (this) {
             BOTH -> SCREEN_ALL
diff --git a/src/com/android/wallpaper/picker/customization/data/repository/WallpaperRepository.kt b/src/com/android/wallpaper/picker/customization/data/repository/WallpaperRepository.kt
index 669f5d0..ec3e58c 100644
--- a/src/com/android/wallpaper/picker/customization/data/repository/WallpaperRepository.kt
+++ b/src/com/android/wallpaper/picker/customization/data/repository/WallpaperRepository.kt
@@ -17,8 +17,10 @@
 
 package com.android.wallpaper.picker.customization.data.repository
 
+import android.app.WallpaperColors
 import android.graphics.Bitmap
 import android.graphics.Point
+import android.graphics.Rect
 import android.util.LruCache
 import com.android.wallpaper.module.WallpaperPreferences
 import com.android.wallpaper.module.logging.UserEventLogger.SetWallpaperEntryPoint
@@ -175,6 +177,9 @@
         }
     }
 
+    suspend fun getWallpaperColors(bitmap: Bitmap, cropHints: Map<Point, Rect>?): WallpaperColors? =
+        withContext(backgroundDispatcher) { client.getWallpaperColors(bitmap, cropHints) }
+
     companion object {
         const val DEFAULT_KEY = "default_missing_key"
         /** The maximum number of options to show, including the currently-selected one. */
diff --git a/src/com/android/wallpaper/picker/preview/domain/interactor/WallpaperPreviewInteractor.kt b/src/com/android/wallpaper/picker/preview/domain/interactor/WallpaperPreviewInteractor.kt
index f243cec..20da4f4 100644
--- a/src/com/android/wallpaper/picker/preview/domain/interactor/WallpaperPreviewInteractor.kt
+++ b/src/com/android/wallpaper/picker/preview/domain/interactor/WallpaperPreviewInteractor.kt
@@ -16,8 +16,10 @@
 
 package com.android.wallpaper.picker.preview.domain.interactor
 
+import android.app.WallpaperColors
 import android.graphics.Bitmap
 import android.graphics.Point
+import android.graphics.Rect
 import com.android.wallpaper.module.logging.UserEventLogger
 import com.android.wallpaper.picker.customization.data.repository.WallpaperRepository
 import com.android.wallpaper.picker.customization.shared.model.WallpaperDestination
@@ -73,4 +75,7 @@
             wallpaperModel,
         )
     }
+
+    suspend fun getWallpaperColors(bitmap: Bitmap, cropHints: Map<Point, Rect>?): WallpaperColors? =
+        wallpaperRepository.getWallpaperColors(bitmap, cropHints)
 }
diff --git a/src/com/android/wallpaper/picker/preview/ui/binder/StaticWallpaperPreviewBinder.kt b/src/com/android/wallpaper/picker/preview/ui/binder/StaticWallpaperPreviewBinder.kt
index bbc5f3a..c481ea7 100644
--- a/src/com/android/wallpaper/picker/preview/ui/binder/StaticWallpaperPreviewBinder.kt
+++ b/src/com/android/wallpaper/picker/preview/ui/binder/StaticWallpaperPreviewBinder.kt
@@ -17,7 +17,6 @@
 
 import android.animation.Animator
 import android.animation.AnimatorListenerAdapter
-import android.graphics.Bitmap
 import android.graphics.Point
 import android.graphics.Rect
 import android.graphics.RenderEffect
@@ -65,7 +64,7 @@
                     viewModel.subsamplingScaleImageViewModel.collect { imageModel ->
                         val cropHint = imageModel.fullPreviewCropModels?.get(displaySize)?.cropHint
                         fullResImageView.setFullResImage(
-                            imageModel.rawWallpaperBitmap,
+                            ImageSource.cachedBitmap(imageModel.rawWallpaperBitmap),
                             imageModel.rawWallpaperSize,
                             displaySize,
                             cropHint,
@@ -120,7 +119,7 @@
      *   this system scale to [SubsamplingScaleImageView].
      */
     private fun SubsamplingScaleImageView.setFullResImage(
-        rawWallpaperBitmap: Bitmap,
+        imageSource: ImageSource,
         rawWallpaperSize: Point,
         displaySize: Point,
         cropHint: Rect?,
@@ -128,7 +127,7 @@
         shouldCalibrateWithSystemScale: Boolean = false,
     ) {
         // Set the full res image
-        setImage(ImageSource.bitmap(rawWallpaperBitmap))
+        setImage(imageSource)
         // Calculate the scale and the center point for the full res image
         FullResImageViewUtil.getScaleAndCenter(
                 Point(measuredWidth, measuredHeight),
diff --git a/src/com/android/wallpaper/picker/preview/ui/binder/WorkspacePreviewBinder.kt b/src/com/android/wallpaper/picker/preview/ui/binder/WorkspacePreviewBinder.kt
index a3c092b..b629d5b 100644
--- a/src/com/android/wallpaper/picker/preview/ui/binder/WorkspacePreviewBinder.kt
+++ b/src/com/android/wallpaper/picker/preview/ui/binder/WorkspacePreviewBinder.kt
@@ -17,6 +17,7 @@
 
 import android.app.WallpaperColors
 import android.os.Bundle
+import android.util.Log
 import android.view.SurfaceHolder
 import android.view.SurfaceView
 import androidx.core.os.bundleOf
@@ -116,13 +117,25 @@
                 object : PreviewUtils.WorkspacePreviewCallback {
                     override fun onPreviewRendered(resultBundle: Bundle?) {
                         if (resultBundle != null) {
-                            surface.setChildSurfacePackage(
-                                SurfaceViewUtils.getSurfacePackage(resultBundle)
-                            )
+                            SurfaceViewUtils.getSurfacePackage(resultBundle).apply {
+                                if (this != null) {
+                                    surface.setChildSurfacePackage(this)
+                                } else {
+                                    Log.w(
+                                        TAG,
+                                        "Result bundle from rendering preview does not contain " +
+                                            "a child surface package."
+                                    )
+                                }
+                            }
+                        } else {
+                            Log.w(TAG, "Result bundle from rendering preview is null.")
                         }
                     }
                 }
             )
         }
     }
+
+    const val TAG = "WorkspacePreviewBinder"
 }
diff --git a/src/com/android/wallpaper/picker/preview/ui/viewmodel/StaticWallpaperPreviewViewModel.kt b/src/com/android/wallpaper/picker/preview/ui/viewmodel/StaticWallpaperPreviewViewModel.kt
index 1e9288a..15f33b2 100644
--- a/src/com/android/wallpaper/picker/preview/ui/viewmodel/StaticWallpaperPreviewViewModel.kt
+++ b/src/com/android/wallpaper/picker/preview/ui/viewmodel/StaticWallpaperPreviewViewModel.kt
@@ -21,6 +21,7 @@
 import android.graphics.BitmapFactory
 import android.graphics.ColorSpace
 import android.graphics.Point
+import android.graphics.Rect
 import com.android.wallpaper.asset.Asset
 import com.android.wallpaper.asset.StreamableAsset
 import com.android.wallpaper.module.WallpaperPreferences
@@ -37,13 +38,16 @@
 import javax.inject.Inject
 import kotlinx.coroutines.CancellableContinuation
 import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.filterNotNull
 import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.shareIn
 import kotlinx.coroutines.suspendCancellableCoroutine
 
 /** View model for static wallpaper preview used in [WallpaperPreviewActivity] and its fragments */
@@ -55,6 +59,7 @@
     @ApplicationContext private val context: Context,
     private val wallpaperPreferences: WallpaperPreferences,
     @BackgroundDispatcher private val bgDispatcher: CoroutineDispatcher,
+    viewModelScope: CoroutineScope,
 ) {
     /**
      * The state of static wallpaper crop in full preview, before user confirmation.
@@ -74,6 +79,11 @@
     private val cropHintsInfo: MutableStateFlow<Map<Point, FullPreviewCropModel>?> =
         MutableStateFlow(null)
 
+    private val cropHints: Flow<Map<Point, Rect>?> =
+        cropHintsInfo.map { cropHintsInfoMap ->
+            cropHintsInfoMap?.map { entry -> entry.key to entry.value.cropHint }?.toMap()
+        }
+
     val staticWallpaperModel: Flow<StaticWallpaperModel> =
         interactor.wallpaperModel.map { it as? StaticWallpaperModel }.filterNotNull()
     val lowResBitmap: Flow<Bitmap> =
@@ -96,6 +106,10 @@
                 }
             }
             .flowOn(bgDispatcher)
+            // We only want to decode bitmap every time when wallpaper model is updated, instead of
+            // a new subscriber listens to this flow. So we need to use shareIn.
+            .shareIn(viewModelScope, SharingStarted.Lazily, 1)
+
     val fullResWallpaperViewModel: Flow<FullResWallpaperViewModel?> =
         combine(assetDetail, cropHintsInfo) { assetDetail, cropHintsInfo ->
                 if (assetDetail == null) {
@@ -110,14 +124,28 @@
             .flowOn(bgDispatcher)
     val subsamplingScaleImageViewModel: Flow<FullResWallpaperViewModel> =
         fullResWallpaperViewModel.filterNotNull()
-    val wallpaperColors: Flow<WallpaperColorsModel> =
+    // TODO (b/315856338): cache wallpaper colors in preferences
+    private val storedWallpaperColors: Flow<WallpaperColors?> =
         staticWallpaperModel
-            .map {
-                WallpaperColorsModel.Loaded(
-                    wallpaperPreferences.getWallpaperColors(it.commonWallpaperData.id.uniqueId)
-                )
-            }
+            .map { wallpaperPreferences.getWallpaperColors(it.commonWallpaperData.id.uniqueId) }
             .distinctUntilChanged()
+    val wallpaperColors: Flow<WallpaperColorsModel> =
+        combine(storedWallpaperColors, subsamplingScaleImageViewModel, cropHints) {
+            storedColors,
+            wallpaperViewModel,
+            cropHints ->
+            WallpaperColorsModel.Loaded(
+                if (cropHints == null) {
+                    storedColors
+                        ?: interactor.getWallpaperColors(
+                            wallpaperViewModel.rawWallpaperBitmap,
+                            null
+                        )
+                } else {
+                    interactor.getWallpaperColors(wallpaperViewModel.rawWallpaperBitmap, cropHints)
+                }
+            )
+        }
 
     /**
      * Updates new cropHints per displaySize that's been confirmed by the user.
@@ -143,7 +171,7 @@
     private suspend fun Asset.decodeBitmap(dimensions: Point): Bitmap? =
         suspendCancellableCoroutine { k: CancellableContinuation<Bitmap?> ->
             val callback = Asset.BitmapReceiver { k.resumeWith(Result.success(it)) }
-            decodeBitmap(dimensions.x, dimensions.y, false, callback)
+            decodeBitmap(dimensions.x, dimensions.y, /* hardwareBitmapAllowed= */ false, callback)
         }
 
     private suspend fun Asset.getStream(): InputStream? =
@@ -176,4 +204,23 @@
         }
         return colors
     }
+
+    class Factory
+    @Inject
+    constructor(
+        private val interactor: WallpaperPreviewInteractor,
+        @ApplicationContext private val context: Context,
+        private val wallpaperPreferences: WallpaperPreferences,
+        @BackgroundDispatcher private val bgDispatcher: CoroutineDispatcher,
+    ) {
+        fun create(viewModelScope: CoroutineScope): StaticWallpaperPreviewViewModel {
+            return StaticWallpaperPreviewViewModel(
+                interactor = interactor,
+                context = context,
+                wallpaperPreferences = wallpaperPreferences,
+                bgDispatcher = bgDispatcher,
+                viewModelScope = viewModelScope,
+            )
+        }
+    }
 }
diff --git a/src/com/android/wallpaper/picker/preview/ui/viewmodel/WallpaperPreviewViewModel.kt b/src/com/android/wallpaper/picker/preview/ui/viewmodel/WallpaperPreviewViewModel.kt
index d185dea..3cd5286 100644
--- a/src/com/android/wallpaper/picker/preview/ui/viewmodel/WallpaperPreviewViewModel.kt
+++ b/src/com/android/wallpaper/picker/preview/ui/viewmodel/WallpaperPreviewViewModel.kt
@@ -57,13 +57,15 @@
 @Inject
 constructor(
     private val interactor: WallpaperPreviewInteractor,
-    val staticWallpaperPreviewViewModel: StaticWallpaperPreviewViewModel,
+    staticWallpaperPreviewViewModelFactory: StaticWallpaperPreviewViewModel.Factory,
     val previewActionsViewModel: PreviewActionsViewModel,
     private val displayUtils: DisplayUtils,
     @HomeScreenPreviewUtils private val homePreviewUtils: PreviewUtils,
     @LockScreenPreviewUtils private val lockPreviewUtils: PreviewUtils,
 ) : ViewModel() {
 
+    val staticWallpaperPreviewViewModel =
+        staticWallpaperPreviewViewModelFactory.create(viewModelScope)
     val smallerDisplaySize = displayUtils.getRealSize(displayUtils.getSmallerDisplay())
     val wallpaperDisplaySize = displayUtils.getRealSize(displayUtils.getWallpaperDisplay())
     var isViewAsHome = false
diff --git a/tests/Android.bp b/tests/Android.bp
index 238d8a2..1004eae 100644
--- a/tests/Android.bp
+++ b/tests/Android.bp
@@ -58,6 +58,7 @@
         "junit",
         "kotlinx_coroutines_test",
         "truth",
+        "flag-junit",
     ],
     libs: [
         "android.test.runner",
diff --git a/tests/common/src/com/android/wallpaper/testing/FakeWallpaperClient.kt b/tests/common/src/com/android/wallpaper/testing/FakeWallpaperClient.kt
index 7999a5c..58d83c6 100644
--- a/tests/common/src/com/android/wallpaper/testing/FakeWallpaperClient.kt
+++ b/tests/common/src/com/android/wallpaper/testing/FakeWallpaperClient.kt
@@ -17,6 +17,7 @@
 
 package com.android.wallpaper.testing
 
+import android.app.WallpaperColors
 import android.graphics.Bitmap
 import android.graphics.Point
 import android.graphics.Rect
@@ -142,6 +143,13 @@
         return emptyMap()
     }
 
+    override suspend fun getWallpaperColors(
+        bitmap: Bitmap,
+        cropHints: Map<Point, Rect>?
+    ): WallpaperColors? {
+        return null
+    }
+
     companion object {
         val INITIAL_RECENT_WALLPAPERS =
             listOf(
diff --git a/tests/src/com/android/wallpaper/picker/preview/ui/WallpaperPreviewActivityTest.kt b/tests/src/com/android/wallpaper/picker/preview/ui/WallpaperPreviewActivityTest.kt
index 91459de..7778b91 100644
--- a/tests/src/com/android/wallpaper/picker/preview/ui/WallpaperPreviewActivityTest.kt
+++ b/tests/src/com/android/wallpaper/picker/preview/ui/WallpaperPreviewActivityTest.kt
@@ -15,20 +15,25 @@
  */
 package com.android.wallpaper.picker.preview.ui
 
+import android.platform.test.annotations.EnableFlags
+import android.platform.test.flag.junit.SetFlagsRule
 import androidx.navigation.fragment.NavHostFragment
 import androidx.test.core.app.ActivityScenario
 import androidx.test.core.app.ApplicationProvider
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.MediumTest
+import com.android.wallpaper.Flags.FLAG_MULTI_CROP_PREVIEW_UI_FLAG
 import com.android.wallpaper.model.WallpaperInfo
 import com.android.wallpaper.module.InjectorProvider
 import com.android.wallpaper.testing.TestInjector
 import com.android.wallpaper.testing.TestStaticWallpaperInfo
+import com.android.window.flags.Flags.FLAG_MULTI_CROP
 import com.google.common.truth.Truth.assertThat
 import dagger.hilt.android.testing.HiltAndroidRule
 import dagger.hilt.android.testing.HiltAndroidTest
 import javax.inject.Inject
 import org.junit.Before
+import org.junit.Ignore
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -39,6 +44,8 @@
 class WallpaperPreviewActivityTest {
     @get:Rule var hiltRule = HiltAndroidRule(this)
 
+    @get:Rule val setFlagsRule = SetFlagsRule()
+
     @Inject lateinit var testInjector: TestInjector
 
     private val testStaticWallpaper =
@@ -59,6 +66,8 @@
     }
 
     @Test
+    @Ignore("b/327241549")
+    @EnableFlags(FLAG_MULTI_CROP_PREVIEW_UI_FLAG, FLAG_MULTI_CROP)
     fun showsNavHostFragment() {
         val scenario: ActivityScenario<WallpaperPreviewActivity> =
             ActivityScenario.launch(activityStartIntent)