Add predictive back support in origin transition backend.

Flag: com.google.android.clockwork.systemui.flags.transitions_enable_origin_transitions_backend
Bug: 381106230
Test: manual test using the test app
Change-Id: I3abc797021de37d6a6596bcc1b79f7d40aa15eac

3 files changed, 161 insertions, 40 deletions

diff --git a/packages/SystemUI/animation/lib/src/com/android/systemui/animation/OriginRemoteTransition.java b/packages/SystemUI/animation/lib/src/com/android/systemui/animation/OriginRemoteTransition.java
index 0f5e3679cc5f..ca2b9578f2be 100644
--- a/packages/SystemUI/animation/lib/src/com/android/systemui/animation/OriginRemoteTransition.java
+++ b/packages/SystemUI/animation/lib/src/com/android/systemui/animation/OriginRemoteTransition.java
@@ -16,6 +16,8 @@
 
 package com.android.systemui.animation;
 
+import static android.view.WindowManager.TRANSIT_CHANGE;
+
 import android.animation.Animator;
 import android.animation.Animator.AnimatorListener;
 import android.animation.ValueAnimator;
@@ -39,6 +41,7 @@ import com.android.internal.policy.ScreenDecorationsUtils;
 import com.android.wm.shell.shared.TransitionUtil;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 
 /**
@@ -90,7 +93,7 @@ public class OriginRemoteTransition extends IRemoteTransition.Stub {
                 () -> {
                     mStartTransaction = t;
                     mFinishCallback = finishCallback;
-                    startAnimationInternal(info);
+                    startAnimationInternal(info, /* states= */ null);
                 });
     }
 
@@ -112,7 +115,13 @@ public class OriginRemoteTransition extends IRemoteTransition.Stub {
             SurfaceControl.Transaction t,
             IRemoteTransitionFinishedCallback finishCallback,
             WindowAnimationState[] states) {
-        logD("takeOverAnimation - " + info);
+        logD("takeOverAnimation - info=" + info + ", states=" + Arrays.toString(states));
+        mHandler.post(
+                () -> {
+                    mStartTransaction = t;
+                    mFinishCallback = finishCallback;
+                    startAnimationInternal(info, states);
+                });
     }
 
     @Override
@@ -121,14 +130,19 @@ public class OriginRemoteTransition extends IRemoteTransition.Stub {
         mHandler.post(this::cancel);
     }
 
-    private void startAnimationInternal(TransitionInfo info) {
+    private void startAnimationInternal(
+            TransitionInfo info, @Nullable WindowAnimationState[] states) {
         if (!prepareUIs(info)) {
             logE("Unable to prepare UI!");
             finishAnimation(/* finished= */ false);
             return;
         }
         // Notify player that we are starting.
-        mPlayer.onStart(info, mStartTransaction, mOrigin, mOriginTransaction);
+        mPlayer.onStart(info, states, mStartTransaction, mOrigin, mOriginTransaction);
+
+        // Apply the initial transactions in case the player forgot to apply them.
+        mOriginTransaction.commit();
+        mStartTransaction.apply();
 
         // Start the animator.
         mAnimator = ValueAnimator.ofFloat(0.0f, 1.0f);
@@ -205,7 +219,8 @@ public class OriginRemoteTransition extends IRemoteTransition.Stub {
                             .setCornerRadius(leash, windowRadius)
                             .setWindowCrop(leash, bounds.width(), bounds.height());
                 }
-            } else if (TransitionUtil.isClosingMode(mode)) {
+            } else if (TransitionUtil.isClosingMode(mode) || mode == TRANSIT_CHANGE) {
+                // TRANSIT_CHANGE refers to the closing window in predictive back animation.
                 closingSurfaces.add(change.getLeash());
                 // For closing surfaces, starting bounds are base bounds. Apply corner radius if
                 // it's full screen.
@@ -236,13 +251,8 @@ public class OriginRemoteTransition extends IRemoteTransition.Stub {
 
         // Attach origin UIComponent to origin leash.
         mOriginTransaction = mOrigin.newTransaction();
-        mOriginTransaction
-                .attachToTransitionLeash(
-                        mOrigin, mOriginLeash, displayBounds.width(), displayBounds.height())
-                .commit();
-
-        // Apply all surface changes.
-        mStartTransaction.apply();
+        mOriginTransaction.attachToTransitionLeash(
+                mOrigin, mOriginLeash, displayBounds.width(), displayBounds.height());
         return true;
     }
 
@@ -328,6 +338,58 @@ public class OriginRemoteTransition extends IRemoteTransition.Stub {
                 /* baseBounds= */ maxBounds);
     }
 
+    private static void applyWindowAnimationStates(
+            TransitionInfo info,
+            @Nullable WindowAnimationState[] states,
+            UIComponent closingApp,
+            UIComponent openingApp) {
+        if (states == null) {
+            // Nothing to apply.
+            return;
+        }
+        // Calculate bounds.
+        Rect maxClosingBounds = new Rect();
+        Rect maxOpeningBounds = new Rect();
+        for (int i = 0; i < info.getChanges().size(); i++) {
+            Rect bound = getBounds(states[i]);
+            if (bound == null) {
+                continue;
+            }
+            int mode = info.getChanges().get(i).getMode();
+            if (TransitionUtil.isOpeningMode(mode)) {
+                maxOpeningBounds.union(bound);
+            } else if (TransitionUtil.isClosingMode(mode) || mode == TRANSIT_CHANGE) {
+                // TRANSIT_CHANGE refers to the closing window in predictive back animation.
+                maxClosingBounds.union(bound);
+            }
+        }
+
+        // Intentionally use a new transaction instead of reusing the existing transaction since we
+        // want to apply window animation states first without committing any other pending changes
+        // in the existing transaction. The existing transaction is expected to be committed by the
+        // onStart() client callback together with client's custom transformation.
+        UIComponent.Transaction transaction = closingApp.newTransaction();
+        if (!maxClosingBounds.isEmpty()) {
+            logD("Applying closing window bounds: " + maxClosingBounds);
+            transaction.setBounds(closingApp, maxClosingBounds);
+        }
+        if (!maxOpeningBounds.isEmpty()) {
+            logD("Applying opening window bounds: " + maxOpeningBounds);
+            transaction.setBounds(openingApp, maxOpeningBounds);
+        }
+        transaction.commit();
+    }
+
+    @Nullable
+    private static Rect getBounds(@Nullable WindowAnimationState state) {
+        if (state == null || state.bounds == null) {
+            return null;
+        }
+        Rect out = new Rect();
+        state.bounds.roundOut(out);
+        return out;
+    }
+
     /**
      * An interface that represents an origin transitions.
      *
@@ -338,9 +400,14 @@ public class OriginRemoteTransition extends IRemoteTransition.Stub {
         /**
          * Called when an origin transition starts. This method exposes the raw {@link
          * TransitionInfo} so that clients can extract more information from it.
+         *
+         * <p>Note: if this transition is taking over a predictive back animation, the {@link
+         * WindowAnimationState} will be passed to this method. The concrete implementation is
+         * expected to apply the {@link WindowAnimationState} before continuing the transition.
          */
         default void onStart(
                 TransitionInfo transitionInfo,
+                @Nullable WindowAnimationState[] states,
                 SurfaceControl.Transaction sfTransaction,
                 UIComponent origin,
                 UIComponent.Transaction uiTransaction) {
@@ -351,12 +418,15 @@ public class OriginRemoteTransition extends IRemoteTransition.Stub {
                             .registerTransactionForClass(
                                     SurfaceUIComponent.class,
                                     new SurfaceUIComponent.Transaction(sfTransaction));
-            // Wrap surfaces and start.
-            onStart(
-                    transactions,
-                    origin,
-                    wrapSurfaces(transitionInfo, /* isOpening= */ false),
-                    wrapSurfaces(transitionInfo, /* isOpening= */ true));
+            // Wrap surfaces.
+            UIComponent closingApp = wrapSurfaces(transitionInfo, /* isOpening= */ false);
+            UIComponent openingApp = wrapSurfaces(transitionInfo, /* isOpening= */ true);
+
+            // Restore the pending animation states coming from predictive back transition.
+            applyWindowAnimationStates(transitionInfo, states, closingApp, openingApp);
+
+            // Start.
+            onStart(transactions, origin, closingApp, openingApp);
         }
 
         /**
diff --git a/packages/SystemUI/animation/lib/src/com/android/systemui/animation/ViewUIComponent.java b/packages/SystemUI/animation/lib/src/com/android/systemui/animation/ViewUIComponent.java
index 9cef43c3deba..cec740ad899b 100644
--- a/packages/SystemUI/animation/lib/src/com/android/systemui/animation/ViewUIComponent.java
+++ b/packages/SystemUI/animation/lib/src/com/android/systemui/animation/ViewUIComponent.java
@@ -89,7 +89,6 @@ public class ViewUIComponent implements UIComponent {
         mSurfaceControl =
                 new SurfaceControl.Builder().setName("ViewUIComponent").setBufferSize(w, h).build();
         mSurface = new Surface(mSurfaceControl);
-        forceDraw();
 
         // Attach surface to transition leash
         SurfaceControl.Transaction t = new SurfaceControl.Transaction();
@@ -99,7 +98,13 @@ public class ViewUIComponent implements UIComponent {
         mView.getViewTreeObserver().addOnDrawListener(mOnDrawListener);
 
         // Make the view invisible AFTER the surface is shown.
-        t.addTransactionCommittedListener(mView::post, () -> mView.setVisibility(View.INVISIBLE))
+        t.addTransactionCommittedListener(
+                        mView::post,
+                        () -> {
+                            logD("Surface attached!");
+                            forceDraw();
+                            mView.setVisibility(View.INVISIBLE);
+                        })
                 .apply();
     }
 
diff --git a/packages/SystemUI/animation/lib/src/com/android/systemui/animation/server/IOriginTransitionsImpl.java b/packages/SystemUI/animation/lib/src/com/android/systemui/animation/server/IOriginTransitionsImpl.java
index 3cbb688ce7ca..6b26ac5f4692 100644
--- a/packages/SystemUI/animation/lib/src/com/android/systemui/animation/server/IOriginTransitionsImpl.java
+++ b/packages/SystemUI/animation/lib/src/com/android/systemui/animation/server/IOriginTransitionsImpl.java
@@ -16,8 +16,10 @@
 
 package com.android.systemui.animation.server;
 
+import static android.view.WindowManager.TRANSIT_CHANGE;
 import static android.view.WindowManager.TRANSIT_CLOSE;
 import static android.view.WindowManager.TRANSIT_OPEN;
+import static android.view.WindowManager.TRANSIT_PREPARE_BACK_NAVIGATION;
 import static android.view.WindowManager.TRANSIT_TO_BACK;
 import static android.view.WindowManager.TRANSIT_TO_FRONT;
 
@@ -26,6 +28,7 @@ import android.annotation.Nullable;
 import android.app.TaskInfo;
 import android.content.ComponentName;
 import android.content.Context;
+import android.os.Build;
 import android.os.IBinder;
 import android.os.RemoteException;
 import android.util.ArrayMap;
@@ -51,8 +54,8 @@ import java.util.function.Predicate;
 
 /** An implementation of the {@link IOriginTransitions}. */
 public class IOriginTransitionsImpl extends IOriginTransitions.Stub {
-    private static final boolean DEBUG = true;
     private static final String TAG = "OriginTransitions";
+    private static final boolean DEBUG = Build.IS_USERDEBUG || Log.isLoggable(TAG, Log.DEBUG);
 
     private final Object mLock = new Object();
     private final ShellTransitions mShellTransitions;
@@ -149,18 +152,7 @@ public class IOriginTransitionsImpl extends IOriginTransitions.Stub {
             if (DEBUG) {
                 Log.d(TAG, "startAnimation: " + info);
             }
-            if (!mOnStarting.test(info)) {
-                Log.w(TAG, "Skipping cancelled transition " + mTransition);
-                t.addTransactionCommittedListener(
-                                mExecutor,
-                                () -> {
-                                    try {
-                                        finishCallback.onTransitionFinished(null, null);
-                                    } catch (RemoteException e) {
-                                        Log.e(TAG, "Unable to report finish.", e);
-                                    }
-                                })
-                        .apply();
+            if (maybeInterceptTransition(info, t, finishCallback)) {
                 return;
             }
             mTransition.startAnimation(token, info, t, finishCallback);
@@ -191,6 +183,9 @@ public class IOriginTransitionsImpl extends IOriginTransitions.Stub {
             if (DEBUG) {
                 Log.d(TAG, "takeOverAnimation: " + info);
             }
+            if (maybeInterceptTransition(info, t, finishCallback)) {
+                return;
+            }
             mTransition.takeOverAnimation(transition, info, t, finishCallback, states);
         }
 
@@ -207,6 +202,27 @@ public class IOriginTransitionsImpl extends IOriginTransitions.Stub {
         public String toString() {
             return "RemoteTransitionDelegate{transition=" + mTransition + "}";
         }
+
+        private boolean maybeInterceptTransition(
+                TransitionInfo info,
+                SurfaceControl.Transaction t,
+                IRemoteTransitionFinishedCallback finishCallback) {
+            if (!mOnStarting.test(info)) {
+                Log.w(TAG, "Intercepting cancelled transition " + mTransition);
+                t.addTransactionCommittedListener(
+                                mExecutor,
+                                () -> {
+                                    try {
+                                        finishCallback.onTransitionFinished(null, null);
+                                    } catch (RemoteException e) {
+                                        Log.e(TAG, "Unable to report finish.", e);
+                                    }
+                                })
+                        .apply();
+                return true;
+            }
+            return false;
+        }
     }
 
     /** A data record containing the origin transition pieces. */
@@ -229,13 +245,25 @@ public class IOriginTransitionsImpl extends IOriginTransitions.Stub {
                 if (mDestroyed) {
                     return false;
                 }
-                TransitionFilter filter = createFilterForReverseTransition(info);
+                TransitionFilter filter =
+                        createFilterForReverseTransition(
+                                info, /* forPredictiveBackTakeover= */ false);
                 if (filter != null) {
                     if (DEBUG) {
                         Log.d(TAG, "Registering filter " + filter);
                     }
                     mShellTransitions.registerRemote(filter, mWrappedReturnTransition);
                 }
+                TransitionFilter takeoverFilter =
+                        createFilterForReverseTransition(
+                                info, /* forPredictiveBackTakeover= */ true);
+                if (takeoverFilter != null) {
+                    if (DEBUG) {
+                        Log.d(TAG, "Registering filter for takeover " + takeoverFilter);
+                    }
+                    mShellTransitions.registerRemoteForTakeover(
+                            takeoverFilter, mWrappedReturnTransition);
+                }
                 return true;
             }
         }
@@ -331,7 +359,8 @@ public class IOriginTransitionsImpl extends IOriginTransitions.Stub {
         }
 
         @Nullable
-        private static TransitionFilter createFilterForReverseTransition(TransitionInfo info) {
+        private static TransitionFilter createFilterForReverseTransition(
+                TransitionInfo info, boolean forPredictiveBackTakeover) {
             TaskInfo launchingTaskInfo = null;
             TaskInfo launchedTaskInfo = null;
             ComponentName launchingActivity = null;
@@ -365,7 +394,9 @@ public class IOriginTransitionsImpl extends IOriginTransitions.Stub {
             if (DEBUG) {
                 Log.d(
                         TAG,
-                        "createFilterForReverseTransition: launchingTaskInfo="
+                        "createFilterForReverseTransition: forPredictiveBackTakeover="
+                                + forPredictiveBackTakeover
+                                + ", launchingTaskInfo="
                                 + launchingTaskInfo
                                 + ", launchedTaskInfo="
                                 + launchedTaskInfo
@@ -395,8 +426,20 @@ public class IOriginTransitionsImpl extends IOriginTransitions.Stub {
                                 + " cookie!");
                 return null;
             }
+            if (forPredictiveBackTakeover && launchedTaskInfo == null) {
+                // Predictive back take over currently only support cross-task transition.
+                Log.d(
+                        TAG,
+                        "createFilterForReverseTransition: skipped - unable to find launched task"
+                                + " for predictive back takeover");
+                return null;
+            }
             TransitionFilter filter = new TransitionFilter();
-            filter.mTypeSet = new int[] {TRANSIT_CLOSE, TRANSIT_TO_BACK};
+            if (forPredictiveBackTakeover) {
+                filter.mTypeSet = new int[] {TRANSIT_PREPARE_BACK_NAVIGATION};
+            } else {
+                filter.mTypeSet = new int[] {TRANSIT_CLOSE, TRANSIT_TO_BACK};
+            }
 
             // The opening activity of the return transition must match the activity we just closed.
             TransitionFilter.Requirement req1 = new TransitionFilter.Requirement();
@@ -405,15 +448,18 @@ public class IOriginTransitionsImpl extends IOriginTransitions.Stub {
                     launchingActivity == null ? launchingTaskInfo.topActivity : launchingActivity;
 
             TransitionFilter.Requirement req2 = new TransitionFilter.Requirement();
-            req2.mModes = new int[] {TRANSIT_CLOSE, TRANSIT_TO_BACK};
+            if (forPredictiveBackTakeover) {
+                req2.mModes = new int[] {TRANSIT_CHANGE};
+            } else {
+                req2.mModes = new int[] {TRANSIT_CLOSE, TRANSIT_TO_BACK};
+            }
             if (launchedTaskInfo != null) {
                 // For task transitions, the closing task's cookie must match the task we just
                 // launched.
                 req2.mLaunchCookie = launchedTaskInfo.launchCookies.get(0);
             } else {
                 // For activity transitions, the closing activity of the return transition must
-                // match
-                // the activity we just launched.
+                // match the activity we just launched.
                 req2.mTopActivity = launchedActivity;
             }