Merge "Fix the regression that closing transition disappeared." into rvc-dev
diff --git a/services/core/java/com/android/server/wm/ActivityRecord.java b/services/core/java/com/android/server/wm/ActivityRecord.java
index 241f5e8..360e40b 100644
--- a/services/core/java/com/android/server/wm/ActivityRecord.java
+++ b/services/core/java/com/android/server/wm/ActivityRecord.java
@@ -2566,7 +2566,7 @@
             if (mayAdjustTop && ((ActivityStack) task).topRunningActivity(true /* focusableOnly */)
                     == null) {
                 task.adjustFocusToNextFocusableTask("finish-top", false /* allowFocusSelf */,
-                        shouldAdjustGlobalFocus);
+                            shouldAdjustGlobalFocus);
             }
 
             finishActivityResults(resultCode, resultData, resultGrants);
@@ -2586,7 +2586,7 @@
                 if (DEBUG_VISIBILITY || DEBUG_TRANSITION) {
                     Slog.v(TAG_TRANSITION, "Prepare close transition: finishing " + this);
                 }
-                getDisplay().mDisplayContent.prepareAppTransition(transit, false);
+                mDisplayContent.prepareAppTransition(transit, false);
 
                 // When finishing the activity preemptively take the snapshot before the app window
                 // is marked as hidden and any configuration changes take place
@@ -2611,6 +2611,13 @@
 
                 if (endTask) {
                     mAtmService.getLockTaskController().clearLockedTask(task);
+                    // This activity was in the top focused stack and this is the last activity in
+                    // that task, give this activity a higher layer so it can stay on top before the
+                    // closing task transition be executed.
+                    if (mayAdjustTop) {
+                        mNeedsZBoost = true;
+                        mDisplayContent.assignWindowLayers(false /* setLayoutNeeded */);
+                    }
                 }
             } else if (!isState(PAUSING)) {
                 if (mVisibleRequested) {
@@ -6111,7 +6118,6 @@
         Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "AR#onAnimationFinished");
         mTransit = TRANSIT_UNSET;
         mTransitFlags = 0;
-        mNeedsZBoost = false;
         mNeedsAnimationBoundsLayer = false;
 
         setAppLayoutChanges(FINISH_LAYOUT_REDO_ANIM | FINISH_LAYOUT_REDO_WALLPAPER,
diff --git a/services/core/java/com/android/server/wm/TaskDisplayArea.java b/services/core/java/com/android/server/wm/TaskDisplayArea.java
index 205a8d2a..c3acb16 100644
--- a/services/core/java/com/android/server/wm/TaskDisplayArea.java
+++ b/services/core/java/com/android/server/wm/TaskDisplayArea.java
@@ -52,6 +52,7 @@
 import android.content.pm.ActivityInfo;
 import android.content.pm.ApplicationInfo;
 import android.os.UserHandle;
+import android.util.IntArray;
 import android.util.Slog;
 import android.view.SurfaceControl;
 import android.window.WindowContainerTransaction;
@@ -106,6 +107,9 @@
     private final ArrayList<ActivityStack> mTmpAlwaysOnTopStacks = new ArrayList<>();
     private final ArrayList<ActivityStack> mTmpNormalStacks = new ArrayList<>();
     private final ArrayList<ActivityStack> mTmpHomeStacks = new ArrayList<>();
+    private final IntArray mTmpNeedsZBoostIndexes = new IntArray();
+    private int mTmpLayerForSplitScreenDividerAnchor;
+    private int mTmpLayerForAnimationLayer;
 
     private ArrayList<Task> mTmpTasks = new ArrayList<>();
 
@@ -633,39 +637,72 @@
 
         int layer = 0;
         // Place home stacks to the bottom.
-        for (int i = 0; i < mTmpHomeStacks.size(); i++) {
-            mTmpHomeStacks.get(i).assignLayer(t, layer++);
-        }
+        layer = adjustRootTaskLayer(t, mTmpHomeStacks, layer, false /* normalStacks */);
         // The home animation layer is between the home stacks and the normal stacks.
         final int layerForHomeAnimationLayer = layer++;
-        int layerForSplitScreenDividerAnchor = layer++;
-        int layerForAnimationLayer = layer++;
-        for (int i = 0; i < mTmpNormalStacks.size(); i++) {
-            final ActivityStack s = mTmpNormalStacks.get(i);
-            s.assignLayer(t, layer++);
-            if (s.inSplitScreenWindowingMode()) {
-                // The split screen divider anchor is located above the split screen window.
-                layerForSplitScreenDividerAnchor = layer++;
-            }
-            if (s.isTaskAnimating() || s.isAppTransitioning()) {
-                // The animation layer is located above the highest animating stack and no
-                // higher.
-                layerForAnimationLayer = layer++;
-            }
-        }
+        mTmpLayerForSplitScreenDividerAnchor = layer++;
+        mTmpLayerForAnimationLayer = layer++;
+        layer = adjustRootTaskLayer(t, mTmpNormalStacks, layer, true /* normalStacks */);
+
         // The boosted animation layer is between the normal stacks and the always on top
         // stacks.
         final int layerForBoostedAnimationLayer = layer++;
-        for (int i = 0; i < mTmpAlwaysOnTopStacks.size(); i++) {
-            mTmpAlwaysOnTopStacks.get(i).assignLayer(t, layer++);
-        }
+        adjustRootTaskLayer(t, mTmpAlwaysOnTopStacks, layer, false /* normalStacks */);
 
         t.setLayer(mHomeAppAnimationLayer, layerForHomeAnimationLayer);
-        t.setLayer(mAppAnimationLayer, layerForAnimationLayer);
-        t.setLayer(mSplitScreenDividerAnchor, layerForSplitScreenDividerAnchor);
+        t.setLayer(mAppAnimationLayer, mTmpLayerForAnimationLayer);
+        t.setLayer(mSplitScreenDividerAnchor, mTmpLayerForSplitScreenDividerAnchor);
         t.setLayer(mBoostedAppAnimationLayer, layerForBoostedAnimationLayer);
     }
 
+    private int adjustNormalStackLayer(ActivityStack s, int layer) {
+        if (s.inSplitScreenWindowingMode()) {
+            // The split screen divider anchor is located above the split screen window.
+            mTmpLayerForSplitScreenDividerAnchor = layer++;
+        }
+        if (s.isTaskAnimating() || s.isAppTransitioning()) {
+            // The animation layer is located above the highest animating stack and no
+            // higher.
+            mTmpLayerForAnimationLayer = layer++;
+        }
+        return layer;
+    }
+
+    /**
+     * Adjusts the layer of the stack which belongs to the same group.
+     * Note that there are three stack groups: home stacks, always on top stacks, and normal stacks.
+     *
+     * @param startLayer The beginning layer of this group of stacks.
+     * @param normalStacks Set {@code true} if this group is neither home nor always on top.
+     * @return The adjusted layer value.
+     */
+    private int adjustRootTaskLayer(SurfaceControl.Transaction t, ArrayList<ActivityStack> stacks,
+            int startLayer, boolean normalStacks) {
+        mTmpNeedsZBoostIndexes.clear();
+        final int stackSize = stacks.size();
+        for (int i = 0; i < stackSize; i++) {
+            final ActivityStack stack = stacks.get(i);
+            if (!stack.needsZBoost()) {
+                stack.assignLayer(t, startLayer++);
+                if (normalStacks) {
+                    startLayer = adjustNormalStackLayer(stack, startLayer);
+                }
+            } else {
+                mTmpNeedsZBoostIndexes.add(i);
+            }
+        }
+
+        final int zBoostSize = mTmpNeedsZBoostIndexes.size();
+        for (int i = 0; i < zBoostSize; i++) {
+            final ActivityStack stack = stacks.get(mTmpNeedsZBoostIndexes.get(i));
+            stack.assignLayer(t, startLayer++);
+            if (normalStacks) {
+                startLayer = adjustNormalStackLayer(stack, startLayer);
+            }
+        }
+        return startLayer;
+    }
+
     @Override
     SurfaceControl getAppAnimationLayer(@AnimationLayer int animationLayer) {
         switch (animationLayer) {
diff --git a/services/core/java/com/android/server/wm/WindowContainer.java b/services/core/java/com/android/server/wm/WindowContainer.java
index 5a73fab..f984a8b 100644
--- a/services/core/java/com/android/server/wm/WindowContainer.java
+++ b/services/core/java/com/android/server/wm/WindowContainer.java
@@ -2440,6 +2440,7 @@
     protected void onAnimationFinished(@AnimationType int type, AnimationAdapter anim) {
         doAnimationFinished(type, anim);
         mWmService.onAnimationFinished();
+        mNeedsZBoost = false;
     }
 
     /**