Merge "[PiP2] Find config-at-end activity in AE case" into main
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTransitionController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTransitionController.java
index cef18f5..c58bb6e 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTransitionController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTransitionController.java
@@ -40,7 +40,6 @@
 import android.view.WindowManager;
 import android.window.TransitionInfo;
 import android.window.TransitionRequestInfo;
-import android.window.WindowContainerToken;
 import android.window.WindowContainerTransaction;
 
 import androidx.annotation.NonNull;
@@ -341,23 +340,6 @@
         return false;
     }
 
-    /**
-     * @return a change representing a config-at-end activity for a given parent.
-     */
-    @Nullable
-    public TransitionInfo.Change getDeferConfigActivityChange(TransitionInfo info,
-            @android.annotation.NonNull WindowContainerToken parent) {
-        for (TransitionInfo.Change change : info.getChanges()) {
-            if (change.getTaskInfo() == null
-                    && change.hasFlags(TransitionInfo.FLAG_CONFIG_AT_END)
-                    && change.getParent() != null && change.getParent().equals(parent)) {
-                return change;
-            }
-        }
-        return null;
-    }
-
-
     /** Whether a particular package is same as current pip package. */
     public boolean isPackageActiveInPip(@Nullable String packageName) {
         // No-op, to be handled differently in PIP1 and PIP2
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTransition.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTransition.java
index cfcd563..5d8d8b6 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTransition.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTransition.java
@@ -76,6 +76,7 @@
 import com.android.wm.shell.pip2.animation.PipAlphaAnimator;
 import com.android.wm.shell.pip2.animation.PipEnterAnimator;
 import com.android.wm.shell.pip2.phone.transition.PipExpandHandler;
+import com.android.wm.shell.pip2.phone.transition.PipTransitionUtils;
 import com.android.wm.shell.shared.TransitionUtil;
 import com.android.wm.shell.splitscreen.SplitScreenController;
 import com.android.wm.shell.sysui.ShellInit;
@@ -387,8 +388,8 @@
         mFinishCallback = finishCallback;
         // We expect the PiP activity as a separate change in a config-at-end transition;
         // only flings are not using config-at-end for resize bounds changes
-        TransitionInfo.Change pipActivityChange = getDeferConfigActivityChange(info,
-                pipChange.getTaskInfo().getToken());
+        TransitionInfo.Change pipActivityChange = PipTransitionUtils.getDeferConfigActivityChange(
+                info, pipChange.getTaskInfo().getToken());
         if (pipActivityChange != null) {
             // Transform calculations use PiP params by default, so make sure they are null to
             // default to using bounds for scaling calculations instead.
@@ -427,8 +428,8 @@
         }
 
         // We expect the PiP activity as a separate change in a config-at-end transition.
-        TransitionInfo.Change pipActivityChange = getDeferConfigActivityChange(info,
-                pipChange.getTaskInfo().getToken());
+        TransitionInfo.Change pipActivityChange = PipTransitionUtils.getDeferConfigActivityChange(
+                info, pipChange.getTaskInfo().getToken());
         if (pipActivityChange == null) {
             return false;
         }
@@ -497,8 +498,8 @@
         }
 
         // We expect the PiP activity as a separate change in a config-at-end transition.
-        TransitionInfo.Change pipActivityChange = getDeferConfigActivityChange(info,
-                pipChange.getTaskInfo().getToken());
+        TransitionInfo.Change pipActivityChange = PipTransitionUtils.getDeferConfigActivityChange(
+                info, pipChange.getTaskInfo().getToken());
         if (pipActivityChange == null) {
             return false;
         }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/transition/PipTransitionUtils.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/transition/PipTransitionUtils.java
index 01cda6c..e562f33 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/transition/PipTransitionUtils.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/transition/PipTransitionUtils.java
@@ -67,6 +67,36 @@
     }
 
     /**
+     * @return a change representing a config-at-end activity for ancestor.
+     */
+    @Nullable
+    public static TransitionInfo.Change getDeferConfigActivityChange(TransitionInfo info,
+            @NonNull WindowContainerToken ancestor) {
+        final TransitionInfo.Change ancestorChange =
+                PipTransitionUtils.getChangeByToken(info, ancestor);
+        if (ancestorChange == null) return null;
+
+        // Iterate through changes bottom-to-top, going up the parent chain starting with ancestor.
+        TransitionInfo.Change lastPipChildChange = ancestorChange;
+        for (int i = info.getChanges().size() - 1; i >= 0; --i) {
+            TransitionInfo.Change change = info.getChanges().get(i);
+            if (change == ancestorChange) continue;
+
+            if (change.getParent() != null
+                    && change.getParent().equals(lastPipChildChange.getContainer())) {
+                // Found a child of the last cached child along the ancestral chain.
+                lastPipChildChange = change;
+                if (change.getTaskInfo() == null
+                        && change.hasFlags(TransitionInfo.FLAG_CONFIG_AT_END)) {
+                    // If this is a config-at-end activity change, then we found the chain leaf.
+                    return change;
+                }
+            }
+        }
+        return null;
+    }
+
+    /**
      * @return the leash to interact with the container this change represents.
      * @throws NullPointerException if the leash is null.
      */
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedTransition.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedTransition.java
index 1853ffa..320a63a 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedTransition.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedTransition.java
@@ -34,6 +34,7 @@
 import com.android.wm.shell.desktopmode.DesktopTasksController;
 import com.android.wm.shell.keyguard.KeyguardTransitionHandler;
 import com.android.wm.shell.pip.PipTransitionController;
+import com.android.wm.shell.pip2.phone.transition.PipTransitionUtils;
 import com.android.wm.shell.protolog.ShellProtoLogGroup;
 import com.android.wm.shell.splitscreen.StageCoordinator;
 import com.android.wm.shell.unfold.UnfoldTransitionHandler;
@@ -132,7 +133,7 @@
 
         TransitionInfo.Change pipActivityChange = null;
         if (pipChange != null) {
-            pipActivityChange = mPipHandler.getDeferConfigActivityChange(
+            pipActivityChange = PipTransitionUtils.getDeferConfigActivityChange(
                     info, pipChange.getContainer());
             everythingElse.getChanges().remove(pipActivityChange);
         }