Setting/getting per-screen wallpaper (1/3).

Necessary changes to support setting and getting wallpapers per-screen,
which is what's needed for the wallpaper quick switcher to apply the
selected wallpaper only to the screen whose tab is currently selected.

Setting the individual per-screen wallpaper is complete but work still
needs to be done to provide separate recents for each screen. That work
is captured in b/265066284.

Bug: 262924643
Test: manually verified that setting per-screen wallpaper works by
switching between the home screen and lock screen tabs and selecting
different wallpapers
Test: manually verified that undo/restore works as expected, especially
given the lack of per-screen recents to serve as a baseline (it always
resets to both having the same wallpaper)
Test: unit tests throughout the layers

Change-Id: I6514247977e497547c4421be6663db84217175cd
diff --git a/src/com/android/wallpaper/module/WallpaperPickerSections.java b/src/com/android/wallpaper/module/WallpaperPickerSections.java
index f06f626..3c41c39 100644
--- a/src/com/android/wallpaper/module/WallpaperPickerSections.java
+++ b/src/com/android/wallpaper/module/WallpaperPickerSections.java
@@ -50,6 +50,7 @@
                         displayUtils));
         sectionControllers.add(
                 new WallpaperQuickSwitchSectionController(
+                        screen,
                         wallpaperQuickSwitchViewModel,
                         lifecycleOwner,
                         sectionNavigationController));
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 2effedc..e849426 100644
--- a/src/com/android/wallpaper/picker/customization/data/content/WallpaperClient.kt
+++ b/src/com/android/wallpaper/picker/customization/data/content/WallpaperClient.kt
@@ -18,6 +18,7 @@
 package com.android.wallpaper.picker.customization.data.content
 
 import android.graphics.Bitmap
+import com.android.wallpaper.picker.customization.shared.model.WallpaperDestination
 import com.android.wallpaper.picker.customization.shared.model.WallpaperModel
 import kotlinx.coroutines.flow.Flow
 
@@ -26,19 +27,27 @@
 
     /** Lists the most recent wallpapers. The first one is the most recent (current) wallpaper. */
     fun recentWallpapers(
+        destination: WallpaperDestination,
         limit: Int,
     ): Flow<List<WallpaperModel>>
 
     /** Returns the selected wallpaper. */
-    suspend fun getCurrentWallpaper(): WallpaperModel
+    suspend fun getCurrentWallpaper(
+        destination: WallpaperDestination,
+    ): WallpaperModel
 
     /**
      * Asynchronously sets the wallpaper to the one with the given ID.
      *
+     * @param destination The screen to set the wallpaper on.
      * @param wallpaperId The ID of the wallpaper to set.
      * @param onDone A callback to invoke when setting is done.
      */
-    suspend fun setWallpaper(wallpaperId: String, onDone: () -> Unit)
+    suspend fun setWallpaper(
+        destination: WallpaperDestination,
+        wallpaperId: String,
+        onDone: () -> Unit
+    )
 
     /** Returns a thumbnail for the wallpaper with the given ID. */
     suspend fun loadThumbnail(wallpaperId: String): Bitmap?
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 aebe836..e85cdfd 100644
--- a/src/com/android/wallpaper/picker/customization/data/content/WallpaperClientImpl.kt
+++ b/src/com/android/wallpaper/picker/customization/data/content/WallpaperClientImpl.kt
@@ -25,6 +25,7 @@
 import android.graphics.BitmapFactory
 import android.net.Uri
 import android.util.Log
+import com.android.wallpaper.picker.customization.shared.model.WallpaperDestination
 import com.android.wallpaper.picker.customization.shared.model.WallpaperModel
 import java.io.IOException
 import kotlinx.coroutines.channels.awaitClose
@@ -37,11 +38,12 @@
 ) : WallpaperClient {
 
     override fun recentWallpapers(
+        destination: WallpaperDestination,
         limit: Int,
     ): Flow<List<WallpaperModel>> {
         return callbackFlow {
             suspend fun queryAndSend(limit: Int) {
-                send(queryRecentWallpapers(limit = limit))
+                send(queryRecentWallpapers(destination = destination, limit = limit))
             }
 
             val contentObserver =
@@ -62,13 +64,20 @@
         }
     }
 
-    override suspend fun getCurrentWallpaper(): WallpaperModel {
-        return queryRecentWallpapers(limit = 1).first()
+    override suspend fun getCurrentWallpaper(
+        destination: WallpaperDestination,
+    ): WallpaperModel {
+        return queryRecentWallpapers(destination = destination, limit = 1).first()
     }
 
-    override suspend fun setWallpaper(wallpaperId: String, onDone: () -> Unit) {
+    override suspend fun setWallpaper(
+        destination: WallpaperDestination,
+        wallpaperId: String,
+        onDone: () -> Unit
+    ) {
         val updateValues = ContentValues()
         updateValues.put(KEY_ID, wallpaperId)
+        updateValues.put(KEY_SCREEN, destination.asString())
         val updatedRowCount = context.contentResolver.update(SET_WALLPAPER_URI, updateValues, null)
         if (updatedRowCount == 0) {
             Log.e(TAG, "Error setting wallpaper: $wallpaperId")
@@ -77,11 +86,12 @@
     }
 
     private suspend fun queryRecentWallpapers(
+        destination: WallpaperDestination,
         limit: Int,
     ): List<WallpaperModel> {
         context.contentResolver
             .query(
-                LIST_RECENTS_URI,
+                LIST_RECENTS_URI.buildUpon().appendPath(destination.asString()).build(),
                 arrayOf(
                     KEY_ID,
                     KEY_PLACEHOLDER_COLOR,
@@ -137,6 +147,14 @@
         return null
     }
 
+    private fun WallpaperDestination.asString(): String {
+        return when (this) {
+            WallpaperDestination.BOTH -> SCREEN_ALL
+            WallpaperDestination.HOME -> SCREEN_HOME
+            WallpaperDestination.LOCK -> SCREEN_LOCK
+        }
+    }
+
     companion object {
         private const val TAG = "WallpaperClientImpl"
         private const val AUTHORITY = "com.google.android.apps.wallpaper.recents"
@@ -162,6 +180,11 @@
 
         /** Key for a parameter used to pass the wallpaper ID to/from the content provider. */
         private const val KEY_ID = "id"
+        /** Key for a parameter used to pass the screen to/from the content provider. */
+        private const val KEY_SCREEN = "screen"
+        private const val SCREEN_ALL = "all_screens"
+        private const val SCREEN_HOME = "home_screen"
+        private const val SCREEN_LOCK = "lock_screen"
         /**
          * Key for a parameter used to get the placeholder color for a wallpaper from the content
          * provider.
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 1a016d7..6234fa5 100644
--- a/src/com/android/wallpaper/picker/customization/data/repository/WallpaperRepository.kt
+++ b/src/com/android/wallpaper/picker/customization/data/repository/WallpaperRepository.kt
@@ -19,6 +19,7 @@
 
 import android.graphics.Bitmap
 import com.android.wallpaper.picker.customization.data.content.WallpaperClient
+import com.android.wallpaper.picker.customization.shared.model.WallpaperDestination
 import com.android.wallpaper.picker.customization.shared.model.WallpaperModel
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
@@ -35,33 +36,44 @@
 
 /** Encapsulates access to wallpaper-related data. */
 class WallpaperRepository(
-    scope: CoroutineScope,
+    private val scope: CoroutineScope,
     private val client: WallpaperClient,
     private val backgroundDispatcher: CoroutineDispatcher,
 ) {
     /** The ID of the currently-selected wallpaper. */
-    val selectedWallpaperId: StateFlow<String> =
-        client
-            .recentWallpapers(limit = 1)
+    fun selectedWallpaperId(
+        destination: WallpaperDestination,
+    ): StateFlow<String> {
+        return client
+            .recentWallpapers(destination = destination, limit = 1)
             .map { previews -> previews.first().wallpaperId }
             .stateIn(
                 scope = scope,
                 started = SharingStarted.WhileSubscribed(),
-                initialValue = runBlocking { client.getCurrentWallpaper().wallpaperId },
+                initialValue =
+                    runBlocking {
+                        client.getCurrentWallpaper(destination = destination).wallpaperId
+                    },
             )
+    }
 
-    private val _selectingWallpaperId = MutableStateFlow<String?>(null)
+    private val _selectingWallpaperId =
+        MutableStateFlow<Map<WallpaperDestination, String?>>(emptyMap())
     /**
      * The ID of the wallpaper that is in the process of becoming the selected wallpaper or `null`
      * if no such transaction is currently taking place.
      */
-    val selectingWallpaperId: StateFlow<String?> = _selectingWallpaperId.asStateFlow()
+    val selectingWallpaperId: StateFlow<Map<WallpaperDestination, String?>> =
+        _selectingWallpaperId.asStateFlow()
 
     /** Lists the most recent wallpapers. The first one is the most recent (current) wallpaper. */
     fun recentWallpapers(
+        destination: WallpaperDestination,
         limit: Int,
     ): Flow<List<WallpaperModel>> {
-        return client.recentWallpapers(limit = limit).flowOn(backgroundDispatcher)
+        return client
+            .recentWallpapers(destination = destination, limit = limit)
+            .flowOn(backgroundDispatcher)
     }
 
     /** Returns a thumbnail for the wallpaper with the given ID. */
@@ -71,14 +83,18 @@
 
     /** Sets the wallpaper to the one with the given ID. */
     suspend fun setWallpaper(
+        destination: WallpaperDestination,
         wallpaperId: String,
     ) {
-        _selectingWallpaperId.value = wallpaperId
+        _selectingWallpaperId.value =
+            _selectingWallpaperId.value.toMutableMap().apply { this[destination] = wallpaperId }
         withContext(backgroundDispatcher) {
             client.setWallpaper(
+                destination = destination,
                 wallpaperId = wallpaperId,
             ) {
-                _selectingWallpaperId.value = null
+                _selectingWallpaperId.value =
+                    _selectingWallpaperId.value.toMutableMap().apply { this[destination] = null }
             }
         }
     }
diff --git a/src/com/android/wallpaper/picker/customization/domain/interactor/WallpaperInteractor.kt b/src/com/android/wallpaper/picker/customization/domain/interactor/WallpaperInteractor.kt
index 6aedc68..6430bee 100644
--- a/src/com/android/wallpaper/picker/customization/domain/interactor/WallpaperInteractor.kt
+++ b/src/com/android/wallpaper/picker/customization/domain/interactor/WallpaperInteractor.kt
@@ -19,6 +19,7 @@
 
 import android.graphics.Bitmap
 import com.android.wallpaper.picker.customization.data.repository.WallpaperRepository
+import com.android.wallpaper.picker.customization.shared.model.WallpaperDestination
 import com.android.wallpaper.picker.customization.shared.model.WallpaperModel
 import javax.inject.Provider
 import kotlinx.coroutines.flow.Flow
@@ -30,13 +31,22 @@
     private val repository: WallpaperRepository,
     private val snapshotRestorer: Provider<WallpaperSnapshotRestorer>,
 ) {
-    /** The ID of the currently-selected wallpaper. */
-    val selectedWallpaperId: StateFlow<String> = repository.selectedWallpaperId
+    /** Returns the ID of the currently-selected wallpaper. */
+    fun selectedWallpaperId(
+        destination: WallpaperDestination,
+    ): StateFlow<String> {
+        return repository.selectedWallpaperId(destination = destination)
+    }
+
     /**
-     * The ID of the wallpaper that is in the process of becoming the selected wallpaper or `null`
-     * if no such transaction is currently taking place.
+     * Returns the ID of the wallpaper that is in the process of becoming the selected wallpaper or
+     * `null` if no such transaction is currently taking place.
      */
-    val selectingWallpaperId: StateFlow<String?> = repository.selectingWallpaperId
+    fun selectingWallpaperId(
+        destination: WallpaperDestination,
+    ): Flow<String?> {
+        return repository.selectingWallpaperId.map { it[destination] }
+    }
 
     /**
      * Lists the [maxResults] most recent wallpapers.
@@ -44,10 +54,12 @@
      * The first one is the most recent (current) wallpaper.
      */
     fun previews(
+        destination: WallpaperDestination,
         maxResults: Int,
     ): Flow<List<WallpaperModel>> {
         return repository
             .recentWallpapers(
+                destination = destination,
                 limit = maxResults,
             )
             .map { previews ->
@@ -61,12 +73,19 @@
 
     /** Sets the wallpaper to the one with the given ID. */
     suspend fun setWallpaper(
+        destination: WallpaperDestination,
         wallpaperId: String,
     ) {
         repository.setWallpaper(
+            destination = destination,
             wallpaperId = wallpaperId,
         )
-        snapshotRestorer.get().storeSnapshot(wallpaperId)
+        snapshotRestorer
+            .get()
+            .storeSnapshot(
+                destination = destination,
+                selectedWallpaperId = wallpaperId,
+            )
     }
 
     /** Returns a thumbnail for the wallpaper with the given ID. */
diff --git a/src/com/android/wallpaper/picker/customization/domain/interactor/WallpaperSnapshotRestorer.kt b/src/com/android/wallpaper/picker/customization/domain/interactor/WallpaperSnapshotRestorer.kt
index 6336ef1..4d6869f 100644
--- a/src/com/android/wallpaper/picker/customization/domain/interactor/WallpaperSnapshotRestorer.kt
+++ b/src/com/android/wallpaper/picker/customization/domain/interactor/WallpaperSnapshotRestorer.kt
@@ -17,7 +17,9 @@
 
 package com.android.wallpaper.picker.customization.domain.interactor
 
+import com.android.wallpaper.picker.customization.shared.model.WallpaperDestination
 import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer
+import com.android.wallpaper.picker.undo.domain.interactor.SnapshotStore
 import com.android.wallpaper.picker.undo.shared.model.RestorableSnapshot
 
 /** Stores and restores undo snapshots for wallpaper state. */
@@ -25,37 +27,75 @@
     private val interactor: WallpaperInteractor,
 ) : SnapshotRestorer {
 
-    private lateinit var updater: (RestorableSnapshot) -> Unit
+    private lateinit var store: SnapshotStore
 
     fun storeSnapshot(
+        destination: WallpaperDestination,
         selectedWallpaperId: String,
     ) {
-        updater(snapshot(selectedWallpaperId))
+        val previousSnapshot = store.retrieve()
+        val nextSnapshot =
+            previousSnapshot.copy { args ->
+                args[destination.toSnapshotKey()] = selectedWallpaperId
+            }
+        store.store(nextSnapshot)
     }
 
     override suspend fun setUpSnapshotRestorer(
-        updater: (RestorableSnapshot) -> Unit,
+        store: SnapshotStore,
     ): RestorableSnapshot {
-        this.updater = updater
-        return snapshot(interactor.selectedWallpaperId.value)
+        this.store = store
+        val snapshot =
+            RestorableSnapshot(
+                args =
+                    buildMap {
+                        put(
+                            SELECTED_HOME_SCREEN_WALLPAPER_ID,
+                            interactor
+                                .selectedWallpaperId(destination = WallpaperDestination.HOME)
+                                .value,
+                        )
+                        put(
+                            SELECTED_LOCK_SCREEN_WALLPAPER_ID,
+                            interactor
+                                .selectedWallpaperId(destination = WallpaperDestination.LOCK)
+                                .value,
+                        )
+                    }
+            )
+        return snapshot
     }
 
     override suspend fun restoreToSnapshot(
         snapshot: RestorableSnapshot,
     ) {
-        val wallpaperId = snapshot.args[SELECTED_WALLPAPER_ID]
-        if (!wallpaperId.isNullOrEmpty()) {
-            interactor.setWallpaper(wallpaperId = wallpaperId)
+        val homeWallpaperId = snapshot.args[SELECTED_HOME_SCREEN_WALLPAPER_ID]
+        if (!homeWallpaperId.isNullOrEmpty()) {
+            interactor.setWallpaper(
+                destination = WallpaperDestination.HOME,
+                wallpaperId = homeWallpaperId
+            )
+        }
+
+        val lockWallpaperId = snapshot.args[SELECTED_LOCK_SCREEN_WALLPAPER_ID]
+        if (!lockWallpaperId.isNullOrEmpty()) {
+            interactor.setWallpaper(
+                destination = WallpaperDestination.LOCK,
+                wallpaperId = lockWallpaperId
+            )
         }
     }
 
-    private fun snapshot(selectedWallpaperId: String): RestorableSnapshot {
-        return RestorableSnapshot(
-            args = buildMap { put(SELECTED_WALLPAPER_ID, selectedWallpaperId) }
-        )
+    private fun WallpaperDestination.toSnapshotKey(): String {
+        return when (this) {
+            WallpaperDestination.HOME -> SELECTED_HOME_SCREEN_WALLPAPER_ID
+            WallpaperDestination.LOCK -> SELECTED_LOCK_SCREEN_WALLPAPER_ID
+            else -> error("Unsupported screen type \"$this\"!")
+        }
     }
 
     companion object {
-        private const val SELECTED_WALLPAPER_ID = "selected_wallpaper_id"
+        private const val SELECTED_HOME_SCREEN_WALLPAPER_ID = "selected_home_screen_wallpaper_id"
+        private const val SELECTED_LOCK_SCREEN_WALLPAPER_ID = "selected_lock_screen_wallpaper_id"
     }
 }
diff --git a/src/com/android/wallpaper/picker/customization/shared/model/WallpaperDestination.kt b/src/com/android/wallpaper/picker/customization/shared/model/WallpaperDestination.kt
new file mode 100644
index 0000000..89d5af4
--- /dev/null
+++ b/src/com/android/wallpaper/picker/customization/shared/model/WallpaperDestination.kt
@@ -0,0 +1,28 @@
+/*
+ * 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.wallpaper.picker.customization.shared.model
+
+/** Enumerates all known wallpaper destinations. */
+enum class WallpaperDestination {
+    /** Both [HOME] and [LOCK] destinations. */
+    BOTH,
+    /** The home screen wallpaper. */
+    HOME,
+    /** The lock screen wallpaper. */
+    LOCK,
+}
diff --git a/src/com/android/wallpaper/picker/customization/ui/section/WallpaperQuickSwitchSectionController.kt b/src/com/android/wallpaper/picker/customization/ui/section/WallpaperQuickSwitchSectionController.kt
index 061a8fe..d5c99aa 100644
--- a/src/com/android/wallpaper/picker/customization/ui/section/WallpaperQuickSwitchSectionController.kt
+++ b/src/com/android/wallpaper/picker/customization/ui/section/WallpaperQuickSwitchSectionController.kt
@@ -23,12 +23,14 @@
 import androidx.lifecycle.LifecycleOwner
 import com.android.wallpaper.R
 import com.android.wallpaper.model.CustomizationSectionController
+import com.android.wallpaper.module.CustomizationSections
 import com.android.wallpaper.picker.CategorySelectorFragment
 import com.android.wallpaper.picker.customization.ui.binder.WallpaperQuickSwitchSectionBinder
 import com.android.wallpaper.picker.customization.ui.viewmodel.WallpaperQuickSwitchViewModel
 
 /** Controls a section that lets the user switch wallpapers quickly. */
 class WallpaperQuickSwitchSectionController(
+    private val screen: CustomizationSections.Screen,
     private val viewModel: WallpaperQuickSwitchViewModel,
     private val lifecycleOwner: LifecycleOwner,
     private val navigator: CustomizationSectionController.CustomizationSectionNavigationController,
@@ -46,6 +48,9 @@
                     R.layout.wallpaper_quick_switch_section,
                     /* parent= */ null,
                 ) as WallpaperQuickSwitchView
+        viewModel.setOnLockScreen(
+            isLockScreenSelected = screen == CustomizationSections.Screen.LOCK_SCREEN,
+        )
         WallpaperQuickSwitchSectionBinder.bind(
             view = view,
             viewModel = viewModel,
@@ -56,4 +61,8 @@
         )
         return view
     }
+
+    override fun onScreenSwitched(isOnLockScreen: Boolean) {
+        viewModel.setOnLockScreen(isLockScreenSelected = isOnLockScreen)
+    }
 }
diff --git a/src/com/android/wallpaper/picker/customization/ui/viewmodel/WallpaperQuickSwitchViewModel.kt b/src/com/android/wallpaper/picker/customization/ui/viewmodel/WallpaperQuickSwitchViewModel.kt
index 8103af7..105dcff 100644
--- a/src/com/android/wallpaper/picker/customization/ui/viewmodel/WallpaperQuickSwitchViewModel.kt
+++ b/src/com/android/wallpaper/picker/customization/ui/viewmodel/WallpaperQuickSwitchViewModel.kt
@@ -25,12 +25,15 @@
 import androidx.lifecycle.viewModelScope
 import androidx.savedstate.SavedStateRegistryOwner
 import com.android.wallpaper.picker.customization.domain.interactor.WallpaperInteractor
+import com.android.wallpaper.picker.customization.shared.model.WallpaperDestination
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 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.distinctUntilChangedBy
+import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.shareIn
 import kotlinx.coroutines.launch
@@ -43,126 +46,167 @@
     private val interactor: WallpaperInteractor,
     maxOptions: Int,
 ) : ViewModel() {
+    private val isLockScreenSelected = MutableStateFlow(false)
 
     private val selectedWallpaperId: Flow<String> =
-        interactor.selectedWallpaperId.shareIn(
-            scope = viewModelScope,
-            started = SharingStarted.WhileSubscribed(),
-            replay = 1,
-        )
+        isLockScreenSelected
+            .flatMapLatest { isOnLockScreen ->
+                interactor.selectedWallpaperId(
+                    destination =
+                        if (isOnLockScreen) {
+                            WallpaperDestination.LOCK
+                        } else {
+                            WallpaperDestination.HOME
+                        },
+                )
+            }
+            .shareIn(
+                scope = viewModelScope,
+                started = SharingStarted.WhileSubscribed(),
+                replay = 1,
+            )
     private val selectingWallpaperId: Flow<String?> =
-        interactor.selectingWallpaperId.shareIn(
-            scope = viewModelScope,
-            started = SharingStarted.WhileSubscribed(),
-            replay = 1,
-        )
+        isLockScreenSelected
+            .flatMapLatest { isOnLockScreen ->
+                interactor.selectingWallpaperId(
+                    destination =
+                        if (isOnLockScreen) {
+                            WallpaperDestination.LOCK
+                        } else {
+                            WallpaperDestination.HOME
+                        },
+                )
+            }
+            .shareIn(
+                scope = viewModelScope,
+                started = SharingStarted.WhileSubscribed(),
+                replay = 1,
+            )
 
     val options: Flow<List<WallpaperQuickSwitchOptionViewModel>> =
-        interactor
-            .previews(
-                maxResults = maxOptions,
-            )
-            .distinctUntilChangedBy { previews ->
-                // Produce a key that's the same if the same set of wallpapers is available, even if
-                // in a different order. This is so that the view can keep from moving the wallpaper
-                // options around when the sort order changes as the user selects different
-                // wallpapers.
-                previews.map { preview -> preview.wallpaperId }.sorted().joinToString(",")
-            }
-            .map { previews ->
-                // True if any option is becoming selected following user click.
-                val isSomethingBecomingSelectedFlow: Flow<Boolean> =
-                    selectingWallpaperId.distinctUntilChanged().map { it != null }
+        isLockScreenSelected
+            .flatMapLatest { isOnLockScreen ->
+                interactor
+                    .previews(
+                        destination =
+                            if (isOnLockScreen) {
+                                WallpaperDestination.LOCK
+                            } else {
+                                WallpaperDestination.HOME
+                            },
+                        maxResults = maxOptions,
+                    )
+                    .distinctUntilChangedBy { previews ->
+                        // Produce a key that's the same if the same set of wallpapers is available,
+                        // even if in a different order. This is so that the view can keep from
+                        // moving the wallpaper options around when the sort order changes as the
+                        // user selects different wallpapers.
+                        previews.map { preview -> preview.wallpaperId }.sorted().joinToString(",")
+                    }
+                    .map { previews ->
+                        // True if any option is becoming selected following user click.
+                        val isSomethingBecomingSelectedFlow: Flow<Boolean> =
+                            selectingWallpaperId.distinctUntilChanged().map { it != null }
 
-                previews.map { preview ->
-                    // True if this option is currently selected.
-                    val isSelectedFlow: Flow<Boolean> =
-                        selectedWallpaperId.distinctUntilChanged().map { it == preview.wallpaperId }
-                    // True if this option is becoming the selected one following user click.
-                    val isBecomingSelectedFlow: Flow<Boolean> =
-                        selectingWallpaperId.distinctUntilChanged().map {
-                            it == preview.wallpaperId
-                        }
-
-                    WallpaperQuickSwitchOptionViewModel(
-                        wallpaperId = preview.wallpaperId,
-                        placeholderColor = preview.placeholderColor,
-                        thumbnail = {
-                            interactor.loadThumbnail(
-                                wallpaperId = preview.wallpaperId,
-                            )
-                        },
-                        isLarge =
-                            combine(
-                                isSelectedFlow,
-                                isBecomingSelectedFlow,
-                                isSomethingBecomingSelectedFlow,
-                            ) { isSelected, isBecomingSelected, isSomethingBecomingSelected ->
-                                // The large option is the one that's currently selected or the one
-                                // that
-                                // is becoming the selected one following user click.
-                                (isSelected && !isSomethingBecomingSelected) || isBecomingSelected
-                            },
-                        // We show the progress indicator if the option is in the process of
-                        // becoming the selected one following user click.
-                        isProgressIndicatorVisible = isBecomingSelectedFlow,
-                        isSelectionBorderVisible =
-                            combine(
-                                isSelectedFlow,
-                                isBecomingSelectedFlow,
-                                isSomethingBecomingSelectedFlow,
-                            ) { isSelected, isBeingSelected, isSomethingBecomingSelected ->
-                                // The selection border is shown for the option that is the one
-                                // that's
-                                // currently selected or the one that is becoming the selected one
-                                // following user click.
-                                (isSelected && !isSomethingBecomingSelected) || isBeingSelected
-                            },
-                        isSelectionIconVisible =
-                            combine(
-                                isSelectedFlow,
-                                isSomethingBecomingSelectedFlow,
-                            ) { isSelected, isSomethingBecomingSelected ->
-                                // The selection icon is shown for the option that is currently
-                                // selected
-                                // but only if nothing else is becoming selected. If anything is
-                                // being
-                                // selected following user click, the selection icon is not shown on
-                                // any
-                                // option.
-                                isSelected && !isSomethingBecomingSelected
-                            },
-                        onSelected =
-                            combine(
-                                    isSelectedFlow,
-                                    isBecomingSelectedFlow,
-                                    isSomethingBecomingSelectedFlow,
-                                ) { isSelected, isBeingSelected, isSomethingBecomingSelected ->
-                                    // An option is selectable if it is not itself becoming selected
-                                    // following user click or if nothing else is becoming selected
-                                    // but this
-                                    // option is not the selected one.
-                                    (isSomethingBecomingSelected && !isBeingSelected) ||
-                                        (!isSomethingBecomingSelected && !isSelected)
+                        previews.map { preview ->
+                            // True if this option is currently selected.
+                            val isSelectedFlow: Flow<Boolean> =
+                                selectedWallpaperId.distinctUntilChanged().map {
+                                    it == preview.wallpaperId
                                 }
-                                .distinctUntilChanged()
-                                .map { isSelectable ->
-                                    if (isSelectable) {
-                                        {
-                                            // A selectable option can become selected.
-                                            viewModelScope.launch {
-                                                interactor.setWallpaper(
-                                                    wallpaperId = preview.wallpaperId,
-                                                )
+                            // True if this option is becoming the selected one following user
+                            // click.
+                            val isBecomingSelectedFlow: Flow<Boolean> =
+                                selectingWallpaperId.distinctUntilChanged().map {
+                                    it == preview.wallpaperId
+                                }
+
+                            WallpaperQuickSwitchOptionViewModel(
+                                wallpaperId = preview.wallpaperId,
+                                placeholderColor = preview.placeholderColor,
+                                thumbnail = {
+                                    interactor.loadThumbnail(
+                                        wallpaperId = preview.wallpaperId,
+                                    )
+                                },
+                                isLarge =
+                                    combine(
+                                        isSelectedFlow,
+                                        isBecomingSelectedFlow,
+                                        isSomethingBecomingSelectedFlow,
+                                    ) { isSelected, isBecomingSelected, isSomethingBecomingSelected
+                                        ->
+                                        // The large option is the one that's currently selected or
+                                        // the one that is becoming the selected one following user
+                                        // click.
+                                        (isSelected && !isSomethingBecomingSelected) ||
+                                            isBecomingSelected
+                                    },
+                                // We show the progress indicator if the option is in the process of
+                                // becoming the selected one following user click.
+                                isProgressIndicatorVisible = isBecomingSelectedFlow,
+                                isSelectionBorderVisible =
+                                    combine(
+                                        isSelectedFlow,
+                                        isBecomingSelectedFlow,
+                                        isSomethingBecomingSelectedFlow,
+                                    ) { isSelected, isBeingSelected, isSomethingBecomingSelected ->
+                                        // The selection border is shown for the option that is the
+                                        // one that's currently selected or the one that is becoming
+                                        // the selected one following user click.
+                                        (isSelected && !isSomethingBecomingSelected) ||
+                                            isBeingSelected
+                                    },
+                                isSelectionIconVisible =
+                                    combine(
+                                        isSelectedFlow,
+                                        isSomethingBecomingSelectedFlow,
+                                    ) { isSelected, isSomethingBecomingSelected ->
+                                        // The selection icon is shown for the option that is
+                                        // currently selected but only if nothing else is becoming
+                                        // selected. If anything is being selected following user
+                                        // click, the selection icon is not shown on any option.
+                                        isSelected && !isSomethingBecomingSelected
+                                    },
+                                onSelected =
+                                    combine(
+                                            isSelectedFlow,
+                                            isBecomingSelectedFlow,
+                                            isSomethingBecomingSelectedFlow,
+                                        ) { isSelected, isBeingSelected, isSomethingBecomingSelected
+                                            ->
+                                            // An option is selectable if it is not itself becoming
+                                            // selected following user click or if nothing else is
+                                            // becoming selected but this option is not the selected
+                                            // one.
+                                            (isSomethingBecomingSelected && !isBeingSelected) ||
+                                                (!isSomethingBecomingSelected && !isSelected)
+                                        }
+                                        .distinctUntilChanged()
+                                        .map { isSelectable ->
+                                            if (isSelectable) {
+                                                {
+                                                    // A selectable option can become selected.
+                                                    viewModelScope.launch {
+                                                        interactor.setWallpaper(
+                                                            destination =
+                                                                if (isOnLockScreen) {
+                                                                    WallpaperDestination.LOCK
+                                                                } else {
+                                                                    WallpaperDestination.HOME
+                                                                },
+                                                            wallpaperId = preview.wallpaperId,
+                                                        )
+                                                    }
+                                                }
+                                            } else {
+                                                // A non-selectable option cannot become selected.
+                                                null
                                             }
                                         }
-                                    } else {
-                                        // A non-selectable option cannot become selected.
-                                        null
-                                    }
-                                }
-                    )
-                }
+                            )
+                        }
+                    }
             }
             .shareIn(
                 scope = viewModelScope,
@@ -170,6 +214,10 @@
                 replay = 1,
             )
 
+    fun setOnLockScreen(isLockScreenSelected: Boolean) {
+        this.isLockScreenSelected.value = isLockScreenSelected
+    }
+
     companion object {
         @JvmStatic
         fun newFactory(
diff --git a/src/com/android/wallpaper/picker/undo/domain/interactor/SnapshotRestorer.kt b/src/com/android/wallpaper/picker/undo/domain/interactor/SnapshotRestorer.kt
index 3047834..a532dfc 100644
--- a/src/com/android/wallpaper/picker/undo/domain/interactor/SnapshotRestorer.kt
+++ b/src/com/android/wallpaper/picker/undo/domain/interactor/SnapshotRestorer.kt
@@ -25,12 +25,12 @@
     /**
      * Sets up the restorer.
      *
-     * @param updater An updater the can be used when a new snapshot should be stored; invoke this
-     * in response to state changes that you wish could be restored when the user asks to reset the
+     * @param store An object the can be used when a new snapshot should be stored; use this in
+     * response to state changes that you wish could be restored when the user asks to reset the
      * changes.
      * @return A snapshot of the initial state as it was at the moment that this method was invoked.
      */
-    suspend fun setUpSnapshotRestorer(updater: (RestorableSnapshot) -> Unit): RestorableSnapshot
+    suspend fun setUpSnapshotRestorer(store: SnapshotStore): RestorableSnapshot
 
     /** Restores the state to what is described in the given snapshot. */
     suspend fun restoreToSnapshot(snapshot: RestorableSnapshot)
diff --git a/src/com/android/wallpaper/picker/undo/domain/interactor/SnapshotStore.kt b/src/com/android/wallpaper/picker/undo/domain/interactor/SnapshotStore.kt
new file mode 100644
index 0000000..39f5009
--- /dev/null
+++ b/src/com/android/wallpaper/picker/undo/domain/interactor/SnapshotStore.kt
@@ -0,0 +1,25 @@
+/*
+ * 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.wallpaper.picker.undo.domain.interactor
+
+import com.android.wallpaper.picker.undo.shared.model.RestorableSnapshot
+
+interface SnapshotStore {
+    fun retrieve(): RestorableSnapshot
+    fun store(snapshot: RestorableSnapshot)
+}
diff --git a/src/com/android/wallpaper/picker/undo/domain/interactor/UndoInteractor.kt b/src/com/android/wallpaper/picker/undo/domain/interactor/UndoInteractor.kt
index 1091f14..981fc03 100644
--- a/src/com/android/wallpaper/picker/undo/domain/interactor/UndoInteractor.kt
+++ b/src/com/android/wallpaper/picker/undo/domain/interactor/UndoInteractor.kt
@@ -18,6 +18,7 @@
 package com.android.wallpaper.picker.undo.domain.interactor
 
 import com.android.wallpaper.picker.undo.data.repository.UndoRepository
+import com.android.wallpaper.picker.undo.shared.model.RestorableSnapshot
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.launch
@@ -55,13 +56,25 @@
         restorerByOwnerId.forEach { (ownerId, restorer) ->
             scope.launch {
                 val initialSnapshot =
-                    restorer.setUpSnapshotRestorer { subsequentSnapshot ->
-                        val initialSnapshot = repository.getSnapshot(ownerId)
-                        repository.putDirty(
-                            ownerId = ownerId,
-                            isDirty = initialSnapshot != subsequentSnapshot
-                        )
-                    }
+                    restorer.setUpSnapshotRestorer(
+                        object : SnapshotStore {
+                            override fun retrieve(): RestorableSnapshot {
+                                return repository.getSnapshot(ownerId)
+                                    ?: error(
+                                        "No snapshot for this owner ID! Did you call this before" +
+                                            " storing a snapshot?"
+                                    )
+                            }
+
+                            override fun store(snapshot: RestorableSnapshot) {
+                                val initialSnapshot = repository.getSnapshot(ownerId)
+                                repository.putDirty(
+                                    ownerId = ownerId,
+                                    isDirty = initialSnapshot != snapshot
+                                )
+                            }
+                        }
+                    )
 
                 repository.putSnapshot(
                     ownerId = ownerId,
diff --git a/src/com/android/wallpaper/picker/undo/shared/model/RestorableSnapshot.kt b/src/com/android/wallpaper/picker/undo/shared/model/RestorableSnapshot.kt
index aac0e22..9ffe9c5 100644
--- a/src/com/android/wallpaper/picker/undo/shared/model/RestorableSnapshot.kt
+++ b/src/com/android/wallpaper/picker/undo/shared/model/RestorableSnapshot.kt
@@ -20,4 +20,32 @@
 /** Models a snapshot of the state of an undo-supporting feature at a given time. */
 data class RestorableSnapshot(
     val args: Map<String, String>,
-)
+) {
+    /**
+     * Returns a copy of the [RestorableSnapshot] but with the [block] applied to its arguments.
+     *
+     * Sample usage:
+     * ```
+     * val previousSnapshot: RestorableSnapshot = ...
+     * val nextSnapshot = previousSnapshot { args ->
+     *     args.put("one", "true")
+     *     args.remove("two")
+     * }
+     *
+     * // Now, nextSnapshot is exactly like previousSnapshot but with its args having "one" mapped
+     * // to "true" and without "two", since it was removed.
+     * ```
+     *
+     * @param block A function that receives the original [args] from the current
+     * [RestorableSnapshot] and can edit them for inclusion into the returned [RestorableSnapshot].
+     */
+    fun copy(
+        block: (MutableMap<String, String>) -> Unit,
+    ): RestorableSnapshot {
+        val mutableArgs = args.toMutableMap()
+        block(mutableArgs)
+        return RestorableSnapshot(
+            args = mutableArgs.toMap(),
+        )
+    }
+}
diff --git a/tests/src/com/android/wallpaper/picker/customization/data/content/FakeWallpaperClient.kt b/tests/src/com/android/wallpaper/picker/customization/data/content/FakeWallpaperClient.kt
index a5afba7..d30d1c9 100644
--- a/tests/src/com/android/wallpaper/picker/customization/data/content/FakeWallpaperClient.kt
+++ b/tests/src/com/android/wallpaper/picker/customization/data/content/FakeWallpaperClient.kt
@@ -18,6 +18,7 @@
 package com.android.wallpaper.picker.customization.data.content
 
 import android.graphics.Bitmap
+import com.android.wallpaper.picker.customization.shared.model.WallpaperDestination
 import com.android.wallpaper.picker.customization.shared.model.WallpaperModel
 import kotlin.math.min
 import kotlinx.coroutines.flow.Flow
@@ -26,12 +27,23 @@
 
 class FakeWallpaperClient : WallpaperClient {
 
-    private val _recentWallpapers = MutableStateFlow(INITIAL_RECENT_WALLPAPERS)
+    private val _recentWallpapers =
+        MutableStateFlow(
+            buildMap {
+                WallpaperDestination.values()
+                    .filter { it != WallpaperDestination.BOTH }
+                    .forEach { screen -> put(screen, INITIAL_RECENT_WALLPAPERS) }
+            }
+        )
     private var isPaused = false
-    private var deferred: (suspend () -> Unit)? = null
+    private var deferred = mutableListOf<(suspend () -> Unit)>()
 
-    fun setRecentWallpapers(recentWallpapers: List<WallpaperModel>) {
-        _recentWallpapers.value = recentWallpapers
+    fun setRecentWallpapers(
+        destination: WallpaperDestination,
+        recentWallpapers: List<WallpaperModel>,
+    ) {
+        _recentWallpapers.value =
+            _recentWallpapers.value.toMutableMap().apply { this[destination] = recentWallpapers }
     }
 
     fun pause() {
@@ -40,14 +52,17 @@
 
     suspend fun unpause() {
         isPaused = false
-        deferred?.invoke()
-        deferred = null
+        deferred.forEach { it.invoke() }
+        deferred.clear()
     }
 
     override fun recentWallpapers(
+        destination: WallpaperDestination,
         limit: Int,
     ): Flow<List<WallpaperModel>> {
-        return _recentWallpapers.map { wallpapers ->
+        return _recentWallpapers.map { wallpapersByScreen ->
+            val wallpapers =
+                wallpapersByScreen[destination] ?: error("No wallpapers for screen $destination")
             if (wallpapers.size > limit) {
                 wallpapers.subList(0, min(limit, wallpapers.size))
             } else {
@@ -56,16 +71,29 @@
         }
     }
 
-    override suspend fun getCurrentWallpaper(): WallpaperModel {
-        return _recentWallpapers.value[0]
+    override suspend fun getCurrentWallpaper(
+        destination: WallpaperDestination,
+    ): WallpaperModel {
+        return _recentWallpapers.value[destination]?.get(0)
+            ?: error("No wallpapers for screen $destination")
     }
 
-    override suspend fun setWallpaper(wallpaperId: String, onDone: () -> Unit) {
+    override suspend fun setWallpaper(
+        destination: WallpaperDestination,
+        wallpaperId: String,
+        onDone: () -> Unit
+    ) {
         if (isPaused) {
-            deferred = { setWallpaper(wallpaperId, onDone) }
+            deferred.add { setWallpaper(destination, wallpaperId, onDone) }
         } else {
             _recentWallpapers.value =
-                _recentWallpapers.value.sortedBy { it.wallpaperId != wallpaperId }
+                _recentWallpapers.value.toMutableMap().apply {
+                    this[destination] =
+                        _recentWallpapers.value[destination]?.sortedBy {
+                            it.wallpaperId != wallpaperId
+                        }
+                            ?: error("No wallpapers for screen $destination")
+                }
             onDone.invoke()
         }
     }
diff --git a/tests/src/com/android/wallpaper/picker/customization/data/repository/WallpaperRepositoryTest.kt b/tests/src/com/android/wallpaper/picker/customization/data/repository/WallpaperRepositoryTest.kt
index 16715c0..7e2c64b 100644
--- a/tests/src/com/android/wallpaper/picker/customization/data/repository/WallpaperRepositoryTest.kt
+++ b/tests/src/com/android/wallpaper/picker/customization/data/repository/WallpaperRepositoryTest.kt
@@ -19,9 +19,11 @@
 
 import androidx.test.filters.SmallTest
 import com.android.wallpaper.picker.customization.data.content.FakeWallpaperClient
+import com.android.wallpaper.picker.customization.shared.model.WallpaperDestination
 import com.android.wallpaper.testing.collectLastValue
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runTest
@@ -57,26 +59,63 @@
     @Test
     fun setWallpaper() =
         testScope.runTest {
-            val recentWallpapers = collectLastValue(underTest.recentWallpapers(limit = 5))
-            val selectedWallpaperId = collectLastValue(underTest.selectedWallpaperId)
-            val selectingWallpaperId = collectLastValue(underTest.selectingWallpaperId)
-            assertThat(recentWallpapers()).isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS)
-            assertThat(selectedWallpaperId())
+            val recentHomeWallpapers =
+                collectLastValue(
+                    underTest.recentWallpapers(destination = WallpaperDestination.HOME, limit = 5)
+                )
+            val recentLockWallpapers =
+                collectLastValue(
+                    underTest.recentWallpapers(destination = WallpaperDestination.LOCK, limit = 5)
+                )
+            val selectedHomeWallpaperId =
+                collectLastValue(underTest.selectedWallpaperId(WallpaperDestination.HOME))
+            val selectedLockWallpaperId =
+                collectLastValue(underTest.selectedWallpaperId(WallpaperDestination.LOCK))
+            val selectingHomeWallpaperId =
+                collectLastValue(
+                    underTest.selectingWallpaperId.map { it[WallpaperDestination.HOME] }
+                )
+            val selectingLockWallpaperId =
+                collectLastValue(
+                    underTest.selectingWallpaperId.map { it[WallpaperDestination.LOCK] }
+                )
+            assertThat(recentHomeWallpapers())
+                .isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS)
+            assertThat(recentLockWallpapers())
+                .isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS)
+            assertThat(selectedHomeWallpaperId())
                 .isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS.first().wallpaperId)
-            assertThat(selectingWallpaperId()).isNull()
+            assertThat(selectedLockWallpaperId())
+                .isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS.first().wallpaperId)
+            assertThat(selectingHomeWallpaperId()).isNull()
+            assertThat(selectingLockWallpaperId()).isNull()
 
             // Pause the client so we can examine the interim state.
             client.pause()
-            underTest.setWallpaper(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[1].wallpaperId)
-            assertThat(recentWallpapers()).isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS)
-            assertThat(selectedWallpaperId())
+            underTest.setWallpaper(
+                WallpaperDestination.HOME,
+                FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[1].wallpaperId,
+            )
+            underTest.setWallpaper(
+                WallpaperDestination.LOCK,
+                FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[2].wallpaperId,
+            )
+            assertThat(recentHomeWallpapers())
+                .isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS)
+            assertThat(recentLockWallpapers())
+                .isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS)
+            assertThat(selectedHomeWallpaperId())
                 .isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS.first().wallpaperId)
-            assertThat(selectingWallpaperId())
+            assertThat(selectedLockWallpaperId())
+                .isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS.first().wallpaperId)
+            assertThat(selectingHomeWallpaperId())
                 .isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[1].wallpaperId)
+            assertThat(selectingLockWallpaperId())
+                .isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[2].wallpaperId)
 
             // Unpause the client so we can examine the final state.
             client.unpause()
-            assertThat(recentWallpapers())
+            assertThat(recentHomeWallpapers())
                 .isEqualTo(
                     listOf(
                         FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[1],
@@ -84,8 +123,19 @@
                         FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[2],
                     )
                 )
-            assertThat(selectedWallpaperId())
+            assertThat(recentLockWallpapers())
+                .isEqualTo(
+                    listOf(
+                        FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[2],
+                        FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[0],
+                        FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[1],
+                    )
+                )
+            assertThat(selectedHomeWallpaperId())
                 .isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[1].wallpaperId)
-            assertThat(selectingWallpaperId()).isNull()
+            assertThat(selectedLockWallpaperId())
+                .isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[2].wallpaperId)
+            assertThat(selectingHomeWallpaperId()).isNull()
+            assertThat(selectingLockWallpaperId()).isNull()
         }
 }
diff --git a/tests/src/com/android/wallpaper/picker/customization/domain/interactor/WallpaperInteractorTest.kt b/tests/src/com/android/wallpaper/picker/customization/domain/interactor/WallpaperInteractorTest.kt
index 161b968..e90a9c3 100644
--- a/tests/src/com/android/wallpaper/picker/customization/domain/interactor/WallpaperInteractorTest.kt
+++ b/tests/src/com/android/wallpaper/picker/customization/domain/interactor/WallpaperInteractorTest.kt
@@ -20,7 +20,9 @@
 import androidx.test.filters.SmallTest
 import com.android.wallpaper.picker.customization.data.content.FakeWallpaperClient
 import com.android.wallpaper.picker.customization.data.repository.WallpaperRepository
+import com.android.wallpaper.picker.customization.shared.model.WallpaperDestination
 import com.android.wallpaper.picker.undo.shared.model.RestorableSnapshot
+import com.android.wallpaper.testing.FakeSnapshotStore
 import com.android.wallpaper.testing.collectLastValue
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -66,9 +68,7 @@
                 interactor = underTest,
             )
         initialSnapshot = runBlocking {
-            snapshotRestorer.setUpSnapshotRestorer {
-                // Do nothing.
-            }
+            snapshotRestorer.setUpSnapshotRestorer(FakeSnapshotStore())
         }
     }
 
@@ -78,6 +78,7 @@
             val limited =
                 collectLastValue(
                     underTest.previews(
+                        destination = WallpaperDestination.HOME,
                         maxResults = FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS.size - 1
                     )
                 )
@@ -94,32 +95,57 @@
     @Test
     fun setWallpaper() =
         testScope.runTest {
-            val previews =
+            val homePreviews =
                 collectLastValue(
                     underTest.previews(
+                        destination = WallpaperDestination.HOME,
                         maxResults = FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS.size
                     )
                 )
-            val selectedWallpaperId = collectLastValue(underTest.selectedWallpaperId)
-            val selectingWallpaperId = collectLastValue(underTest.selectingWallpaperId)
-            assertThat(previews()).isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS)
-            assertThat(selectedWallpaperId())
+            val lockPreviews =
+                collectLastValue(
+                    underTest.previews(
+                        destination = WallpaperDestination.LOCK,
+                        maxResults = FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS.size
+                    )
+                )
+            val selectedHomeWallpaperId =
+                collectLastValue(underTest.selectedWallpaperId(WallpaperDestination.HOME))
+            val selectedLockWallpaperId =
+                collectLastValue(underTest.selectedWallpaperId(WallpaperDestination.LOCK))
+            val selectingHomeWallpaperId =
+                collectLastValue(underTest.selectingWallpaperId(WallpaperDestination.HOME))
+            val selectingLockWallpaperId =
+                collectLastValue(underTest.selectingWallpaperId(WallpaperDestination.LOCK))
+            assertThat(homePreviews()).isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS)
+            assertThat(lockPreviews()).isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS)
+            assertThat(selectedLockWallpaperId())
                 .isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[0].wallpaperId)
-            assertThat(selectingWallpaperId()).isNull()
-            val wallpaperId = FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[1].wallpaperId
+            assertThat(selectedLockWallpaperId())
+                .isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[0].wallpaperId)
+            assertThat(selectingHomeWallpaperId()).isNull()
+            assertThat(selectingLockWallpaperId()).isNull()
+            val homeWallpaperId = FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[1].wallpaperId
+            val lockWallpaperId = FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[2].wallpaperId
 
             // Pause the client so we can examine the interim state.
             client.pause()
-            underTest.setWallpaper(wallpaperId)
-            assertThat(previews()).isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS)
-            assertThat(selectedWallpaperId())
+            underTest.setWallpaper(WallpaperDestination.HOME, homeWallpaperId)
+            underTest.setWallpaper(WallpaperDestination.LOCK, lockWallpaperId)
+            assertThat(homePreviews()).isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS)
+            assertThat(lockPreviews()).isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS)
+            assertThat(selectedHomeWallpaperId())
                 .isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[0].wallpaperId)
-            assertThat(selectingWallpaperId())
+            assertThat(selectedLockWallpaperId())
+                .isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[0].wallpaperId)
+            assertThat(selectingHomeWallpaperId())
                 .isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[1].wallpaperId)
+            assertThat(selectingLockWallpaperId())
+                .isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[2].wallpaperId)
 
             // Unpause the client so we can examine the final state.
             client.unpause()
-            assertThat(previews())
+            assertThat(homePreviews())
                 .isEqualTo(
                     listOf(
                         FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[1],
@@ -127,32 +153,56 @@
                         FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[2],
                     )
                 )
-            assertThat(selectedWallpaperId()).isEqualTo(wallpaperId)
-            assertThat(selectingWallpaperId()).isNull()
+            assertThat(lockPreviews())
+                .isEqualTo(
+                    listOf(
+                        FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[2],
+                        FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[0],
+                        FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[1],
+                    )
+                )
+            assertThat(selectedHomeWallpaperId()).isEqualTo(homeWallpaperId)
+            assertThat(selectedLockWallpaperId()).isEqualTo(lockWallpaperId)
+            assertThat(selectingHomeWallpaperId()).isNull()
+            assertThat(selectingLockWallpaperId()).isNull()
         }
 
     @Test
     fun restore() =
         testScope.runTest {
-            val previews =
+            val homePreviews =
                 collectLastValue(
                     underTest.previews(
+                        destination = WallpaperDestination.HOME,
                         maxResults = FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS.size
                     )
                 )
-            val selectedWallpaperId = collectLastValue(underTest.selectedWallpaperId)
-            val selectingWallpaperId = collectLastValue(underTest.selectingWallpaperId)
-            assertThat(previews()).isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS)
-            assertThat(selectedWallpaperId())
+            val lockPreviews =
+                collectLastValue(
+                    underTest.previews(
+                        destination = WallpaperDestination.LOCK,
+                        maxResults = FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS.size
+                    )
+                )
+            val selectedHomeWallpaperId =
+                collectLastValue(underTest.selectedWallpaperId(WallpaperDestination.HOME))
+            val selectedLockWallpaperId =
+                collectLastValue(underTest.selectedWallpaperId(WallpaperDestination.LOCK))
+            assertThat(homePreviews()).isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS)
+            assertThat(lockPreviews()).isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS)
+            assertThat(selectedHomeWallpaperId())
                 .isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[0].wallpaperId)
-            assertThat(selectingWallpaperId()).isNull()
-            val wallpaperId = FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[1].wallpaperId
-            underTest.setWallpaper(wallpaperId)
+            assertThat(selectedLockWallpaperId())
+                .isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[0].wallpaperId)
+            val homeWallpaperId = FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[1].wallpaperId
+            val lockWallpaperId = FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[2].wallpaperId
+            underTest.setWallpaper(WallpaperDestination.HOME, homeWallpaperId)
+            underTest.setWallpaper(WallpaperDestination.LOCK, lockWallpaperId)
 
             // Pause the client so we can examine the interim state.
             client.pause()
             snapshotRestorer.restoreToSnapshot(initialSnapshot)
-            assertThat(previews())
+            assertThat(homePreviews())
                 .isEqualTo(
                     listOf(
                         FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[1],
@@ -160,15 +210,31 @@
                         FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[2],
                     )
                 )
-            assertThat(selectedWallpaperId()).isEqualTo(wallpaperId)
-            assertThat(selectingWallpaperId())
-                .isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[0].wallpaperId)
+            assertThat(lockPreviews())
+                .isEqualTo(
+                    listOf(
+                        FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[2],
+                        FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[0],
+                        FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[1],
+                    )
+                )
+            assertThat(selectedHomeWallpaperId()).isEqualTo(homeWallpaperId)
+            assertThat(selectedLockWallpaperId()).isEqualTo(lockWallpaperId)
 
             // Unpause the client so we can examine the final state.
             client.unpause()
-            assertThat(previews()).isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS)
-            assertThat(selectedWallpaperId())
+            assertThat(homePreviews()).isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS)
+            assertThat(lockPreviews())
+                .isEqualTo(
+                    listOf(
+                        FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[0],
+                        FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[2],
+                        FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[1],
+                    )
+                )
+            assertThat(selectedHomeWallpaperId())
                 .isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[0].wallpaperId)
-            assertThat(selectingWallpaperId()).isNull()
+            assertThat(selectedLockWallpaperId())
+                .isEqualTo(FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[0].wallpaperId)
         }
 }
diff --git a/tests/src/com/android/wallpaper/picker/customization/ui/viewmodel/WallpaperQuickSwitchViewModelTest.kt b/tests/src/com/android/wallpaper/picker/customization/ui/viewmodel/WallpaperQuickSwitchViewModelTest.kt
index ec86331..6b8f7df 100644
--- a/tests/src/com/android/wallpaper/picker/customization/ui/viewmodel/WallpaperQuickSwitchViewModelTest.kt
+++ b/tests/src/com/android/wallpaper/picker/customization/ui/viewmodel/WallpaperQuickSwitchViewModelTest.kt
@@ -22,7 +22,9 @@
 import com.android.wallpaper.picker.customization.data.repository.WallpaperRepository
 import com.android.wallpaper.picker.customization.domain.interactor.WallpaperInteractor
 import com.android.wallpaper.picker.customization.domain.interactor.WallpaperSnapshotRestorer
+import com.android.wallpaper.picker.customization.shared.model.WallpaperDestination
 import com.android.wallpaper.picker.customization.shared.model.WallpaperModel
+import com.android.wallpaper.testing.FakeSnapshotStore
 import com.android.wallpaper.testing.collectLastValue
 import com.google.common.truth.Truth.assertThat
 import com.google.common.truth.Truth.assertWithMessage
@@ -78,11 +80,7 @@
             WallpaperSnapshotRestorer(
                 interactor = interactor,
             )
-        runBlocking {
-            snapshotRestorer.setUpSnapshotRestorer {
-                // Do nothing.
-            }
-        }
+        runBlocking { snapshotRestorer.setUpSnapshotRestorer(FakeSnapshotStore()) }
     }
 
     @After
@@ -120,7 +118,7 @@
                         placeholderColor = 1400,
                     ),
                 )
-            client.setRecentWallpapers(models)
+            client.setRecentWallpapers(WallpaperDestination.HOME, models)
 
             assertOptions(
                 observed = options(),
@@ -163,6 +161,51 @@
             )
         }
 
+    @Test
+    fun `switches between screens`() =
+        testScope.runTest {
+            val options = collectLastValue(underTest.options)
+
+            // We begin on the home screen by default.
+            // Select option at index 2 on the home screen.
+            val selectedIndex = 2
+            val optionToSelect = checkNotNull(options()?.get(selectedIndex))
+            val onSelected = collectLastValue(optionToSelect.onSelected)
+            onSelected()?.invoke()
+            runCurrent()
+            assertOptions(
+                observed = options(),
+                expected =
+                    expectations(
+                        selectedIndex = selectedIndex,
+                    ),
+            )
+
+            // Switch to the lock screen, it should still have the original option selected.
+            underTest.setOnLockScreen(isLockScreenSelected = true)
+            runCurrent()
+            assertOptions(
+                observed = options(),
+                expected = expectations(),
+            )
+
+            // Switch back to the home screen, it should still have option at index 2 selected.
+            underTest.setOnLockScreen(isLockScreenSelected = false)
+            runCurrent()
+            assertOptions(
+                observed = options(),
+                expected =
+                    expectations(
+                        models =
+                            listOf(
+                                FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[2],
+                                FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[0],
+                                FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS[1],
+                            ),
+                    ),
+            )
+        }
+
     private fun expectations(
         models: List<WallpaperModel> = FakeWallpaperClient.INITIAL_RECENT_WALLPAPERS,
         selectedIndex: Int = 0,
diff --git a/tests/src/com/android/wallpaper/testing/FakeSnapshotStore.kt b/tests/src/com/android/wallpaper/testing/FakeSnapshotStore.kt
new file mode 100644
index 0000000..a2a8b7c
--- /dev/null
+++ b/tests/src/com/android/wallpaper/testing/FakeSnapshotStore.kt
@@ -0,0 +1,35 @@
+/*
+ * 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.wallpaper.testing
+
+import com.android.wallpaper.picker.undo.domain.interactor.SnapshotStore
+import com.android.wallpaper.picker.undo.shared.model.RestorableSnapshot
+
+class FakeSnapshotStore(
+    initialSnapshot: RestorableSnapshot = RestorableSnapshot(emptyMap()),
+) : SnapshotStore {
+    private var snapshot = initialSnapshot
+
+    override fun retrieve(): RestorableSnapshot {
+        return snapshot
+    }
+
+    override fun store(snapshot: RestorableSnapshot) {
+        this.snapshot = snapshot
+    }
+}
diff --git a/tests/src/com/android/wallpaper/testing/UndoTestUtil.kt b/tests/src/com/android/wallpaper/testing/UndoTestUtil.kt
index 0f30db2..04cd756 100644
--- a/tests/src/com/android/wallpaper/testing/UndoTestUtil.kt
+++ b/tests/src/com/android/wallpaper/testing/UndoTestUtil.kt
@@ -18,6 +18,7 @@
 package com.android.wallpaper.testing
 
 import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer
+import com.android.wallpaper.picker.undo.domain.interactor.SnapshotStore
 import com.android.wallpaper.picker.undo.shared.model.RestorableSnapshot
 
 val FAKE_RESTORERS =
@@ -39,18 +40,18 @@
 class FakeSnapshotRestorer(
     private val ownerId: Int,
 ) : SnapshotRestorer {
-    private lateinit var updater: (RestorableSnapshot) -> Unit
+    private lateinit var store: SnapshotStore
     var restored: RestorableSnapshot? = null
         private set
 
     fun update(version: Int) {
-        updater(snapshot(ownerId, version))
+        store.store(snapshot(ownerId, version))
     }
 
     override suspend fun setUpSnapshotRestorer(
-        updater: (RestorableSnapshot) -> Unit,
+        store: SnapshotStore,
     ): RestorableSnapshot {
-        this.updater = updater
+        this.store = store
         return snapshot(
             ownerId = ownerId,
             version = 0,