Extend recents mixer to handle finishing with a bookend transition

- Currently, splitscreen will listen for the recents transition to start
  (to include the divider in the transition) and end (to restore the
  divider visibility if returning to the app).  It is notified of the
  recents events by RecentsMixedTransition, which wraps the finish
  callback, and passes the existing finishWCT (for the transition that
  starts the recents transition), which split will then query for changes
  that reorder the split stages to the top.

  However, when the recents bookend transitions are enabled, the
  restoring of the transient order is done in a separate transition,
  and the existing logic of looking at the finishWCT fails.

  Instead, we can extend the RecentsMixedHandler callbacks to allow
  RecentsMixedTransition to be notified when a finish is about to
  happen with direct information about whether the finish will return
  to the app.  Using this, split can then modify the finishWCT to
  include any necessary cleanup.

Bug: 346588978
Bug: 391486839
Flag: com.android.wm.shell.enable_recents_bookend_transition
Test: atest WMShellUnitTests
Change-Id: Ifebd4fbde6c7871832309b362553286f017bf2c7

5 files changed, 102 insertions, 19 deletions

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentsTransitionHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentsTransitionHandler.java
index afc6fee2eca3..55133780f517 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentsTransitionHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentsTransitionHandler.java
@@ -222,7 +222,7 @@ public class RecentsTransitionHandler implements Transitions.TransitionHandler,
         RecentsMixedHandler mixer = null;
         Consumer<IBinder> setTransitionForMixer = null;
         for (int i = 0; i < mMixers.size(); ++i) {
-            setTransitionForMixer = mMixers.get(i).handleRecentsRequest(wct);
+            setTransitionForMixer = mMixers.get(i).handleRecentsRequest();
             if (setTransitionForMixer != null) {
                 mixer = mMixers.get(i);
                 break;
@@ -1455,6 +1455,11 @@ public class RecentsTransitionHandler implements Transitions.TransitionHandler,
                 }
             }
 
+            // Notify the mixers of the pending finish
+            for (int i = 0; i < mMixers.size(); ++i) {
+                mMixers.get(i).handleFinishRecents(returningToApp, wct, t);
+            }
+
             if (Flags.enableRecentsBookendTransition()) {
                 if (!wct.isEmpty()) {
                     ProtoLog.v(ShellProtoLogGroup.WM_SHELL_RECENTS_TRANSITION,
@@ -1653,15 +1658,22 @@ public class RecentsTransitionHandler implements Transitions.TransitionHandler,
      */
     public interface RecentsMixedHandler extends Transitions.TransitionHandler {
         /**
-         * Called when a recents request comes in. The handler can add operations to outWCT. If
-         * the handler wants to "accept" the transition, it should return a Consumer accepting the
-         * IBinder for the transition. If not, it should return `null`.
+         * Called when a recents request comes in. If the handler wants to "accept" the transition,
+         * it should return a Consumer accepting the IBinder for the transition. If not, it should
+         * return `null`.
          *
          * If a mixed-handler accepts this recents, it will be the de-facto handler for this
          * transition and is required to call the associated {@link #startAnimation},
          * {@link #mergeAnimation}, and {@link #onTransitionConsumed} methods.
          */
         @Nullable
-        Consumer<IBinder> handleRecentsRequest(WindowContainerTransaction outWCT);
+        Consumer<IBinder> handleRecentsRequest();
+
+        /**
+         * Called when a recents transition has finished, with a WCT and SurfaceControl Transaction
+         * that can be used to add to any changes needed to restore the state.
+         */
+        void handleFinishRecents(boolean returnToApp, @NonNull WindowContainerTransaction finishWct,
+                @NonNull SurfaceControl.Transaction finishT);
     }
 }
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 511e426cc681..722494c05e32 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
@@ -129,6 +129,7 @@ import com.android.internal.logging.InstanceId;
 import com.android.internal.policy.FoldLockSettingsObserver;
 import com.android.internal.protolog.ProtoLog;
 import com.android.launcher3.icons.IconProvider;
+import com.android.wm.shell.Flags;
 import com.android.wm.shell.R;
 import com.android.wm.shell.RootTaskDisplayAreaOrganizer;
 import com.android.wm.shell.ShellTaskOrganizer;
@@ -3766,13 +3767,31 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         mTaskOrganizer.applyTransaction(wct);
     }
 
+    public void onRecentsInSplitAnimationFinishing(boolean returnToApp,
+            @NonNull WindowContainerTransaction finishWct,
+            @NonNull SurfaceControl.Transaction finishT) {
+        if (!Flags.enableRecentsBookendTransition()) {
+            // The non-bookend recents transition case will be handled by
+            // RecentsMixedTransition wrapping the finish callback and calling
+            // onRecentsInSplitAnimationFinish()
+            return;
+        }
+
+        onRecentsInSplitAnimationFinishInner(returnToApp, finishWct, finishT);
+    }
+
     /** Call this when the recents animation during split-screen finishes. */
-    public void onRecentsInSplitAnimationFinish(WindowContainerTransaction finishWct,
-            SurfaceControl.Transaction finishT) {
-        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "onRecentsInSplitAnimationFinish");
-        mPausingTasks.clear();
+    public void onRecentsInSplitAnimationFinish(@NonNull WindowContainerTransaction finishWct,
+            @NonNull SurfaceControl.Transaction finishT) {
+        if (Flags.enableRecentsBookendTransition()) {
+            // The bookend recents transition case will be handled by
+            // onRecentsInSplitAnimationFinishing above
+            return;
+        }
+
         // Check if the recent transition is finished by returning to the current
         // split, so we can restore the divider bar.
+        boolean returnToApp = false;
         for (int i = 0; i < finishWct.getHierarchyOps().size(); ++i) {
             final WindowContainerTransaction.HierarchyOp op =
                     finishWct.getHierarchyOps().get(i);
@@ -3787,13 +3806,26 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
             }
             if (op.getType() == HIERARCHY_OP_TYPE_REORDER && op.getToTop()
                     && anyStageContainsContainer) {
-                updateSurfaceBounds(mSplitLayout, finishT,
-                        false /* applyResizingOffset */);
-                finishT.reparent(mSplitLayout.getDividerLeash(), mRootTaskLeash);
-                setDividerVisibility(true, finishT);
-                return;
+                returnToApp = true;
             }
         }
+        onRecentsInSplitAnimationFinishInner(returnToApp, finishWct, finishT);
+    }
+
+    /** Call this when the recents animation during split-screen finishes. */
+    public void onRecentsInSplitAnimationFinishInner(boolean returnToApp,
+            @NonNull WindowContainerTransaction finishWct,
+            @NonNull SurfaceControl.Transaction finishT) {
+        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "onRecentsInSplitAnimationFinish: returnToApp=%b",
+                returnToApp);
+        mPausingTasks.clear();
+        if (returnToApp) {
+            updateSurfaceBounds(mSplitLayout, finishT,
+                    false /* applyResizingOffset */);
+            finishT.reparent(mSplitLayout.getDividerLeash(), mRootTaskLeash);
+            setDividerVisibility(true, finishT);
+            return;
+        }
 
         setSplitsVisible(false);
         finishWct.setReparentLeafTaskIfRelaunch(mRootTaskInfo.token,
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 2177986bccd5..d8e7c2ccb15f 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
@@ -367,7 +367,7 @@ public class DefaultMixedHandler implements MixedTransitionHandler,
     }
 
     @Override
-    public Consumer<IBinder> handleRecentsRequest(WindowContainerTransaction outWCT) {
+    public Consumer<IBinder> handleRecentsRequest() {
         if (mRecentsHandler != null) {
             if (mSplitHandler.isSplitScreenVisible()) {
                 return this::setRecentsTransitionDuringSplit;
@@ -383,6 +383,21 @@ public class DefaultMixedHandler implements MixedTransitionHandler,
         return null;
     }
 
+    @Override
+    public void handleFinishRecents(boolean returnToApp,
+            @NonNull WindowContainerTransaction finishWct,
+            @NonNull SurfaceControl.Transaction finishT) {
+        if (mRecentsHandler != null) {
+            for (int i = mActiveTransitions.size() - 1; i >= 0; --i) {
+                final MixedTransition mixed = mActiveTransitions.get(i);
+                if (mixed.mType == MixedTransition.TYPE_RECENTS_DURING_SPLIT) {
+                    ((RecentsMixedTransition) mixed).onAnimateRecentsDuringSplitFinishing(
+                            returnToApp, finishWct, finishT);
+                }
+            }
+        }
+    }
+
     private void setRecentsTransitionDuringSplit(IBinder transition) {
         ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS, " Got a recents request while "
                 + "Split-Screen is foreground, so treat it as Mixed.");
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/RecentsMixedTransition.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/RecentsMixedTransition.java
index 8cdbe26a2c76..1847af07f275 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/RecentsMixedTransition.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/RecentsMixedTransition.java
@@ -159,6 +159,8 @@ class RecentsMixedTransition extends DefaultMixedHandler.MixedTransition {
             // If pair-to-pair switching, the post-recents clean-up isn't needed.
             wct = wct != null ? wct : new WindowContainerTransaction();
             if (mAnimType != ANIM_TYPE_PAIR_TO_PAIR) {
+                // TODO(b/346588978): Only called if !enableRecentsBookendTransition(), can remove
+                // once that rolls out
                 mSplitHandler.onRecentsInSplitAnimationFinish(wct, finishTransaction);
             } else {
                 // notify pair-to-pair recents animation finish
@@ -177,6 +179,17 @@ class RecentsMixedTransition extends DefaultMixedHandler.MixedTransition {
         return handled;
     }
 
+    /**
+     * Called when the recents animation during split is about to finish.
+     */
+    void onAnimateRecentsDuringSplitFinishing(boolean returnToApp,
+            @NonNull WindowContainerTransaction finishWct,
+            @NonNull SurfaceControl.Transaction finishT) {
+        if (mAnimType != ANIM_TYPE_PAIR_TO_PAIR) {
+            mSplitHandler.onRecentsInSplitAnimationFinishing(returnToApp, finishWct, finishT);
+        }
+    }
+
     @Override
     void mergeAnimation(
             @NonNull IBinder transition, @NonNull TransitionInfo info,
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTransitionTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTransitionTests.java
index 4211e4682810..b9d6a454694d 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTransitionTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTransitionTests.java
@@ -67,6 +67,7 @@ import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.launcher3.icons.IconProvider;
+import com.android.wm.shell.Flags;
 import com.android.wm.shell.MockToken;
 import com.android.wm.shell.RootTaskDisplayAreaOrganizer;
 import com.android.wm.shell.ShellTaskOrganizer;
@@ -355,8 +356,13 @@ public class SplitTransitionTests extends ShellTestCase {
 
         // Make sure it cleans-up if recents doesn't restore
         WindowContainerTransaction commitWCT = new WindowContainerTransaction();
-        mStageCoordinator.onRecentsInSplitAnimationFinish(commitWCT,
-                mock(SurfaceControl.Transaction.class));
+        if (Flags.enableRecentsBookendTransition()) {
+            mStageCoordinator.onRecentsInSplitAnimationFinishing(false /* returnToApp */, commitWCT,
+                    mock(SurfaceControl.Transaction.class));
+        } else {
+            mStageCoordinator.onRecentsInSplitAnimationFinish(commitWCT,
+                    mock(SurfaceControl.Transaction.class));
+        }
         assertFalse(mStageCoordinator.isSplitScreenVisible());
     }
 
@@ -420,8 +426,13 @@ public class SplitTransitionTests extends ShellTestCase {
         // simulate the restoreWCT being applied:
         mMainStage.onTaskAppeared(mMainChild, mock(SurfaceControl.class));
         mSideStage.onTaskAppeared(mSideChild, mock(SurfaceControl.class));
-        mStageCoordinator.onRecentsInSplitAnimationFinish(restoreWCT,
-                mock(SurfaceControl.Transaction.class));
+        if (Flags.enableRecentsBookendTransition()) {
+            mStageCoordinator.onRecentsInSplitAnimationFinishing(true /* returnToApp */, restoreWCT,
+                    mock(SurfaceControl.Transaction.class));
+        } else {
+            mStageCoordinator.onRecentsInSplitAnimationFinish(restoreWCT,
+                    mock(SurfaceControl.Transaction.class));
+        }
         assertTrue(mStageCoordinator.isSplitScreenVisible());
     }