backnav2: fixes to cross-activity flow in BackNavigationController

- Minor refactor to move screenshotting logic to prepare for a cache
  implementation.
- Prepare animation for CROSS_ACTIVITY type.

Test: atest
WmTests:com.android.server.wm.BackNavigationControllerTests#backTypeCrossActivityWhenBackToPreviousActivity
also tested locally with other changes

Bug: 207481538
Change-Id: I5b5cd35fd0a80f4ac0e03e07f86e15619db092e6

services/core/java/com/android/server/wm/BackNavigationController.java

@@ -24,7 +24,6 @@ import static com.android.internal.protolog.ProtoLogGroup.WM_DEBUG_BACK_PREVIEW;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
-import android.content.ComponentName;
 import android.graphics.Point;
 import android.graphics.Rect;
 import android.hardware.HardwareBuffer;
@@ -41,7 +40,6 @@ import android.window.BackAnimationAdapter;
 import android.window.BackNavigationInfo;
 import android.window.IBackAnimationFinishedCallback;
 import android.window.OnBackInvokedCallbackInfo;
-import android.window.ScreenCapture;
 import android.window.TaskSnapshot;
 import android.window.WindowContainerToken;
 
@@ -269,10 +267,11 @@ class BackNavigationController {
                     removedWindowContainer,
                     BackNavigationInfo.typeToString(backType));
 
-            // For now, we only animate when going home and cross task.
+            // For now, we only animate when going home, cross task or cross-activity.
             boolean prepareAnimation =
                     (backType == BackNavigationInfo.TYPE_RETURN_TO_HOME
-                            || backType == BackNavigationInfo.TYPE_CROSS_TASK)
+                            || backType == BackNavigationInfo.TYPE_CROSS_TASK
+                            || backType == BackNavigationInfo.TYPE_CROSS_ACTIVITY)
                     && adapter != null;
 
             // Only prepare animation if no leash has been created (no animation is running).
@@ -336,6 +335,7 @@ class BackNavigationController {
         RemoteAnimationTarget behindAppTarget = null;
         if (needsScreenshot(backType)) {
             HardwareBuffer screenshotBuffer = null;
+            Task backTargetTask = prevTask;
             switch(backType) {
                 case BackNavigationInfo.TYPE_CROSS_TASK:
                     int prevTaskId = prevTask != null ? prevTask.mTaskId : 0;
@@ -343,14 +343,10 @@ class BackNavigationController {
                     screenshotBuffer = getTaskSnapshot(prevTaskId, prevUserId);
                     break;
                 case BackNavigationInfo.TYPE_CROSS_ACTIVITY:
-                    //TODO(207481538) Remove once the infrastructure to support per-activity
-                    // screenshot is implemented. For now we simply have the mBackScreenshots hash
-                    // map that dumbly saves the screenshots.
-                    if (prevActivity != null
-                            && prevActivity.mActivityComponent != null) {
-                        screenshotBuffer =
-                                getActivitySnapshot(currentTask, prevActivity.mActivityComponent);
+                    if (prevActivity != null && prevActivity.mActivityComponent != null) {
+                        screenshotBuffer = getActivitySnapshot(currentTask, prevActivity);
                     }
+                    backTargetTask = currentTask;
                     break;
             }
 
@@ -369,8 +365,9 @@ class BackNavigationController {
                 // leash needs to be added before to be in the synchronized block.
                 startedTransaction.setLayer(topAppTarget.leash, 1);
 
-                behindAppTarget = createRemoteAnimationTargetLocked(
-                        prevTask, screenshotSurface, MODE_OPENING);
+                behindAppTarget =
+                        createRemoteAnimationTargetLocked(
+                                backTargetTask, screenshotSurface, MODE_OPENING);
 
                 // reset leash after animation finished.
                 leashes.add(screenshotSurface);
@@ -543,14 +540,8 @@ class BackNavigationController {
         mShowWallpaper = false;
     }
 
-    private HardwareBuffer getActivitySnapshot(@NonNull Task task,
-            ComponentName activityComponent) {
-        // Check if we have a screenshot of the previous activity, indexed by its
-        // component name.
-        ScreenCapture.ScreenshotHardwareBuffer backBuffer = task.mBackScreenshots
-                .get(activityComponent.flattenToString());
-        return backBuffer != null ? backBuffer.getHardwareBuffer() : null;
-
+    private HardwareBuffer getActivitySnapshot(@NonNull Task task, ActivityRecord r) {
+        return task.getSnapshotForActivityRecord(r);
     }
 
     private HardwareBuffer getTaskSnapshot(int taskId, int userId) {

services/core/java/com/android/server/wm/TaskFragment.java

@@ -79,6 +79,7 @@ import android.app.servertransaction.ResumeActivityItem;
 import android.content.res.Configuration;
 import android.graphics.Point;
 import android.graphics.Rect;
+import android.hardware.HardwareBuffer;
 import android.os.IBinder;
 import android.os.RemoteException;
 import android.os.UserHandle;
@@ -1859,7 +1860,6 @@ class TaskFragment extends WindowContainer<WindowContainer> {
         super.addChild(child, index);
 
         if (isAddingActivity && task != null) {
-
             // TODO(b/207481538): temporary per-activity screenshoting
             if (r != null && BackNavigationController.isScreenshotEnabled()) {
                 ProtoLog.v(WM_DEBUG_BACK_PREVIEW, "Screenshotting Activity %s",
@@ -2528,6 +2528,19 @@ class TaskFragment extends WindowContainer<WindowContainer> {
         return !mCreatedByOrganizer || mIsRemovalRequested;
     }
 
+    @Nullable
+    HardwareBuffer getSnapshotForActivityRecord(@Nullable ActivityRecord r) {
+        if (!BackNavigationController.isScreenshotEnabled()) {
+            return null;
+        }
+        if (r != null && r.mActivityComponent != null) {
+            ScreenCapture.ScreenshotHardwareBuffer backBuffer =
+                    mBackScreenshots.get(r.mActivityComponent.flattenToString());
+            return backBuffer != null ? backBuffer.getHardwareBuffer() : null;
+        }
+        return null;
+    }
+
     @Override
     void removeChild(WindowContainer child) {
         removeChild(child, true /* removeSelfIfPossible */);

services/tests/wmtests/src/com/android/server/wm/BackNavigationControllerTests.java

@@ -115,18 +115,14 @@ public class BackNavigationControllerTests extends WindowTestsBase {
 
     @Test
     public void backTypeCrossActivityWhenBackToPreviousActivity() {
-        Task task = createTopTaskWithActivity();
-        WindowState window = createAppWindow(task, FIRST_APPLICATION_WINDOW, "window");
-        addToWindowMap(window, true);
-        IOnBackInvokedCallback callback = createOnBackInvokedCallback();
-        window.setOnBackInvokedCallbackInfo(
-                new OnBackInvokedCallbackInfo(callback, OnBackInvokedDispatcher.PRIORITY_SYSTEM));
+        CrossActivityTestCase testCase = createTopTaskWithTwoActivities();
+        IOnBackInvokedCallback callback = withSystemCallback(testCase.task);
+
         BackNavigationInfo backNavigationInfo = startBackNavigation();
         assertWithMessage("BackNavigationInfo").that(backNavigationInfo).isNotNull();
+        assertThat(backNavigationInfo.getOnBackInvokedCallback()).isEqualTo(callback);
         assertThat(typeToString(backNavigationInfo.getType()))
                 .isEqualTo(typeToString(BackNavigationInfo.TYPE_CROSS_ACTIVITY));
-        assertWithMessage("Activity callback").that(
-                backNavigationInfo.getOnBackInvokedCallback()).isEqualTo(callback);
     }
 
     @Test
@@ -302,6 +298,34 @@ public class BackNavigationControllerTests extends WindowTestsBase {
         return task;
     }
 
+    @NonNull
+    private CrossActivityTestCase createTopTaskWithTwoActivities() {
+        Task task = createTask(mDefaultDisplay);
+        ActivityRecord record1 = createActivityRecord(task);
+        ActivityRecord record2 = createActivityRecord(task);
+        // enable OnBackInvokedCallbacks
+        record2.info.applicationInfo.privateFlagsExt |=
+                PRIVATE_FLAG_EXT_ENABLE_ON_BACK_INVOKED_CALLBACK;
+        WindowState window1 = createWindow(null, FIRST_APPLICATION_WINDOW, record1, "window1");
+        WindowState window2 = createWindow(null, FIRST_APPLICATION_WINDOW, record2, "window2");
+        when(task.mSurfaceControl.isValid()).thenReturn(true);
+        when(record1.mSurfaceControl.isValid()).thenReturn(true);
+        when(record2.mSurfaceControl.isValid()).thenReturn(true);
+        Mockito.doNothing().when(task).reparentSurfaceControl(any(), any());
+        Mockito.doNothing().when(record1).reparentSurfaceControl(any(), any());
+        Mockito.doNothing().when(record2).reparentSurfaceControl(any(), any());
+        mAtm.setFocusedTask(task.mTaskId, record1);
+        mAtm.setFocusedTask(task.mTaskId, record2);
+        addToWindowMap(window1, true);
+        addToWindowMap(window2, true);
+
+        CrossActivityTestCase testCase = new CrossActivityTestCase();
+        testCase.task = task;
+        testCase.recordBack = record1;
+        testCase.recordFront = record2;
+        return testCase;
+    }
+
     private void addToWindowMap(WindowState window, boolean focus) {
         mWm.mWindowMap.put(window.mClient.asBinder(), window);
         if (focus) {
@@ -310,4 +334,10 @@ public class BackNavigationControllerTests extends WindowTestsBase {
             doReturn(window).when(mWm).getFocusedWindowLocked();
         }
     }
+
+    private class CrossActivityTestCase {
+        public Task task;
+        public ActivityRecord recordBack;
+        public ActivityRecord recordFront;
+    }
 }