Use the proper transaction when finishing a split screen transition

It will restore the shifted split layout after entered overview panel,
this makes sure to update the proper finish transaction so the split
layout will be updated properly after animated the transition.

Also support to modify the finish transaction when a transition got
merged, so we can append clean-up operations for the cases like merging
an enter split screen transition to another handler.

Fix: 238856352
Test: enabled shell transition, verified split layout has been
      shifted/restored properly after re-entering a split pair.

Change-Id: I630c54ed7c25246634c314eb9b2561a8556d8570
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitLayout.java b/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitLayout.java
index 40cf9a3..85c8ebf 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitLayout.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitLayout.java
@@ -1130,16 +1130,16 @@
             boolean adjusted = false;
             if (mYOffsetForIme != 0) {
                 if (dividerLeash != null) {
-                    mTempRect.set(mDividerBounds);
+                    getRefDividerBounds(mTempRect);
                     mTempRect.offset(0, mYOffsetForIme);
                     t.setPosition(dividerLeash, mTempRect.left, mTempRect.top);
                 }
 
-                mTempRect.set(mBounds1);
+                getRefBounds1(mTempRect);
                 mTempRect.offset(0, mYOffsetForIme);
                 t.setPosition(leash1, mTempRect.left, mTempRect.top);
 
-                mTempRect.set(mBounds2);
+                getRefBounds2(mTempRect);
                 mTempRect.offset(0, mYOffsetForIme);
                 t.setPosition(leash2, mTempRect.left, mTempRect.top);
                 adjusted = true;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTransition.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTransition.java
index 51be2a5..dd09b2c 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTransition.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTransition.java
@@ -315,7 +315,8 @@
     }
 
     @Override
-    public void onTransitionConsumed(@NonNull IBinder transition, boolean aborted) {
+    public void onTransitionConsumed(@NonNull IBinder transition, boolean aborted,
+            @Nullable SurfaceControl.Transaction finishT) {
         if (transition != mExitTransition) {
             return;
         }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenTransitions.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenTransitions.java
index 83bdf8b..d7ca791 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenTransitions.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenTransitions.java
@@ -94,6 +94,7 @@
             @NonNull WindowContainerToken topRoot) {
         mFinishCallback = finishCallback;
         mAnimatingTransition = transition;
+        mFinishTransaction = finishTransaction;
         if (mPendingRemoteHandler != null) {
             mPendingRemoteHandler.startAnimation(transition, info, startTransaction,
                     finishTransaction, mRemoteFinishCB);
@@ -107,8 +108,6 @@
     private void playInternalAnimation(@NonNull IBinder transition, @NonNull TransitionInfo info,
             @NonNull SurfaceControl.Transaction t, @NonNull WindowContainerToken mainRoot,
             @NonNull WindowContainerToken sideRoot, @NonNull WindowContainerToken topRoot) {
-        mFinishTransaction = mTransactionPool.acquire();
-
         // Play some place-holder fade animations
         for (int i = info.getChanges().size() - 1; i >= 0; --i) {
             final TransitionInfo.Change change = info.getChanges().get(i);
@@ -287,16 +286,14 @@
         return true;
     }
 
-    void onTransitionConsumed(@NonNull IBinder transition, boolean aborted) {
+    void onTransitionConsumed(@NonNull IBinder transition, boolean aborted,
+            @Nullable SurfaceControl.Transaction finishT) {
         if (isPendingEnter(transition)) {
             if (!aborted) {
                 // An enter transition got merged, appends the rest operations to finish entering
                 // split screen.
-                // TODO (b/238856352): Passed-in the proper finish transition to merge instead.
-                if (mFinishTransaction == null) {
-                    mFinishTransaction = mTransactionPool.acquire();
-                }
-                mStageCoordinator.finishEnterSplitScreen(mFinishTransaction);
+                mStageCoordinator.finishEnterSplitScreen(finishT);
+                mPendingRemoteHandler = null;
             }
 
             mPendingEnter.mCallback.onTransitionConsumed(aborted);
@@ -339,11 +336,6 @@
         mAnimatingTransition = null;
 
         mOnFinish.run();
-        if (mFinishTransaction != null) {
-            mFinishTransaction.apply();
-            mTransactionPool.release(mFinishTransaction);
-            mFinishTransaction = null;
-        }
         if (mFinishCallback != null) {
             mFinishCallback.onTransitionFinished(wct /* wct */, wctCB /* wctCB */);
             mFinishCallback = null;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
index 3c7db33..530f479 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
@@ -229,6 +229,7 @@
                 if (op.getType() == HIERARCHY_OP_TYPE_REORDER && op.getToTop()
                         && (mMainStage.containsContainer(container)
                         || mSideStage.containsContainer(container))) {
+                    updateSurfaceBounds(mSplitLayout, finishT, false /* applyResizingOffset */);
                     setDividerVisibility(true, finishT);
                     return;
                 }
@@ -1677,8 +1678,9 @@
     }
 
     @Override
-    public void onTransitionConsumed(@NonNull IBinder transition, boolean aborted) {
-        mSplitTransitions.onTransitionConsumed(transition, aborted);
+    public void onTransitionConsumed(@NonNull IBinder transition, boolean aborted,
+            @Nullable SurfaceControl.Transaction finishT) {
+        mSplitTransitions.onTransitionConsumed(transition, aborted, finishT);
     }
 
     @Override
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedHandler.java
index 11b453c..5cce6b9 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedHandler.java
@@ -274,7 +274,8 @@
     }
 
     @Override
-    public void onTransitionConsumed(@NonNull IBinder transition, boolean aborted) {
+    public void onTransitionConsumed(@NonNull IBinder transition, boolean aborted,
+            @Nullable SurfaceControl.Transaction finishT) {
         MixedTransition mixed = null;
         for (int i = mActiveTransitions.size() - 1; i >= 0; --i) {
             if (mActiveTransitions.get(i).mTransition != transition) continue;
@@ -283,7 +284,7 @@
         }
         if (mixed == null) return;
         if (mixed.mType == MixedTransition.TYPE_ENTER_PIP_FROM_SPLIT) {
-            mPipHandler.onTransitionConsumed(transition, aborted);
+            mPipHandler.onTransitionConsumed(transition, aborted, finishT);
         }
     }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/RemoteTransitionHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/RemoteTransitionHandler.java
index b15c48c..bbb0c9d 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/RemoteTransitionHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/RemoteTransitionHandler.java
@@ -83,7 +83,8 @@
     }
 
     @Override
-    public void onTransitionConsumed(@NonNull IBinder transition, boolean aborted) {
+    public void onTransitionConsumed(@NonNull IBinder transition, boolean aborted,
+            @Nullable SurfaceControl.Transaction finishT) {
         mRequestedRemotes.remove(transition);
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/Transitions.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/Transitions.java
index fa22c7c..0ef6200 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/Transitions.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/Transitions.java
@@ -516,7 +516,8 @@
             active.mMerged = true;
             active.mAborted = abort;
             if (active.mHandler != null) {
-                active.mHandler.onTransitionConsumed(active.mToken, abort);
+                active.mHandler.onTransitionConsumed(
+                        active.mToken, abort, abort ? null : active.mFinishT);
             }
             return;
         }
@@ -524,7 +525,8 @@
         active.mAborted = abort;
         if (active.mAborted && active.mHandler != null) {
             // Notifies to clean-up the aborted transition.
-            active.mHandler.onTransitionConsumed(transition, true /* aborted */);
+            active.mHandler.onTransitionConsumed(
+                    transition, true /* aborted */, null /* finishTransaction */);
         }
         ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS,
                 "Transition animation finished (abort=%b), notifying core %s", abort, transition);
@@ -560,7 +562,8 @@
             ActiveTransition aborted = mActiveTransitions.remove(activeIdx);
             // Notifies to clean-up the aborted transition.
             if (aborted.mHandler != null) {
-                aborted.mHandler.onTransitionConsumed(transition, true /* aborted */);
+                aborted.mHandler.onTransitionConsumed(
+                        transition, true /* aborted */, null /* finishTransaction */);
             }
             mOrganizer.finishTransition(aborted.mToken, null /* wct */, null /* wctCB */);
         }
@@ -746,8 +749,13 @@
          * Called when a transition which was already "claimed" by this handler has been merged
          * into another animation or has been aborted. Gives this handler a chance to clean-up any
          * expectations.
+         *
+         * @param transition The transition been consumed.
+         * @param aborted Whether the transition is aborted or not.
+         * @param finishTransaction The transaction to be applied after the transition animated.
          */
-        default void onTransitionConsumed(@NonNull IBinder transition, boolean aborted) { }
+        default void onTransitionConsumed(@NonNull IBinder transition, boolean aborted,
+                @Nullable SurfaceControl.Transaction finishTransaction) { }
 
         /**
          * Sets transition animation scale settings value to handler.