Merge "Implement new Layout manager APIs"
diff --git a/core/java/android/hardware/input/KeyboardLayout.java b/core/java/android/hardware/input/KeyboardLayout.java
index 58f7759..0311da4 100644
--- a/core/java/android/hardware/input/KeyboardLayout.java
+++ b/core/java/android/hardware/input/KeyboardLayout.java
@@ -23,6 +23,7 @@
 
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Objects;
 
 /**
  * Describes a keyboard layout.
@@ -30,6 +31,37 @@
  * @hide
  */
 public final class KeyboardLayout implements Parcelable, Comparable<KeyboardLayout> {
+
+    /** Undefined keyboard layout */
+    public static final String LAYOUT_TYPE_UNDEFINED = "undefined";
+
+    /** Qwerty-based keyboard layout */
+    public static final String LAYOUT_TYPE_QWERTY = "qwerty";
+
+    /** Qwertz-based keyboard layout */
+    public static final String LAYOUT_TYPE_QWERTZ = "qwertz";
+
+    /** Azerty-based keyboard layout */
+    public static final String LAYOUT_TYPE_AZERTY = "azerty";
+
+    /** Dvorak keyboard layout */
+    public static final String LAYOUT_TYPE_DVORAK = "dvorak";
+
+    /** Colemak keyboard layout */
+    public static final String LAYOUT_TYPE_COLEMAK = "colemak";
+
+    /** Workman keyboard layout */
+    public static final String LAYOUT_TYPE_WORKMAN = "workman";
+
+    /** Turkish-F keyboard layout */
+    public static final String LAYOUT_TYPE_TURKISH_F = "turkish_f";
+
+    /** Turkish-Q keyboard layout */
+    public static final String LAYOUT_TYPE_TURKISH_Q = "turkish_q";
+
+    /** Keyboard layout that has been enhanced with a large number of extra characters */
+    public static final String LAYOUT_TYPE_EXTENDED = "extended";
+
     private final String mDescriptor;
     private final String mLabel;
     private final String mCollection;
@@ -42,31 +74,24 @@
 
     /** Currently supported Layout types in the KCM files */
     private enum LayoutType {
-        UNDEFINED(0, "undefined"),
-        QWERTY(1, "qwerty"),
-        QWERTZ(2, "qwertz"),
-        AZERTY(3, "azerty"),
-        DVORAK(4, "dvorak"),
-        COLEMAK(5, "colemak"),
-        WORKMAN(6, "workman"),
-        TURKISH_F(7, "turkish_f"),
-        TURKISH_Q(8, "turkish_q"),
-        EXTENDED(9, "extended");
+        UNDEFINED(0, LAYOUT_TYPE_UNDEFINED),
+        QWERTY(1, LAYOUT_TYPE_QWERTY),
+        QWERTZ(2, LAYOUT_TYPE_QWERTZ),
+        AZERTY(3, LAYOUT_TYPE_AZERTY),
+        DVORAK(4, LAYOUT_TYPE_DVORAK),
+        COLEMAK(5, LAYOUT_TYPE_COLEMAK),
+        WORKMAN(6, LAYOUT_TYPE_WORKMAN),
+        TURKISH_F(7, LAYOUT_TYPE_TURKISH_F),
+        TURKISH_Q(8, LAYOUT_TYPE_TURKISH_Q),
+        EXTENDED(9, LAYOUT_TYPE_EXTENDED);
 
         private final int mValue;
         private final String mName;
         private static final Map<Integer, LayoutType> VALUE_TO_ENUM_MAP = new HashMap<>();
         static {
-            VALUE_TO_ENUM_MAP.put(UNDEFINED.mValue, UNDEFINED);
-            VALUE_TO_ENUM_MAP.put(QWERTY.mValue, QWERTY);
-            VALUE_TO_ENUM_MAP.put(QWERTZ.mValue, QWERTZ);
-            VALUE_TO_ENUM_MAP.put(AZERTY.mValue, AZERTY);
-            VALUE_TO_ENUM_MAP.put(DVORAK.mValue, DVORAK);
-            VALUE_TO_ENUM_MAP.put(COLEMAK.mValue, COLEMAK);
-            VALUE_TO_ENUM_MAP.put(WORKMAN.mValue, WORKMAN);
-            VALUE_TO_ENUM_MAP.put(TURKISH_F.mValue, TURKISH_F);
-            VALUE_TO_ENUM_MAP.put(TURKISH_Q.mValue, TURKISH_Q);
-            VALUE_TO_ENUM_MAP.put(EXTENDED.mValue, EXTENDED);
+            for (LayoutType type : LayoutType.values()) {
+                VALUE_TO_ENUM_MAP.put(type.mValue, type);
+            }
         }
 
         private static LayoutType of(int value) {
@@ -207,6 +232,9 @@
         // keyboards to be listed before lower priority keyboards.
         int result = Integer.compare(another.mPriority, mPriority);
         if (result == 0) {
+            result = Integer.compare(mLayoutType.mValue, another.mLayoutType.mValue);
+        }
+        if (result == 0) {
             result = mLabel.compareToIgnoreCase(another.mLabel);
         }
         if (result == 0) {
@@ -226,4 +254,21 @@
                 + ", vendorId: " + mVendorId
                 + ", productId: " + mProductId;
     }
+
+    /**
+     * Check if the provided layout type is supported/valid.
+     *
+     * @param layoutName name of layout type
+     * @return {@code true} if the provided layout type is supported/valid.
+     */
+    public static boolean isLayoutTypeValid(@NonNull String layoutName) {
+        Objects.requireNonNull(layoutName, "Provided layout name should not be null");
+        for (LayoutType layoutType : LayoutType.values()) {
+            if (layoutName.equals(layoutType.getName())) {
+                return true;
+            }
+        }
+        // Layout doesn't match any supported layout types
+        return false;
+    }
 }
diff --git a/services/core/java/com/android/server/input/InputManagerService.java b/services/core/java/com/android/server/input/InputManagerService.java
index ea7f0bb..e6bd493 100644
--- a/services/core/java/com/android/server/input/InputManagerService.java
+++ b/services/core/java/com/android/server/input/InputManagerService.java
@@ -3346,10 +3346,7 @@
         public void onInputMethodSubtypeChangedForKeyboardLayoutMapping(@UserIdInt int userId,
                 @Nullable InputMethodSubtypeHandle subtypeHandle,
                 @Nullable InputMethodSubtype subtype) {
-            if (DEBUG) {
-                Slog.i(TAG, "InputMethodSubtype changed: userId=" + userId
-                        + " subtypeHandle=" + subtypeHandle);
-            }
+            mKeyboardLayoutManager.onInputMethodSubtypeChanged(userId, subtypeHandle, subtype);
         }
 
         @Override
diff --git a/services/core/java/com/android/server/input/KeyboardLayoutManager.java b/services/core/java/com/android/server/input/KeyboardLayoutManager.java
index b8eb901..9e8b9f1 100644
--- a/services/core/java/com/android/server/input/KeyboardLayoutManager.java
+++ b/services/core/java/com/android/server/input/KeyboardLayoutManager.java
@@ -37,6 +37,8 @@
 import android.hardware.input.InputDeviceIdentifier;
 import android.hardware.input.InputManager;
 import android.hardware.input.KeyboardLayout;
+import android.icu.lang.UScript;
+import android.icu.util.ULocale;
 import android.os.Bundle;
 import android.os.Handler;
 import android.os.LocaleList;
@@ -45,6 +47,8 @@
 import android.os.UserHandle;
 import android.provider.Settings;
 import android.text.TextUtils;
+import android.util.ArrayMap;
+import android.util.FeatureFlagUtils;
 import android.util.Log;
 import android.util.Slog;
 import android.view.InputDevice;
@@ -54,6 +58,7 @@
 
 import com.android.internal.R;
 import com.android.internal.annotations.GuardedBy;
+import com.android.internal.inputmethod.InputMethodSubtypeHandle;
 import com.android.internal.messages.nano.SystemMessageProto;
 import com.android.internal.notification.SystemNotificationChannels;
 import com.android.internal.util.XmlUtils;
@@ -63,10 +68,12 @@
 import java.io.IOException;
 import java.io.InputStreamReader;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Locale;
+import java.util.Map;
 import java.util.Objects;
 import java.util.stream.Stream;
 
@@ -94,10 +101,18 @@
     @GuardedBy("mDataStore")
     private final PersistentDataStore mDataStore;
     private final Handler mHandler;
+
     private final List<InputDevice> mKeyboardsWithMissingLayouts = new ArrayList<>();
     private boolean mKeyboardLayoutNotificationShown = false;
     private Toast mSwitchedKeyboardLayoutToast;
 
+    // This cache stores "best-matched" layouts so that we don't need to run the matching
+    // algorithm repeatedly.
+    @GuardedBy("mKeyboardLayoutCache")
+    private final Map<String, String> mKeyboardLayoutCache = new ArrayMap<>();
+    @Nullable
+    private ImeInfo mCurrentImeInfo;
+
     KeyboardLayoutManager(Context context, NativeInputManagerService nativeService,
             PersistentDataStore dataStore, Looper looper) {
         mContext = context;
@@ -139,8 +154,10 @@
 
     @Override
     public void onInputDeviceRemoved(int deviceId) {
-        mKeyboardsWithMissingLayouts.removeIf(device -> device.getId() == deviceId);
-        maybeUpdateNotification();
+        if (!useNewSettingsUi()) {
+            mKeyboardsWithMissingLayouts.removeIf(device -> device.getId() == deviceId);
+            maybeUpdateNotification();
+        }
     }
 
     @Override
@@ -149,18 +166,21 @@
         if (inputDevice == null || inputDevice.isVirtual() || !inputDevice.isFullKeyboard()) {
             return;
         }
-        synchronized (mDataStore) {
-            String layout = getCurrentKeyboardLayoutForInputDevice(inputDevice.getIdentifier());
-            if (layout == null) {
-                layout = getDefaultKeyboardLayout(inputDevice);
-                if (layout != null) {
-                    setCurrentKeyboardLayoutForInputDevice(inputDevice.getIdentifier(), layout);
-                } else {
-                    mKeyboardsWithMissingLayouts.add(inputDevice);
+        if (!useNewSettingsUi()) {
+            synchronized (mDataStore) {
+                String layout = getCurrentKeyboardLayoutForInputDevice(inputDevice.getIdentifier());
+                if (layout == null) {
+                    layout = getDefaultKeyboardLayout(inputDevice);
+                    if (layout != null) {
+                        setCurrentKeyboardLayoutForInputDevice(inputDevice.getIdentifier(), layout);
+                    } else {
+                        mKeyboardsWithMissingLayouts.add(inputDevice);
+                    }
                 }
+                maybeUpdateNotification();
             }
-            maybeUpdateNotification();
         }
+        // TODO(b/259530132): Show notification for new Settings UI
     }
 
     private String getDefaultKeyboardLayout(final InputDevice inputDevice) {
@@ -244,6 +264,12 @@
             }
         }
 
+        synchronized (mKeyboardLayoutCache) {
+            // Invalidate the cache: With packages being installed/removed, existing cache of
+            // auto-selected layout might not be the best layouts anymore.
+            mKeyboardLayoutCache.clear();
+        }
+
         // Reload keyboard layouts.
         reloadKeyboardLayouts();
     }
@@ -256,6 +282,9 @@
 
     public KeyboardLayout[] getKeyboardLayoutsForInputDevice(
             final InputDeviceIdentifier identifier) {
+        if (useNewSettingsUi()) {
+            return new KeyboardLayout[0];
+        }
         final String[] enabledLayoutDescriptors =
                 getEnabledKeyboardLayoutsForInputDevice(identifier);
         final ArrayList<KeyboardLayout> enabledLayouts =
@@ -296,6 +325,7 @@
                 KeyboardLayout[]::new);
     }
 
+    @Nullable
     public KeyboardLayout getKeyboardLayout(String keyboardLayoutDescriptor) {
         Objects.requireNonNull(keyboardLayoutDescriptor,
                 "keyboardLayoutDescriptor must not be null");
@@ -434,21 +464,25 @@
         return LocaleList.forLanguageTags(languageTags.replace('|', ','));
     }
 
-    /**
-     * Builds a layout descriptor for the vendor/product. This returns the
-     * descriptor for ids that aren't useful (such as the default 0, 0).
-     */
-    private String getLayoutDescriptor(InputDeviceIdentifier identifier) {
+    private static String getLayoutDescriptor(@NonNull InputDeviceIdentifier identifier) {
         Objects.requireNonNull(identifier, "identifier must not be null");
         Objects.requireNonNull(identifier.getDescriptor(), "descriptor must not be null");
 
         if (identifier.getVendorId() == 0 && identifier.getProductId() == 0) {
             return identifier.getDescriptor();
         }
+        // If vendor id and product id is available, use it as keys. This allows us to have the
+        // same setup for all keyboards with same product and vendor id. i.e. User can swap 2
+        // identical keyboards and still get the same setup.
         return "vendor:" + identifier.getVendorId() + ",product:" + identifier.getProductId();
     }
 
+    @Nullable
     public String getCurrentKeyboardLayoutForInputDevice(InputDeviceIdentifier identifier) {
+        if (useNewSettingsUi()) {
+            Slog.e(TAG, "getCurrentKeyboardLayoutForInputDevice API not supported");
+            return null;
+        }
         String key = getLayoutDescriptor(identifier);
         synchronized (mDataStore) {
             String layout;
@@ -468,9 +502,13 @@
 
     public void setCurrentKeyboardLayoutForInputDevice(InputDeviceIdentifier identifier,
             String keyboardLayoutDescriptor) {
+        if (useNewSettingsUi()) {
+            Slog.e(TAG, "setCurrentKeyboardLayoutForInputDevice API not supported");
+            return;
+        }
+
         Objects.requireNonNull(keyboardLayoutDescriptor,
                 "keyboardLayoutDescriptor must not be null");
-
         String key = getLayoutDescriptor(identifier);
         synchronized (mDataStore) {
             try {
@@ -489,6 +527,10 @@
     }
 
     public String[] getEnabledKeyboardLayoutsForInputDevice(InputDeviceIdentifier identifier) {
+        if (useNewSettingsUi()) {
+            Slog.e(TAG, "getEnabledKeyboardLayoutsForInputDevice API not supported");
+            return new String[0];
+        }
         String key = getLayoutDescriptor(identifier);
         synchronized (mDataStore) {
             String[] layouts = mDataStore.getKeyboardLayouts(key);
@@ -502,6 +544,10 @@
 
     public void addKeyboardLayoutForInputDevice(InputDeviceIdentifier identifier,
             String keyboardLayoutDescriptor) {
+        if (useNewSettingsUi()) {
+            Slog.e(TAG, "addKeyboardLayoutForInputDevice API not supported");
+            return;
+        }
         Objects.requireNonNull(keyboardLayoutDescriptor,
                 "keyboardLayoutDescriptor must not be null");
 
@@ -525,6 +571,10 @@
 
     public void removeKeyboardLayoutForInputDevice(InputDeviceIdentifier identifier,
             String keyboardLayoutDescriptor) {
+        if (useNewSettingsUi()) {
+            Slog.e(TAG, "removeKeyboardLayoutForInputDevice API not supported");
+            return;
+        }
         Objects.requireNonNull(keyboardLayoutDescriptor,
                 "keyboardLayoutDescriptor must not be null");
 
@@ -551,31 +601,11 @@
         }
     }
 
-    public String getKeyboardLayoutForInputDevice(InputDeviceIdentifier identifier,
-            @UserIdInt int userId, @NonNull InputMethodInfo imeInfo,
-            @Nullable InputMethodSubtype imeSubtype) {
-        // TODO(b/259530132): Implement the new keyboard layout API: Returning non-IME specific
-        //  layout for now.
-        return getCurrentKeyboardLayoutForInputDevice(identifier);
-    }
-
-    public void setKeyboardLayoutForInputDevice(InputDeviceIdentifier identifier,
-            @UserIdInt int userId, @NonNull InputMethodInfo imeInfo,
-            @Nullable InputMethodSubtype imeSubtype, String keyboardLayoutDescriptor) {
-        // TODO(b/259530132): Implement the new keyboard layout API: setting non-IME specific
-        //  layout for now.
-        setCurrentKeyboardLayoutForInputDevice(identifier, keyboardLayoutDescriptor);
-    }
-
-    public KeyboardLayout[] getKeyboardLayoutListForInputDevice(InputDeviceIdentifier identifier,
-            @UserIdInt int userId, @NonNull InputMethodInfo imeInfo,
-            @Nullable InputMethodSubtype imeSubtype) {
-        // TODO(b/259530132): Implement the new keyboard layout API: Returning list of all
-        //  layouts for now.
-        return getKeyboardLayouts();
-    }
-
     public void switchKeyboardLayout(int deviceId, int direction) {
+        if (useNewSettingsUi()) {
+            Slog.e(TAG, "switchKeyboardLayout API not supported");
+            return;
+        }
         mHandler.obtainMessage(MSG_SWITCH_KEYBOARD_LAYOUT, deviceId, direction).sendToTarget();
     }
 
@@ -616,8 +646,21 @@
         }
     }
 
+    @Nullable
     public String[] getKeyboardLayoutOverlay(InputDeviceIdentifier identifier) {
-        String keyboardLayoutDescriptor = getCurrentKeyboardLayoutForInputDevice(identifier);
+        String keyboardLayoutDescriptor;
+        if (useNewSettingsUi()) {
+            if (mCurrentImeInfo == null) {
+                // Haven't received onInputMethodSubtypeChanged() callback from IMMS. Will reload
+                // keyboard layouts once we receive the callback.
+                return null;
+            }
+
+            keyboardLayoutDescriptor = getKeyboardLayoutForInputDeviceInternal(identifier,
+                    mCurrentImeInfo);
+        } else {
+            keyboardLayoutDescriptor = getCurrentKeyboardLayoutForInputDevice(identifier);
+        }
         if (keyboardLayoutDescriptor == null) {
             return null;
         }
@@ -640,6 +683,287 @@
         return result;
     }
 
+    @Nullable
+    public String getKeyboardLayoutForInputDevice(InputDeviceIdentifier identifier,
+            @UserIdInt int userId, @NonNull InputMethodInfo imeInfo,
+            @Nullable InputMethodSubtype imeSubtype) {
+        if (!useNewSettingsUi()) {
+            Slog.e(TAG, "getKeyboardLayoutForInputDevice() API not supported");
+            return null;
+        }
+        InputMethodSubtypeHandle subtypeHandle = InputMethodSubtypeHandle.of(imeInfo, imeSubtype);
+        String layout = getKeyboardLayoutForInputDeviceInternal(identifier,
+                new ImeInfo(userId, subtypeHandle, imeSubtype));
+        if (DEBUG) {
+            Slog.d(TAG, "getKeyboardLayoutForInputDevice() " + identifier.toString() + ", userId : "
+                    + userId + ", subtypeHandle = " + subtypeHandle + " -> " + layout);
+        }
+        return layout;
+    }
+
+    public void setKeyboardLayoutForInputDevice(InputDeviceIdentifier identifier,
+            @UserIdInt int userId, @NonNull InputMethodInfo imeInfo,
+            @Nullable InputMethodSubtype imeSubtype,
+            String keyboardLayoutDescriptor) {
+        if (!useNewSettingsUi()) {
+            Slog.e(TAG, "setKeyboardLayoutForInputDevice() API not supported");
+            return;
+        }
+        Objects.requireNonNull(keyboardLayoutDescriptor,
+                "keyboardLayoutDescriptor must not be null");
+        String key = createLayoutKey(identifier, userId,
+                InputMethodSubtypeHandle.of(imeInfo, imeSubtype));
+        synchronized (mDataStore) {
+            try {
+                // Key for storing into data store = <device descriptor>,<userId>,<subtypeHandle>
+                if (mDataStore.setKeyboardLayout(getLayoutDescriptor(identifier), key,
+                        keyboardLayoutDescriptor)) {
+                    if (DEBUG) {
+                        Slog.d(TAG, "setKeyboardLayoutForInputDevice() " + identifier
+                                + " key: " + key
+                                + " keyboardLayoutDescriptor: " + keyboardLayoutDescriptor);
+                    }
+                    mHandler.sendEmptyMessage(MSG_RELOAD_KEYBOARD_LAYOUTS);
+                }
+            } finally {
+                mDataStore.saveIfNeeded();
+            }
+        }
+    }
+
+    public KeyboardLayout[] getKeyboardLayoutListForInputDevice(InputDeviceIdentifier identifier,
+            @UserIdInt int userId, @NonNull InputMethodInfo imeInfo,
+            @Nullable InputMethodSubtype imeSubtype) {
+        if (!useNewSettingsUi()) {
+            Slog.e(TAG, "getKeyboardLayoutListForInputDevice() API not supported");
+            return new KeyboardLayout[0];
+        }
+        return getKeyboardLayoutListForInputDeviceInternal(identifier, new ImeInfo(userId,
+                InputMethodSubtypeHandle.of(imeInfo, imeSubtype), imeSubtype));
+    }
+
+    private KeyboardLayout[] getKeyboardLayoutListForInputDeviceInternal(
+            InputDeviceIdentifier identifier, ImeInfo imeInfo) {
+        String key = createLayoutKey(identifier, imeInfo.mUserId, imeInfo.mImeSubtypeHandle);
+
+        // Fetch user selected layout and always include it in layout list.
+        String userSelectedLayout;
+        synchronized (mDataStore) {
+            userSelectedLayout = mDataStore.getKeyboardLayout(getLayoutDescriptor(identifier), key);
+        }
+
+        final ArrayList<KeyboardLayout> potentialLayouts = new ArrayList<>();
+        String imeLanguageTag;
+        if (imeInfo.mImeSubtype == null) {
+            imeLanguageTag = "";
+        } else {
+            ULocale imeLocale = imeInfo.mImeSubtype.getPhysicalKeyboardHintLanguageTag();
+            imeLanguageTag = imeLocale != null ? imeLocale.toLanguageTag()
+                    : imeInfo.mImeSubtype.getCanonicalizedLanguageTag();
+        }
+
+        visitAllKeyboardLayouts(new KeyboardLayoutVisitor() {
+            boolean mDeviceSpecificLayoutAvailable;
+
+            @Override
+            public void visitKeyboardLayout(Resources resources,
+                    int keyboardLayoutResId, KeyboardLayout layout) {
+                // Next find any potential layouts that aren't yet enabled for the device. For
+                // devices that have special layouts we assume there's a reason that the generic
+                // layouts don't work for them, so we don't want to return them since it's likely
+                // to result in a poor user experience.
+                if (layout.getVendorId() == identifier.getVendorId()
+                        && layout.getProductId() == identifier.getProductId()) {
+                    if (!mDeviceSpecificLayoutAvailable) {
+                        mDeviceSpecificLayoutAvailable = true;
+                        potentialLayouts.clear();
+                    }
+                    potentialLayouts.add(layout);
+                } else if (layout.getVendorId() == -1 && layout.getProductId() == -1
+                        && !mDeviceSpecificLayoutAvailable && isLayoutCompatibleWithLanguageTag(
+                        layout, imeLanguageTag)) {
+                    potentialLayouts.add(layout);
+                } else if (layout.getDescriptor().equals(userSelectedLayout)) {
+                    potentialLayouts.add(layout);
+                }
+            }
+        });
+        // Sort the Keyboard layouts. This is done first by priority then by label. So, system
+        // layouts will come above 3rd party layouts.
+        Collections.sort(potentialLayouts);
+        return potentialLayouts.toArray(new KeyboardLayout[0]);
+    }
+
+    public void onInputMethodSubtypeChanged(@UserIdInt int userId,
+            @Nullable InputMethodSubtypeHandle subtypeHandle,
+            @Nullable InputMethodSubtype subtype) {
+        if (!useNewSettingsUi()) {
+            Slog.e(TAG, "onInputMethodSubtypeChanged() API not supported");
+            return;
+        }
+        if (subtypeHandle == null) {
+            if (DEBUG) {
+                Slog.d(TAG, "No InputMethod is running, ignoring change");
+            }
+            return;
+        }
+        if (mCurrentImeInfo == null || !subtypeHandle.equals(mCurrentImeInfo.mImeSubtypeHandle)
+                || mCurrentImeInfo.mUserId != userId) {
+            mCurrentImeInfo = new ImeInfo(userId, subtypeHandle, subtype);
+            mHandler.sendEmptyMessage(MSG_RELOAD_KEYBOARD_LAYOUTS);
+            if (DEBUG) {
+                Slog.d(TAG, "InputMethodSubtype changed: userId=" + userId
+                        + " subtypeHandle=" + subtypeHandle);
+            }
+        }
+    }
+
+    @Nullable
+    private String getKeyboardLayoutForInputDeviceInternal(InputDeviceIdentifier identifier,
+            ImeInfo imeInfo) {
+        InputDevice inputDevice = getInputDevice(identifier);
+        if (inputDevice == null || inputDevice.isVirtual() || !inputDevice.isFullKeyboard()) {
+            return null;
+        }
+        String key = createLayoutKey(identifier, imeInfo.mUserId, imeInfo.mImeSubtypeHandle);
+        String layout;
+        synchronized (mDataStore) {
+            layout = mDataStore.getKeyboardLayout(getLayoutDescriptor(identifier), key);
+        }
+        if (layout == null) {
+            synchronized (mKeyboardLayoutCache) {
+                // Check Auto-selected layout cache to see if layout had been previously selected
+                if (mKeyboardLayoutCache.containsKey(key)) {
+                    layout = mKeyboardLayoutCache.get(key);
+                } else {
+                    // NOTE: This list is already filtered based on IME Script code
+                    KeyboardLayout[] layoutList = getKeyboardLayoutListForInputDeviceInternal(
+                            identifier, imeInfo);
+                    // Call auto-matching algorithm to find the best matching layout
+                    layout = getDefaultKeyboardLayoutBasedOnImeInfo(inputDevice, imeInfo,
+                            layoutList);
+                    mKeyboardLayoutCache.put(key, layout);
+                }
+            }
+        }
+        return layout;
+    }
+
+    @Nullable
+    private static String getDefaultKeyboardLayoutBasedOnImeInfo(InputDevice inputDevice,
+            ImeInfo imeInfo, KeyboardLayout[] layoutList) {
+        if (imeInfo.mImeSubtypeHandle == null) {
+            return null;
+        }
+
+        Arrays.sort(layoutList);
+
+        // Check <VendorID, ProductID> matching for explicitly declared custom KCM files.
+        for (KeyboardLayout layout : layoutList) {
+            if (layout.getVendorId() == inputDevice.getVendorId()
+                    && layout.getProductId() == inputDevice.getProductId()) {
+                if (DEBUG) {
+                    Slog.d(TAG,
+                            "getDefaultKeyboardLayoutBasedOnImeInfo() : Layout found based on "
+                                    + "vendor and product Ids. " + inputDevice.getIdentifier()
+                                    + " : " + layout.getDescriptor());
+                }
+                return layout.getDescriptor();
+            }
+        }
+
+        // Check layout type, language tag information from InputDevice for matching
+        String inputLanguageTag = inputDevice.getKeyboardLanguageTag();
+        if (inputLanguageTag != null) {
+            String layoutDesc = getMatchingLayoutForProvidedLanguageTagAndLayoutType(layoutList,
+                    inputLanguageTag, inputDevice.getKeyboardLayoutType());
+
+            if (layoutDesc != null) {
+                if (DEBUG) {
+                    Slog.d(TAG,
+                            "getDefaultKeyboardLayoutBasedOnImeInfo() : Layout found based on "
+                                    + "HW information (Language tag and Layout type). "
+                                    + inputDevice.getIdentifier() + " : " + layoutDesc);
+                }
+                return layoutDesc;
+            }
+        }
+
+        InputMethodSubtype subtype = imeInfo.mImeSubtype;
+        // Can't auto select layout based on IME if subtype or language tag is null
+        if (subtype == null) {
+            return null;
+        }
+
+        // Check layout type, language tag information from IME for matching
+        ULocale pkLocale = subtype.getPhysicalKeyboardHintLanguageTag();
+        String pkLanguageTag =
+                pkLocale != null ? pkLocale.toLanguageTag() : subtype.getCanonicalizedLanguageTag();
+        String layoutDesc = getMatchingLayoutForProvidedLanguageTagAndLayoutType(layoutList,
+                pkLanguageTag, subtype.getPhysicalKeyboardHintLayoutType());
+        if (DEBUG) {
+            Slog.d(TAG,
+                    "getDefaultKeyboardLayoutBasedOnImeInfo() : Layout found based on "
+                            + "IME locale matching. " + inputDevice.getIdentifier() + " : "
+                            + layoutDesc);
+        }
+        return layoutDesc;
+    }
+
+    @Nullable
+    private static String getMatchingLayoutForProvidedLanguageTagAndLayoutType(
+            KeyboardLayout[] layoutList, @NonNull String languageTag, @Nullable String layoutType) {
+        if (layoutType == null || !KeyboardLayout.isLayoutTypeValid(layoutType)) {
+            layoutType = KeyboardLayout.LAYOUT_TYPE_UNDEFINED;
+        }
+        List<KeyboardLayout> layoutsFilteredByLayoutType = new ArrayList<>();
+        for (KeyboardLayout layout : layoutList) {
+            if (layout.getLayoutType().equals(layoutType)) {
+                layoutsFilteredByLayoutType.add(layout);
+            }
+        }
+        String layoutDesc = getMatchingLayoutForProvidedLanguageTag(layoutsFilteredByLayoutType,
+                languageTag);
+        if (layoutDesc != null) {
+            return layoutDesc;
+        }
+
+        return getMatchingLayoutForProvidedLanguageTag(Arrays.asList(layoutList), languageTag);
+    }
+
+    @Nullable
+    private static String getMatchingLayoutForProvidedLanguageTag(List<KeyboardLayout> layoutList,
+            @NonNull String languageTag) {
+        Locale locale = Locale.forLanguageTag(languageTag);
+        String layoutMatchingLanguage = null;
+        String layoutMatchingLanguageAndCountry = null;
+
+        for (KeyboardLayout layout : layoutList) {
+            final LocaleList locales = layout.getLocales();
+            for (int i = 0; i < locales.size(); i++) {
+                final Locale l = locales.get(i);
+                if (l == null) {
+                    continue;
+                }
+                if (l.getLanguage().equals(locale.getLanguage())) {
+                    if (layoutMatchingLanguage == null) {
+                        layoutMatchingLanguage = layout.getDescriptor();
+                    }
+                    if (l.getCountry().equals(locale.getCountry())) {
+                        if (layoutMatchingLanguageAndCountry == null) {
+                            layoutMatchingLanguageAndCountry = layout.getDescriptor();
+                        }
+                        if (l.getVariant().equals(locale.getVariant())) {
+                            return layout.getDescriptor();
+                        }
+                    }
+                }
+            }
+        }
+        return layoutMatchingLanguageAndCountry != null
+                    ? layoutMatchingLanguageAndCountry : layoutMatchingLanguage;
+    }
+
     private void reloadKeyboardLayouts() {
         if (DEBUG) {
             Slog.d(TAG, "Reloading keyboard layouts.");
@@ -734,11 +1058,65 @@
         }
     }
 
+    private boolean useNewSettingsUi() {
+        return FeatureFlagUtils.isEnabled(mContext, FeatureFlagUtils.SETTINGS_NEW_KEYBOARD_UI);
+    }
+
+    @Nullable
     private InputDevice getInputDevice(int deviceId) {
         InputManager inputManager = mContext.getSystemService(InputManager.class);
         return inputManager != null ? inputManager.getInputDevice(deviceId) : null;
     }
 
+    @Nullable
+    private InputDevice getInputDevice(InputDeviceIdentifier identifier) {
+        InputManager inputManager = mContext.getSystemService(InputManager.class);
+        return inputManager != null ? inputManager.getInputDeviceByDescriptor(
+                identifier.getDescriptor()) : null;
+    }
+
+    private static String createLayoutKey(InputDeviceIdentifier identifier, int userId,
+            @NonNull InputMethodSubtypeHandle subtypeHandle) {
+        Objects.requireNonNull(subtypeHandle, "subtypeHandle must not be null");
+        return "layoutDescriptor:" + getLayoutDescriptor(identifier) + ",userId:" + userId
+                + ",subtypeHandle:" + subtypeHandle.toStringHandle();
+    }
+
+    private static boolean isLayoutCompatibleWithLanguageTag(KeyboardLayout layout,
+            @NonNull String languageTag) {
+        final int[] scriptsFromLanguageTag = UScript.getCode(Locale.forLanguageTag(languageTag));
+        if (scriptsFromLanguageTag.length == 0) {
+            // If no scripts inferred from languageTag then allowing the layout
+            return true;
+        }
+        LocaleList locales = layout.getLocales();
+        if (locales.isEmpty()) {
+            // KCM file doesn't have an associated language tag. This can be from
+            // a 3rd party app so need to include it as a potential layout.
+            return true;
+        }
+        for (int i = 0; i < locales.size(); i++) {
+            final Locale locale = locales.get(i);
+            if (locale == null) {
+                continue;
+            }
+            int[] scripts = UScript.getCode(locale);
+            if (scripts != null && haveCommonValue(scripts, scriptsFromLanguageTag)) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    private static boolean haveCommonValue(int[] arr1, int[] arr2) {
+        for (int a1 : arr1) {
+            for (int a2 : arr2) {
+                if (a1 == a2) return true;
+            }
+        }
+        return false;
+    }
+
     private static final class KeyboardLayoutDescriptor {
         public String packageName;
         public String receiverName;
@@ -767,6 +1145,19 @@
         }
     }
 
+    private static class ImeInfo {
+        @UserIdInt int mUserId;
+        @NonNull InputMethodSubtypeHandle mImeSubtypeHandle;
+        @Nullable InputMethodSubtype mImeSubtype;
+
+        ImeInfo(@UserIdInt int userId, @NonNull InputMethodSubtypeHandle imeSubtypeHandle,
+                @Nullable InputMethodSubtype imeSubtype) {
+            mUserId = userId;
+            mImeSubtypeHandle = imeSubtypeHandle;
+            mImeSubtype = imeSubtype;
+        }
+    }
+
     private interface KeyboardLayoutVisitor {
         void visitKeyboardLayout(Resources resources,
                 int keyboardLayoutResId, KeyboardLayout layout);
diff --git a/services/core/java/com/android/server/input/PersistentDataStore.java b/services/core/java/com/android/server/input/PersistentDataStore.java
index 375377a7..a2b18362 100644
--- a/services/core/java/com/android/server/input/PersistentDataStore.java
+++ b/services/core/java/com/android/server/input/PersistentDataStore.java
@@ -18,6 +18,7 @@
 
 import android.annotation.Nullable;
 import android.hardware.input.TouchCalibration;
+import android.util.ArrayMap;
 import android.util.AtomicFile;
 import android.util.Slog;
 import android.util.SparseIntArray;
@@ -42,6 +43,7 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.OptionalInt;
@@ -121,6 +123,7 @@
         return false;
     }
 
+    @Nullable
     public String getCurrentKeyboardLayout(String inputDeviceDescriptor) {
         InputDeviceState state = getInputDeviceState(inputDeviceDescriptor);
         return state != null ? state.getCurrentKeyboardLayout() : null;
@@ -136,6 +139,22 @@
         return false;
     }
 
+    @Nullable
+    public String getKeyboardLayout(String inputDeviceDescriptor, String key) {
+        InputDeviceState state = getInputDeviceState(inputDeviceDescriptor);
+        return state != null ? state.getKeyboardLayout(key) : null;
+    }
+
+    public boolean setKeyboardLayout(String inputDeviceDescriptor, String key,
+            String keyboardLayoutDescriptor) {
+        InputDeviceState state = getOrCreateInputDeviceState(inputDeviceDescriptor);
+        if (state.setKeyboardLayout(key, keyboardLayoutDescriptor)) {
+            setDirty();
+            return true;
+        }
+        return false;
+    }
+
     public String[] getKeyboardLayouts(String inputDeviceDescriptor) {
         InputDeviceState state = getInputDeviceState(inputDeviceDescriptor);
         if (state == null) {
@@ -387,6 +406,8 @@
         private final ArrayList<String> mKeyboardLayouts = new ArrayList<String>();
         private final SparseIntArray mKeyboardBacklightBrightnessMap = new SparseIntArray();
 
+        private final Map<String, String> mKeyboardLayoutMap = new ArrayMap<>();
+
         public TouchCalibration getTouchCalibration(int surfaceRotation) {
             try {
                 return mTouchCalibration[surfaceRotation];
@@ -410,6 +431,15 @@
         }
 
         @Nullable
+        public String getKeyboardLayout(String key) {
+            return mKeyboardLayoutMap.get(key);
+        }
+
+        public boolean setKeyboardLayout(String key, String keyboardLayout) {
+            return !Objects.equals(mKeyboardLayoutMap.put(key, keyboardLayout), keyboardLayout);
+        }
+
+        @Nullable
         public String getCurrentKeyboardLayout() {
             return mCurrentKeyboardLayout;
         }
@@ -507,6 +537,18 @@
                     changed = true;
                 }
             }
+            List<String> removedEntries = new ArrayList<>();
+            for (String key : mKeyboardLayoutMap.keySet()) {
+                if (!availableKeyboardLayouts.contains(mKeyboardLayoutMap.get(key))) {
+                    removedEntries.add(key);
+                }
+            }
+            if (!removedEntries.isEmpty()) {
+                for (String key : removedEntries) {
+                    mKeyboardLayoutMap.remove(key);
+                }
+                changed = true;
+            }
             return changed;
         }
 
@@ -534,6 +576,18 @@
                         }
                         mCurrentKeyboardLayout = descriptor;
                     }
+                } else if (parser.getName().equals("keyed-keyboard-layout")) {
+                    String key = parser.getAttributeValue(null, "key");
+                    if (key == null) {
+                        throw new XmlPullParserException(
+                                "Missing key attribute on keyed-keyboard-layout.");
+                    }
+                    String layout = parser.getAttributeValue(null, "layout");
+                    if (layout == null) {
+                        throw new XmlPullParserException(
+                                "Missing layout attribute on keyed-keyboard-layout.");
+                    }
+                    mKeyboardLayoutMap.put(key, layout);
                 } else if (parser.getName().equals("light-info")) {
                     int lightId = parser.getAttributeInt(null, "light-id");
                     int lightBrightness = parser.getAttributeInt(null, "light-brightness");
@@ -607,6 +661,13 @@
                 serializer.endTag(null, "keyboard-layout");
             }
 
+            for (String key : mKeyboardLayoutMap.keySet()) {
+                serializer.startTag(null, "keyed-keyboard-layout");
+                serializer.attribute(null, "key", key);
+                serializer.attribute(null, "layout", mKeyboardLayoutMap.get(key));
+                serializer.endTag(null, "keyed-keyboard-layout");
+            }
+
             for (int i = 0; i < mKeyboardBacklightBrightnessMap.size(); i++) {
                 serializer.startTag(null, "light-info");
                 serializer.attributeInt(null, "light-id", mKeyboardBacklightBrightnessMap.keyAt(i));
diff --git a/services/tests/servicestests/res/raw/dummy_keyboard_layout.kcm b/services/tests/servicestests/res/raw/dummy_keyboard_layout.kcm
new file mode 100644
index 0000000..ea6bc98
--- /dev/null
+++ b/services/tests/servicestests/res/raw/dummy_keyboard_layout.kcm
@@ -0,0 +1,311 @@
+# Copyright (C) 2023 The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+#
+# English (US) keyboard layout.
+# Unlike the default (generic) keyboard layout, English (US) does not contain any
+# special ALT characters.
+#
+
+type OVERLAY
+
+### ROW 1
+
+key GRAVE {
+    label:                              '`'
+    base:                               '`'
+    shift:                              '~'
+}
+
+key 1 {
+    label:                              '1'
+    base:                               '1'
+    shift:                              '!'
+}
+
+key 2 {
+    label:                              '2'
+    base:                               '2'
+    shift:                              '@'
+}
+
+key 3 {
+    label:                              '3'
+    base:                               '3'
+    shift:                              '#'
+}
+
+key 4 {
+    label:                              '4'
+    base:                               '4'
+    shift:                              '$'
+}
+
+key 5 {
+    label:                              '5'
+    base:                               '5'
+    shift:                              '%'
+}
+
+key 6 {
+    label:                              '6'
+    base:                               '6'
+    shift:                              '^'
+}
+
+key 7 {
+    label:                              '7'
+    base:                               '7'
+    shift:                              '&'
+}
+
+key 8 {
+    label:                              '8'
+    base:                               '8'
+    shift:                              '*'
+}
+
+key 9 {
+    label:                              '9'
+    base:                               '9'
+    shift:                              '('
+}
+
+key 0 {
+    label:                              '0'
+    base:                               '0'
+    shift:                              ')'
+}
+
+key MINUS {
+    label:                              '-'
+    base:                               '-'
+    shift:                              '_'
+}
+
+key EQUALS {
+    label:                              '='
+    base:                               '='
+    shift:                              '+'
+}
+
+### ROW 2
+
+key Q {
+    label:                              'Q'
+    base:                               'q'
+    shift, capslock:                    'Q'
+}
+
+key W {
+    label:                              'W'
+    base:                               'w'
+    shift, capslock:                    'W'
+}
+
+key E {
+    label:                              'E'
+    base:                               'e'
+    shift, capslock:                    'E'
+}
+
+key R {
+    label:                              'R'
+    base:                               'r'
+    shift, capslock:                    'R'
+}
+
+key T {
+    label:                              'T'
+    base:                               't'
+    shift, capslock:                    'T'
+}
+
+key Y {
+    label:                              'Y'
+    base:                               'y'
+    shift, capslock:                    'Y'
+}
+
+key U {
+    label:                              'U'
+    base:                               'u'
+    shift, capslock:                    'U'
+}
+
+key I {
+    label:                              'I'
+    base:                               'i'
+    shift, capslock:                    'I'
+}
+
+key O {
+    label:                              'O'
+    base:                               'o'
+    shift, capslock:                    'O'
+}
+
+key P {
+    label:                              'P'
+    base:                               'p'
+    shift, capslock:                    'P'
+}
+
+key LEFT_BRACKET {
+    label:                              '['
+    base:                               '['
+    shift:                              '{'
+}
+
+key RIGHT_BRACKET {
+    label:                              ']'
+    base:                               ']'
+    shift:                              '}'
+}
+
+key BACKSLASH {
+    label:                              '\\'
+    base:                               '\\'
+    shift:                              '|'
+}
+
+### ROW 3
+
+key A {
+    label:                              'A'
+    base:                               'a'
+    shift, capslock:                    'A'
+}
+
+key S {
+    label:                              'S'
+    base:                               's'
+    shift, capslock:                    'S'
+}
+
+key D {
+    label:                              'D'
+    base:                               'd'
+    shift, capslock:                    'D'
+}
+
+key F {
+    label:                              'F'
+    base:                               'f'
+    shift, capslock:                    'F'
+}
+
+key G {
+    label:                              'G'
+    base:                               'g'
+    shift, capslock:                    'G'
+}
+
+key H {
+    label:                              'H'
+    base:                               'h'
+    shift, capslock:                    'H'
+}
+
+key J {
+    label:                              'J'
+    base:                               'j'
+    shift, capslock:                    'J'
+}
+
+key K {
+    label:                              'K'
+    base:                               'k'
+    shift, capslock:                    'K'
+}
+
+key L {
+    label:                              'L'
+    base:                               'l'
+    shift, capslock:                    'L'
+}
+
+key SEMICOLON {
+    label:                              ';'
+    base:                               ';'
+    shift:                              ':'
+}
+
+key APOSTROPHE {
+    label:                              '\''
+    base:                               '\''
+    shift:                              '"'
+}
+
+### ROW 4
+
+key Z {
+    label:                              'Z'
+    base:                               'z'
+    shift, capslock:                    'Z'
+}
+
+key X {
+    label:                              'X'
+    base:                               'x'
+    shift, capslock:                    'X'
+}
+
+key C {
+    label:                              'C'
+    base:                               'c'
+    shift, capslock:                    'C'
+}
+
+key V {
+    label:                              'V'
+    base:                               'v'
+    shift, capslock:                    'V'
+}
+
+key B {
+    label:                              'B'
+    base:                               'b'
+    shift, capslock:                    'B'
+}
+
+key N {
+    label:                              'N'
+    base:                               'n'
+    shift, capslock:                    'N'
+}
+
+key M {
+    label:                              'M'
+    base:                               'm'
+    shift, capslock:                    'M'
+}
+
+key COMMA {
+    label:                              ','
+    base:                               ','
+    shift:                              '<'
+}
+
+key PERIOD {
+    label:                              '.'
+    base:                               '.'
+    shift:                              '>'
+}
+
+key SLASH {
+    label:                              '/'
+    base:                               '/'
+    shift:                              '?'
+}
diff --git a/services/tests/servicestests/res/xml/keyboard_layouts.xml b/services/tests/servicestests/res/xml/keyboard_layouts.xml
new file mode 100644
index 0000000..b5a05fc
--- /dev/null
+++ b/services/tests/servicestests/res/xml/keyboard_layouts.xml
@@ -0,0 +1,81 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  ~ Copyright (C) 2023 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+
+<keyboard-layouts xmlns:android="http://schemas.android.com/apk/res/android"
+                  xmlns:androidprv="http://schemas.android.com/apk/prv/res/android">
+    <keyboard-layout
+        android:name="keyboard_layout_english_uk"
+        android:label="English (UK)"
+        android:keyboardLayout="@raw/dummy_keyboard_layout"
+        android:keyboardLocale="en-Latn-GB"
+        android:keyboardLayoutType="qwerty" />
+
+    <keyboard-layout
+        android:name="keyboard_layout_english_us"
+        android:label="English (US)"
+        android:keyboardLayout="@raw/dummy_keyboard_layout"
+        android:keyboardLocale="en-Latn,en-Latn-US"
+        android:keyboardLayoutType="qwerty" />
+
+    <keyboard-layout
+        android:name="keyboard_layout_english_us_intl"
+        android:label="English (International)"
+        android:keyboardLayout="@raw/dummy_keyboard_layout"
+        android:keyboardLocale="en-Latn-US"
+        android:keyboardLayoutType="extended" />
+
+    <keyboard-layout
+        android:name="keyboard_layout_english_us_dvorak"
+        android:label="English (Dvorak)"
+        android:keyboardLayout="@raw/dummy_keyboard_layout"
+        android:keyboardLocale="en-Latn-US"
+        android:keyboardLayoutType="dvorak" />
+
+    <keyboard-layout
+        android:name="keyboard_layout_german"
+        android:label="German"
+        android:keyboardLayout="@raw/dummy_keyboard_layout"
+        android:keyboardLocale="de-Latn"
+        android:keyboardLayoutType="qwertz" />
+
+    <keyboard-layout
+        android:name="keyboard_layout_french"
+        android:label="French"
+        android:keyboardLayout="@raw/dummy_keyboard_layout"
+        android:keyboardLocale="fr-Latn-FR"
+        android:keyboardLayoutType="azerty" />
+
+    <keyboard-layout
+        android:name="keyboard_layout_russian_qwerty"
+        android:label="Russian"
+        android:keyboardLayout="@raw/dummy_keyboard_layout"
+        android:keyboardLocale="ru-Cyrl"
+        android:keyboardLayoutType="qwerty" />
+
+    <keyboard-layout
+        android:name="keyboard_layout_russian"
+        android:label="Russian"
+        android:keyboardLayout="@raw/dummy_keyboard_layout"
+        android:keyboardLocale="ru-Cyrl" />
+
+    <keyboard-layout
+        android:name="keyboard_layout_vendorId:1,productId:1"
+        android:label="vendorId:1,productId:1"
+        android:keyboardLayout="@raw/dummy_keyboard_layout"
+        androidprv:vendorId="1"
+        androidprv:productId="1" />
+</keyboard-layouts>
diff --git a/services/tests/servicestests/src/com/android/server/input/KeyboardLayoutManagerTests.kt b/services/tests/servicestests/src/com/android/server/input/KeyboardLayoutManagerTests.kt
new file mode 100644
index 0000000..34540c3
--- /dev/null
+++ b/services/tests/servicestests/src/com/android/server/input/KeyboardLayoutManagerTests.kt
@@ -0,0 +1,820 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.input
+
+import android.content.Context
+import android.content.ContextWrapper
+import android.content.pm.ActivityInfo
+import android.content.pm.ApplicationInfo
+import android.content.pm.PackageManager
+import android.content.pm.ResolveInfo
+import android.content.pm.ServiceInfo
+import android.hardware.input.IInputManager
+import android.hardware.input.InputManager
+import android.hardware.input.KeyboardLayout
+import android.icu.lang.UScript
+import android.icu.util.ULocale
+import android.os.Bundle
+import android.os.test.TestLooper
+import android.platform.test.annotations.Presubmit
+import android.provider.Settings
+import android.view.InputDevice
+import android.view.inputmethod.InputMethodInfo
+import android.view.inputmethod.InputMethodSubtype
+import androidx.test.core.R
+import androidx.test.core.app.ApplicationProvider
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertNotEquals
+import org.junit.Assert.assertNull
+import org.junit.Assert.assertTrue
+import org.junit.Assert.assertThrows
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.mockito.Mock
+import org.mockito.Mockito
+import org.mockito.junit.MockitoJUnit
+import java.io.FileNotFoundException
+import java.io.FileOutputStream
+import java.io.IOException
+import java.io.InputStream
+import java.util.Locale
+
+private fun createKeyboard(
+    deviceId: Int,
+    vendorId: Int,
+    productId: Int,
+    languageTag: String,
+    layoutType: String
+): InputDevice =
+    InputDevice.Builder()
+        .setId(deviceId)
+        .setName("Device $deviceId")
+        .setDescriptor("descriptor $deviceId")
+        .setSources(InputDevice.SOURCE_KEYBOARD)
+        .setKeyboardType(InputDevice.KEYBOARD_TYPE_ALPHABETIC)
+        .setExternal(true)
+        .setVendorId(vendorId)
+        .setProductId(productId)
+        .setKeyboardLanguageTag(languageTag)
+        .setKeyboardLayoutType(layoutType)
+        .build()
+
+/**
+ * Tests for {@link Default UI} and {@link New UI}.
+ *
+ * Build/Install/Run:
+ * atest FrameworksServicesTests:KeyboardLayoutManagerTests
+ */
+@Presubmit
+class KeyboardLayoutManagerTests {
+    companion object {
+        const val DEVICE_ID = 1
+        const val VENDOR_SPECIFIC_DEVICE_ID = 2
+        const val ENGLISH_DVORAK_DEVICE_ID = 3
+        const val USER_ID = 4
+        const val IME_ID = "ime_id"
+        const val PACKAGE_NAME = "KeyboardLayoutManagerTests"
+        const val RECEIVER_NAME = "DummyReceiver"
+        private const val ENGLISH_US_LAYOUT_NAME = "keyboard_layout_english_us"
+        private const val ENGLISH_UK_LAYOUT_NAME = "keyboard_layout_english_uk"
+        private const val VENDOR_SPECIFIC_LAYOUT_NAME = "keyboard_layout_vendorId:1,productId:1"
+    }
+
+    private val ENGLISH_US_LAYOUT_DESCRIPTOR = createLayoutDescriptor(ENGLISH_US_LAYOUT_NAME)
+    private val ENGLISH_UK_LAYOUT_DESCRIPTOR = createLayoutDescriptor(ENGLISH_UK_LAYOUT_NAME)
+    private val VENDOR_SPECIFIC_LAYOUT_DESCRIPTOR =
+        createLayoutDescriptor(VENDOR_SPECIFIC_LAYOUT_NAME)
+
+    @get:Rule
+    val rule = MockitoJUnit.rule()!!
+
+    @Mock
+    private lateinit var iInputManager: IInputManager
+
+    @Mock
+    private lateinit var native: NativeInputManagerService
+
+    @Mock
+    private lateinit var packageManager: PackageManager
+    private lateinit var keyboardLayoutManager: KeyboardLayoutManager
+
+    private lateinit var imeInfo: InputMethodInfo
+    private var nextImeSubtypeId = 0
+    private lateinit var context: Context
+    private lateinit var dataStore: PersistentDataStore
+    private lateinit var testLooper: TestLooper
+
+    // Devices
+    private lateinit var keyboardDevice: InputDevice
+    private lateinit var vendorSpecificKeyboardDevice: InputDevice
+    private lateinit var englishDvorakKeyboardDevice: InputDevice
+
+    @Before
+    fun setup() {
+        context = Mockito.spy(ContextWrapper(ApplicationProvider.getApplicationContext()))
+        dataStore = PersistentDataStore(object : PersistentDataStore.Injector() {
+            override fun openRead(): InputStream? {
+                throw FileNotFoundException()
+            }
+
+            override fun startWrite(): FileOutputStream? {
+                throw IOException()
+            }
+
+            override fun finishWrite(fos: FileOutputStream?, success: Boolean) {}
+        })
+        testLooper = TestLooper()
+        keyboardLayoutManager = KeyboardLayoutManager(context, native, dataStore, testLooper.looper)
+        setupInputDevices()
+        setupBroadcastReceiver()
+        setupIme()
+    }
+
+    private fun setupInputDevices() {
+        val inputManager = InputManager.resetInstance(iInputManager)
+        Mockito.`when`(context.getSystemService(Mockito.eq(Context.INPUT_SERVICE)))
+            .thenReturn(inputManager)
+
+        keyboardDevice = createKeyboard(DEVICE_ID, 0, 0, "", "")
+        vendorSpecificKeyboardDevice = createKeyboard(VENDOR_SPECIFIC_DEVICE_ID, 1, 1, "", "")
+        englishDvorakKeyboardDevice =
+            createKeyboard(ENGLISH_DVORAK_DEVICE_ID, 0, 0, "en", "dvorak")
+        Mockito.`when`(iInputManager.inputDeviceIds)
+            .thenReturn(intArrayOf(DEVICE_ID, VENDOR_SPECIFIC_DEVICE_ID, ENGLISH_DVORAK_DEVICE_ID))
+        Mockito.`when`(iInputManager.getInputDevice(DEVICE_ID)).thenReturn(keyboardDevice)
+        Mockito.`when`(iInputManager.getInputDevice(VENDOR_SPECIFIC_DEVICE_ID))
+            .thenReturn(vendorSpecificKeyboardDevice)
+        Mockito.`when`(iInputManager.getInputDevice(ENGLISH_DVORAK_DEVICE_ID))
+            .thenReturn(englishDvorakKeyboardDevice)
+    }
+
+    private fun setupBroadcastReceiver() {
+        Mockito.`when`(context.packageManager).thenReturn(packageManager)
+
+        val info = createMockReceiver()
+        Mockito.`when`(packageManager.queryBroadcastReceivers(Mockito.any(), Mockito.anyInt()))
+            .thenReturn(listOf(info))
+        Mockito.`when`(packageManager.getReceiverInfo(Mockito.any(), Mockito.anyInt()))
+            .thenReturn(info.activityInfo)
+
+        val resources = context.resources
+        Mockito.`when`(
+            packageManager.getResourcesForApplication(
+                Mockito.any(
+                    ApplicationInfo::class.java
+                )
+            )
+        ).thenReturn(resources)
+    }
+
+    private fun setupIme() {
+        imeInfo = InputMethodInfo(PACKAGE_NAME, RECEIVER_NAME, "", "", 0)
+    }
+
+    @Test
+    fun testDefaultUi_getKeyboardLayouts() {
+        NewSettingsApiFlag(false).use {
+            val keyboardLayouts = keyboardLayoutManager.keyboardLayouts
+            assertNotEquals(
+                "Default UI: Keyboard layout API should not return empty array",
+                0,
+                keyboardLayouts.size
+            )
+            assertTrue(
+                "Default UI: Keyboard layout API should provide English(US) layout",
+                hasLayout(keyboardLayouts, ENGLISH_US_LAYOUT_DESCRIPTOR)
+            )
+        }
+    }
+
+    @Test
+    fun testNewUi_getKeyboardLayouts() {
+        NewSettingsApiFlag(true).use {
+            val keyboardLayouts = keyboardLayoutManager.keyboardLayouts
+            assertNotEquals(
+                "New UI: Keyboard layout API should not return empty array",
+                0,
+                keyboardLayouts.size
+            )
+            assertTrue(
+                "New UI: Keyboard layout API should provide English(US) layout",
+                hasLayout(keyboardLayouts, ENGLISH_US_LAYOUT_DESCRIPTOR)
+            )
+        }
+    }
+
+    @Test
+    fun testDefaultUi_getKeyboardLayoutsForInputDevice() {
+        NewSettingsApiFlag(false).use {
+            val keyboardLayouts =
+                keyboardLayoutManager.getKeyboardLayoutsForInputDevice(keyboardDevice.identifier)
+            assertNotEquals(
+                "Default UI: getKeyboardLayoutsForInputDevice API should not return empty array",
+                0,
+                keyboardLayouts.size
+            )
+            assertTrue(
+                "Default UI: getKeyboardLayoutsForInputDevice API should provide English(US) " +
+                        "layout",
+                hasLayout(keyboardLayouts, ENGLISH_US_LAYOUT_DESCRIPTOR)
+            )
+
+            val vendorSpecificKeyboardLayouts =
+                keyboardLayoutManager.getKeyboardLayoutsForInputDevice(
+                    vendorSpecificKeyboardDevice.identifier
+                )
+            assertEquals(
+                "Default UI: getKeyboardLayoutsForInputDevice API should return only vendor " +
+                        "specific layout",
+                1,
+                vendorSpecificKeyboardLayouts.size
+            )
+            assertEquals(
+                "Default UI: getKeyboardLayoutsForInputDevice API should return vendor specific " +
+                        "layout",
+                VENDOR_SPECIFIC_LAYOUT_DESCRIPTOR,
+                vendorSpecificKeyboardLayouts[0].descriptor
+            )
+        }
+    }
+
+    @Test
+    fun testNewUi_getKeyboardLayoutsForInputDevice() {
+        NewSettingsApiFlag(true).use {
+            val keyboardLayouts =
+                keyboardLayoutManager.getKeyboardLayoutsForInputDevice(keyboardDevice.identifier)
+            assertEquals(
+                "New UI: getKeyboardLayoutsForInputDevice API should always return empty array",
+                0,
+                keyboardLayouts.size
+            )
+        }
+    }
+
+    @Test
+    fun testDefaultUi_getSetCurrentKeyboardLayoutForInputDevice() {
+        NewSettingsApiFlag(false).use {
+            assertNull(
+                "Default UI: getCurrentKeyboardLayoutForInputDevice API should return null if " +
+                        "nothing was set",
+                keyboardLayoutManager.getCurrentKeyboardLayoutForInputDevice(
+                    keyboardDevice.identifier
+                )
+            )
+
+            keyboardLayoutManager.setCurrentKeyboardLayoutForInputDevice(
+                keyboardDevice.identifier,
+                ENGLISH_US_LAYOUT_DESCRIPTOR
+            )
+            val keyboardLayout =
+                keyboardLayoutManager.getCurrentKeyboardLayoutForInputDevice(
+                    keyboardDevice.identifier
+                )
+            assertEquals(
+                "Default UI: getCurrentKeyboardLayoutForInputDevice API should return the set " +
+                        "layout",
+                ENGLISH_US_LAYOUT_DESCRIPTOR,
+                keyboardLayout
+            )
+        }
+    }
+
+    @Test
+    fun testNewUi_getSetCurrentKeyboardLayoutForInputDevice() {
+        NewSettingsApiFlag(true).use {
+            keyboardLayoutManager.setCurrentKeyboardLayoutForInputDevice(
+                keyboardDevice.identifier,
+                ENGLISH_US_LAYOUT_DESCRIPTOR
+            )
+            assertNull(
+                "New UI: getCurrentKeyboardLayoutForInputDevice API should always return null " +
+                        "even after setCurrentKeyboardLayoutForInputDevice",
+                keyboardLayoutManager.getCurrentKeyboardLayoutForInputDevice(
+                    keyboardDevice.identifier
+                )
+            )
+        }
+    }
+
+    @Test
+    fun testDefaultUi_getEnabledKeyboardLayoutsForInputDevice() {
+        NewSettingsApiFlag(false).use {
+            keyboardLayoutManager.addKeyboardLayoutForInputDevice(
+                keyboardDevice.identifier, ENGLISH_US_LAYOUT_DESCRIPTOR
+            )
+
+            val keyboardLayouts =
+                keyboardLayoutManager.getEnabledKeyboardLayoutsForInputDevice(
+                    keyboardDevice.identifier
+                )
+            assertEquals(
+                "Default UI: getEnabledKeyboardLayoutsForInputDevice API should return added " +
+                        "layout",
+                1,
+                keyboardLayouts.size
+            )
+            assertEquals(
+                "Default UI: getEnabledKeyboardLayoutsForInputDevice API should return " +
+                        "English(US) layout",
+                ENGLISH_US_LAYOUT_DESCRIPTOR,
+                keyboardLayouts[0]
+            )
+            assertEquals(
+                "Default UI: getCurrentKeyboardLayoutForInputDevice API should return " +
+                        "English(US) layout (Auto select the first enabled layout)",
+                ENGLISH_US_LAYOUT_DESCRIPTOR,
+                keyboardLayoutManager.getCurrentKeyboardLayoutForInputDevice(
+                    keyboardDevice.identifier
+                )
+            )
+
+            keyboardLayoutManager.removeKeyboardLayoutForInputDevice(
+                keyboardDevice.identifier, ENGLISH_US_LAYOUT_DESCRIPTOR
+            )
+            assertEquals(
+                "Default UI: getKeyboardLayoutsForInputDevice API should return 0 layouts",
+                0,
+                keyboardLayoutManager.getEnabledKeyboardLayoutsForInputDevice(
+                    keyboardDevice.identifier
+                ).size
+            )
+            assertNull(
+                "Default UI: getCurrentKeyboardLayoutForInputDevice API should return null after " +
+                        "the enabled layout is removed",
+                keyboardLayoutManager.getCurrentKeyboardLayoutForInputDevice(
+                    keyboardDevice.identifier
+                )
+            )
+        }
+    }
+
+    @Test
+    fun testNewUi_getEnabledKeyboardLayoutsForInputDevice() {
+        NewSettingsApiFlag(true).use {
+            keyboardLayoutManager.addKeyboardLayoutForInputDevice(
+                keyboardDevice.identifier, ENGLISH_US_LAYOUT_DESCRIPTOR
+            )
+
+            assertEquals(
+                "New UI: getEnabledKeyboardLayoutsForInputDevice API should return always return " +
+                        "an empty array",
+                0,
+                keyboardLayoutManager.getEnabledKeyboardLayoutsForInputDevice(
+                    keyboardDevice.identifier
+                ).size
+            )
+            assertNull(
+                "New UI: getCurrentKeyboardLayoutForInputDevice API should always return null",
+                keyboardLayoutManager.getCurrentKeyboardLayoutForInputDevice(
+                    keyboardDevice.identifier
+                )
+            )
+        }
+    }
+
+    @Test
+    fun testDefaultUi_switchKeyboardLayout() {
+        NewSettingsApiFlag(false).use {
+            keyboardLayoutManager.addKeyboardLayoutForInputDevice(
+                keyboardDevice.identifier, ENGLISH_US_LAYOUT_DESCRIPTOR
+            )
+            keyboardLayoutManager.addKeyboardLayoutForInputDevice(
+                keyboardDevice.identifier, ENGLISH_UK_LAYOUT_DESCRIPTOR
+            )
+            assertEquals(
+                "Default UI: getCurrentKeyboardLayoutForInputDevice API should return " +
+                        "English(US) layout",
+                ENGLISH_US_LAYOUT_DESCRIPTOR,
+                keyboardLayoutManager.getCurrentKeyboardLayoutForInputDevice(
+                    keyboardDevice.identifier
+                )
+            )
+
+            keyboardLayoutManager.switchKeyboardLayout(DEVICE_ID, 1)
+
+            // Throws null pointer because trying to show toast using TestLooper
+            assertThrows(NullPointerException::class.java) { testLooper.dispatchAll() }
+            assertEquals("Default UI: getCurrentKeyboardLayoutForInputDevice API should return " +
+                    "English(UK) layout",
+                ENGLISH_UK_LAYOUT_DESCRIPTOR,
+                keyboardLayoutManager.getCurrentKeyboardLayoutForInputDevice(
+                    keyboardDevice.identifier
+                )
+            )
+        }
+    }
+
+    @Test
+    fun testNewUi_switchKeyboardLayout() {
+        NewSettingsApiFlag(true).use {
+            keyboardLayoutManager.addKeyboardLayoutForInputDevice(
+                keyboardDevice.identifier, ENGLISH_US_LAYOUT_DESCRIPTOR
+            )
+            keyboardLayoutManager.addKeyboardLayoutForInputDevice(
+                keyboardDevice.identifier, ENGLISH_UK_LAYOUT_DESCRIPTOR
+            )
+
+            keyboardLayoutManager.switchKeyboardLayout(DEVICE_ID, 1)
+            testLooper.dispatchAll()
+
+            assertNull("New UI: getCurrentKeyboardLayoutForInputDevice API should always return " +
+                    "null",
+                keyboardLayoutManager.getCurrentKeyboardLayoutForInputDevice(
+                    keyboardDevice.identifier
+                )
+            )
+        }
+    }
+
+    @Test
+    fun testDefaultUi_getKeyboardLayout() {
+        NewSettingsApiFlag(false).use {
+            val keyboardLayout =
+                keyboardLayoutManager.getKeyboardLayout(ENGLISH_US_LAYOUT_DESCRIPTOR)
+            assertEquals("Default UI: getKeyboardLayout API should return correct Layout from " +
+                    "available layouts",
+                ENGLISH_US_LAYOUT_DESCRIPTOR,
+                keyboardLayout!!.descriptor
+            )
+        }
+    }
+
+    @Test
+    fun testNewUi_getKeyboardLayout() {
+        NewSettingsApiFlag(true).use {
+            val keyboardLayout =
+                keyboardLayoutManager.getKeyboardLayout(ENGLISH_US_LAYOUT_DESCRIPTOR)
+            assertEquals("New UI: getKeyboardLayout API should return correct Layout from " +
+                    "available layouts",
+                ENGLISH_US_LAYOUT_DESCRIPTOR,
+                keyboardLayout!!.descriptor
+            )
+        }
+    }
+
+    @Test
+    fun testDefaultUi_getSetKeyboardLayoutForInputDevice_WithImeInfo() {
+        NewSettingsApiFlag(false).use {
+            val imeSubtype = createImeSubtype()
+            keyboardLayoutManager.setKeyboardLayoutForInputDevice(
+                keyboardDevice.identifier, USER_ID, imeInfo, imeSubtype,
+                ENGLISH_UK_LAYOUT_DESCRIPTOR
+            )
+            val keyboardLayout =
+                keyboardLayoutManager.getKeyboardLayoutForInputDevice(
+                    keyboardDevice.identifier, USER_ID, imeInfo, imeSubtype
+                )
+            assertNull(
+                "Default UI: getKeyboardLayoutForInputDevice API should always return null",
+                keyboardLayout
+            )
+        }
+    }
+
+    @Test
+    fun testNewUi_getSetKeyboardLayoutForInputDevice_withImeInfo() {
+        NewSettingsApiFlag(true).use {
+            val imeSubtype = createImeSubtype()
+
+            keyboardLayoutManager.setKeyboardLayoutForInputDevice(
+                keyboardDevice.identifier, USER_ID, imeInfo, imeSubtype,
+                ENGLISH_UK_LAYOUT_DESCRIPTOR
+            )
+            assertEquals(
+                "New UI: getKeyboardLayoutForInputDevice API should return the set layout",
+                ENGLISH_UK_LAYOUT_DESCRIPTOR,
+                keyboardLayoutManager.getKeyboardLayoutForInputDevice(
+                    keyboardDevice.identifier, USER_ID, imeInfo, imeSubtype
+                )
+            )
+
+            // This should replace previously set layout
+            keyboardLayoutManager.setKeyboardLayoutForInputDevice(
+                keyboardDevice.identifier, USER_ID, imeInfo, imeSubtype,
+                ENGLISH_US_LAYOUT_DESCRIPTOR
+            )
+            assertEquals(
+                "New UI: getKeyboardLayoutForInputDevice API should return the last set layout",
+                ENGLISH_US_LAYOUT_DESCRIPTOR,
+                keyboardLayoutManager.getKeyboardLayoutForInputDevice(
+                    keyboardDevice.identifier, USER_ID, imeInfo, imeSubtype
+                )
+            )
+        }
+    }
+
+    @Test
+    fun testDefaultUi_getKeyboardLayoutListForInputDevice() {
+        NewSettingsApiFlag(false).use {
+            val keyboardLayouts =
+                keyboardLayoutManager.getKeyboardLayoutListForInputDevice(
+                    keyboardDevice.identifier, USER_ID, imeInfo,
+                    createImeSubtype()
+                )
+            assertEquals("Default UI: getKeyboardLayoutListForInputDevice API should always " +
+                    "return empty array",
+                0,
+                keyboardLayouts.size
+            )
+        }
+    }
+
+    @Test
+    fun testNewUi_getKeyboardLayoutListForInputDevice() {
+        NewSettingsApiFlag(true).use {
+            // Check Layouts for "hi-Latn". It should return all 'Latn' keyboard layouts
+            var keyboardLayouts =
+                keyboardLayoutManager.getKeyboardLayoutListForInputDevice(
+                    keyboardDevice.identifier, USER_ID, imeInfo,
+                    createImeSubtypeForLanguageTag("hi-Latn")
+                )
+            assertNotEquals(
+                "New UI: getKeyboardLayoutListForInputDevice API should return the list of " +
+                        "supported layouts with matching script code",
+                0,
+                keyboardLayouts.size
+            )
+
+            val englishScripts = UScript.getCode(Locale.forLanguageTag("hi-Latn"))
+            for (kl in keyboardLayouts) {
+                var isCompatible = false
+                for (i in 0 until kl.locales.size()) {
+                    val locale: Locale = kl.locales.get(i) ?: continue
+                    val scripts = UScript.getCode(locale)
+                    if (scripts != null && areScriptsCompatible(scripts, englishScripts)) {
+                        isCompatible = true
+                        break
+                    }
+                }
+                assertTrue(
+                    "New UI: getKeyboardLayoutListForInputDevice API should only return " +
+                            "compatible layouts but found " + kl.descriptor,
+                    isCompatible
+                )
+            }
+
+            // Check Layouts for "hi" which by default uses 'Deva' script.
+            keyboardLayouts =
+                keyboardLayoutManager.getKeyboardLayoutListForInputDevice(
+                    keyboardDevice.identifier, USER_ID, imeInfo,
+                    createImeSubtypeForLanguageTag("hi")
+                )
+            assertEquals("New UI: getKeyboardLayoutListForInputDevice API should return empty " +
+                    "list if no supported layouts available",
+                0,
+                keyboardLayouts.size
+            )
+
+            // If user manually selected some layout, always provide it in the layout list
+            val imeSubtype = createImeSubtypeForLanguageTag("hi")
+            keyboardLayoutManager.setKeyboardLayoutForInputDevice(
+                keyboardDevice.identifier, USER_ID, imeInfo, imeSubtype,
+                ENGLISH_US_LAYOUT_DESCRIPTOR
+            )
+            keyboardLayouts =
+                keyboardLayoutManager.getKeyboardLayoutListForInputDevice(
+                    keyboardDevice.identifier, USER_ID, imeInfo,
+                    imeSubtype
+                )
+            assertEquals("New UI: getKeyboardLayoutListForInputDevice API should return user " +
+                    "selected layout even if the script is incompatible with IME",
+                    1,
+                keyboardLayouts.size
+            )
+        }
+    }
+
+    @Test
+    fun testNewUi_getDefaultKeyboardLayoutForInputDevice_withImeLanguageTag() {
+        NewSettingsApiFlag(true).use {
+            assertCorrectLayout(
+                keyboardDevice,
+                createImeSubtypeForLanguageTag("en-US"),
+                ENGLISH_US_LAYOUT_DESCRIPTOR
+            )
+            assertCorrectLayout(
+                keyboardDevice,
+                createImeSubtypeForLanguageTag("en-GB"),
+                ENGLISH_UK_LAYOUT_DESCRIPTOR
+            )
+            assertCorrectLayout(
+                keyboardDevice,
+                createImeSubtypeForLanguageTag("de"),
+                createLayoutDescriptor("keyboard_layout_german")
+            )
+            assertCorrectLayout(
+                keyboardDevice,
+                createImeSubtypeForLanguageTag("fr-FR"),
+                createLayoutDescriptor("keyboard_layout_french")
+            )
+            assertCorrectLayout(
+                keyboardDevice,
+                createImeSubtypeForLanguageTag("ru"),
+                createLayoutDescriptor("keyboard_layout_russian")
+            )
+            assertNull(
+                "New UI: getDefaultKeyboardLayoutForInputDevice should return null when no " +
+                        "layout available",
+                keyboardLayoutManager.getKeyboardLayoutForInputDevice(
+                    keyboardDevice.identifier, USER_ID, imeInfo,
+                    createImeSubtypeForLanguageTag("it")
+                )
+            )
+            assertNull(
+                "New UI: getDefaultKeyboardLayoutForInputDevice should return null when no " +
+                        "layout for script code is available",
+                keyboardLayoutManager.getKeyboardLayoutForInputDevice(
+                    keyboardDevice.identifier, USER_ID, imeInfo,
+                    createImeSubtypeForLanguageTag("en-Deva")
+                )
+            )
+        }
+    }
+
+    @Test
+    fun testNewUi_getDefaultKeyboardLayoutForInputDevice_withImeLanguageTagAndLayoutType() {
+        NewSettingsApiFlag(true).use {
+            assertCorrectLayout(
+                keyboardDevice,
+                createImeSubtypeForLanguageTagAndLayoutType("en-US", "qwerty"),
+                ENGLISH_US_LAYOUT_DESCRIPTOR
+            )
+            assertCorrectLayout(
+                keyboardDevice,
+                createImeSubtypeForLanguageTagAndLayoutType("en-US", "dvorak"),
+                createLayoutDescriptor("keyboard_layout_english_us_dvorak")
+            )
+            // Try to match layout type even if country doesn't match
+            assertCorrectLayout(
+                keyboardDevice,
+                createImeSubtypeForLanguageTagAndLayoutType("en-GB", "dvorak"),
+                createLayoutDescriptor("keyboard_layout_english_us_dvorak")
+            )
+            // Choose layout based on layout type priority, if layout type is not provided by IME
+            // (Qwerty > Dvorak > Extended)
+            assertCorrectLayout(
+                keyboardDevice,
+                createImeSubtypeForLanguageTagAndLayoutType("en-US", ""),
+                ENGLISH_US_LAYOUT_DESCRIPTOR
+            )
+            assertCorrectLayout(
+                keyboardDevice,
+                createImeSubtypeForLanguageTagAndLayoutType("en-GB", "qwerty"),
+                ENGLISH_UK_LAYOUT_DESCRIPTOR
+            )
+            assertCorrectLayout(
+                keyboardDevice,
+                createImeSubtypeForLanguageTagAndLayoutType("de", "qwertz"),
+                createLayoutDescriptor("keyboard_layout_german")
+            )
+            // Wrong layout type should match with language if provided layout type not available
+            assertCorrectLayout(
+                keyboardDevice,
+                createImeSubtypeForLanguageTagAndLayoutType("de", "qwerty"),
+                createLayoutDescriptor("keyboard_layout_german")
+            )
+            assertCorrectLayout(
+                keyboardDevice,
+                createImeSubtypeForLanguageTagAndLayoutType("fr-FR", "azerty"),
+                createLayoutDescriptor("keyboard_layout_french")
+            )
+            assertCorrectLayout(
+                keyboardDevice,
+                createImeSubtypeForLanguageTagAndLayoutType("ru", "qwerty"),
+                createLayoutDescriptor("keyboard_layout_russian_qwerty")
+            )
+            // If layout type is empty then prioritize KCM with empty layout type
+            assertCorrectLayout(
+                keyboardDevice,
+                createImeSubtypeForLanguageTagAndLayoutType("ru", ""),
+                createLayoutDescriptor("keyboard_layout_russian")
+            )
+            assertNull("New UI: getDefaultKeyboardLayoutForInputDevice should return null when " +
+                    "no layout for script code is available",
+                keyboardLayoutManager.getKeyboardLayoutForInputDevice(
+                    keyboardDevice.identifier, USER_ID, imeInfo,
+                    createImeSubtypeForLanguageTagAndLayoutType("en-Deva-US", "")
+                )
+            )
+        }
+    }
+
+    @Test
+    fun testNewUi_getDefaultKeyboardLayoutForInputDevice_withHwLanguageTagAndLayoutType() {
+        NewSettingsApiFlag(true).use {
+            // Should return English dvorak even if IME current layout is qwerty, since HW says the
+            // keyboard is a Dvorak keyboard
+            assertCorrectLayout(
+                englishDvorakKeyboardDevice,
+                createImeSubtypeForLanguageTagAndLayoutType("en", "qwerty"),
+                createLayoutDescriptor("keyboard_layout_english_us_dvorak")
+            )
+
+            // Fallback to IME information if the HW provided layout script is incompatible with the
+            // provided IME subtype
+            assertCorrectLayout(
+                englishDvorakKeyboardDevice,
+                createImeSubtypeForLanguageTagAndLayoutType("ru", ""),
+                createLayoutDescriptor("keyboard_layout_russian")
+            )
+        }
+    }
+
+    private fun assertCorrectLayout(
+        device: InputDevice,
+        imeSubtype: InputMethodSubtype,
+        expectedLayout: String
+    ) {
+        assertEquals(
+            "New UI: getDefaultKeyboardLayoutForInputDevice should return $expectedLayout",
+            expectedLayout,
+            keyboardLayoutManager.getKeyboardLayoutForInputDevice(
+                device.identifier, USER_ID, imeInfo, imeSubtype
+            )
+        )
+    }
+
+    private fun createImeSubtype(): InputMethodSubtype =
+        InputMethodSubtype.InputMethodSubtypeBuilder().setSubtypeId(nextImeSubtypeId++).build()
+
+    private fun createImeSubtypeForLanguageTag(languageTag: String): InputMethodSubtype =
+        InputMethodSubtype.InputMethodSubtypeBuilder().setSubtypeId(nextImeSubtypeId++)
+            .setLanguageTag(languageTag).build()
+
+    private fun createImeSubtypeForLanguageTagAndLayoutType(
+        languageTag: String,
+        layoutType: String
+    ): InputMethodSubtype =
+        InputMethodSubtype.InputMethodSubtypeBuilder().setSubtypeId(nextImeSubtypeId++)
+            .setPhysicalKeyboardHint(ULocale.forLanguageTag(languageTag), layoutType).build()
+
+    private fun hasLayout(layoutList: Array<KeyboardLayout>, layoutDesc: String): Boolean {
+        for (kl in layoutList) {
+            if (kl.descriptor == layoutDesc) {
+                return true
+            }
+        }
+        return false
+    }
+
+    private fun createLayoutDescriptor(keyboardName: String): String =
+        "$PACKAGE_NAME/$RECEIVER_NAME/$keyboardName"
+
+    private fun areScriptsCompatible(scriptList1: IntArray, scriptList2: IntArray): Boolean {
+        for (s1 in scriptList1) {
+            for (s2 in scriptList2) {
+                if (s1 == s2) return true
+            }
+        }
+        return false
+    }
+
+    private fun createMockReceiver(): ResolveInfo {
+        val info = ResolveInfo()
+        info.activityInfo = ActivityInfo()
+        info.activityInfo.packageName = PACKAGE_NAME
+        info.activityInfo.name = RECEIVER_NAME
+        info.activityInfo.applicationInfo = ApplicationInfo()
+        info.activityInfo.metaData = Bundle()
+        info.activityInfo.metaData.putInt(
+            InputManager.META_DATA_KEYBOARD_LAYOUTS,
+            R.xml.keyboard_layouts
+        )
+        info.serviceInfo = ServiceInfo()
+        info.serviceInfo.packageName = PACKAGE_NAME
+        info.serviceInfo.name = RECEIVER_NAME
+        return info
+    }
+
+    private inner class NewSettingsApiFlag constructor(enabled: Boolean) : AutoCloseable {
+        init {
+            Settings.Global.putString(
+                context.contentResolver,
+                "settings_new_keyboard_ui", enabled.toString()
+            )
+        }
+
+        override fun close() {
+            Settings.Global.putString(
+                context.contentResolver,
+                "settings_new_keyboard_ui",
+                ""
+            )
+        }
+    }
+}
\ No newline at end of file