Fix wrong ime parent in embedded activity

We have CL[1] to polish the ime transition. But in
ActivityEmbedding, the ime layering target may higher
than input target. And the ime parent didn't be updated
since input target and layering target are unsync.

Update ime parent for this case.

[1]: I332c0e4fff62df5d7b793eda2767bb58fe85a938

Bug: 268102890
Bug: 260387203
Test: Manual test with test apk in the bug
Test: atest TaskFragmentTest#testUpdateImeParentForActivityEmbedding
Merged-In: I937560f7479535f0e2dc77ebe0f5056fa85bbb04
Change-Id: I78c82a0677d2ab7a1e247e91e5dfa1e5235e95bb

2 files changed, 103 insertions, 18 deletions

diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index b0eecebc7d79..ec355b73a7be 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -4505,24 +4505,8 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
      */
     @VisibleForTesting
     SurfaceControl computeImeParent() {
-        if (mImeLayeringTarget != null) {
-            // Ensure changing the IME parent when the layering target that may use IME has
-            // became to the input target for preventing IME flickers.
-            // Note that:
-            // 1) For the imeLayeringTarget that may not use IME but requires IME on top
-            // of it (e.g. an overlay window with NOT_FOCUSABLE|ALT_FOCUSABLE_IM flags), we allow
-            // it to re-parent the IME on top the display to keep the legacy behavior.
-            // 2) Even though the starting window won't use IME, the associated activity
-            // behind the starting window may request the input. If so, then we should still hold
-            // the IME parent change until the activity started the input.
-            boolean imeLayeringTargetMayUseIme =
-                    LayoutParams.mayUseInputMethod(mImeLayeringTarget.mAttrs.flags)
-                    || mImeLayeringTarget.mAttrs.type == TYPE_APPLICATION_STARTING;
-            if (imeLayeringTargetMayUseIme && (mImeInputTarget == null
-                    || mImeLayeringTarget.mActivityRecord != mImeInputTarget.getActivityRecord())) {
-                // Do not change parent if the window hasn't requested IME.
-                return null;
-            }
+        if (!canComputeImeParent(mImeLayeringTarget, mImeInputTarget)) {
+            return null;
         }
         // Attach it to app if the target is part of an app and such app is covering the entire
         // screen. If it's not covering the entire screen the IME might extend beyond the apps
@@ -4535,6 +4519,76 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
                 ? mImeWindowsContainer.getParent().getSurfaceControl() : null;
     }
 
+    private static boolean canComputeImeParent(@Nullable WindowState imeLayeringTarget,
+            @Nullable InputTarget imeInputTarget) {
+        if (imeLayeringTarget == null) {
+            return false;
+        }
+        if (shouldComputeImeParentForEmbeddedActivity(imeLayeringTarget, imeInputTarget)) {
+            return true;
+        }
+        // Ensure changing the IME parent when the layering target that may use IME has
+        // became to the input target for preventing IME flickers.
+        // Note that:
+        // 1) For the imeLayeringTarget that may not use IME but requires IME on top
+        // of it (e.g. an overlay window with NOT_FOCUSABLE|ALT_FOCUSABLE_IM flags), we allow
+        // it to re-parent the IME on top the display to keep the legacy behavior.
+        // 2) Even though the starting window won't use IME, the associated activity
+        // behind the starting window may request the input. If so, then we should still hold
+        // the IME parent change until the activity started the input.
+        boolean imeLayeringTargetMayUseIme =
+                LayoutParams.mayUseInputMethod(imeLayeringTarget.mAttrs.flags)
+                        || imeLayeringTarget.mAttrs.type == TYPE_APPLICATION_STARTING;
+
+        if (imeLayeringTargetMayUseIme && (imeInputTarget == null
+                || imeLayeringTarget.mActivityRecord != imeInputTarget.getActivityRecord())) {
+            // Do not change parent if the window hasn't requested IME.
+            return false;
+        }
+        return true;
+    }
+
+    /**
+     * Called from {@link #computeImeParent()} to check the given IME targets if the IME surface
+     * parent should be updated in ActivityEmbeddings.
+     *
+     * As the IME layering target is calculated according to the window hierarchy by
+     * {@link #computeImeTarget}, the layering target and input target may be different when the
+     * window hasn't started input connection, WindowManagerService hasn't yet received the
+     * input target which reported from InputMethodManagerService. To make the IME surface will be
+     * shown on the best fit IME layering target, we basically won't update IME parent until both
+     * IME input and layering target updated for better IME transition.
+     *
+     * However, in activity embedding, tapping a window won't update it to the top window so the
+     * calculated IME layering target may higher than input target. Update IME parent for this case.
+     *
+     * @return {@code true} means the layer of IME layering target is higher than the input target
+     * and {@link #computeImeParent()} should keep progressing to update the IME
+     * surface parent on the display in case the IME surface left behind.
+     */
+    private static boolean shouldComputeImeParentForEmbeddedActivity(
+            @Nullable WindowState imeLayeringTarget, @Nullable InputTarget imeInputTarget) {
+        if (imeInputTarget == null || imeLayeringTarget == null) {
+            return false;
+        }
+        final WindowState inputTargetWindow = imeInputTarget.getWindowState();
+        if (inputTargetWindow == null || !imeLayeringTarget.isAttached()
+                || !inputTargetWindow.isAttached()) {
+            return false;
+        }
+
+        final ActivityRecord inputTargetRecord = imeInputTarget.getActivityRecord();
+        final ActivityRecord layeringTargetRecord = imeLayeringTarget.getActivityRecord();
+        if (inputTargetRecord == null || layeringTargetRecord == null
+                || inputTargetRecord == layeringTargetRecord
+                || inputTargetRecord.getTask() != layeringTargetRecord.getTask()
+                || !inputTargetRecord.isEmbedded() || !layeringTargetRecord.isEmbedded()) {
+            // Check whether the input target and layering target are embedded in the same Task.
+            return false;
+        }
+        return imeLayeringTarget.compareTo(inputTargetWindow) > 0;
+    }
+
     void setLayoutNeeded() {
         if (DEBUG_LAYOUT) Slog.w(TAG_WM, "setLayoutNeeded: callers=" + Debug.getCallers(3));
         mLayoutNeeded = true;
diff --git a/services/tests/wmtests/src/com/android/server/wm/TaskFragmentTest.java b/services/tests/wmtests/src/com/android/server/wm/TaskFragmentTest.java
index c8932550d877..132aa90f1367 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TaskFragmentTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TaskFragmentTest.java
@@ -24,6 +24,7 @@ import static android.content.pm.ActivityInfo.SCREEN_ORIENTATION_LANDSCAPE;
 import static android.content.pm.ActivityInfo.SCREEN_ORIENTATION_PORTRAIT;
 import static android.content.pm.ActivityInfo.SCREEN_ORIENTATION_UNSET;
 import static android.content.pm.ActivityInfo.SCREEN_ORIENTATION_UNSPECIFIED;
+import static android.view.WindowManager.LayoutParams.TYPE_BASE_APPLICATION;
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.any;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.doNothing;
@@ -603,4 +604,34 @@ public class TaskFragmentTest extends WindowTestsBase {
         assertEquals(SCREEN_ORIENTATION_UNSPECIFIED, tf1.getOrientation(SCREEN_ORIENTATION_UNSET));
         assertEquals(SCREEN_ORIENTATION_UNSPECIFIED, task.getOrientation(SCREEN_ORIENTATION_UNSET));
     }
+
+    @Test
+    public void testUpdateImeParentForActivityEmbedding() {
+        // Setup two activities in ActivityEmbedding.
+        final Task task = createTask(mDisplayContent);
+        final TaskFragment tf0 = new TaskFragmentBuilder(mAtm)
+                .setParentTask(task)
+                .createActivityCount(1)
+                .setOrganizer(mOrganizer)
+                .setFragmentToken(new Binder())
+                .build();
+        final TaskFragment tf1 = new TaskFragmentBuilder(mAtm)
+                .setParentTask(task)
+                .createActivityCount(1)
+                .setOrganizer(mOrganizer)
+                .setFragmentToken(new Binder())
+                .build();
+        final ActivityRecord activity0 = tf0.getTopMostActivity();
+        final ActivityRecord activity1 = tf1.getTopMostActivity();
+        final WindowState win0 = createWindow(null, TYPE_BASE_APPLICATION, activity0, "win0");
+        final WindowState win1 = createWindow(null, TYPE_BASE_APPLICATION, activity1, "win1");
+        doReturn(false).when(mDisplayContent).shouldImeAttachedToApp();
+
+        mDisplayContent.setImeInputTarget(win0);
+        mDisplayContent.setImeLayeringTarget(win1);
+
+        // The ImeParent should be the display.
+        assertEquals(mDisplayContent.getImeContainer().getParent().getSurfaceControl(),
+                mDisplayContent.computeImeParent());
+    }
 }