Fixing ProxyPrefs not removing values correctly

Bug: 394600803
Flag: EXEMPT bugfix
Test: atest ProxyPrefsTest
Change-Id: Icb116df1826e9538c929a08f4e8afc32d1b1f183
diff --git a/src/com/android/launcher3/LauncherPrefs.kt b/src/com/android/launcher3/LauncherPrefs.kt
index 484cef4..1120ec8 100644
--- a/src/com/android/launcher3/LauncherPrefs.kt
+++ b/src/com/android/launcher3/LauncherPrefs.kt
@@ -34,6 +34,7 @@
 import com.android.launcher3.states.RotationHelper
 import com.android.launcher3.util.DaggerSingletonObject
 import com.android.launcher3.util.DisplayController
+import java.util.concurrent.ConcurrentHashMap
 import javax.inject.Inject
 
 /**
@@ -52,17 +53,18 @@
             .getSharedPreferences(BOOT_AWARE_PREFS_KEY, MODE_PRIVATE)
     }
 
-    open val Item.sharedPrefs: SharedPreferences
-        get() =
+    open protected fun getSharedPrefs(item: Item): SharedPreferences =
+        item.run {
             if (encryptionType == EncryptionType.DEVICE_PROTECTED) deviceProtectedSharedPrefs
             else encryptedContext.getSharedPreferences(sharedPrefFile, MODE_PRIVATE)
+        }
 
     /** Returns the value with type [T] for [item]. */
-    open fun <T> get(item: ContextualItem<T>): T =
+    fun <T> get(item: ContextualItem<T>): T =
         getInner(item, item.defaultValueFromContext(encryptedContext))
 
     /** Returns the value with type [T] for [item]. */
-    open fun <T> get(item: ConstantItem<T>): T = getInner(item, item.defaultValue)
+    fun <T> get(item: ConstantItem<T>): T = getInner(item, item.defaultValue)
 
     /**
      * Retrieves the value for an [Item] from [SharedPreferences]. It handles method typing via the
@@ -71,7 +73,7 @@
      */
     @Suppress("IMPLICIT_CAST_TO_ANY", "UNCHECKED_CAST")
     private fun <T> getInner(item: Item, default: T): T {
-        val sp = item.sharedPrefs
+        val sp = getSharedPrefs(item)
 
         return when (item.type) {
             String::class.java -> sp.getString(item.sharedPrefKey, default as? String)
@@ -127,7 +129,7 @@
     private fun prepareToPutValues(
         updates: Array<out Pair<Item, Any>>
     ): List<SharedPreferences.Editor> {
-        val updatesPerPrefFile = updates.groupBy { it.first.sharedPrefs }.toMap()
+        val updatesPerPrefFile = updates.groupBy { getSharedPrefs(it.first) }.toMap()
 
         return updatesPerPrefFile.map { (sharedPref, itemList) ->
             sharedPref.edit().apply { itemList.forEach { (item, value) -> putValue(item, value) } }
@@ -140,7 +142,7 @@
      * types of Item values.
      */
     @Suppress("UNCHECKED_CAST")
-    private fun SharedPreferences.Editor.putValue(
+    internal fun SharedPreferences.Editor.putValue(
         item: Item,
         value: Any?,
     ): SharedPreferences.Editor =
@@ -168,7 +170,7 @@
      */
     fun addListener(listener: LauncherPrefChangeListener, vararg items: Item) {
         items
-            .map { it.sharedPrefs }
+            .map { getSharedPrefs(it) }
             .distinct()
             .forEach { it.registerOnSharedPreferenceChangeListener(listener) }
     }
@@ -180,7 +182,7 @@
     fun removeListener(listener: LauncherPrefChangeListener, vararg items: Item) {
         // If a listener is not registered to a SharedPreference, unregistering it does nothing
         items
-            .map { it.sharedPrefs }
+            .map { getSharedPrefs(it) }
             .distinct()
             .forEach { it.unregisterOnSharedPreferenceChangeListener(listener) }
     }
@@ -191,7 +193,7 @@
      */
     fun has(vararg items: Item): Boolean {
         items
-            .groupBy { it.sharedPrefs }
+            .groupBy { getSharedPrefs(it) }
             .forEach { (prefs, itemsSublist) ->
                 if (!itemsSublist.none { !prefs.contains(it.sharedPrefKey) }) return false
             }
@@ -215,7 +217,7 @@
      *   .apply() or .commit()
      */
     private fun prepareToRemove(items: Array<out Item>): List<SharedPreferences.Editor> {
-        val itemsPerFile = items.groupBy { it.sharedPrefs }.toMap()
+        val itemsPerFile = items.groupBy { getSharedPrefs(it) }.toMap()
 
         return itemsPerFile.map { (prefs, items) ->
             prefs.edit().also { editor ->
@@ -412,14 +414,20 @@
  */
 class ProxyPrefs(context: Context, private val prefs: SharedPreferences) : LauncherPrefs(context) {
 
-    private val realPrefs = LauncherPrefs(context)
+    private val copiedPrefs = ConcurrentHashMap<SharedPreferences, Boolean>()
 
-    override val Item.sharedPrefs: SharedPreferences
-        get() = prefs
-
-    override fun <T> get(item: ConstantItem<T>) =
-        super.get(backedUpItem(item.sharedPrefKey, realPrefs.get(item)))
-
-    override fun <T> get(item: ContextualItem<T>) =
-        super.get(backedUpItem(item.sharedPrefKey, realPrefs.get(item)))
+    override fun getSharedPrefs(item: Item): SharedPreferences {
+        val originalPrefs = super.getSharedPrefs(item)
+        // Copy all existing values, when the pref is accessed for the first time
+        copiedPrefs.computeIfAbsent(originalPrefs) { op ->
+            val editor = prefs.edit()
+            op.all.forEach { (key, value) ->
+                if (value != null) {
+                    editor.putValue(backedUpItem(key, value), value)
+                }
+            }
+            editor.commit()
+        }
+        return prefs
+    }
 }
diff --git a/tests/multivalentTests/src/com/android/launcher3/FakeLauncherPrefs.kt b/tests/multivalentTests/src/com/android/launcher3/FakeLauncherPrefs.kt
index 7573d2f..5e1e548 100644
--- a/tests/multivalentTests/src/com/android/launcher3/FakeLauncherPrefs.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/FakeLauncherPrefs.kt
@@ -35,6 +35,5 @@
             MODE_PRIVATE,
         )
 
-    override val Item.sharedPrefs: SharedPreferences
-        get() = backingPrefs
+    override fun getSharedPrefs(item: Item): SharedPreferences = backingPrefs
 }
diff --git a/tests/multivalentTests/src/com/android/launcher3/ProxyPrefsTest.kt b/tests/multivalentTests/src/com/android/launcher3/ProxyPrefsTest.kt
new file mode 100644
index 0000000..54f6f63
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/ProxyPrefsTest.kt
@@ -0,0 +1,99 @@
+/*
+ * Copyright (C) 2025 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.launcher3
+
+import android.content.Context.MODE_PRIVATE
+import android.platform.uiautomatorhelpers.DeviceHelpers.context
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.launcher3.LauncherPrefs.Companion.backedUpItem
+import com.google.common.truth.Truth.assertThat
+import java.util.UUID
+import org.junit.After
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class ProxyPrefsTest {
+
+    private val prefName = "pref-test-" + UUID.randomUUID().toString()
+
+    private val proxyPrefs by lazy {
+        ProxyPrefs(
+            context,
+            context.getSharedPreferences(prefName, MODE_PRIVATE).apply { edit().clear().commit() },
+        )
+    }
+    private val launcherPrefs by lazy { LauncherPrefs(context) }
+
+    @After
+    fun tearDown() {
+        context.deleteSharedPreferences(prefName)
+    }
+
+    @Test
+    fun `returns fallback value if present`() {
+        launcherPrefs.putSync(TEST_ENTRY.to("new_value"))
+        assertThat(proxyPrefs.get(TEST_ENTRY)).isEqualTo("new_value")
+    }
+
+    @Test
+    fun `returns default value if not present`() {
+        launcherPrefs.removeSync(TEST_ENTRY)
+        assertThat(proxyPrefs.get(TEST_ENTRY)).isEqualTo("default_value")
+    }
+
+    @Test
+    fun `returns overridden value if present`() {
+        launcherPrefs.putSync(TEST_ENTRY.to("new_value"))
+        proxyPrefs.putSync(TEST_ENTRY.to("overridden_value"))
+        assertThat(proxyPrefs.get(TEST_ENTRY)).isEqualTo("overridden_value")
+    }
+
+    @Test
+    fun `value not present when removed`() {
+        launcherPrefs.putSync(TEST_ENTRY.to("new_value"))
+        proxyPrefs.removeSync(TEST_ENTRY)
+        assertThat(proxyPrefs.has(TEST_ENTRY)).isFalse()
+    }
+
+    @Test
+    fun `returns default if removed`() {
+        launcherPrefs.putSync(TEST_ENTRY.to("new_value"))
+        proxyPrefs.removeSync(TEST_ENTRY)
+        assertThat(proxyPrefs.get(TEST_ENTRY)).isEqualTo("default_value")
+    }
+
+    @Test
+    fun `value present on init`() {
+        launcherPrefs.putSync(TEST_ENTRY.to("new_value"))
+        assertThat(proxyPrefs.has(TEST_ENTRY)).isTrue()
+    }
+
+    @Test
+    fun `value absent on init`() {
+        launcherPrefs.removeSync(TEST_ENTRY)
+        assertThat(proxyPrefs.has(TEST_ENTRY)).isFalse()
+    }
+
+    companion object {
+
+        val TEST_ENTRY =
+            backedUpItem("test_prefs", "default_value", EncryptionType.DEVICE_PROTECTED)
+    }
+}