Merge changes I243f1bd9,I977831b2,Ia6b0f35b,Icb3fa967 into main

* changes:
  Update IME Switcher Menu spacing
  Enable omitting IME Switcher Menu header
  Split IME Switcher menu header and item
  Move IME MenuItem list creation in Menu Controller

9 files changed, 687 insertions, 234 deletions

diff --git a/core/res/res/layout/input_method_switch_dialog_new.xml b/core/res/res/layout/input_method_switch_dialog_new.xml
index 058fe3fe3076..118f93b89237 100644
--- a/core/res/res/layout/input_method_switch_dialog_new.xml
+++ b/core/res/res/layout/input_method_switch_dialog_new.xml
@@ -39,7 +39,7 @@
                 android:id="@+id/list"
                 android:layout_width="match_parent"
                 android:layout_height="match_parent"
-                android:paddingVertical="8dp"
+                android:paddingTop="8dp"
                 android:clipToPadding="false"
                 android:layoutManager="com.android.internal.widget.LinearLayoutManager"/>
 
@@ -74,8 +74,7 @@
             android:text="@string/input_method_switcher_settings_button"
             android:fontFamily="google-sans-text"
             android:textAppearance="?attr/textAppearance"
-            android:contentDescription="@string/input_method_language_settings"
-            android:visibility="gone"/>
+            android:contentDescription="@string/input_method_language_settings"/>
 
     </LinearLayout>
 
diff --git a/core/res/res/layout/input_method_switch_item_divider.xml b/core/res/res/layout/input_method_switch_item_divider.xml
new file mode 100644
index 000000000000..4f8c963ff8cf
--- /dev/null
+++ b/core/res/res/layout/input_method_switch_item_divider.xml
@@ -0,0 +1,34 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  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.
+-->
+
+<LinearLayout xmlns:android="http://schemas.android.com/apk/res/android"
+    android:layout_width="match_parent"
+    android:layout_height="wrap_content"
+    android:orientation="vertical"
+    android:layout_marginHorizontal="16dp"
+    android:layout_marginTop="8dp"
+    android:layout_marginBottom="16dp">
+
+    <View
+        android:layout_width="match_parent"
+        android:layout_height="1dp"
+        android:background="?attr/materialColorOutlineVariant"
+        android:layout_marginStart="20dp"
+        android:layout_marginEnd="24dp"
+        android:importantForAccessibility="no"/>
+
+</LinearLayout>
diff --git a/core/res/res/layout/input_method_switch_item_header.xml b/core/res/res/layout/input_method_switch_item_header.xml
new file mode 100644
index 000000000000..f0080a630025
--- /dev/null
+++ b/core/res/res/layout/input_method_switch_item_header.xml
@@ -0,0 +1,38 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  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.
+-->
+
+<LinearLayout xmlns:android="http://schemas.android.com/apk/res/android"
+    android:layout_width="match_parent"
+    android:layout_height="wrap_content"
+    android:orientation="vertical"
+    android:layout_marginHorizontal="16dp"
+    android:layout_marginTop="4dp"
+    android:layout_marginBottom="16dp">
+
+    <TextView
+        android:id="@+id/header_text"
+        android:layout_width="wrap_content"
+        android:layout_height="wrap_content"
+        android:layout_marginHorizontal="8dp"
+        android:ellipsize="end"
+        android:singleLine="true"
+        android:fontFamily="google-sans-text"
+        android:textAppearance="?attr/textAppearance"
+        android:accessibilityHeading="true"
+        android:textColor="?attr/materialColorPrimary"/>
+
+</LinearLayout>
diff --git a/core/res/res/layout/input_method_switch_item_new.xml b/core/res/res/layout/input_method_switch_item_new.xml
index 10d938c71ea4..f8710cc45358 100644
--- a/core/res/res/layout/input_method_switch_item_new.xml
+++ b/core/res/res/layout/input_method_switch_item_new.xml
@@ -16,76 +16,45 @@
 -->
 
 <LinearLayout xmlns:android="http://schemas.android.com/apk/res/android"
+    android:id="@+id/list_item"
     android:layout_width="match_parent"
-    android:layout_height="wrap_content"
-    android:orientation="vertical"
-    android:paddingHorizontal="16dp"
-    android:paddingBottom="8dp">
-
-    <View
-        android:id="@+id/divider"
-        android:layout_width="match_parent"
-        android:layout_height="1dp"
-        android:background="?attr/materialColorSurfaceVariant"
-        android:layout_marginStart="20dp"
-        android:layout_marginTop="8dp"
-        android:layout_marginEnd="24dp"
-        android:layout_marginBottom="12dp"
-        android:visibility="gone"
-        android:importantForAccessibility="no"/>
-
-    <TextView
-        android:id="@+id/header_text"
-        android:layout_width="wrap_content"
-        android:layout_height="wrap_content"
-        android:padding="8dp"
-        android:ellipsize="end"
-        android:singleLine="true"
-        android:fontFamily="google-sans-text"
-        android:textAppearance="?attr/textAppearance"
-        android:textColor="?attr/materialColorPrimary"
-        android:visibility="gone"/>
+    android:layout_height="72dp"
+    android:background="@drawable/input_method_switch_item_background"
+    android:gravity="center_vertical"
+    android:orientation="horizontal"
+    android:layout_marginHorizontal="16dp"
+    android:layout_marginBottom="8dp"
+    android:paddingStart="20dp"
+    android:paddingEnd="24dp">
 
     <LinearLayout
-        android:id="@+id/list_item"
-        android:layout_width="match_parent"
-        android:layout_height="72dp"
-        android:background="@drawable/input_method_switch_item_background"
-        android:gravity="center_vertical"
-        android:orientation="horizontal"
-        android:paddingStart="20dp"
-        android:paddingEnd="24dp">
-
-        <LinearLayout
-            android:layout_width="0dp"
-            android:layout_height="match_parent"
-            android:layout_weight="1"
-            android:gravity="start|center_vertical"
-            android:orientation="vertical">
-
-            <TextView
-                android:id="@+id/text"
-                android:layout_width="wrap_content"
-                android:layout_height="wrap_content"
-                android:ellipsize="marquee"
-                android:singleLine="true"
-                android:fontFamily="google-sans-text"
-                android:textColor="@color/input_method_switch_on_item"
-                android:textAppearance="?attr/textAppearanceListItem"/>
-
-        </LinearLayout>
-
-        <ImageView
-            android:id="@+id/image"
-            android:layout_width="24dp"
-            android:layout_height="24dp"
-            android:gravity="center_vertical"
-            android:layout_marginStart="12dp"
-            android:src="@drawable/ic_check_24dp"
-            android:tint="@color/input_method_switch_on_item"
-            android:visibility="gone"
-            android:importantForAccessibility="no"/>
+        android:layout_width="0dp"
+        android:layout_height="match_parent"
+        android:layout_weight="1"
+        android:gravity="start|center_vertical"
+        android:orientation="vertical">
+
+        <TextView
+            android:id="@+id/text"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:ellipsize="marquee"
+            android:singleLine="true"
+            android:fontFamily="google-sans-text"
+            android:textColor="@color/input_method_switch_on_item"
+            android:textAppearance="?attr/textAppearanceListItem"/>
 
     </LinearLayout>
 
+    <ImageView
+        android:id="@+id/image"
+        android:layout_width="24dp"
+        android:layout_height="24dp"
+        android:gravity="center_vertical"
+        android:layout_marginStart="12dp"
+        android:src="@drawable/ic_check_24dp"
+        android:tint="@color/input_method_switch_on_item"
+        android:visibility="gone"
+        android:importantForAccessibility="no"/>
+
 </LinearLayout>
diff --git a/core/res/res/values/symbols.xml b/core/res/res/values/symbols.xml
index c1893ab85d2b..4f63fac96a8d 100644
--- a/core/res/res/values/symbols.xml
+++ b/core/res/res/values/symbols.xml
@@ -1582,6 +1582,8 @@
   <java-symbol type="layout" name="input_method" />
   <java-symbol type="layout" name="input_method_extract_view" />
   <java-symbol type="layout" name="input_method_switch_item" />
+  <java-symbol type="layout" name="input_method_switch_item_divider" />
+  <java-symbol type="layout" name="input_method_switch_item_header" />
   <java-symbol type="layout" name="input_method_switch_item_new" />
   <java-symbol type="layout" name="input_method_switch_dialog_new" />
   <java-symbol type="layout" name="input_method_switch_dialog_title" />
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
index 35b517118aab..939aad469bd8 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
@@ -187,7 +187,6 @@ import com.android.server.SystemService;
 import com.android.server.companion.virtual.VirtualDeviceManagerInternal;
 import com.android.server.input.InputManagerInternal;
 import com.android.server.inputmethod.InputMethodManagerInternal.InputMethodListListener;
-import com.android.server.inputmethod.InputMethodMenuControllerNew.MenuItem;
 import com.android.server.inputmethod.InputMethodSubtypeSwitchingController.ImeSubtypeListItem;
 import com.android.server.pm.UserManagerInternal;
 import com.android.server.statusbar.StatusBarManagerInternal;
@@ -4078,65 +4077,6 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
         }
     }
 
-    /**
-     * Gets the list of Input Method Switcher Menu items and the index of the selected item.
-     *
-     * @param items                the list of input method and subtype items.
-     * @param selectedImeId        the ID of the selected input method.
-     * @param selectedSubtypeIndex the index of the selected subtype in the input method's array of
-     *                             subtypes, or {@link InputMethodUtils#NOT_A_SUBTYPE_INDEX} if no
-     *                             subtype is selected.
-     * @param userId               the ID of the user for which to get the menu items.
-     * @return the list of menu items, and the index of the selected item,
-     * or {@code -1} if no item is selected.
-     */
-    @GuardedBy("ImfLock.class")
-    @NonNull
-    private Pair<List<MenuItem>, Integer> getInputMethodPickerItems(
-            @NonNull List<ImeSubtypeListItem> items, @Nullable String selectedImeId,
-            int selectedSubtypeIndex, @UserIdInt int userId) {
-        final var bindingController = getInputMethodBindingController(userId);
-        final var settings = InputMethodSettingsRepository.get(userId);
-
-        if (selectedSubtypeIndex == NOT_A_SUBTYPE_INDEX) {
-            // TODO(b/351124299): Check if this fallback logic is still necessary.
-            final var curSubtype = bindingController.getCurrentInputMethodSubtype();
-            if (curSubtype != null) {
-                final var curMethodId = bindingController.getSelectedMethodId();
-                final var curImi = settings.getMethodMap().get(curMethodId);
-                selectedSubtypeIndex = SubtypeUtils.getSubtypeIndexFromHashCode(
-                        curImi, curSubtype.hashCode());
-            }
-        }
-
-        // No item is selected by default. When we have a list of explicitly enabled
-        // subtypes, the implicit subtype is no longer listed. If the implicit one
-        // is still selected, no items will be shown as selected.
-        int selectedIndex = -1;
-        String prevImeId = null;
-        final var menuItems = new ArrayList<MenuItem>();
-        for (int i = 0; i < items.size(); i++) {
-            final var item = items.get(i);
-            final var imeId = item.mImi.getId();
-            if (imeId.equals(selectedImeId)) {
-                final int subtypeIndex = item.mSubtypeIndex;
-                // Check if this is the selected IME-subtype pair.
-                if ((subtypeIndex == 0 && selectedSubtypeIndex == NOT_A_SUBTYPE_INDEX)
-                        || subtypeIndex == NOT_A_SUBTYPE_INDEX
-                        || subtypeIndex == selectedSubtypeIndex) {
-                    selectedIndex = i;
-                }
-            }
-            final boolean hasHeader = !imeId.equals(prevImeId);
-            final boolean hasDivider = hasHeader && prevImeId != null;
-            prevImeId = imeId;
-            menuItems.add(new MenuItem(item.mImeName, item.mSubtypeName, item.mImi,
-                    item.mSubtypeIndex, hasHeader, hasDivider));
-        }
-
-        return new Pair<>(menuItems, selectedIndex);
-    }
-
     @IInputMethodManagerImpl.PermissionVerified(allOf = {
             Manifest.permission.INTERACT_ACROSS_USERS_FULL,
             Manifest.permission.WRITE_SECURE_SETTINGS})
@@ -4973,18 +4913,21 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
                         + " preferredInputMethodSubtypeIndex=" + lastInputMethodSubtypeIndex);
             }
 
-            final var itemsAndIndex = getInputMethodPickerItems(imList,
-                    lastInputMethodId, lastInputMethodSubtypeIndex, userId);
-            final var menuItems = itemsAndIndex.first;
-            final int selectedIndex = itemsAndIndex.second;
-
-            if (selectedIndex == -1) {
-                Slog.w(TAG, "Switching menu shown with no item selected"
-                        + ", IME id: " + lastInputMethodId
-                        + ", subtype index: " + lastInputMethodSubtypeIndex);
+            int selectedSubtypeIndex = lastInputMethodSubtypeIndex;
+            if (selectedSubtypeIndex == NOT_A_SUBTYPE_INDEX) {
+                // TODO(b/351124299): Check if this fallback logic is still necessary.
+                final var bindingController = getInputMethodBindingController(userId);
+                final var curSubtype = bindingController.getCurrentInputMethodSubtype();
+                if (curSubtype != null) {
+                    final var curMethodId = bindingController.getSelectedMethodId();
+                    final var curImi = settings.getMethodMap().get(curMethodId);
+                    selectedSubtypeIndex = SubtypeUtils.getSubtypeIndexFromHashCode(
+                            curImi, curSubtype.hashCode());
+                }
             }
 
-            mMenuControllerNew.show(menuItems, selectedIndex, displayId, userId);
+            mMenuControllerNew.show(imList, lastInputMethodId, selectedSubtypeIndex, displayId,
+                    userId);
         } else {
             mMenuController.showInputMethodMenuLocked(showAuxSubtypes, displayId,
                     lastInputMethodId, lastInputMethodSubtypeIndex, imList, userId);
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodMenuControllerNew.java b/services/core/java/com/android/server/inputmethod/InputMethodMenuControllerNew.java
index cf2cdc1500f8..1d0e3c639c3f 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodMenuControllerNew.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodMenuControllerNew.java
@@ -30,7 +30,6 @@ import android.annotation.RequiresPermission;
 import android.annotation.UserIdInt;
 import android.app.AlertDialog;
 import android.content.Context;
-import android.content.DialogInterface;
 import android.content.Intent;
 import android.os.UserHandle;
 import android.provider.Settings;
@@ -48,8 +47,11 @@ import android.widget.Button;
 import android.widget.ImageView;
 import android.widget.TextView;
 
+import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.widget.RecyclerView;
+import com.android.server.inputmethod.InputMethodSubtypeSwitchingController.ImeSubtypeListItem;
 
+import java.util.ArrayList;
 import java.util.List;
 
 /**
@@ -80,18 +82,27 @@ final class InputMethodMenuControllerNew {
     /**
      * Shows the Input Method Switcher Menu, with a list of IMEs and their subtypes.
      *
-     * @param items         the list of menu items.
-     * @param selectedIndex the index of the menu item that is selected.
-     *                      If no other IMEs are enabled, this index will be out of reach.
-     * @param displayId     the ID of the display where the menu was requested.
-     * @param userId        the ID of the user that requested the menu.
+     * @param items                the list of input method and subtype items.
+     * @param selectedImeId        the ID of the selected input method.
+     * @param selectedSubtypeIndex the index of the selected subtype in the input method's array of
+     *                             subtypes, or {@link InputMethodUtils#NOT_A_SUBTYPE_INDEX} if no
+     *                             subtype is selected.
+     * @param displayId            the ID of the display where the menu was requested.
+     * @param userId               the ID of the user that requested the menu.
      */
     @RequiresPermission(allOf = {INTERACT_ACROSS_USERS, HIDE_OVERLAY_WINDOWS})
-    void show(@NonNull List<MenuItem> items, int selectedIndex, int displayId,
-            @UserIdInt int userId) {
+    void show(@NonNull List<ImeSubtypeListItem> items, @Nullable String selectedImeId,
+            int selectedSubtypeIndex, int displayId, @UserIdInt int userId) {
         // Hide the menu in case it was already showing.
         hide(displayId, userId);
 
+        final var menuItems = getMenuItems(items);
+        final int selectedIndex = getSelectedIndex(menuItems, selectedImeId, selectedSubtypeIndex);
+        if (selectedIndex == -1) {
+            Slog.w(TAG, "Switching menu shown with no item selected, IME id: " + selectedImeId
+                    + ", subtype index: " + selectedSubtypeIndex);
+        }
+
         final Context dialogWindowContext = mDialogWindowContext.get(displayId);
         final var builder = new AlertDialog.Builder(dialogWindowContext,
                 com.android.internal.R.style.Theme_DeviceDefault_InputMethodSwitcherDialog);
@@ -104,52 +115,28 @@ final class InputMethodMenuControllerNew {
                 dialogWindowContext.getText(com.android.internal.R.string.select_input_method));
         builder.setView(contentView);
 
-        final DialogInterface.OnClickListener onClickListener = (dialog, which) -> {
-            if (which != selectedIndex) {
-                final var item = items.get(which);
+        final OnClickListener onClickListener = (item, isSelected) -> {
+            if (!isSelected) {
                 InputMethodManagerInternal.get()
                         .switchToInputMethod(item.mImi.getId(), item.mSubtypeIndex, userId);
             }
             hide(displayId, userId);
         };
 
-        final var selectedImi = selectedIndex >= 0 ? items.get(selectedIndex).mImi : null;
-        final var languageSettingsIntent = selectedImi != null
-                ? selectedImi.createImeLanguageSettingsActivityIntent() : null;
-        final boolean isDeviceProvisioned = Settings.Global.getInt(
-                dialogWindowContext.getContentResolver(), Settings.Global.DEVICE_PROVISIONED,
-                0) != 0;
-        final boolean hasLanguageSettingsButton = languageSettingsIntent != null
-                && isDeviceProvisioned;
-        if (hasLanguageSettingsButton) {
-            final View buttonBar = contentView
-                    .requireViewById(com.android.internal.R.id.button_bar);
-            buttonBar.setVisibility(View.VISIBLE);
-
-            languageSettingsIntent.setFlags(Intent.FLAG_ACTIVITY_NEW_TASK);
-            final Button languageSettingsButton = contentView
-                    .requireViewById(com.android.internal.R.id.button1);
-            languageSettingsButton.setVisibility(View.VISIBLE);
-            languageSettingsButton.setOnClickListener(v -> {
-                v.getContext().startActivityAsUser(languageSettingsIntent, UserHandle.of(userId));
-                hide(displayId, userId);
-            });
-        }
-
         // Create the current IME subtypes list.
         final RecyclerView recyclerView = contentView
                 .requireViewById(com.android.internal.R.id.list);
-        recyclerView.setAdapter(new Adapter(items, selectedIndex, inflater, onClickListener));
+        recyclerView.setAdapter(new Adapter(menuItems, selectedIndex, inflater, onClickListener));
         // Scroll to the currently selected IME. This must run after the recycler view is laid out.
         recyclerView.post(() -> recyclerView.scrollToPosition(selectedIndex));
-        // Indicate that the list can be scrolled.
-        recyclerView.setScrollIndicators(
-                hasLanguageSettingsButton ? View.SCROLL_INDICATOR_BOTTOM : 0);
         // Request focus to enable rotary scrolling on watches.
         recyclerView.requestFocus();
 
+        final var selectedItem = selectedIndex > -1 ? menuItems.get(selectedIndex) : null;
+        updateLanguageSettingsButton(selectedItem, contentView, displayId, userId);
+
         builder.setOnCancelListener(dialog -> hide(displayId, userId));
-        mMenuItems = items;
+        mMenuItems = menuItems;
         mDialog = builder.create();
         mDialog.setCanceledOnTouchOutside(true);
         final Window w = mDialog.getWindow();
@@ -208,98 +195,303 @@ final class InputMethodMenuControllerNew {
     }
 
     /**
-     * Item to be shown in the Input Method Switcher Menu, containing an input method and
-     * optionally an input method subtype.
+     * Creates the list of menu items from the given list of input methods and subtypes. This
+     * handles adding headers and dividers between groups of items from different input methods
+     * as follows:
+     *
+     * <li>If there is only one group, no divider or header will be added.</li>
+     * <li>A divider is added before each group, except the first one.</li>
+     * <li>A header is added before each group (after the divider, if it exists) if the group has
+     * at least two items, or a single item with a subtype name.</li>
+     *
+     * @param items the list of input method and subtype items.
      */
-    static class MenuItem {
+    @VisibleForTesting
+    @NonNull
+    static List<MenuItem> getMenuItems(@NonNull List<ImeSubtypeListItem> items) {
+        final var menuItems = new ArrayList<MenuItem>();
+        if (items.isEmpty()) {
+            return menuItems;
+        }
+
+        final var itemsArray = (ArrayList<ImeSubtypeListItem>) items;
+        final int numItems = itemsArray.size();
+        // Initialize to the last IME id to avoid headers if there is only a single IME.
+        String prevImeId = itemsArray.getLast().mImi.getId();
+        boolean firstGroup = true;
+        for (int i = 0; i < numItems; i++) {
+            final var item = itemsArray.get(i);
+
+            final var imeId = item.mImi.getId();
+            final boolean groupChange = !imeId.equals(prevImeId);
+            if (groupChange) {
+                if (!firstGroup) {
+                    menuItems.add(DividerItem.getInstance());
+                }
+                // Add a header if we have at least two items, or a single item with a subtype name.
+                final var nextItemId = i + 1 < numItems ? itemsArray.get(i + 1).mImi.getId() : null;
+                final boolean addHeader = item.mSubtypeName != null || imeId.equals(nextItemId);
+                if (addHeader) {
+                    menuItems.add(new HeaderItem(item.mImeName));
+                }
+                firstGroup = false;
+                prevImeId = imeId;
+            }
+
+            menuItems.add(new SubtypeItem(item.mImeName, item.mSubtypeName, item.mImi,
+                    item.mSubtypeIndex));
+        }
+
+        return menuItems;
+    }
+
+    /**
+     * Gets the index of the selected item.
+     *
+     * @param items                the list of menu items.
+     * @param selectedImeId        the ID of the selected input method.
+     * @param selectedSubtypeIndex the index of the selected subtype in the input method's array of
+     *                             subtypes, or {@link InputMethodUtils#NOT_A_SUBTYPE_INDEX} if no
+     *                             subtype is selected.
+     * @return the index of the selected item, or {@code -1} if no item is selected.
+     */
+    @VisibleForTesting
+    @IntRange(from = -1)
+    static int getSelectedIndex(@NonNull List<MenuItem> items, @Nullable String selectedImeId,
+            int selectedSubtypeIndex) {
+        for (int i = 0; i < items.size(); i++) {
+            final var item = items.get(i);
+            if (item instanceof SubtypeItem subtypeItem) {
+                final var imeId = subtypeItem.mImi.getId();
+                final int subtypeIndex = subtypeItem.mSubtypeIndex;
+                if (imeId.equals(selectedImeId)
+                        && ((subtypeIndex == 0 && selectedSubtypeIndex == NOT_A_SUBTYPE_INDEX)
+                            || subtypeIndex == NOT_A_SUBTYPE_INDEX
+                            || subtypeIndex == selectedSubtypeIndex)) {
+                    return i;
+                }
+            }
+        }
+        // Either there is no selected IME, or the selected subtype is enabled but not in the list.
+        // This can happen if an implicit subtype is selected, but we got a list of explicit
+        // subtypes. In this case, the implicit subtype will no longer be included in the list.
+        return -1;
+    }
+
+    /**
+     * Updates the visibility of the Language Settings button to visible if the currently selected
+     * item specifies a (language) settings activity and the device is provisioned. Otherwise,
+     * the button won't be shown.
+     *
+     * @param selectedItem the currently selected item, or {@code null} if no item is selected.
+     * @param view         the menu dialog view.
+     * @param displayId    the ID of the display where the menu was requested.
+     * @param userId       the ID of the user that requested the menu.
+     */
+    @RequiresPermission(allOf = {INTERACT_ACROSS_USERS})
+    private void updateLanguageSettingsButton(@Nullable MenuItem selectedItem, @NonNull View view,
+            int displayId, @UserIdInt int userId) {
+        final var settingsIntent = (selectedItem instanceof SubtypeItem selectedSubtypeItem)
+                ? selectedSubtypeItem.mImi.createImeLanguageSettingsActivityIntent() : null;
+        final boolean isDeviceProvisioned = Settings.Global.getInt(
+                view.getContext().getContentResolver(), Settings.Global.DEVICE_PROVISIONED,
+                0) != 0;
+        final boolean hasButton = settingsIntent != null && isDeviceProvisioned;
+        final View buttonBar = view.requireViewById(com.android.internal.R.id.button_bar);
+        final Button button = view.requireViewById(com.android.internal.R.id.button1);
+        final RecyclerView recyclerView = view.requireViewById(com.android.internal.R.id.list);
+        if (hasButton) {
+            settingsIntent.setFlags(Intent.FLAG_ACTIVITY_NEW_TASK);
+            buttonBar.setVisibility(View.VISIBLE);
+            button.setOnClickListener(v -> {
+                v.getContext().startActivityAsUser(settingsIntent, UserHandle.of(userId));
+                hide(displayId, userId);
+            });
+            // Indicate that the list can be scrolled.
+            recyclerView.setScrollIndicators(View.SCROLL_INDICATOR_BOTTOM);
+        } else {
+            buttonBar.setVisibility(View.GONE);
+            button.setOnClickListener(null);
+            // Remove scroll indicator as there is nothing drawn below the list.
+            recyclerView.setScrollIndicators(0 /* indicators */);
+        }
+    }
+
+    /**
+     * Interface definition for callbacks to be invoked when a {@link SubtypeItem} is clicked.
+     */
+    private interface OnClickListener {
+
+        /**
+         * Called when an item is clicked.
+         *
+         * @param item       The item that was clicked.
+         * @param isSelected Whether the item is the currently selected one.
+         */
+        void onClick(@NonNull SubtypeItem item, boolean isSelected);
+    }
+
+    /** Item to be displayed in the menu. */
+    sealed interface MenuItem {}
+
+    /** Subtype item containing an input method and optionally an input method subtype. */
+    static final class SubtypeItem implements MenuItem {
 
         /** The name of the input method. */
         @NonNull
-        private final CharSequence mImeName;
+        final CharSequence mImeName;
 
         /**
          * The name of the input method subtype, or {@code null} if this item doesn't have a
          * subtype.
          */
         @Nullable
-        private final CharSequence mSubtypeName;
+        final CharSequence mSubtypeName;
 
         /** The info of the input method. */
         @NonNull
-        private final InputMethodInfo mImi;
+        final InputMethodInfo mImi;
 
         /**
          * The index of the subtype in the input method's array of subtypes,
          * or {@link InputMethodUtils#NOT_A_SUBTYPE_INDEX} if this item doesn't have a subtype.
          */
         @IntRange(from = NOT_A_SUBTYPE_INDEX)
-        private final int mSubtypeIndex;
+        final int mSubtypeIndex;
 
-        /** Whether this item has a group header (only the first item of each input method). */
-        private final boolean mHasHeader;
-
-        /**
-         * Whether this item should has a group divider (same as {@link #mHasHeader},
-         * excluding the first IME).
-         */
-        private final boolean mHasDivider;
-
-        MenuItem(@NonNull CharSequence imeName, @Nullable CharSequence subtypeName,
+        SubtypeItem(@NonNull CharSequence imeName, @Nullable CharSequence subtypeName,
                 @NonNull InputMethodInfo imi,
-                @IntRange(from = NOT_A_SUBTYPE_INDEX) int subtypeIndex, boolean hasHeader,
-                boolean hasDivider) {
+                @IntRange(from = NOT_A_SUBTYPE_INDEX) int subtypeIndex) {
             mImeName = imeName;
             mSubtypeName = subtypeName;
             mImi = imi;
             mSubtypeIndex = subtypeIndex;
-            mHasHeader = hasHeader;
-            mHasDivider = hasDivider;
         }
 
         @Override
         public String toString() {
-            return "MenuItem{"
+            return "SubtypeItem{"
                     + "mImeName=" + mImeName
                     + " mSubtypeName=" + mSubtypeName
                     + " mSubtypeIndex=" + mSubtypeIndex
-                    + " mHasHeader=" + mHasHeader
-                    + " mHasDivider=" + mHasDivider
                     + "}";
         }
     }
 
-    private static class Adapter extends RecyclerView.Adapter<Adapter.ViewHolder> {
+    /** Header item displayed before a group of {@link SubtypeItem} of the same input method. */
+    static final class HeaderItem implements MenuItem {
+
+        /** The header title. */
+        @NonNull
+        final CharSequence mTitle;
+
+        HeaderItem(@NonNull CharSequence title) {
+            mTitle = title;
+        }
+
+        @Override
+        public String toString() {
+            return "HeaderItem{"
+                    + "mTitle=" + mTitle
+                    + "}";
+        }
+    }
+
+    /** Divider item displayed before a {@link HeaderItem}. */
+    static final class DividerItem implements MenuItem {
+
+        private static DividerItem sInstance;
+
+        /** Gets a singleton instance of DividerItem. */
+        @NonNull
+        static DividerItem getInstance() {
+            if (sInstance == null) {
+                sInstance = new DividerItem();
+            }
+            return sInstance;
+        }
+
+        @Override
+        public String toString() {
+            return "DividerItem{}";
+        }
+    }
+
+    private static final class Adapter extends RecyclerView.Adapter<RecyclerView.ViewHolder> {
+
+        /** View type for unknown item. */
+        private static final int TYPE_UNKNOWN = -1;
+
+        /** View type for {@link SubtypeItem}. */
+        private static final int TYPE_SUBTYPE = 0;
+
+        /** View type for {@link HeaderItem}. */
+        private static final int TYPE_HEADER = 1;
+
+        /** View type for {@link DividerItem}. */
+        private static final int TYPE_DIVIDER = 2;
 
         /** The list of items to show. */
         @NonNull
         private final List<MenuItem> mItems;
         /** The index of the selected item. */
+        @IntRange(from = -1)
         private final int mSelectedIndex;
         @NonNull
         private final LayoutInflater mInflater;
+        /** The listener used to handle clicks on {@link SubtypeViewHolder} items. */
         @NonNull
-        private final DialogInterface.OnClickListener mOnClickListener;
+        private final OnClickListener mListener;
 
-        Adapter(@NonNull List<MenuItem> items, int selectedIndex,
+        Adapter(@NonNull List<MenuItem> items, @IntRange(from = -1) int selectedIndex,
                 @NonNull LayoutInflater inflater,
-                @NonNull DialogInterface.OnClickListener onClickListener) {
+                @NonNull OnClickListener listener) {
             mItems = items;
             mSelectedIndex = selectedIndex;
             mInflater = inflater;
-            mOnClickListener = onClickListener;
+            mListener = listener;
         }
 
         @Override
-        public ViewHolder onCreateViewHolder(ViewGroup parent, int viewType) {
-            final View view = mInflater.inflate(
-                    com.android.internal.R.layout.input_method_switch_item_new, parent, false);
-
-            return new ViewHolder(view, mOnClickListener);
+        public RecyclerView.ViewHolder onCreateViewHolder(ViewGroup parent, int viewType) {
+            switch (viewType) {
+                case TYPE_SUBTYPE -> {
+                    final View view = mInflater.inflate(
+                            com.android.internal.R.layout.input_method_switch_item_new, parent,
+                            false);
+                    return new SubtypeViewHolder(view, mListener);
+                }
+                case TYPE_HEADER -> {
+                    final View view = mInflater.inflate(
+                            com.android.internal.R.layout.input_method_switch_item_header, parent,
+                            false);
+                    return new HeaderViewHolder(view);
+                }
+                case TYPE_DIVIDER -> {
+                    final View view = mInflater.inflate(
+                            com.android.internal.R.layout.input_method_switch_item_divider, parent,
+                            false);
+                    return new DividerViewHolder(view);
+                }
+                default -> throw new IllegalArgumentException("Unknown viewType: " + viewType);
+            }
         }
 
         @Override
-        public void onBindViewHolder(ViewHolder holder, int position) {
-            holder.bind(mItems.get(position), position == mSelectedIndex /* isSelected */);
+        public void onBindViewHolder(RecyclerView.ViewHolder holder, int position) {
+            final var item = mItems.get(position);
+            if (holder instanceof SubtypeViewHolder subtypeHolder
+                    && item instanceof SubtypeItem subtypeItem) {
+                subtypeHolder.bind(subtypeItem, position == mSelectedIndex /* isSelected */);
+            } else if (holder instanceof HeaderViewHolder headerHolder
+                    && item instanceof HeaderItem headerItem) {
+                headerHolder.bind(headerItem);
+            } else if (holder instanceof DividerViewHolder && item instanceof DividerItem) {
+                // Nothing to bind for dividers.
+                return;
+            } else {
+                Slog.w(TAG, "Holder type: " + holder + " doesn't match item type: " + item);
+            }
         }
 
         @Override
@@ -307,7 +499,21 @@ final class InputMethodMenuControllerNew {
             return mItems.size();
         }
 
-        private static class ViewHolder extends RecyclerView.ViewHolder {
+        @Override
+        public int getItemViewType(int position) {
+            final var item = mItems.get(position);
+            if (item instanceof SubtypeItem) {
+                return TYPE_SUBTYPE;
+            } else if (item instanceof HeaderItem) {
+                return TYPE_HEADER;
+            } else if (item instanceof DividerItem) {
+                return TYPE_DIVIDER;
+            } else {
+                return TYPE_UNKNOWN;
+            }
+        }
+
+        private static final class SubtypeViewHolder extends RecyclerView.ViewHolder {
 
             /** The container of the item. */
             @NonNull
@@ -318,46 +524,74 @@ final class InputMethodMenuControllerNew {
             /** Indicator for the selected status of the item. */
             @NonNull
             private final ImageView mCheckmark;
-            /** The group header optionally drawn above the item. */
-            @NonNull
-            private final TextView mHeader;
-            /** The group divider optionally drawn above the item. */
-            @NonNull
-            private final View mDivider;
 
-            private ViewHolder(@NonNull View itemView,
-                    @NonNull DialogInterface.OnClickListener onClickListener) {
+            /** The bound item data, or {@code null} if no item was bound yet. */
+            @Nullable
+            private SubtypeItem mItem;
+            /** Whether this item is the currently selected one. */
+            private boolean mIsSelected;
+
+            SubtypeViewHolder(@NonNull View itemView, @NonNull OnClickListener listener) {
                 super(itemView);
 
-                mContainer = itemView.requireViewById(com.android.internal.R.id.list_item);
+                mContainer = itemView;
                 mName = itemView.requireViewById(com.android.internal.R.id.text);
                 mCheckmark = itemView.requireViewById(com.android.internal.R.id.image);
-                mHeader = itemView.requireViewById(com.android.internal.R.id.header_text);
-                mDivider = itemView.requireViewById(com.android.internal.R.id.divider);
 
-                mContainer.setOnClickListener((v) ->
-                        onClickListener.onClick(null /* dialog */, getAdapterPosition()));
+                mContainer.setOnClickListener((v) -> {
+                    if (mItem != null) {
+                        listener.onClick(mItem, mIsSelected);
+                    }
+                });
             }
 
             /**
              * Binds the given item to the current view.
              *
              * @param item       the item to bind.
-             * @param isSelected whether this is selected.
+             * @param isSelected whether the item is selected.
              */
-            private void bind(@NonNull MenuItem item, boolean isSelected) {
+            void bind(@NonNull SubtypeItem item, boolean isSelected) {
+                mItem = item;
+                mIsSelected = isSelected;
                 // Use the IME name for subtypes with an empty subtype name.
                 final var name = TextUtils.isEmpty(item.mSubtypeName)
                         ? item.mImeName : item.mSubtypeName;
                 mContainer.setActivated(isSelected);
                 // Activated is the correct state, but we also set selected for accessibility info.
                 mContainer.setSelected(isSelected);
+                // Trigger the ellipsize marquee behaviour by selecting the name.
                 mName.setSelected(isSelected);
                 mName.setText(name);
                 mCheckmark.setVisibility(isSelected ? View.VISIBLE : View.GONE);
-                mHeader.setText(item.mImeName);
-                mHeader.setVisibility(item.mHasHeader ? View.VISIBLE : View.GONE);
-                mDivider.setVisibility(item.mHasDivider ? View.VISIBLE : View.GONE);
+            }
+        }
+
+        private static final class HeaderViewHolder extends RecyclerView.ViewHolder {
+
+            /** The title view, only visible if the bound item has a title. */
+            private final TextView mTitle;
+
+            HeaderViewHolder(@NonNull View itemView) {
+                super(itemView);
+
+                mTitle = itemView.requireViewById(com.android.internal.R.id.header_text);
+            }
+
+            /**
+             * Binds the given item to the current view.
+             *
+             * @param item the item to bind.
+             */
+            void bind(@NonNull HeaderItem item) {
+                mTitle.setText(item.mTitle);
+            }
+        }
+
+        private static final class DividerViewHolder extends RecyclerView.ViewHolder {
+
+            DividerViewHolder(@NonNull View itemView) {
+                super(itemView);
             }
         }
     }
diff --git a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodMenuControllerTest.java b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodMenuControllerTest.java
new file mode 100644
index 000000000000..02dc86bffe2d
--- /dev/null
+++ b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodMenuControllerTest.java
@@ -0,0 +1,234 @@
+/*
+ * 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.android.server.inputmethod.InputMethodMenuControllerNew.getMenuItems;
+import static com.android.server.inputmethod.InputMethodMenuControllerNew.getSelectedIndex;
+import static com.android.server.inputmethod.InputMethodSubtypeSwitchingControllerTest.addTestImeSubtypeListItems;
+import static com.android.server.inputmethod.InputMethodUtils.NOT_A_SUBTYPE_INDEX;
+
+import static com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.Truth.assertWithMessage;
+
+import android.platform.test.annotations.RequiresFlagsEnabled;
+import android.platform.test.flag.junit.CheckFlagsRule;
+import android.platform.test.flag.junit.DeviceFlagsValueProvider;
+import android.view.inputmethod.Flags;
+
+import com.android.server.inputmethod.InputMethodMenuControllerNew.DividerItem;
+import com.android.server.inputmethod.InputMethodMenuControllerNew.HeaderItem;
+import com.android.server.inputmethod.InputMethodMenuControllerNew.SubtypeItem;
+import com.android.server.inputmethod.InputMethodSubtypeSwitchingController.ImeSubtypeListItem;
+
+import org.junit.Rule;
+import org.junit.Test;
+
+import java.util.ArrayList;
+import java.util.List;
+
+@RequiresFlagsEnabled(Flags.FLAG_IME_SWITCHER_REVAMP)
+public class InputMethodMenuControllerTest {
+
+    @Rule
+    public final CheckFlagsRule mCheckFlagsRule = DeviceFlagsValueProvider.createCheckFlagsRule();
+
+    /** Verifies that getMenuItems maintains the same order and information from the given items. */
+    @Test
+    public void testGetMenuItems() {
+        final var items = new ArrayList<ImeSubtypeListItem>();
+        addTestImeSubtypeListItems(items, "LatinIme", "LatinIme",
+                List.of("en", "fr"), true /* supportsSwitchingToNextInputMethod */);
+        addTestImeSubtypeListItems(items, "SimpleIme", "SimpleIme",
+                null, true /* supportsSwitchingToNextInputMethod */);
+
+        final var menuItems = getMenuItems(items);
+
+        int itemsIndex = 0;
+
+        for (int i = 0; i < menuItems.size(); i++) {
+            final var menuItem = menuItems.get(i);
+            if (menuItem instanceof SubtypeItem subtypeItem) {
+                final var item = items.get(itemsIndex);
+
+                assertWithMessage("IME name does not match").that(subtypeItem.mImeName)
+                        .isEqualTo(item.mImeName);
+                assertWithMessage("Subtype name does not match").that(subtypeItem.mSubtypeName)
+                        .isEqualTo(item.mSubtypeName);
+                assertWithMessage("InputMethodInfo does not match").that(subtypeItem.mImi)
+                        .isEqualTo(item.mImi);
+                assertWithMessage("Subtype index does not match").that(subtypeItem.mSubtypeIndex)
+                        .isEqualTo(item.mSubtypeIndex);
+
+                itemsIndex++;
+            }
+        }
+
+        assertWithMessage("Items list was not fully traversed").that(itemsIndex)
+                .isEqualTo(items.size());
+    }
+
+    /**
+     * Verifies that getMenuItems does not add a header or divider if all the items belong to
+     * a single input method.
+     */
+    @Test
+    public void testGetMenuItemsNoHeaderOrDividerForSingleInputMethod() {
+        final var items = new ArrayList<ImeSubtypeListItem>();
+        addTestImeSubtypeListItems(items, "LatinIme", "LatinIme",
+                List.of("en", "fr"), true /* supportsSwitchingToNextInputMethod */);
+
+        final var menuItems = getMenuItems(items);
+
+        assertThat(menuItems.stream()
+                .filter(item -> item instanceof HeaderItem || item instanceof DividerItem).toList())
+                .isEmpty();
+    }
+
+    /**
+     * Verifies that getMenuItems only adds headers for item groups with at least two items,
+     * or with a single item with a subtype name.
+     */
+    @Test
+    public void testGetMenuItemsHeaders() {
+        final var items = new ArrayList<ImeSubtypeListItem>();
+        addTestImeSubtypeListItems(items, "DefaultIme", "DefaultIme",
+                null, true /* supportsSwitchingToNextInputMethod */);
+        addTestImeSubtypeListItems(items, "LatinIme", "LatinIme",
+                List.of("en", "fr"), true /* supportsSwitchingToNextInputMethod */);
+        addTestImeSubtypeListItems(items, "ItalianIme", "ItalianIme",
+                List.of("it"), true /* supportsSwitchingToNextInputMethod */);
+        addTestImeSubtypeListItems(items, "SimpleIme", "SimpleIme",
+                null, true /* supportsSwitchingToNextInputMethod */);
+
+        final var menuItems = getMenuItems(items);
+
+        assertWithMessage("Must have menu items").that(menuItems).isNotEmpty();
+
+        final var headersAndDividers = menuItems.stream()
+                .filter(item -> item instanceof HeaderItem || item instanceof DividerItem)
+                .toList();
+
+        assertWithMessage("Must have header and divider items").that(headersAndDividers).hasSize(5);
+
+        assertWithMessage("First group has no header")
+                .that(menuItems.getFirst()).isInstanceOf(SubtypeItem.class);
+        assertWithMessage("Group with multiple items has divider")
+                .that(headersAndDividers.get(0)).isInstanceOf(DividerItem.class);
+        assertWithMessage("Group with multiple items has header")
+                .that(headersAndDividers.get(1)).isInstanceOf(HeaderItem.class);
+        assertWithMessage("Group with single item with subtype name has divider")
+                .that(headersAndDividers.get(2)).isInstanceOf(DividerItem.class);
+        assertWithMessage("Group with single item with subtype name has header")
+                .that(headersAndDividers.get(3)).isInstanceOf(HeaderItem.class);
+        assertWithMessage("Group with single item without subtype name has divider only")
+                .that(headersAndDividers.get(4)).isInstanceOf(DividerItem.class);
+    }
+
+    /** Verifies that getMenuItems adds a divider before every header except the first one. */
+    @Test
+    public void testGetMenuItemsDivider() {
+        final var items = new ArrayList<ImeSubtypeListItem>();
+        addTestImeSubtypeListItems(items, "LatinIme", "LatinIme",
+                List.of("en", "fr"), true /* supportsSwitchingToNextInputMethod */);
+        addTestImeSubtypeListItems(items, "ItalianIme", "ItalianIme",
+                List.of("it"), true /* supportsSwitchingToNextInputMethod */);
+        addTestImeSubtypeListItems(items, "SimpleIme", "SimpleIme",
+                null, true /* supportsSwitchingToNextInputMethod */);
+
+        final var menuItems = getMenuItems(items);
+
+        assertWithMessage("First item is a header")
+                .that(menuItems.getFirst()).isInstanceOf(HeaderItem.class);
+        assertWithMessage("Last item is a subtype")
+                .that(menuItems.getLast()).isInstanceOf(SubtypeItem.class);
+
+        for (int i = 0; i < menuItems.size(); i++) {
+            final var item = menuItems.get(i);
+            if (item instanceof HeaderItem && i > 0) {
+                final var prevItem = menuItems.get(i - 1);
+                assertWithMessage("The item before a header should be a divider")
+                        .that(prevItem).isInstanceOf(DividerItem.class);
+            } else if (item instanceof DividerItem && i < menuItems.size() - 1) {
+                final var nextItem = menuItems.get(i + 1);
+                assertWithMessage("The item after a divider should be a header or subtype")
+                        .that(nextItem instanceof HeaderItem || nextItem instanceof SubtypeItem)
+                        .isTrue();
+            }
+        }
+    }
+
+    /**
+     * Verifies that getSelectedIndex returns the matching item when the selected subtype is given.
+     */
+    @Test
+    public void testGetSelectedIndexWithSelectedSubtype() {
+        final var items = new ArrayList<ImeSubtypeListItem>();
+        addTestImeSubtypeListItems(items, "LatinIme", "LatinIme",
+                List.of("en", "fr"), true /* supportsSwitchingToNextInputMethod */);
+        addTestImeSubtypeListItems(items, "SimpleIme", "SimpleIme",
+                List.of("it", "jp", "pt"),  true /* supportsSwitchingToNextInputMethod */);
+
+        final var simpleImeId = items.get(2).mImi.getId();
+        final var menuItems = getMenuItems(items);
+
+        final int selectedIndex = getSelectedIndex(menuItems, simpleImeId, 1);
+        // Two headers + one divider + three items
+        assertThat(selectedIndex).isEqualTo(6);
+    }
+
+    /**
+     * Verifies that getSelectedIndex returns the first item of the selected input method,
+     * when no selected subtype is given.
+     */
+    @Test
+    public void testGetSelectedIndexWithoutSelectedSubtype() {
+        final var items = new ArrayList<ImeSubtypeListItem>();
+        addTestImeSubtypeListItems(items, "LatinIme", "LatinIme",
+                List.of("en", "fr"), true /* supportsSwitchingToNextInputMethod */);
+        addTestImeSubtypeListItems(items, "SimpleIme", "SimpleIme",
+                List.of("it", "jp", "pt"),  true /* supportsSwitchingToNextInputMethod */);
+
+        final var simpleImeId = items.get(2).mImi.getId();
+        final var menuItems = getMenuItems(items);
+
+        final int selectedIndex = getSelectedIndex(menuItems, simpleImeId, NOT_A_SUBTYPE_INDEX);
+
+        // Two headers + one divider + two items
+        assertThat(selectedIndex).isEqualTo(5);
+    }
+
+    /**
+     * Verifies that getSelectedIndex will return the item of the selected input method that has
+     * no subtype, when this is the first one reached, regardless of the given selected subtype.
+     */
+    @Test
+    public void getSelectedIndexNoSubtype() {
+        final var items = new ArrayList<ImeSubtypeListItem>();
+        addTestImeSubtypeListItems(items, "LatinIme", "LatinIme",
+                List.of("en", "fr"), true /* supportsSwitchingToNextInputMethod */);
+        addTestImeSubtypeListItems(items, "SimpleIme", "SimpleIme",
+                null,  true /* supportsSwitchingToNextInputMethod */);
+
+        final var simpleImeId = items.get(2).mImi.getId();
+        final var menuItems = getMenuItems(items);
+
+        final int selectedIndex = getSelectedIndex(menuItems, simpleImeId, 1);
+
+        // One header + one divider + two items
+        assertThat(selectedIndex).isEqualTo(4);
+    }
+}
diff --git a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodSubtypeSwitchingControllerTest.java b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodSubtypeSwitchingControllerTest.java
index 770451cc838d..a804f24acc8f 100644
--- a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodSubtypeSwitchingControllerTest.java
+++ b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodSubtypeSwitchingControllerTest.java
@@ -75,7 +75,7 @@ public final class InputMethodSubtypeSwitchingControllerTest {
                 .build();
     }
 
-    private static void addTestImeSubtypeListItems(@NonNull List<ImeSubtypeListItem> items,
+    static void addTestImeSubtypeListItems(@NonNull List<ImeSubtypeListItem> items,
             @NonNull String imeName, @NonNull String imeLabel,
             @Nullable List<String> subtypeLocales, boolean supportsSwitchingToNextInputMethod) {
         final ApplicationInfo ai = new ApplicationInfo();