diff --git a/services/core/java/com/android/server/am/ActivityTaskManagerService.java b/services/core/java/com/android/server/am/ActivityTaskManagerService.java
index f79d9aa..4897c05 100644
--- a/services/core/java/com/android/server/am/ActivityTaskManagerService.java
+++ b/services/core/java/com/android/server/am/ActivityTaskManagerService.java
@@ -19,7 +19,6 @@
 import static android.Manifest.permission.BIND_VOICE_INTERACTION;
 import static android.Manifest.permission.CHANGE_CONFIGURATION;
 import static android.Manifest.permission.CONTROL_REMOTE_APP_TRANSITION_ANIMATIONS;
-import static android.Manifest.permission.FILTER_EVENTS;
 import static android.Manifest.permission.INTERNAL_SYSTEM_WINDOW;
 import static android.Manifest.permission.MANAGE_ACTIVITY_STACKS;
 import static android.Manifest.permission.READ_FRAME_BUFFER;
@@ -28,9 +27,9 @@
 import static android.Manifest.permission.STOP_APP_SWITCHES;
 import static android.app.ActivityManager.LOCK_TASK_MODE_NONE;
 import static android.app.ActivityManagerInternal.ALLOW_FULL_ONLY;
+import static android.app.ActivityTaskManager.INVALID_TASK_ID;
 import static android.app.ActivityTaskManager.RESIZE_MODE_PRESERVE_WINDOW;
 import static android.app.ActivityTaskManager.SPLIT_SCREEN_CREATE_MODE_TOP_OR_LEFT;
-import static android.app.AppOpsManager.OP_NONE;
 import static android.app.WindowConfiguration.ACTIVITY_TYPE_UNDEFINED;
 import static android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM;
 import static android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN;
@@ -71,6 +70,10 @@
 import static android.view.WindowManager.TRANSIT_TASK_OPEN;
 import static android.view.WindowManager.TRANSIT_TASK_TO_FRONT;
 
+import static com.android.server.am.ActivityManagerService.ANR_TRACE_DIR;
+import static com.android.server.am.ActivityManagerService.MY_PID;
+import static com.android.server.am.ActivityManagerService.STOCK_PM_FLAGS;
+import static com.android.server.am.ActivityManagerService.dumpStackTraces;
 import static com.android.server.am.ActivityManagerServiceDumpProcessesProto.CONFIG_WILL_CHANGE;
 import static com.android.server.am.ActivityManagerServiceDumpProcessesProto.CONTROLLER;
 import static com.android.server.am.ActivityManagerServiceDumpProcessesProto.CURRENT_TRACKER;
@@ -81,12 +84,21 @@
 import static com.android.server.am.ActivityManagerServiceDumpProcessesProto.HOME_PROC;
 import static com.android.server.am.ActivityManagerServiceDumpProcessesProto.LAUNCHING_ACTIVITY;
 import static com.android.server.am.ActivityManagerServiceDumpProcessesProto.PREVIOUS_PROC;
-import static com.android.server.am.ActivityManagerServiceDumpProcessesProto.PREVIOUS_PROC_VISIBLE_TIME_MS;
+import static com.android.server.am.ActivityManagerServiceDumpProcessesProto
+        .PREVIOUS_PROC_VISIBLE_TIME_MS;
 import static com.android.server.am.ActivityManagerServiceDumpProcessesProto.SCREEN_COMPAT_PACKAGES;
-import static com.android.server.am.ActivityManagerServiceDumpProcessesProto.ScreenCompatPackage.MODE;
-import static com.android.server.am.ActivityManagerServiceDumpProcessesProto.ScreenCompatPackage.PACKAGE;
+import static com.android.server.am.ActivityManagerServiceDumpProcessesProto.ScreenCompatPackage
+        .MODE;
+import static com.android.server.am.ActivityManagerServiceDumpProcessesProto.ScreenCompatPackage
+        .PACKAGE;
+import static com.android.server.am.ActivityStack.REMOVE_TASK_MODE_DESTROYING;
+import static com.android.server.am.ActivityStackSupervisor.DEFER_RESUME;
+import static com.android.server.am.ActivityStackSupervisor.MATCH_TASK_IN_STACKS_ONLY;
+import static com.android.server.am.ActivityStackSupervisor.MATCH_TASK_IN_STACKS_OR_RECENT_TASKS;
+import static com.android.server.am.ActivityStackSupervisor.ON_TOP;
+import static com.android.server.am.ActivityStackSupervisor.PRESERVE_WINDOWS;
+import static com.android.server.am.ActivityStackSupervisor.REMOVE_FROM_RECENTS;
 import static com.android.server.am.ActivityTaskManagerDebugConfig.DEBUG_ALL;
-import static com.android.server.am.ActivityManagerService.ANR_TRACE_DIR;
 import static com.android.server.am.ActivityTaskManagerDebugConfig.DEBUG_CONFIGURATION;
 import static com.android.server.am.ActivityTaskManagerDebugConfig.DEBUG_FOCUS;
 import static com.android.server.am.ActivityTaskManagerDebugConfig.DEBUG_IMMERSIVE;
@@ -104,19 +116,8 @@
 import static com.android.server.am.ActivityTaskManagerDebugConfig.POSTFIX_VISIBILITY;
 import static com.android.server.am.ActivityTaskManagerDebugConfig.TAG_ATM;
 import static com.android.server.am.ActivityTaskManagerDebugConfig.TAG_WITH_CLASS_NAME;
-import static com.android.server.am.ActivityManagerService.MY_PID;
-import static com.android.server.am.ActivityManagerService.STOCK_PM_FLAGS;
-import static com.android.server.am.ActivityManagerService.dumpStackTraces;
-import static com.android.server.am.ActivityStack.REMOVE_TASK_MODE_DESTROYING;
-import static com.android.server.am.ActivityStackSupervisor.DEFER_RESUME;
-import static com.android.server.am.ActivityStackSupervisor.MATCH_TASK_IN_STACKS_ONLY;
-import static com.android.server.am.ActivityStackSupervisor.MATCH_TASK_IN_STACKS_OR_RECENT_TASKS;
-import static com.android.server.am.ActivityStackSupervisor.ON_TOP;
-import static com.android.server.am.ActivityStackSupervisor.PRESERVE_WINDOWS;
-import static com.android.server.am.ActivityStackSupervisor.REMOVE_FROM_RECENTS;
 import static com.android.server.am.ActivityTaskManagerService.H.REPORT_TIME_TRACKER_MSG;
 import static com.android.server.am.ActivityTaskManagerService.UiHandler.DISMISS_DIALOG_UI_MSG;
-import static android.app.ActivityTaskManager.INVALID_TASK_ID;
 import static com.android.server.am.TaskRecord.LOCK_TASK_AUTH_DONT_LOCK;
 import static com.android.server.am.TaskRecord.REPARENT_KEEP_STACK_AT_FRONT;
 import static com.android.server.am.TaskRecord.REPARENT_LEAVE_STACK_IN_PLACE;
@@ -4761,7 +4762,8 @@
         }
 
         for (int i = mPidMap.size() - 1; i >= 0; i--) {
-            WindowProcessController app = mPidMap.get(mPidMap.keyAt(i));
+            final int pid = mPidMap.keyAt(i);
+            final WindowProcessController app = mPidMap.get(pid);
             if (DEBUG_CONFIGURATION) {
                 Slog.v(TAG_CONFIGURATION, "Update process config of "
                         + app.mName + " to new config " + configCopy);
@@ -6015,7 +6017,7 @@
          * @param displayId The ID of the display showing the IME.
          */
         @Override
-        public void onImeWindowSetOnDisplay(int pid, int displayId) {
+        public void onImeWindowSetOnDisplay(final int pid, final int displayId) {
             if (pid == MY_PID || pid < 0) {
                 if (DEBUG_CONFIGURATION) {
                     Slog.w(TAG,
@@ -6025,29 +6027,28 @@
             }
             mH.post(() -> {
                 synchronized (mGlobalLock) {
-                    // Check if display is initialized in AM.
-                    if (!mStackSupervisor.isDisplayAdded(displayId)) {
-                        // Call come when display is not yet added or has already been removed.
+                    final ActivityDisplay activityDisplay =
+                            mStackSupervisor.getActivityDisplay(displayId);
+                    if (activityDisplay == null) {
+                        // Call might come when display is not yet added or has been removed.
                         if (DEBUG_CONFIGURATION) {
                             Slog.w(TAG, "Trying to update display configuration for non-existing "
-                                            + "displayId=" + displayId);
+                                    + "displayId=" + displayId);
                         }
                         return;
                     }
-                    final WindowProcessController imeProcess = mPidMap.get(pid);
-                    if (imeProcess == null) {
+                    final WindowProcessController process = mPidMap.get(pid);
+                    if (process == null) {
                         if (DEBUG_CONFIGURATION) {
-                            Slog.w(TAG, "Trying to update display configuration for invalid pid: "
-                                            + pid);
+                            Slog.w(TAG, "Trying to update display configuration for invalid "
+                                    + "process, pid=" + pid);
                         }
                         return;
                     }
-                    // Fetch the current override configuration of the display and set it to the
-                    // process global configuration.
-                    imeProcess.onConfigurationChanged(
-                            mStackSupervisor.getDisplayOverrideConfiguration(displayId));
+                    process.registerDisplayConfigurationListenerLocked(activityDisplay);
                 }
             });
+
         }
 
         @Override