Support to get the exact task which IME is attached to

To distinguish which side of the split that IME panel is attached to
after having a single-top root task for split screen, update
getImeTarget API to return the exact task that IME panel is attached to.

Bug: 207185041
Test: WMShellUnitTests
Test: enter split in portrait, toggle IME panel on the bottom split and
      verified split layout will be push up as expectedly.
Change-Id: Ib6ba74ecb1b844bb408475d5b6c28dd5816a8e29

core/java/android/window/ITaskOrganizerController.aidl

@@ -52,7 +52,7 @@ interface ITaskOrganizerController {
     /** Gets all root tasks on a display (ordered from top-to-bottom) */
     List<ActivityManager.RunningTaskInfo> getRootTasks(int displayId, in int[] activityTypes);
 
-    /** Get the root task which contains the current ime target */
+    /** Get the {@link WindowContainerToken} of the task which contains the current ime target */
     WindowContainerToken getImeTarget(int display);
 
     /**

core/java/android/window/TaskOrganizer.java

@@ -199,7 +199,7 @@ public class TaskOrganizer extends WindowOrganizer {
         }
     }
 
-    /** Get the root task which contains the current ime target */
+    /** Get the {@link WindowContainerToken} of the task which contains the current ime target */
     @RequiresPermission(android.Manifest.permission.MANAGE_ACTIVITY_TASKS)
     @Nullable
     public WindowContainerToken getImeTarget(int display) {

libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java

@@ -1115,9 +1115,9 @@ class StageCoordinator implements SplitLayout.SplitLayoutHandler,
             return SPLIT_POSITION_UNDEFINED;
         }
 
-        if (token.equals(mMainStage.mRootTaskInfo.getToken())) {
+        if (mMainStage.containsToken(token)) {
             return getMainStagePosition();
-        } else if (token.equals(mSideStage.mRootTaskInfo.getToken())) {
+        } else if (mSideStage.containsToken(token)) {
             return getSideStagePosition();
         }
 

libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java

@@ -124,6 +124,20 @@ class StageTaskListener implements ShellTaskOrganizer.TaskListener {
         return mChildrenTaskInfo.contains(taskId);
     }
 
+    boolean containsToken(WindowContainerToken token) {
+        if (token.equals(mRootTaskInfo.token)) {
+            return true;
+        }
+
+        for (int i = mChildrenTaskInfo.size() - 1; i >= 0; --i) {
+            if (token.equals(mChildrenTaskInfo.valueAt(i).token)) {
+                return true;
+            }
+        }
+
+        return false;
+    }
+
     /**
      * Returns the top visible child task's id.
      */

services/core/java/com/android/server/wm/TaskOrganizerController.java

@@ -790,17 +790,24 @@ class TaskOrganizerController extends ITaskOrganizerController.Stub {
         final long origId = Binder.clearCallingIdentity();
         try {
             synchronized (mGlobalLock) {
-                DisplayContent dc = mService.mWindowManager.mRoot
+                final DisplayContent dc = mService.mWindowManager.mRoot
                         .getDisplayContent(displayId);
-                if (dc == null || dc.getImeTarget(IME_TARGET_LAYERING) == null) {
+                if (dc == null) {
                     return null;
                 }
+
+                final InsetsControlTarget imeLayeringTarget = dc.getImeTarget(IME_TARGET_LAYERING);
+                if (imeLayeringTarget == null || imeLayeringTarget.getWindow() == null) {
+                    return null;
+                }
+
                 // Avoid WindowState#getRootTask() so we don't attribute system windows to a task.
-                final Task task = dc.getImeTarget(IME_TARGET_LAYERING).getWindow().getTask();
+                final Task task = imeLayeringTarget.getWindow().asTask();
                 if (task == null) {
                     return null;
                 }
-                return task.getRootTask().mRemoteToken.toWindowContainerToken();
+
+                return task.mRemoteToken.toWindowContainerToken();
             }
         } finally {
             Binder.restoreCallingIdentity(origId);