Handle z-layering in animation layer

We use the prefix order to determine base layer within the
animation layer. This ensure that for the animation windows, the
z-ordering during animations doesn't change.

We then boost anything that needs a z-boost to 800570000 + prefix
order, such order within the boosted layers is preserved as well.

Also fix an issue where the thumbnail wasn't attached to the
animation layer.

Test: WindowContainerTests
Test: go/wm-smoke
Bug: 64674361
Change-Id: If5909bd87a12f1d8920c7232acab0f3d17be0f6c
diff --git a/services/core/java/com/android/server/wm/AppWindowThumbnail.java b/services/core/java/com/android/server/wm/AppWindowThumbnail.java
index 65abcf5..b01b268 100644
--- a/services/core/java/com/android/server/wm/AppWindowThumbnail.java
+++ b/services/core/java/com/android/server/wm/AppWindowThumbnail.java
@@ -141,7 +141,7 @@
 
     @Override
     public Builder makeAnimationLeash() {
-        return mAppToken.makeSurface();
+        return mAppToken.makeSurface().setParent(mAppToken.getAppAnimationLayer());
     }
 
     @Override
diff --git a/services/core/java/com/android/server/wm/AppWindowToken.java b/services/core/java/com/android/server/wm/AppWindowToken.java
index 444a05d..af7523c 100644
--- a/services/core/java/com/android/server/wm/AppWindowToken.java
+++ b/services/core/java/com/android/server/wm/AppWindowToken.java
@@ -65,6 +65,7 @@
 import android.util.Slog;
 import android.util.proto.ProtoOutputStream;
 import android.view.DisplayInfo;
+import android.view.SurfaceControl.Transaction;
 import android.view.animation.Animation;
 import android.view.IApplicationToken;
 import android.view.SurfaceControl;
@@ -91,6 +92,11 @@
 class AppWindowToken extends WindowToken implements WindowManagerService.AppFreezeListener {
     private static final String TAG = TAG_WITH_CLASS_NAME ? "AppWindowToken" : TAG_WM;
 
+    /**
+     * Value to increment the z-layer when boosting a layer during animations. BOOST in l33tsp34k.
+     */
+    private static final int Z_BOOST_BASE = 800570000;
+
     // Non-null only for application tokens.
     final IApplicationToken appToken;
 
@@ -1537,11 +1543,10 @@
                         new WindowAnimationSpec(a, mTmpPoint,
                                 mService.mAppTransition.canSkipFirstFrame()),
                         mService.mSurfaceAnimationRunner);
-                startAnimation(getPendingTransaction(), adapter, !isVisible());
                 if (a.getZAdjustment() == Animation.ZORDER_TOP) {
                     mNeedsZBoost = true;
-                    getDisplayContent().assignWindowLayers(false /* setLayoutNeeded */);
                 }
+                startAnimation(getPendingTransaction(), adapter, !isVisible());
                 mTransit = transit;
                 mTransitFlags = mService.mAppTransition.getTransitFlags();
                 // TODO: Skip first frame and app stack clip mode.
@@ -1610,6 +1615,39 @@
         return a;
     }
 
+    @Override
+    protected void setLayer(Transaction t, int layer) {
+        if (!mSurfaceAnimator.hasLeash()) {
+            t.setLayer(mSurfaceControl, layer);
+        }
+    }
+
+    @Override
+    protected void setRelativeLayer(Transaction t, SurfaceControl relativeTo, int layer) {
+        if (!mSurfaceAnimator.hasLeash()) {
+            t.setRelativeLayer(mSurfaceControl, relativeTo, layer);
+        }
+    }
+
+    @Override
+    protected void reparentSurfaceControl(Transaction t, SurfaceControl newParent) {
+        if (!mSurfaceAnimator.hasLeash()) {
+            t.reparent(mSurfaceControl, newParent.getHandle());
+        }
+    }
+
+    @Override
+    public void onAnimationLeashCreated(Transaction t, SurfaceControl leash) {
+
+        // The leash is parented to the animation layer. We need to preserve the z-order by using
+        // the prefix order index, but we boost if necessary.
+        int layer = getPrefixOrderIndex();
+        if (mNeedsZBoost) {
+            layer += Z_BOOST_BASE;
+        }
+        leash.setLayer(layer);
+    }
+
     /**
      * This must be called while inside a transaction.
      */
diff --git a/services/core/java/com/android/server/wm/WindowContainer.java b/services/core/java/com/android/server/wm/WindowContainer.java
index 7641cbc..b251b53 100644
--- a/services/core/java/com/android/server/wm/WindowContainer.java
+++ b/services/core/java/com/android/server/wm/WindowContainer.java
@@ -96,6 +96,9 @@
 
     private final Point mTmpPos = new Point();
 
+    /** Total number of elements in this subtree, including our own hierarchy element. */
+    private int mTreeWeight = 1;
+
     WindowContainer(WindowManagerService service) {
         mService = service;
         mPendingTransaction = service.mTransactionFactory.make();
@@ -157,7 +160,7 @@
             // surface animator such that hierarchy is preserved when animating, i.e.
             // mSurfaceControl stays attached to the leash and we just reparent the leash to the
             // new parent.
-            mSurfaceAnimator.reparent(getPendingTransaction(), mParent.mSurfaceControl);
+            reparentSurfaceControl(getPendingTransaction(), mParent.mSurfaceControl);
         }
 
         // Either way we need to ask the parent to assign us a Z-order.
@@ -200,6 +203,8 @@
         } else {
             mChildren.add(positionToAdd, child);
         }
+        onChildAdded(child);
+
         // Set the parent after we've actually added a child in case a subclass depends on this.
         child.setParent(this);
     }
@@ -213,10 +218,21 @@
                     + " can't add to container=" + getName());
         }
         mChildren.add(index, child);
+        onChildAdded(child);
+
         // Set the parent after we've actually added a child in case a subclass depends on this.
         child.setParent(this);
     }
 
+    private void onChildAdded(WindowContainer child) {
+        mTreeWeight += child.mTreeWeight;
+        WindowContainer parent = getParent();
+        while (parent != null) {
+            parent.mTreeWeight += child.mTreeWeight;
+            parent = parent.getParent();
+        }
+    }
+
     /**
      * Removes the input child container from this container which is its parent.
      *
@@ -225,6 +241,7 @@
     @CallSuper
     void removeChild(E child) {
         if (mChildren.remove(child)) {
+            onChildRemoved(child);
             child.setParent(null);
         } else {
             throw new IllegalArgumentException("removeChild: container=" + child.getName()
@@ -232,6 +249,15 @@
         }
     }
 
+    private void onChildRemoved(WindowContainer child) {
+        mTreeWeight -= child.mTreeWeight;
+        WindowContainer parent = getParent();
+        while (parent != null) {
+            parent.mTreeWeight -= child.mTreeWeight;
+            parent = parent.getParent();
+        }
+    }
+
     /**
      * Removes this window container and its children with no regard for what else might be going on
      * in the system. For example, the container will be removed during animation if this method is
@@ -246,7 +272,9 @@
             // Need to do this after calling remove on the child because the child might try to
             // remove/detach itself from its parent which will cause an exception if we remove
             // it before calling remove on the child.
-            mChildren.remove(child);
+            if (mChildren.remove(child)) {
+                onChildRemoved(child);
+            }
         }
 
         if (mSurfaceControl != null) {
@@ -265,6 +293,34 @@
     }
 
     /**
+     * @return The index of this element in the hierarchy tree in prefix order.
+     */
+    int getPrefixOrderIndex() {
+        if (mParent == null) {
+            return 0;
+        }
+        return mParent.getPrefixOrderIndex(this);
+    }
+
+    private int getPrefixOrderIndex(WindowContainer child) {
+        int order = 0;
+        for (int i = 0; i < mChildren.size(); i++) {
+            final WindowContainer childI = mChildren.get(i);
+            if (child == childI) {
+                break;
+            }
+            order += childI.mTreeWeight;
+        }
+        if (mParent != null) {
+            order += mParent.getPrefixOrderIndex(this);
+        }
+
+        // We also need to count ourselves.
+        order++;
+        return order;
+    }
+
+    /**
      * Removes this window container and its children taking care not to remove them during a
      * critical stage in the system. For example, some containers will not be removed during
      * animation if this method is called.
@@ -831,10 +887,7 @@
     void assignLayer(Transaction t, int layer) {
         final boolean changed = layer != mLastLayer || mLastRelativeToLayer != null;
         if (mSurfaceControl != null && changed) {
-
-            // Route through surface animator to accommodate that our surface control might be
-            // attached to the leash, and leash is attached to parent container.
-            mSurfaceAnimator.setLayer(t, layer);
+            setLayer(t, layer);
             mLastLayer = layer;
             mLastRelativeToLayer = null;
         }
@@ -843,15 +896,30 @@
     void assignRelativeLayer(Transaction t, SurfaceControl relativeTo, int layer) {
         final boolean changed = layer != mLastLayer || mLastRelativeToLayer != relativeTo;
         if (mSurfaceControl != null && changed) {
-
-            // Route through surface animator to accommodate that our surface control might be
-            // attached to the leash, and leash is attached to parent container.
-            mSurfaceAnimator.setRelativeLayer(t, relativeTo, layer);
+            setRelativeLayer(t, relativeTo, layer);
             mLastLayer = layer;
             mLastRelativeToLayer = relativeTo;
         }
     }
 
+    protected void setLayer(Transaction t, int layer) {
+
+        // Route through surface animator to accommodate that our surface control might be
+        // attached to the leash, and leash is attached to parent container.
+        mSurfaceAnimator.setLayer(t, layer);
+    }
+
+    protected void setRelativeLayer(Transaction t, SurfaceControl relativeTo, int layer) {
+
+        // Route through surface animator to accommodate that our surface control might be
+        // attached to the leash, and leash is attached to parent container.
+        mSurfaceAnimator.setRelativeLayer(t, relativeTo, layer);
+    }
+
+    protected void reparentSurfaceControl(Transaction t, SurfaceControl newParent) {
+        mSurfaceAnimator.reparent(t, newParent);
+    }
+
     void assignChildLayers(Transaction t) {
         int layer = 0;