Merge "Add an immutable AdditionalSubtypeMap class" into main

7 files changed, 343 insertions, 61 deletions

diff --git a/services/core/java/com/android/server/inputmethod/AdditionalSubtypeMap.java b/services/core/java/com/android/server/inputmethod/AdditionalSubtypeMap.java
new file mode 100644
index 000000000000..f2185163824f
--- /dev/null
+++ b/services/core/java/com/android/server/inputmethod/AdditionalSubtypeMap.java
@@ -0,0 +1,155 @@
+/*
+ * 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.ArrayMap;
+import android.view.inputmethod.InputMethodSubtype;
+
+import java.util.Collection;
+import java.util.List;
+
+/**
+ * An on-memory immutable data representation of subtype.xml, which contains so-called additional
+ * {@link InputMethodSubtype}.
+ *
+ * <p>While the data structure could be also used for general purpose map from IME ID to
+ * a list of {@link InputMethodSubtype}, unlike {@link InputMethodMap} this particular data
+ * structure is currently used only around additional {@link InputMethodSubtype}, which is why this
+ * class is (still) called {@code AdditionalSubtypeMap} rather than {@code InputMethodSubtypeMap}.
+ * </p>
+ */
+final class AdditionalSubtypeMap {
+    /**
+     * An empty {@link AdditionalSubtypeMap}.
+     */
+    static final AdditionalSubtypeMap EMPTY_MAP = new AdditionalSubtypeMap(new ArrayMap<>());
+
+    @NonNull
+    private final ArrayMap<String, List<InputMethodSubtype>> mMap;
+
+    @AnyThread
+    @NonNull
+    private static AdditionalSubtypeMap createOrEmpty(
+            @NonNull ArrayMap<String, List<InputMethodSubtype>> map) {
+        return map.isEmpty() ? EMPTY_MAP : new AdditionalSubtypeMap(map);
+    }
+
+    /**
+     * Create a new instance from the given {@link ArrayMap}.
+     *
+     * <p>This method effectively creates a new copy of map.</p>
+     *
+     * @param map An {@link ArrayMap} from which {@link AdditionalSubtypeMap} is to be created.
+     * @return A {@link AdditionalSubtypeMap} that contains a new copy of {@code map}.
+     */
+    @AnyThread
+    @NonNull
+    static AdditionalSubtypeMap of(@NonNull ArrayMap<String, List<InputMethodSubtype>> map) {
+        return createOrEmpty(map);
+    }
+
+    /**
+     * Create a new instance of {@link AdditionalSubtypeMap} from an existing
+     * {@link AdditionalSubtypeMap} by removing {@code key}, or return {@code map} itself if it does
+     * not contain an entry of {@code key}.
+     *
+     * @param key The key to be removed from {@code map}.
+     * @return A new instance of {@link AdditionalSubtypeMap}, which is guaranteed to not contain
+     *         {@code key}, or {@code map} itself if it does not contain an entry of {@code key}.
+     */
+    @AnyThread
+    @NonNull
+    AdditionalSubtypeMap cloneWithRemoveOrSelf(@NonNull String key) {
+        if (isEmpty() || !containsKey(key)) {
+            return this;
+        }
+        final ArrayMap<String, List<InputMethodSubtype>> newMap = new ArrayMap<>(mMap);
+        newMap.remove(key);
+        return createOrEmpty(newMap);
+    }
+
+    /**
+     * Create a new instance of {@link AdditionalSubtypeMap} from an existing
+     * {@link AdditionalSubtypeMap} by removing {@code keys} or return {@code map} itself if it does
+     * not contain any entry for {@code keys}.
+     *
+     * @param keys Keys to be removed from {@code map}.
+     * @return A new instance of {@link AdditionalSubtypeMap}, which is guaranteed to not contain
+     *         {@code keys}, or {@code map} itself if it does not contain any entry of {@code keys}.
+     */
+    @AnyThread
+    @NonNull
+    AdditionalSubtypeMap cloneWithRemoveOrSelf(@NonNull Collection<String> keys) {
+        if (isEmpty()) {
+            return this;
+        }
+        final ArrayMap<String, List<InputMethodSubtype>> newMap = new ArrayMap<>(mMap);
+        return newMap.removeAll(keys) ? createOrEmpty(newMap) : this;
+    }
+
+    /**
+     * Create a new instance of {@link AdditionalSubtypeMap} from an existing
+     * {@link AdditionalSubtypeMap} by putting {@code key} and {@code value}.
+     *
+     * @param key Key to be put into {@code map}.
+     * @param value Value to be put into {@code map}.
+     * @return A new instance of {@link AdditionalSubtypeMap}, which is guaranteed to contain the
+     *         pair of {@code key} and {@code value}.
+     */
+    @AnyThread
+    @NonNull
+    AdditionalSubtypeMap cloneWithPut(
+            @Nullable String key, @NonNull List<InputMethodSubtype> value) {
+        final ArrayMap<String, List<InputMethodSubtype>> newMap = new ArrayMap<>(mMap);
+        newMap.put(key, value);
+        return new AdditionalSubtypeMap(newMap);
+    }
+
+    private AdditionalSubtypeMap(@NonNull ArrayMap<String, List<InputMethodSubtype>> map) {
+        mMap = map;
+    }
+
+    @AnyThread
+    @Nullable
+    List<InputMethodSubtype> get(@Nullable String key) {
+        return mMap.get(key);
+    }
+
+    @AnyThread
+    boolean containsKey(@Nullable String key) {
+        return mMap.containsKey(key);
+    }
+
+    @AnyThread
+    boolean isEmpty() {
+        return mMap.isEmpty();
+    }
+
+    @AnyThread
+    @NonNull
+    Collection<String> keySet() {
+        return mMap.keySet();
+    }
+
+    @AnyThread
+    int size() {
+        return mMap.size();
+    }
+}
diff --git a/services/core/java/com/android/server/inputmethod/AdditionalSubtypeUtils.java b/services/core/java/com/android/server/inputmethod/AdditionalSubtypeUtils.java
index fba71fd0ff9e..146ce1732070 100644
--- a/services/core/java/com/android/server/inputmethod/AdditionalSubtypeUtils.java
+++ b/services/core/java/com/android/server/inputmethod/AdditionalSubtypeUtils.java
@@ -108,12 +108,12 @@ final class AdditionalSubtypeUtils {
      * multiple threads are not calling this method at the same time for the same {@code userId}.
      * </p>
      *
-     * @param allSubtypes {@link ArrayMap} from IME ID to additional subtype list. Passing an empty
-     *                    map deletes the file.
+     * @param allSubtypes {@link AdditionalSubtypeMap} from IME ID to additional subtype list.
+     *                    Passing an empty map deletes the file.
      * @param methodMap   {@link ArrayMap} from IME ID to {@link InputMethodInfo}.
      * @param userId      The user ID to be associated with.
      */
-    static void save(ArrayMap<String, List<InputMethodSubtype>> allSubtypes,
+    static void save(AdditionalSubtypeMap allSubtypes,
             InputMethodMap methodMap, @UserIdInt int userId) {
         final File inputMethodDir = getInputMethodDir(userId);
 
@@ -142,7 +142,7 @@ final class AdditionalSubtypeUtils {
     }
 
     @VisibleForTesting
-    static void saveToFile(ArrayMap<String, List<InputMethodSubtype>> allSubtypes,
+    static void saveToFile(AdditionalSubtypeMap allSubtypes,
             InputMethodMap methodMap, AtomicFile subtypesFile) {
         // Safety net for the case that this function is called before methodMap is set.
         final boolean isSetMethodMap = methodMap != null && methodMap.size() > 0;
@@ -212,24 +212,21 @@ final class AdditionalSubtypeUtils {
      * multiple threads are not calling this method at the same time for the same {@code userId}.
      * </p>
      *
-     * @param allSubtypes {@link ArrayMap} from IME ID to additional subtype list. This parameter
-     *                    will be used to return the result.
-     * @param userId      The user ID to be associated with.
+     * @param userId The user ID to be associated with.
+     * @return {@link AdditionalSubtypeMap} that contains the additional {@link InputMethodSubtype}.
      */
-    static void load(@NonNull ArrayMap<String, List<InputMethodSubtype>> allSubtypes,
-            @UserIdInt int userId) {
-        allSubtypes.clear();
-
+    static AdditionalSubtypeMap load(@UserIdInt int userId) {
         final AtomicFile subtypesFile = getAdditionalSubtypeFile(getInputMethodDir(userId));
         // Not having the file means there is no additional subtype.
         if (subtypesFile.exists()) {
-            loadFromFile(allSubtypes, subtypesFile);
+            return loadFromFile(subtypesFile);
         }
+        return AdditionalSubtypeMap.EMPTY_MAP;
     }
 
     @VisibleForTesting
-    static void loadFromFile(@NonNull ArrayMap<String, List<InputMethodSubtype>> allSubtypes,
-            AtomicFile subtypesFile) {
+    static AdditionalSubtypeMap loadFromFile(AtomicFile subtypesFile) {
+        final ArrayMap<String, List<InputMethodSubtype>> allSubtypes = new ArrayMap<>();
         try (FileInputStream fis = subtypesFile.openRead()) {
             final TypedXmlPullParser parser = Xml.resolvePullParser(fis);
             int type = parser.next();
@@ -310,5 +307,6 @@ final class AdditionalSubtypeUtils {
         } catch (XmlPullParserException | IOException | NumberFormatException e) {
             Slog.w(TAG, "Error reading subtypes", e);
         }
+        return AdditionalSubtypeMap.of(allSubtypes);
     }
 }
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
index ffdc9347369e..b8a63cd7941c 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
@@ -286,8 +286,10 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
     final InputManagerInternal mInputManagerInternal;
     final ImePlatformCompatUtils mImePlatformCompatUtils;
     final InputMethodDeviceConfigs mInputMethodDeviceConfigs;
-    private final ArrayMap<String, List<InputMethodSubtype>> mAdditionalSubtypeMap =
-            new ArrayMap<>();
+
+    @GuardedBy("ImfLock.class")
+    @NonNull
+    private AdditionalSubtypeMap mAdditionalSubtypeMap = AdditionalSubtypeMap.EMPTY_MAP;
     private final UserManagerInternal mUserManagerInternal;
     private final InputMethodMenuController mMenuController;
     @NonNull private final InputMethodBindingController mBindingController;
@@ -1332,16 +1334,21 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
         @Override
         public void onPackageDataCleared(String packageName, int uid) {
             synchronized (ImfLock.class) {
-                boolean changed = false;
+                // Note that one package may implement multiple IMEs.
+                final ArrayList<String> changedImes = new ArrayList<>();
                 for (InputMethodInfo imi : mSettings.getMethodList()) {
                     if (imi.getPackageName().equals(packageName)) {
-                        mAdditionalSubtypeMap.remove(imi.getId());
-                        changed = true;
+                        changedImes.add(imi.getId());
                     }
                 }
-                if (changed) {
+                final AdditionalSubtypeMap newMap =
+                        mAdditionalSubtypeMap.cloneWithRemoveOrSelf(changedImes);
+                if (newMap != mAdditionalSubtypeMap) {
+                    mAdditionalSubtypeMap = newMap;
                     AdditionalSubtypeUtils.save(
                             mAdditionalSubtypeMap, mSettings.getMethodMap(), mSettings.getUserId());
+                }
+                if (!changedImes.isEmpty()) {
                     mChangedPackages.add(packageName);
                 }
             }
@@ -1413,7 +1420,8 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
                             Slog.i(TAG,
                                     "Input method reinstalling, clearing additional subtypes: "
                                             + imi.getComponent());
-                            mAdditionalSubtypeMap.remove(imi.getId());
+                            mAdditionalSubtypeMap =
+                                    mAdditionalSubtypeMap.cloneWithRemoveOrSelf(imi.getId());
                             AdditionalSubtypeUtils.save(mAdditionalSubtypeMap,
                                     mSettings.getMethodMap(), mSettings.getUserId());
                         }
@@ -1648,7 +1656,7 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
         // mSettings should be created before buildInputMethodListLocked
         mSettings = InputMethodSettings.createEmptyMap(userId);
 
-        AdditionalSubtypeUtils.load(mAdditionalSubtypeMap, userId);
+        mAdditionalSubtypeMap = AdditionalSubtypeUtils.load(userId);
         mSwitchingController =
                 InputMethodSubtypeSwitchingController.createInstanceLocked(context,
                         mSettings.getMethodMap(), userId);
@@ -1783,7 +1791,7 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
 
         mSettings = InputMethodSettings.createEmptyMap(newUserId);
         // Additional subtypes should be reset when the user is changed
-        AdditionalSubtypeUtils.load(mAdditionalSubtypeMap, newUserId);
+        mAdditionalSubtypeMap = AdditionalSubtypeUtils.load(newUserId);
         final String defaultImiId = mSettings.getSelectedInputMethod();
 
         if (DEBUG) {
@@ -2016,9 +2024,7 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
                 && directBootAwareness == DirectBootAwareness.AUTO) {
             settings = mSettings;
         } else {
-            final ArrayMap<String, List<InputMethodSubtype>> additionalSubtypeMap =
-                    new ArrayMap<>();
-            AdditionalSubtypeUtils.load(additionalSubtypeMap, userId);
+            final AdditionalSubtypeMap additionalSubtypeMap = AdditionalSubtypeUtils.load(userId);
             settings = queryInputMethodServicesInternal(mContext, userId, additionalSubtypeMap,
                     directBootAwareness);
         }
@@ -4218,10 +4224,15 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
             }
 
             if (mSettings.getUserId() == userId) {
-                if (!mSettings.setAdditionalInputMethodSubtypes(imiId, toBeAdded,
-                        mAdditionalSubtypeMap, mPackageManagerInternal, callingUid)) {
+                final var newAdditionalSubtypeMap = mSettings.getNewAdditionalSubtypeMap(
+                        imiId, toBeAdded, mAdditionalSubtypeMap, mPackageManagerInternal,
+                        callingUid);
+                if (mAdditionalSubtypeMap == newAdditionalSubtypeMap) {
                     return;
                 }
+                AdditionalSubtypeUtils.save(newAdditionalSubtypeMap, mSettings.getMethodMap(),
+                        mSettings.getUserId());
+                mAdditionalSubtypeMap = newAdditionalSubtypeMap;
                 final long ident = Binder.clearCallingIdentity();
                 try {
                     buildInputMethodListLocked(false /* resetDefaultEnabledIme */);
@@ -4231,13 +4242,15 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
                 return;
             }
 
-            final ArrayMap<String, List<InputMethodSubtype>> additionalSubtypeMap =
-                    new ArrayMap<>();
-            AdditionalSubtypeUtils.load(additionalSubtypeMap, userId);
+            final AdditionalSubtypeMap additionalSubtypeMap = AdditionalSubtypeUtils.load(userId);
             final InputMethodSettings settings = queryInputMethodServicesInternal(mContext, userId,
                     additionalSubtypeMap, DirectBootAwareness.AUTO);
-            settings.setAdditionalInputMethodSubtypes(imiId, toBeAdded, additionalSubtypeMap,
-                    mPackageManagerInternal, callingUid);
+            final var newAdditionalSubtypeMap = settings.getNewAdditionalSubtypeMap(
+                    imiId, toBeAdded, additionalSubtypeMap, mPackageManagerInternal, callingUid);
+            if (additionalSubtypeMap != newAdditionalSubtypeMap) {
+                AdditionalSubtypeUtils.save(newAdditionalSubtypeMap, settings.getMethodMap(),
+                        settings.getUserId());
+            }
         }
     }
 
@@ -5072,7 +5085,7 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
 
     @NonNull
     static InputMethodSettings queryInputMethodServicesInternal(Context context,
-            @UserIdInt int userId, ArrayMap<String, List<InputMethodSubtype>> additionalSubtypeMap,
+            @UserIdInt int userId, @NonNull AdditionalSubtypeMap additionalSubtypeMap,
             @DirectBootAwareness int directBootAwareness) {
         final Context userAwareContext = context.getUserId() == userId
                 ? context
@@ -5112,7 +5125,7 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
 
     @NonNull
     static InputMethodMap filterInputMethodServices(
-            ArrayMap<String, List<InputMethodSubtype>> additionalSubtypeMap,
+            @NonNull AdditionalSubtypeMap additionalSubtypeMap,
             List<String> enabledInputMethodList, Context userAwareContext,
             List<ResolveInfo> services) {
         final ArrayMap<String, Integer> imiPackageCount = new ArrayMap<>();
@@ -5512,9 +5525,7 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
         if (userId == mSettings.getUserId()) {
             settings = mSettings;
         } else {
-            final ArrayMap<String, List<InputMethodSubtype>> additionalSubtypeMap =
-                    new ArrayMap<>();
-            AdditionalSubtypeUtils.load(additionalSubtypeMap, userId);
+            final AdditionalSubtypeMap additionalSubtypeMap = AdditionalSubtypeUtils.load(userId);
             settings = queryInputMethodServicesInternal(mContext, userId,
                     additionalSubtypeMap, DirectBootAwareness.AUTO);
         }
@@ -5522,9 +5533,7 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
     }
 
     private InputMethodSettings queryMethodMapForUser(@UserIdInt int userId) {
-        final ArrayMap<String, List<InputMethodSubtype>> additionalSubtypeMap =
-                new ArrayMap<>();
-        AdditionalSubtypeUtils.load(additionalSubtypeMap, userId);
+        final AdditionalSubtypeMap additionalSubtypeMap = AdditionalSubtypeUtils.load(userId);
         return queryInputMethodServicesInternal(mContext, userId, additionalSubtypeMap,
                 DirectBootAwareness.AUTO);
     }
@@ -6572,9 +6581,8 @@ public final class InputMethodManagerService extends IInputMethodManager.Stub
                         nextIme = mSettings.getSelectedInputMethod();
                         nextEnabledImes = mSettings.getEnabledInputMethodList();
                     } else {
-                        final ArrayMap<String, List<InputMethodSubtype>> additionalSubtypeMap =
-                                new ArrayMap<>();
-                        AdditionalSubtypeUtils.load(additionalSubtypeMap, userId);
+                        final AdditionalSubtypeMap additionalSubtypeMap =
+                                AdditionalSubtypeUtils.load(userId);
                         final InputMethodSettings settings = queryInputMethodServicesInternal(
                                 mContext, userId, additionalSubtypeMap, DirectBootAwareness.AUTO);
 
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodSettings.java b/services/core/java/com/android/server/inputmethod/InputMethodSettings.java
index e444db1b79e8..a558838172f8 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodSettings.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodSettings.java
@@ -24,7 +24,6 @@ import android.content.pm.PackageManagerInternal;
 import android.os.LocaleList;
 import android.provider.Settings;
 import android.text.TextUtils;
-import android.util.ArrayMap;
 import android.util.IntArray;
 import android.util.Pair;
 import android.util.Printer;
@@ -614,26 +613,27 @@ final class InputMethodSettings {
                 explicitlyOrImplicitlyEnabledSubtypes, null, locale, true);
     }
 
-    boolean setAdditionalInputMethodSubtypes(@NonNull String imeId,
+    @NonNull
+    AdditionalSubtypeMap getNewAdditionalSubtypeMap(@NonNull String imeId,
             @NonNull ArrayList<InputMethodSubtype> subtypes,
-            @NonNull ArrayMap<String, List<InputMethodSubtype>> additionalSubtypeMap,
+            @NonNull AdditionalSubtypeMap additionalSubtypeMap,
             @NonNull PackageManagerInternal packageManagerInternal, int callingUid) {
         final InputMethodInfo imi = mMethodMap.get(imeId);
         if (imi == null) {
-            return false;
+            return additionalSubtypeMap;
         }
         if (!InputMethodUtils.checkIfPackageBelongsToUid(packageManagerInternal, callingUid,
                 imi.getPackageName())) {
-            return false;
+            return additionalSubtypeMap;
         }
 
+        final AdditionalSubtypeMap newMap;
         if (subtypes.isEmpty()) {
-            additionalSubtypeMap.remove(imi.getId());
+            newMap = additionalSubtypeMap.cloneWithRemoveOrSelf(imi.getId());
         } else {
-            additionalSubtypeMap.put(imi.getId(), subtypes);
+            newMap = additionalSubtypeMap.cloneWithPut(imi.getId(), subtypes);
         }
-        AdditionalSubtypeUtils.save(additionalSubtypeMap, mMethodMap, getUserId());
-        return true;
+        return newMap;
     }
 
     boolean setEnabledInputMethodSubtypes(@NonNull String imeId,
diff --git a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/AdditionalSubtypeMapTest.java b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/AdditionalSubtypeMapTest.java
new file mode 100644
index 000000000000..3bb6712a34f1
--- /dev/null
+++ b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/AdditionalSubtypeMapTest.java
@@ -0,0 +1,125 @@
+/*
+ * 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 android.util.ArrayMap;
+import android.view.inputmethod.InputMethodSubtype;
+
+import androidx.annotation.NonNull;
+
+import org.junit.Test;
+
+import java.util.List;
+public final class AdditionalSubtypeMapTest {
+
+    private static final String TEST_IME1_ID = "com.android.test.inputmethod/.TestIme1";
+    private static final String TEST_IME2_ID = "com.android.test.inputmethod/.TestIme2";
+
+    private static InputMethodSubtype createTestSubtype(String locale) {
+        return new InputMethodSubtype
+                .InputMethodSubtypeBuilder()
+                .setSubtypeNameResId(0)
+                .setSubtypeIconResId(0)
+                .setSubtypeLocale(locale)
+                .setIsAsciiCapable(true)
+                .build();
+    }
+
+    private static final InputMethodSubtype TEST_SUBTYPE_EN_US = createTestSubtype("en_US");
+    private static final InputMethodSubtype TEST_SUBTYPE_JA_JP = createTestSubtype("ja_JP");
+
+    private static final List<InputMethodSubtype> TEST_SUBTYPE_LIST1 = List.of(TEST_SUBTYPE_EN_US);
+    private static final List<InputMethodSubtype> TEST_SUBTYPE_LIST2 = List.of(TEST_SUBTYPE_JA_JP);
+
+    private static ArrayMap<String, List<InputMethodSubtype>> mapOf(
+            @NonNull String key1, @NonNull List<InputMethodSubtype> value1) {
+        final ArrayMap<String, List<InputMethodSubtype>> map = new ArrayMap<>();
+        map.put(key1, value1);
+        return map;
+    }
+
+    private static ArrayMap<String, List<InputMethodSubtype>> mapOf(
+            @NonNull String key1, @NonNull List<InputMethodSubtype> value1,
+            @NonNull String key2, @NonNull List<InputMethodSubtype> value2) {
+        final ArrayMap<String, List<InputMethodSubtype>> map = new ArrayMap<>();
+        map.put(key1, value1);
+        map.put(key2, value2);
+        return map;
+    }
+
+    @Test
+    public void testOfReturnsEmptyInstance() {
+        assertThat(AdditionalSubtypeMap.of(new ArrayMap<>()))
+                .isSameInstanceAs(AdditionalSubtypeMap.EMPTY_MAP);
+    }
+
+    @Test
+    public void testOfReturnsNewInstance() {
+        final AdditionalSubtypeMap instance = AdditionalSubtypeMap.of(
+                mapOf(TEST_IME1_ID, TEST_SUBTYPE_LIST1));
+        assertThat(instance.keySet()).containsExactly(TEST_IME1_ID);
+        assertThat(instance.get(TEST_IME1_ID)).containsExactlyElementsIn(TEST_SUBTYPE_LIST1);
+    }
+
+    @Test
+    public void testCloneWithRemoveOrSelfReturnsEmptyInstance() {
+        final AdditionalSubtypeMap original = AdditionalSubtypeMap.of(
+                mapOf(TEST_IME1_ID, TEST_SUBTYPE_LIST1));
+        final AdditionalSubtypeMap result = original.cloneWithRemoveOrSelf(TEST_IME1_ID);
+        assertThat(result).isSameInstanceAs(AdditionalSubtypeMap.EMPTY_MAP);
+    }
+
+    @Test
+    public void testCloneWithRemoveOrSelfWithMultipleKeysReturnsEmptyInstance() {
+        final AdditionalSubtypeMap original = AdditionalSubtypeMap.of(
+                mapOf(TEST_IME1_ID, TEST_SUBTYPE_LIST1, TEST_IME2_ID, TEST_SUBTYPE_LIST2));
+        final AdditionalSubtypeMap result = original.cloneWithRemoveOrSelf(
+                List.of(TEST_IME1_ID, TEST_IME2_ID));
+        assertThat(result).isSameInstanceAs(AdditionalSubtypeMap.EMPTY_MAP);
+    }
+
+    @Test
+    public void testCloneWithRemoveOrSelfReturnsNewInstance() {
+        final AdditionalSubtypeMap original = AdditionalSubtypeMap.of(
+                mapOf(TEST_IME1_ID, TEST_SUBTYPE_LIST1, TEST_IME2_ID, TEST_SUBTYPE_LIST2));
+        final AdditionalSubtypeMap result = original.cloneWithRemoveOrSelf(TEST_IME1_ID);
+        assertThat(result.keySet()).containsExactly(TEST_IME2_ID);
+        assertThat(result.get(TEST_IME2_ID)).containsExactlyElementsIn(TEST_SUBTYPE_LIST2);
+    }
+
+    @Test
+    public void testCloneWithPutWithNewKey() {
+        final AdditionalSubtypeMap original = AdditionalSubtypeMap.of(
+                mapOf(TEST_IME1_ID, TEST_SUBTYPE_LIST1));
+        final AdditionalSubtypeMap result = original.cloneWithPut(TEST_IME2_ID, TEST_SUBTYPE_LIST2);
+        assertThat(result.keySet()).containsExactly(TEST_IME1_ID, TEST_IME2_ID);
+        assertThat(result.get(TEST_IME1_ID)).containsExactlyElementsIn(TEST_SUBTYPE_LIST1);
+        assertThat(result.get(TEST_IME2_ID)).containsExactlyElementsIn(TEST_SUBTYPE_LIST2);
+    }
+
+    @Test
+    public void testCloneWithPutWithExistingKey() {
+        final AdditionalSubtypeMap original = AdditionalSubtypeMap.of(
+                mapOf(TEST_IME1_ID, TEST_SUBTYPE_LIST1, TEST_IME2_ID, TEST_SUBTYPE_LIST2));
+        final AdditionalSubtypeMap result = original.cloneWithPut(TEST_IME2_ID, TEST_SUBTYPE_LIST1);
+        assertThat(result.keySet()).containsExactly(TEST_IME1_ID, TEST_IME2_ID);
+        assertThat(result.get(TEST_IME1_ID)).containsExactlyElementsIn(TEST_SUBTYPE_LIST1);
+        assertThat(result.get(TEST_IME2_ID)).containsExactlyElementsIn(TEST_SUBTYPE_LIST1);
+    }
+}
diff --git a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/AdditionalSubtypeUtilsTest.java b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/AdditionalSubtypeUtilsTest.java
index 0edb3dfc0bc0..63224bb2aa3f 100644
--- a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/AdditionalSubtypeUtilsTest.java
+++ b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/AdditionalSubtypeUtilsTest.java
@@ -55,9 +55,9 @@ public final class AdditionalSubtypeUtilsTest {
         // Save & load.
         AtomicFile atomicFile = new AtomicFile(
                 new File(InstrumentationRegistry.getContext().getCacheDir(), "subtypes.xml"));
-        AdditionalSubtypeUtils.saveToFile(allSubtypes, InputMethodMap.of(methodMap), atomicFile);
-        ArrayMap<String, List<InputMethodSubtype>> loadedSubtypes = new ArrayMap<>();
-        AdditionalSubtypeUtils.loadFromFile(loadedSubtypes, atomicFile);
+        AdditionalSubtypeUtils.saveToFile(AdditionalSubtypeMap.of(allSubtypes),
+                InputMethodMap.of(methodMap), atomicFile);
+        AdditionalSubtypeMap loadedSubtypes = AdditionalSubtypeUtils.loadFromFile(atomicFile);
 
         // Verifies the loaded data.
         assertEquals(1, loadedSubtypes.size());
diff --git a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodManagerServiceRestrictImeAmountTest.java b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodManagerServiceRestrictImeAmountTest.java
index 71752ba3b393..2ea2e2264bec 100644
--- a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodManagerServiceRestrictImeAmountTest.java
+++ b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodManagerServiceRestrictImeAmountTest.java
@@ -25,10 +25,8 @@ import android.content.Intent;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageManager;
 import android.content.pm.ResolveInfo;
-import android.util.ArrayMap;
 import android.view.inputmethod.InputMethod;
 import android.view.inputmethod.InputMethodInfo;
-import android.view.inputmethod.InputMethodSubtype;
 
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 
@@ -124,10 +122,8 @@ public class InputMethodManagerServiceRestrictImeAmountTest extends
 
     private List<InputMethodInfo> filterInputMethodServices(List<ResolveInfo> resolveInfoList,
             List<String> enabledComponents) {
-        final ArrayMap<String, List<InputMethodSubtype>> emptyAdditionalSubtypeMap =
-                new ArrayMap<>();
         final InputMethodMap methodMap = InputMethodManagerService.filterInputMethodServices(
-                emptyAdditionalSubtypeMap, enabledComponents, mContext, resolveInfoList);
+                AdditionalSubtypeMap.EMPTY_MAP, enabledComponents, mContext, resolveInfoList);
         return methodMap.values();
     }