Merge "Make *Repository#get(int userId) lock-free" into main
diff --git a/services/core/java/com/android/server/inputmethod/AdditionalSubtypeMapRepository.java b/services/core/java/com/android/server/inputmethod/AdditionalSubtypeMapRepository.java
index 99f4747..b08f917 100644
--- a/services/core/java/com/android/server/inputmethod/AdditionalSubtypeMapRepository.java
+++ b/services/core/java/com/android/server/inputmethod/AdditionalSubtypeMapRepository.java
@@ -38,10 +38,11 @@
 final class AdditionalSubtypeMapRepository {
     private static final String TAG = "AdditionalSubtypeMapRepository";
 
-    // TODO(b/352594784): Should we user other lock primitives?
-    @GuardedBy("sPerUserMap")
+    private static final Object sMutationLock = new Object();
+
     @NonNull
-    private static final SparseArray<AdditionalSubtypeMap> sPerUserMap = new SparseArray<>();
+    private static volatile ImmutableSparseArray<AdditionalSubtypeMap> sPerUserMap =
+            ImmutableSparseArray.empty();
 
     record WriteTask(@UserIdInt int userId, @NonNull AdditionalSubtypeMap subtypeMap,
                      @NonNull InputMethodMap inputMethodMap) {
@@ -198,7 +199,7 @@
     /**
      * Returns {@link AdditionalSubtypeMap} for the given user.
      *
-     * <p>This method is expected be called after {@link #ensureInitializedAndGet(int)}. Otherwise
+     * <p>This method is expected be called after {@link #initializeIfNecessary(int)}. Otherwise
      * {@link AdditionalSubtypeMap#EMPTY_MAP} will be returned.</p>
      *
      * @param userId the user to be queried about
@@ -207,10 +208,7 @@
     @AnyThread
     @NonNull
     static AdditionalSubtypeMap get(@UserIdInt int userId) {
-        final AdditionalSubtypeMap map;
-        synchronized (sPerUserMap) {
-            map = sPerUserMap.get(userId);
-        }
+        final AdditionalSubtypeMap map = sPerUserMap.get(userId);
         if (map == null) {
             Slog.e(TAG, "get(userId=" + userId + ") is called before loadInitialDataAndGet()."
                     + " Returning an empty map");
@@ -220,28 +218,24 @@
     }
 
     /**
-     * Ensures that {@link AdditionalSubtypeMap} is initialized for the given user. Load it from
-     * the persistent storage if {@link #putAndSave(int, AdditionalSubtypeMap, InputMethodMap)} has
-     * not been called yet.
+     * Ensures that {@link AdditionalSubtypeMap} is initialized for the given user.
      *
      * @param userId the user to be initialized
-     * @return {@link AdditionalSubtypeMap} that is associated with the given user. If
-     *         {@link #putAndSave(int, AdditionalSubtypeMap, InputMethodMap)} is already called
-     *         then the given {@link AdditionalSubtypeMap}.
      */
     @AnyThread
     @NonNull
-    static AdditionalSubtypeMap ensureInitializedAndGet(@UserIdInt int userId) {
-        final var map = AdditionalSubtypeUtils.load(userId);
-        synchronized (sPerUserMap) {
-            final AdditionalSubtypeMap previous = sPerUserMap.get(userId);
-            // If putAndSave() has already been called, then use it.
-            if (previous != null) {
-                return previous;
-            }
-            sPerUserMap.put(userId, map);
+    static void initializeIfNecessary(@UserIdInt int userId) {
+        if (sPerUserMap.contains(userId)) {
+            // Fast-pass. If putAndSave() is already called, then do nothing.
+            return;
         }
-        return map;
+        final var map = AdditionalSubtypeUtils.load(userId);
+        synchronized (sMutationLock) {
+            // Check the condition again.
+            if (!sPerUserMap.contains(userId)) {
+                sPerUserMap = sPerUserMap.cloneWithPutOrSelf(userId, map);
+            }
+        }
     }
 
     /**
@@ -255,12 +249,8 @@
     @AnyThread
     static void putAndSave(@UserIdInt int userId, @NonNull AdditionalSubtypeMap map,
             @NonNull InputMethodMap inputMethodMap) {
-        synchronized (sPerUserMap) {
-            final AdditionalSubtypeMap previous = sPerUserMap.get(userId);
-            if (previous == map) {
-                return;
-            }
-            sPerUserMap.put(userId, map);
+        synchronized (sMutationLock) {
+            sPerUserMap = sPerUserMap.cloneWithPutOrSelf(userId, map);
             sWriter.scheduleWriteTask(userId, map, inputMethodMap);
         }
     }
@@ -277,9 +267,9 @@
 
     @AnyThread
     static void remove(@UserIdInt int userId) {
-        synchronized (sPerUserMap) {
+        synchronized (sMutationLock) {
             sWriter.onUserRemoved(userId);
-            sPerUserMap.remove(userId);
+            sPerUserMap = sPerUserMap.cloneWithRemoveOrSelf(userId);
         }
     }
 }
diff --git a/services/core/java/com/android/server/inputmethod/ImmutableSparseArray.java b/services/core/java/com/android/server/inputmethod/ImmutableSparseArray.java
new file mode 100644
index 0000000..382aa8a
--- /dev/null
+++ b/services/core/java/com/android/server/inputmethod/ImmutableSparseArray.java
@@ -0,0 +1,183 @@
+/*
+ * Copyright (C) 2024 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.server.inputmethod;
+
+
+import android.annotation.AnyThread;
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.util.SparseArray;
+
+import java.util.function.Consumer;
+
+/**
+ * A holder object to expose {@link SparseArray} to multiple threads in a thread-safe manner through
+ * "final Field Semantics" defined in JLS 17.5, with only exposing thread-safe methods such as
+ * {@link SparseArray#get(int)} and {@link SparseArray#size()} from {@link SparseArray}, and with
+ * adding clone-with-update style methods {@link #cloneWithPutOrSelf(int, Object)} and
+ * {@link #cloneWithRemoveOrSelf(int)} instead of exposing mutation methods.
+ *
+ * @param <E> Type of the element
+ */
+final class ImmutableSparseArray<E> {
+    @NonNull
+    private final SparseArray<E> mArray;
+
+    private static final ImmutableSparseArray<Object> EMPTY =
+            new ImmutableSparseArray<>(new SparseArray<>());
+
+    /**
+     * Returns an empty {@link ImmutableSparseArray} instance.
+     *
+     * @return An empty {@link ImmutableSparseArray} instance.
+     * @param <T> Type of the element
+     */
+    @SuppressWarnings("unchecked")
+    @AnyThread
+    @NonNull
+    static <T> ImmutableSparseArray<T> empty() {
+        return (ImmutableSparseArray<T>) EMPTY;
+    }
+
+    private ImmutableSparseArray(@NonNull SparseArray<E> array) {
+        mArray = array;
+    }
+
+    /**
+     * @return the size of this array
+     */
+    @AnyThread
+    int size() {
+        return mArray.size();
+    }
+
+    /**
+     * Returns the key of the specified index.
+     *
+     * @return the key of the specified index
+     * @throws ArrayIndexOutOfBoundsException when the index is out of range
+     */
+    @AnyThread
+    int keyAt(int index) {
+        return mArray.keyAt(index);
+    }
+
+    /**
+     * Returns the value of the specified index.
+     *
+     * @return the value of the specified index
+     * @throws ArrayIndexOutOfBoundsException when the index is out of range
+     */
+    @AnyThread
+    @Nullable
+    public E valueAt(int index) {
+        return mArray.valueAt(index);
+    }
+
+    /**
+     * Returns the index of the specified key.
+     *
+     * @return the index of the specified key if exists. Otherwise {@code -1}
+     */
+    @AnyThread
+    int indexOfKey(int key) {
+        return mArray.indexOfKey(key);
+    }
+
+    /**
+     * Returns {@code true} if the given {@code key} exists.
+     *
+     * @param key the key to be queried
+     * @return    {@code true} if the given {@code key} exists
+     */
+    @AnyThread
+    boolean contains(int key) {
+        return mArray.contains(key);
+    }
+
+    /**
+     * Returns the value associated with the {@code key}.
+     *
+     * @param key the key to be queried
+     * @return    the value associated with the {@code key} if exists. Otherwise {@code null}
+     */
+    @AnyThread
+    @Nullable
+    E get(int key) {
+        return mArray.get(key);
+    }
+
+    /**
+     * Run {@link Consumer} for each value.
+     *
+     * @param consumer {@link Consumer} to be called back
+     */
+    @AnyThread
+    void forEach(@NonNull Consumer<E> consumer) {
+        final int size = mArray.size();
+        for (int i = 0; i < size; ++i) {
+            consumer.accept(mArray.valueAt(i));
+        }
+    }
+
+    /**
+     * Returns an instance of {@link ImmutableSparseArray} that has the given key and value on top
+     * of items cloned from this instance.
+     *
+     * @param key   the key to be added
+     * @param value the value to be added
+     * @return      the same {@link ImmutableSparseArray} instance if there is actually no update.
+     *              Otherwise, a new instance of {@link ImmutableSparseArray}
+     */
+    @AnyThread
+    @NonNull
+    ImmutableSparseArray<E> cloneWithPutOrSelf(int key, @Nullable E value) {
+        final var prevKeyIndex = mArray.indexOfKey(key);
+        if (prevKeyIndex >= 0) {
+            final var prevValue = mArray.valueAt(prevKeyIndex);
+            if (prevValue == value) {
+                return this;
+            }
+        }
+        final var clone = mArray.clone();
+        clone.put(key, value);
+        return new ImmutableSparseArray<>(clone);
+    }
+
+    /**
+     * Returns an instance of {@link ImmutableSparseArray} that does not have the given key on top
+     * of items cloned from this instance.
+     *
+     * @param key the key to be removed
+     * @return    the same {@link ImmutableSparseArray} instance if there is actually no update.
+     *            Otherwise, a new instance of {@link ImmutableSparseArray}
+     */
+    @AnyThread
+    @NonNull
+    ImmutableSparseArray<E> cloneWithRemoveOrSelf(int key) {
+        final int index = indexOfKey(key);
+        if (index < 0) {
+            return this;
+        }
+        if (mArray.size() == 1) {
+            return empty();
+        }
+        final var clone = mArray.clone();
+        clone.remove(key);
+        return new ImmutableSparseArray<>(clone);
+    }
+}
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
index 9d007c6..9de9c2d 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
@@ -1063,8 +1063,8 @@
 
                 for (int userId : userIds) {
                     Slog.d(TAG, "Start initialization for user=" + userId);
-                    final var additionalSubtypeMap =
-                            AdditionalSubtypeMapRepository.ensureInitializedAndGet(userId);
+                    AdditionalSubtypeMapRepository.initializeIfNecessary(userId);
+                    final var additionalSubtypeMap = AdditionalSubtypeMapRepository.get(userId);
                     final var settings = InputMethodManagerService.queryInputMethodServicesInternal(
                             context, userId, additionalSubtypeMap,
                             DirectBootAwareness.AUTO).getMethodMap();
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodSettingsRepository.java b/services/core/java/com/android/server/inputmethod/InputMethodSettingsRepository.java
index 1b84036..4f5af63 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodSettingsRepository.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodSettingsRepository.java
@@ -19,15 +19,13 @@
 import android.annotation.AnyThread;
 import android.annotation.NonNull;
 import android.annotation.UserIdInt;
-import android.util.SparseArray;
-
-import com.android.internal.annotations.GuardedBy;
 
 final class InputMethodSettingsRepository {
-    // TODO(b/352594784): Should we user other lock primitives?
-    @GuardedBy("sPerUserMap")
+    private static final Object sMutationLock = new Object();
+
     @NonNull
-    private static final SparseArray<InputMethodSettings> sPerUserMap = new SparseArray<>();
+    private static volatile ImmutableSparseArray<InputMethodSettings> sPerUserMap =
+            ImmutableSparseArray.empty();
 
     /**
      * Not intended to be instantiated.
@@ -38,10 +36,7 @@
     @NonNull
     @AnyThread
     static InputMethodSettings get(@UserIdInt int userId) {
-        final InputMethodSettings obj;
-        synchronized (sPerUserMap) {
-            obj = sPerUserMap.get(userId);
-        }
+        final InputMethodSettings obj = sPerUserMap.get(userId);
         if (obj != null) {
             return obj;
         }
@@ -50,15 +45,15 @@
 
     @AnyThread
     static void put(@UserIdInt int userId, @NonNull InputMethodSettings obj) {
-        synchronized (sPerUserMap) {
-            sPerUserMap.put(userId, obj);
+        synchronized (sMutationLock) {
+            sPerUserMap = sPerUserMap.cloneWithPutOrSelf(userId, obj);
         }
     }
 
     @AnyThread
     static void remove(@UserIdInt int userId) {
-        synchronized (sPerUserMap) {
-            sPerUserMap.remove(userId);
+        synchronized (sMutationLock) {
+            sPerUserMap = sPerUserMap.cloneWithRemoveOrSelf(userId);
         }
     }
 }
diff --git a/services/core/java/com/android/server/inputmethod/SecureSettingsWrapper.java b/services/core/java/com/android/server/inputmethod/SecureSettingsWrapper.java
index e7cff20..b3500be 100644
--- a/services/core/java/com/android/server/inputmethod/SecureSettingsWrapper.java
+++ b/services/core/java/com/android/server/inputmethod/SecureSettingsWrapper.java
@@ -24,7 +24,6 @@
 import android.provider.Settings;
 import android.util.ArrayMap;
 import android.util.ArraySet;
-import android.util.SparseArray;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.server.LocalServices;
@@ -38,6 +37,13 @@
  * to the persistent value when the user storage is unlocked.</p>
  */
 final class SecureSettingsWrapper {
+
+    private static final Object sMutationLock = new Object();
+
+    @NonNull
+    private static volatile ImmutableSparseArray<ReaderWriter> sUserMap =
+            ImmutableSparseArray.empty();
+
     @Nullable
     private static volatile ContentResolver sContentResolver = null;
 
@@ -61,8 +67,8 @@
      */
     @AnyThread
     static void endTestMode() {
-        synchronized (sUserMap) {
-            sUserMap.clear();
+        synchronized (sMutationLock) {
+            sUserMap = ImmutableSparseArray.empty();
         }
         sTestMode = false;
     }
@@ -243,10 +249,6 @@
         }
     }
 
-    @GuardedBy("sUserMap")
-    @NonNull
-    private static final SparseArray<ReaderWriter> sUserMap = new SparseArray<>();
-
     private static final ReaderWriter NOOP = new ReaderWriter() {
         @Override
         public void putString(String key, String str) {
@@ -282,15 +284,15 @@
     private static ReaderWriter putOrGet(@UserIdInt int userId,
             @NonNull ReaderWriter readerWriter) {
         final boolean isUnlockedUserImpl = readerWriter instanceof UnlockedUserImpl;
-        synchronized (sUserMap) {
+        synchronized (sMutationLock) {
             final ReaderWriter current = sUserMap.get(userId);
             if (current == null) {
-                sUserMap.put(userId, readerWriter);
+                sUserMap = sUserMap.cloneWithPutOrSelf(userId, readerWriter);
                 return readerWriter;
             }
             // Upgrading from CopyOnWriteImpl to DirectImpl is allowed.
             if (current instanceof LockedUserImpl && isUnlockedUserImpl) {
-                sUserMap.put(userId, readerWriter);
+                sUserMap = sUserMap.cloneWithPutOrSelf(userId, readerWriter);
                 return readerWriter;
             }
             return current;
@@ -300,11 +302,9 @@
     @NonNull
     @AnyThread
     private static ReaderWriter get(@UserIdInt int userId) {
-        synchronized (sUserMap) {
-            final ReaderWriter readerWriter = sUserMap.get(userId);
-            if (readerWriter != null) {
-                return readerWriter;
-            }
+        final ReaderWriter readerWriter = sUserMap.get(userId);
+        if (readerWriter != null) {
+            return readerWriter;
         }
         if (sTestMode) {
             return putOrGet(userId, new FakeReaderWriterImpl());
@@ -363,8 +363,8 @@
      */
     @AnyThread
     static void onUserRemoved(@UserIdInt int userId) {
-        synchronized (sUserMap) {
-            sUserMap.remove(userId);
+        synchronized (sMutationLock) {
+            sUserMap = sUserMap.cloneWithRemoveOrSelf(userId);
         }
     }
 
diff --git a/services/core/java/com/android/server/inputmethod/UserDataRepository.java b/services/core/java/com/android/server/inputmethod/UserDataRepository.java
index 6f831cc..e3524b1 100644
--- a/services/core/java/com/android/server/inputmethod/UserDataRepository.java
+++ b/services/core/java/com/android/server/inputmethod/UserDataRepository.java
@@ -19,51 +19,39 @@
 import android.annotation.AnyThread;
 import android.annotation.NonNull;
 import android.annotation.UserIdInt;
-import android.util.SparseArray;
 
-import com.android.internal.annotations.GuardedBy;
-
-import java.util.concurrent.locks.ReentrantReadWriteLock;
 import java.util.function.Consumer;
 import java.util.function.IntFunction;
 
 final class UserDataRepository {
 
-    private final ReentrantReadWriteLock mUserDataLock = new ReentrantReadWriteLock();
+    private final Object mMutationLock = new Object();
 
-    @GuardedBy("mUserDataLock")
-    private final SparseArray<UserData> mUserData = new SparseArray<>();
+    @NonNull
+    private volatile ImmutableSparseArray<UserData> mUserData = ImmutableSparseArray.empty();
 
     private final IntFunction<InputMethodBindingController> mBindingControllerFactory;
 
     @AnyThread
     @NonNull
     UserData getOrCreate(@UserIdInt int userId) {
-        mUserDataLock.writeLock().lock();
-        try {
-            UserData userData = mUserData.get(userId);
-            if (userData == null) {
-                userData = new UserData(userId, mBindingControllerFactory.apply(userId));
-                mUserData.put(userId, userData);
-            }
+        // Do optimistic read first for optimization.
+        final var userData = mUserData.get(userId);
+        if (userData != null) {
             return userData;
-        } finally {
-            mUserDataLock.writeLock().unlock();
+        }
+        // Note that the below line can be called concurrently. Here we assume that
+        // instantiating UserData for the same user multiple times would have no side effect.
+        final var newUserData = new UserData(userId, mBindingControllerFactory.apply(userId));
+        synchronized (mMutationLock) {
+            mUserData = mUserData.cloneWithPutOrSelf(userId, newUserData);
+            return newUserData;
         }
     }
 
     @AnyThread
     void forAllUserData(Consumer<UserData> consumer) {
-        final SparseArray<UserData> copiedArray;
-        mUserDataLock.readLock().lock();
-        try {
-            copiedArray = mUserData.clone();
-        } finally {
-            mUserDataLock.readLock().unlock();
-        }
-        for (int i = 0; i < copiedArray.size(); i++) {
-            consumer.accept(copiedArray.valueAt(i));
-        }
+        mUserData.forEach(consumer);
     }
 
     UserDataRepository(
@@ -73,11 +61,8 @@
 
     @AnyThread
     void remove(@UserIdInt int userId) {
-        mUserDataLock.writeLock().lock();
-        try {
-            mUserData.remove(userId);
-        } finally {
-            mUserDataLock.writeLock().unlock();
+        synchronized (mMutationLock) {
+            mUserData = mUserData.cloneWithRemoveOrSelf(userId);
         }
     }
 }
diff --git a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/ImmutableSparseArrayTest.java b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/ImmutableSparseArrayTest.java
new file mode 100644
index 0000000..944b7c6
--- /dev/null
+++ b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/ImmutableSparseArrayTest.java
@@ -0,0 +1,278 @@
+/*
+ * Copyright (C) 2024 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.server.inputmethod;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.junit.Assert.assertThrows;
+
+import android.annotation.NonNull;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.util.ArrayList;
+
+@RunWith(AndroidJUnit4.class)
+public final class ImmutableSparseArrayTest {
+
+    @Test
+    public void testEmptyObject() {
+        final ImmutableSparseArray<Object> empty = ImmutableSparseArray.empty();
+
+        assertThat(empty.size()).isEqualTo(0);
+        verifyCommonBehaviors(empty);
+    }
+
+    @Test
+    public void testEmptyMethod() {
+        assertThat(ImmutableSparseArray.empty()).isSameInstanceAs(ImmutableSparseArray.empty());
+    }
+
+    @Test
+    public void testCloneWithPutOrSelf_appendingFromEmpty() {
+        final int key1 = 1;
+        final Object value1 = new Object();
+        final int key2 = -2;  // intentionally negative
+        final Object value2 = new Object();
+        final int key3 = -3;  // intentionally negative
+        final Object value3 = new Object();
+        final int key4 = 4;
+        final Object value4 = new Object();
+
+        final ImmutableSparseArray<Object> oneItemArray = ImmutableSparseArray.empty()
+                .cloneWithPutOrSelf(key1, value1);
+        verifyCommonBehaviors(oneItemArray);
+        assertThat(oneItemArray.size()).isEqualTo(1);
+        assertThat(oneItemArray.get(key1)).isSameInstanceAs(value1);
+
+        final ImmutableSparseArray<Object> twoItemArray =
+                oneItemArray.cloneWithPutOrSelf(key2, value2);
+        assertThat(twoItemArray).isNotSameInstanceAs(oneItemArray);
+        verifyCommonBehaviors(twoItemArray);
+        assertThat(twoItemArray.size()).isEqualTo(2);
+        assertThat(twoItemArray.get(key1)).isSameInstanceAs(value1);
+        assertThat(twoItemArray.get(key2)).isSameInstanceAs(value2);
+
+        final ImmutableSparseArray<Object> threeItemArray =
+                twoItemArray.cloneWithPutOrSelf(key3, value3);
+        assertThat(threeItemArray).isNotSameInstanceAs(twoItemArray);
+        verifyCommonBehaviors(threeItemArray);
+        assertThat(threeItemArray.size()).isEqualTo(3);
+        assertThat(threeItemArray.get(key1)).isSameInstanceAs(value1);
+        assertThat(threeItemArray.get(key2)).isSameInstanceAs(value2);
+        assertThat(threeItemArray.get(key3)).isSameInstanceAs(value3);
+
+        final ImmutableSparseArray<Object> fourItemArray =
+                threeItemArray.cloneWithPutOrSelf(key4, value4);
+        assertThat(fourItemArray).isNotSameInstanceAs(threeItemArray);
+        verifyCommonBehaviors(fourItemArray);
+        assertThat(fourItemArray.size()).isEqualTo(4);
+        assertThat(fourItemArray.get(key1)).isSameInstanceAs(value1);
+        assertThat(fourItemArray.get(key2)).isSameInstanceAs(value2);
+        assertThat(fourItemArray.get(key3)).isSameInstanceAs(value3);
+        assertThat(fourItemArray.get(key4)).isSameInstanceAs(value4);
+    }
+
+    @Test
+    public void testCloneWithPutOrSelf_returnSelf() {
+        final int key1 = 1;
+        final Object value1 = new Object();
+        final ImmutableSparseArray<Object> array = ImmutableSparseArray
+                .empty()
+                .cloneWithPutOrSelf(key1, value1);
+        assertThat(array.cloneWithPutOrSelf(key1, value1)).isSameInstanceAs(array);
+    }
+
+    @Test
+    public void testCloneWithPutOrSelf_updateExistingValue() {
+        final int key1 = 1;
+        final Object value1 = new Object();
+        final int key2 = 2;
+        final Object value2 = new Object();
+        final Object value2updated = new Object();
+        final int key3 = 3;
+        final Object value3 = new Object();
+
+        final ImmutableSparseArray<Object> array = ImmutableSparseArray
+                .empty()
+                .cloneWithPutOrSelf(key1, value1)
+                .cloneWithPutOrSelf(key2, value2)
+                .cloneWithPutOrSelf(key3, value3);
+
+        final var updatedArray = array.cloneWithPutOrSelf(key2, value2updated);
+        verifyCommonBehaviors(updatedArray);
+
+        assertThat(updatedArray.size()).isEqualTo(3);
+        assertThat(updatedArray.get(key1)).isSameInstanceAs(value1);
+        assertThat(updatedArray.get(key2)).isSameInstanceAs(value2updated);
+        assertThat(updatedArray.get(key3)).isSameInstanceAs(value3);
+    }
+
+    @Test
+    public void testCloneWithRemoveOrSelf_empty() {
+        final ImmutableSparseArray<Object> empty = ImmutableSparseArray.empty();
+        assertThat(empty.cloneWithRemoveOrSelf(0)).isSameInstanceAs(empty);
+    }
+
+    @Test
+    public void testCloneWithRemoveOrSelf_singleInstance() {
+        final int key = 1;
+        final Object value = new Object();
+        final ImmutableSparseArray<Object> array = ImmutableSparseArray
+                .empty()
+                .cloneWithPutOrSelf(key, value);
+        assertThat(array.cloneWithRemoveOrSelf(key)).isSameInstanceAs(ImmutableSparseArray.empty());
+    }
+
+    @Test
+    public void testCloneWithRemoveOrSelf_firstItem() {
+        final int key1 = 1;
+        final Object value1 = new Object();
+        final int key2 = 2;
+        final Object value2 = new Object();
+        final int key3 = 3;
+        final Object value3 = new Object();
+
+        final ImmutableSparseArray<Object> array = ImmutableSparseArray
+                .empty()
+                .cloneWithPutOrSelf(key1, value1)
+                .cloneWithPutOrSelf(key2, value2)
+                .cloneWithPutOrSelf(key3, value3)
+                .cloneWithRemoveOrSelf(key1);
+        verifyCommonBehaviors(array);
+
+        assertThat(array.size()).isEqualTo(2);
+        assertThat(array.get(key1)).isNull();
+        assertThat(array.get(key2)).isSameInstanceAs(value2);
+        assertThat(array.get(key3)).isSameInstanceAs(value3);
+        assertThat(array.keyAt(0)).isEqualTo(key2);
+        assertThat(array.keyAt(1)).isEqualTo(key3);
+    }
+
+    @Test
+    public void testCloneWithRemoveOrSelf_lastItem() {
+        final int key1 = 1;
+        final Object value1 = new Object();
+        final int key2 = 2;
+        final Object value2 = new Object();
+        final int key3 = 3;
+        final Object value3 = new Object();
+
+        final ImmutableSparseArray<Object> array = ImmutableSparseArray
+                .empty()
+                .cloneWithPutOrSelf(key1, value1)
+                .cloneWithPutOrSelf(key2, value2)
+                .cloneWithPutOrSelf(key3, value3)
+                .cloneWithRemoveOrSelf(key3);
+        verifyCommonBehaviors(array);
+
+        assertThat(array.size()).isEqualTo(2);
+        assertThat(array.get(key1)).isSameInstanceAs(value1);
+        assertThat(array.get(key2)).isSameInstanceAs(value2);
+        assertThat(array.get(key3)).isNull();
+    }
+
+    @Test
+    public void testCloneWithRemoveOrSelf_middleItem() {
+        final int key1 = 1;
+        final Object value1 = new Object();
+        final int key2 = 2;
+        final Object value2 = new Object();
+        final int key3 = 3;
+        final Object value3 = new Object();
+
+        final ImmutableSparseArray<Object> array = ImmutableSparseArray
+                .empty()
+                .cloneWithPutOrSelf(key1, value1)
+                .cloneWithPutOrSelf(key2, value2)
+                .cloneWithPutOrSelf(key3, value3)
+                .cloneWithRemoveOrSelf(key2);
+        verifyCommonBehaviors(array);
+
+        assertThat(array.size()).isEqualTo(2);
+        assertThat(array.get(key1)).isSameInstanceAs(value1);
+        assertThat(array.get(key2)).isNull();
+        assertThat(array.get(key3)).isSameInstanceAs(value3);
+    }
+
+    @Test
+    public void testCloneWithRemoveOrSelf_nonExistentItem() {
+        final int key1 = 1;
+        final Object value1 = new Object();
+        final int key2 = 2;
+        final Object value2 = new Object();
+        final int key3 = 3;
+        final Object value3 = new Object();
+        final int key4 = 4;
+
+        final ImmutableSparseArray<Object> array = ImmutableSparseArray
+                .empty()
+                .cloneWithPutOrSelf(key1, value1)
+                .cloneWithPutOrSelf(key2, value2)
+                .cloneWithPutOrSelf(key3, value3);
+
+        assertThat(array.cloneWithRemoveOrSelf(key4)).isSameInstanceAs(array);
+    }
+
+    @Test
+    public void testForEach() {
+        final int key1 = 1;
+        final Object value1 = new Object();
+        final int key2 = 2;
+        final Object value2 = new Object();
+        final int key3 = 3;
+        final Object value3 = new Object();
+
+        final ImmutableSparseArray<Object> array = ImmutableSparseArray
+                .empty()
+                .cloneWithPutOrSelf(key1, value1)
+                .cloneWithPutOrSelf(key2, value2)
+                .cloneWithPutOrSelf(key3, value3);
+
+        final ArrayList<Object> list = new ArrayList<>();
+        array.forEach(list::add);
+        assertThat(list).containsExactlyElementsIn(new Object[]{ value1, value2, value3 })
+                .inOrder();
+    }
+
+
+    private void verifyCommonBehaviors(@NonNull ImmutableSparseArray<Object> sparseArray) {
+        verifyInvalidKeyBehaviors(sparseArray);
+        verifyOutOfBoundsBehaviors(sparseArray);
+    }
+
+    private void verifyInvalidKeyBehaviors(@NonNull ImmutableSparseArray<Object> sparseArray) {
+        final int invalid_key = -123456678;
+        assertThat(sparseArray.get(invalid_key)).isNull();
+        assertThat(sparseArray.indexOfKey(invalid_key)).isEqualTo(-1);
+    }
+
+    private void verifyOutOfBoundsBehaviors(@NonNull ImmutableSparseArray<Object> sparseArray) {
+        final int size = sparseArray.size();
+        assertThrows(ArrayIndexOutOfBoundsException.class,
+                () -> sparseArray.keyAt(size));
+        assertThrows(ArrayIndexOutOfBoundsException.class,
+                () -> sparseArray.valueAt(size));
+        assertThrows(ArrayIndexOutOfBoundsException.class,
+                () -> sparseArray.keyAt(-1));
+        assertThrows(ArrayIndexOutOfBoundsException.class,
+                () -> sparseArray.valueAt(-1));
+    }
+}
diff --git a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodManagerServiceTestBase.java b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodManagerServiceTestBase.java
index ec9bfa7..d427a6d 100644
--- a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodManagerServiceTestBase.java
+++ b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodManagerServiceTestBase.java
@@ -272,7 +272,7 @@
 
         // Certain tests rely on TEST_IME_ID that is installed with AndroidTest.xml.
         // TODO(b/352615651): Consider just synthesizing test InputMethodInfo then injecting it.
-        AdditionalSubtypeMapRepository.ensureInitializedAndGet(mCallingUserId);
+        AdditionalSubtypeMapRepository.initializeIfNecessary(mCallingUserId);
         final var settings = InputMethodManagerService.queryInputMethodServicesInternal(mContext,
                 mCallingUserId, AdditionalSubtypeMapRepository.get(mCallingUserId),
                 DirectBootAwareness.AUTO);