Merge "Correct record activity snapshot condition." into main
diff --git a/core/java/android/window/WindowOnBackInvokedDispatcher.java b/core/java/android/window/WindowOnBackInvokedDispatcher.java
index 3d4bc2f..6a8ca33 100644
--- a/core/java/android/window/WindowOnBackInvokedDispatcher.java
+++ b/core/java/android/window/WindowOnBackInvokedDispatcher.java
@@ -41,6 +41,7 @@
 import java.util.HashMap;
 import java.util.Objects;
 import java.util.TreeMap;
+import java.util.function.Supplier;
 
 /**
  * Provides window based implementation of {@link OnBackInvokedDispatcher}.
@@ -271,7 +272,7 @@
      * Returns false if the legacy back behavior should be used.
      */
     public boolean isOnBackInvokedCallbackEnabled() {
-        return Checker.isOnBackInvokedCallbackEnabled(mChecker.getContext());
+        return isOnBackInvokedCallbackEnabled(mChecker.getContext());
     }
 
     /**
@@ -394,7 +395,18 @@
      * {@link OnBackInvokedCallback}.
      */
     public static boolean isOnBackInvokedCallbackEnabled(@NonNull Context context) {
-        return Checker.isOnBackInvokedCallbackEnabled(context);
+        final Context originalContext = context;
+        while ((context instanceof ContextWrapper) && !(context instanceof Activity)) {
+            context = ((ContextWrapper) context).getBaseContext();
+        }
+        final ActivityInfo activityInfo = (context instanceof Activity)
+                ? ((Activity) context).getActivityInfo() : null;
+        final ApplicationInfo applicationInfo = context.getApplicationInfo();
+
+        return WindowOnBackInvokedDispatcher
+                .isOnBackInvokedCallbackEnabled(activityInfo, applicationInfo,
+                        () -> originalContext.obtainStyledAttributes(
+                                new int[] {android.R.attr.windowSwipeToDismiss}), true);
     }
 
     @Override
@@ -426,7 +438,7 @@
          */
         public boolean checkApplicationCallbackRegistration(int priority,
                 OnBackInvokedCallback callback) {
-            if (!isOnBackInvokedCallbackEnabled(getContext())
+            if (!WindowOnBackInvokedDispatcher.isOnBackInvokedCallbackEnabled(getContext())
                     && !(callback instanceof CompatOnBackInvokedCallback)) {
                 Log.w(TAG,
                         "OnBackInvokedCallback is not enabled for the application."
@@ -445,97 +457,76 @@
         private Context getContext() {
             return mContext.get();
         }
+    }
 
-        private static boolean isOnBackInvokedCallbackEnabled(@Nullable Context context) {
-            // new back is enabled if the feature flag is enabled AND the app does not explicitly
-            // request legacy back.
-            boolean featureFlagEnabled = ENABLE_PREDICTIVE_BACK;
-            if (!featureFlagEnabled) {
-                return false;
+    /**
+     * @hide
+     */
+    public static boolean isOnBackInvokedCallbackEnabled(@Nullable ActivityInfo activityInfo,
+            @NonNull ApplicationInfo applicationInfo,
+            @NonNull Supplier<TypedArray> windowAttrSupplier, boolean recycleTypedArray) {
+        // new back is enabled if the feature flag is enabled AND the app does not explicitly
+        // request legacy back.
+        if (!ENABLE_PREDICTIVE_BACK) {
+            return false;
+        }
+
+        if (ALWAYS_ENFORCE_PREDICTIVE_BACK) {
+            return true;
+        }
+
+        boolean requestsPredictiveBack;
+        // Activity
+        if (activityInfo != null && activityInfo.hasOnBackInvokedCallbackEnabled()) {
+            requestsPredictiveBack = activityInfo.isOnBackInvokedCallbackEnabled();
+            if (DEBUG) {
+                Log.d(TAG, TextUtils.formatSimple(
+                        "Activity: %s isPredictiveBackEnabled=%s",
+                        activityInfo.getComponentName(),
+                        requestsPredictiveBack));
             }
-
-            if (ALWAYS_ENFORCE_PREDICTIVE_BACK) {
-                return true;
-            }
-
-            // If the context is null, return false to use legacy back.
-            if (context == null) {
-                Log.w(TAG, "OnBackInvokedCallback is not enabled because context is null.");
-                return false;
-            }
-
-            boolean requestsPredictiveBack = false;
-
-            // Check if the context is from an activity.
-            Context originalContext = context;
-            while ((context instanceof ContextWrapper) && !(context instanceof Activity)) {
-                context = ((ContextWrapper) context).getBaseContext();
-            }
-
-            boolean shouldCheckActivity = false;
-
-            if (context instanceof Activity) {
-                final Activity activity = (Activity) context;
-
-                final ActivityInfo activityInfo = activity.getActivityInfo();
-                if (activityInfo != null) {
-                    if (activityInfo.hasOnBackInvokedCallbackEnabled()) {
-                        shouldCheckActivity = true;
-                        requestsPredictiveBack = activityInfo.isOnBackInvokedCallbackEnabled();
-
-                        if (DEBUG) {
-                            Log.d(TAG, TextUtils.formatSimple(
-                                    "Activity: %s isPredictiveBackEnabled=%s",
-                                    activity.getComponentName(),
-                                    requestsPredictiveBack));
-                        }
-                    }
-                } else {
-                    Log.w(TAG, "The ActivityInfo is null, so we cannot verify if this Activity"
-                            + " has the 'android:enableOnBackInvokedCallback' attribute."
-                            + " The application attribute will be used as a fallback.");
-                }
-            }
-
-            if (!shouldCheckActivity) {
-                final ApplicationInfo applicationInfo = context.getApplicationInfo();
-                requestsPredictiveBack = applicationInfo.isOnBackInvokedCallbackEnabled();
-
-                if (DEBUG) {
-                    Log.d(TAG, TextUtils.formatSimple("App: %s requestsPredictiveBack=%s",
-                            applicationInfo.packageName,
-                            requestsPredictiveBack));
-                }
-
-                if (PREDICTIVE_BACK_FALLBACK_WINDOW_ATTRIBUTE && !requestsPredictiveBack) {
-                    // Compatibility check for legacy window style flag used by Wear OS.
-                    // Note on compatibility behavior:
-                    // 1. windowSwipeToDismiss should be respected for all apps not opted in.
-                    // 2. windowSwipeToDismiss should be true for all apps not opted in, which
-                    //    enables the PB animation for them.
-                    // 3. windowSwipeToDismiss=false should be respected for apps not opted in,
-                    //    which disables PB & onBackPressed caused by BackAnimController's
-                    //    setTrigger(true)
-                    // Use the original context to resolve the styled attribute so that they stay
-                    // true to the window.
-                    TypedArray windowAttr =
-                            originalContext.obtainStyledAttributes(
-                                    new int[] {android.R.attr.windowSwipeToDismiss});
-                    boolean windowSwipeToDismiss = true;
-                    if (windowAttr.getIndexCount() > 0) {
-                        windowSwipeToDismiss = windowAttr.getBoolean(0, true);
-                    }
-                    windowAttr.recycle();
-
-                    if (DEBUG) {
-                        Log.i(TAG, "falling back to windowSwipeToDismiss: " + windowSwipeToDismiss);
-                    }
-
-                    requestsPredictiveBack = windowSwipeToDismiss;
-                }
-            }
-
             return requestsPredictiveBack;
         }
+
+        // Application
+        requestsPredictiveBack = applicationInfo.isOnBackInvokedCallbackEnabled();
+        if (DEBUG) {
+            Log.d(TAG, TextUtils.formatSimple("App: %s requestsPredictiveBack=%s",
+                    applicationInfo.packageName,
+                    requestsPredictiveBack));
+        }
+        if (requestsPredictiveBack) {
+            return true;
+        }
+
+        if (PREDICTIVE_BACK_FALLBACK_WINDOW_ATTRIBUTE) {
+            // Compatibility check for legacy window style flag used by Wear OS.
+            // Note on compatibility behavior:
+            // 1. windowSwipeToDismiss should be respected for all apps not opted in.
+            // 2. windowSwipeToDismiss should be true for all apps not opted in, which
+            //    enables the PB animation for them.
+            // 3. windowSwipeToDismiss=false should be respected for apps not opted in,
+            //    which disables PB & onBackPressed caused by BackAnimController's
+            //    setTrigger(true)
+            // Use the original context to resolve the styled attribute so that they stay
+            // true to the window.
+            TypedArray windowAttr = windowAttrSupplier.get();
+            boolean windowSwipeToDismiss = true;
+            if (windowAttr != null) {
+                if (windowAttr.getIndexCount() > 0) {
+                    windowSwipeToDismiss = windowAttr.getBoolean(0, true);
+                }
+                if (recycleTypedArray) {
+                    windowAttr.recycle();
+                }
+            }
+
+            if (DEBUG) {
+                Log.i(TAG, "falling back to windowSwipeToDismiss: " + windowSwipeToDismiss);
+            }
+
+            requestsPredictiveBack = windowSwipeToDismiss;
+        }
+        return requestsPredictiveBack;
     }
 }
diff --git a/services/core/java/com/android/server/wm/ActivityRecord.java b/services/core/java/com/android/server/wm/ActivityRecord.java
index 04ac9fb..2df965c 100644
--- a/services/core/java/com/android/server/wm/ActivityRecord.java
+++ b/services/core/java/com/android/server/wm/ActivityRecord.java
@@ -353,6 +353,7 @@
 import android.window.TaskSnapshot;
 import android.window.TransitionInfo.AnimationOptions;
 import android.window.WindowContainerToken;
+import android.window.WindowOnBackInvokedDispatcher;
 
 import com.android.internal.R;
 import com.android.internal.annotations.GuardedBy;
@@ -978,6 +979,8 @@
     int mAllowedTouchUid;
     // Whether client has requested a scene transition when exiting.
     final boolean mHasSceneTransition;
+    // Whether the app has opt-in enableOnBackInvokedCallback.
+    final boolean mOptInOnBackInvoked;
 
     // Whether the ActivityEmbedding is enabled on the app.
     private final boolean mAppActivityEmbeddingSplitsEnabled;
@@ -2231,6 +2234,10 @@
             // No such property name.
         }
         mAppActivityEmbeddingSplitsEnabled = appActivityEmbeddingEnabled;
+
+        mOptInOnBackInvoked = WindowOnBackInvokedDispatcher
+                .isOnBackInvokedCallbackEnabled(info, info.applicationInfo,
+                        () -> ent != null ? ent.array : null, false);
     }
 
     /**
diff --git a/services/core/java/com/android/server/wm/SnapshotController.java b/services/core/java/com/android/server/wm/SnapshotController.java
index 01fa39b..b6f040a 100644
--- a/services/core/java/com/android/server/wm/SnapshotController.java
+++ b/services/core/java/com/android/server/wm/SnapshotController.java
@@ -25,6 +25,7 @@
 import static android.view.WindowManager.TRANSIT_TO_FRONT;
 
 import android.os.Trace;
+import android.util.ArrayMap;
 import android.view.WindowManager;
 import android.window.TaskSnapshot;
 
@@ -80,6 +81,7 @@
         if (!isTransitionOpen && !isTransitionClose && type < TRANSIT_FIRST_CUSTOM) {
             return;
         }
+        ActivitiesByTask activityTargets = null;
         for (int i = changeInfos.size() - 1; i >= 0; --i) {
             Transition.ChangeInfo info = changeInfos.get(i);
             // Intentionally skip record snapshot for changes originated from PiP.
@@ -98,14 +100,68 @@
                 final TaskFragment tf = info.mContainer.asTaskFragment();
                 final ActivityRecord ar = tf != null ? tf.getTopMostActivity()
                         : info.mContainer.asActivityRecord();
-                if (ar != null && !ar.isVisibleRequested() && ar.getTask().isVisibleRequested()) {
-                    final WindowState mainWindow = ar.findMainWindow(false);
-                    // Only capture activity snapshot if this app has adapted to back predict
-                    if (mainWindow != null
-                            && mainWindow.getOnBackInvokedCallbackInfo() != null
-                            && mainWindow.getOnBackInvokedCallbackInfo().isSystemCallback()) {
-                        mActivitySnapshotController.recordSnapshot(ar);
+                if (ar != null && ar.getTask().isVisibleRequested()) {
+                    if (activityTargets == null) {
+                        activityTargets = new ActivitiesByTask();
                     }
+                    activityTargets.put(ar);
+                }
+            }
+        }
+        if (activityTargets != null) {
+            activityTargets.recordSnapshot(mActivitySnapshotController);
+        }
+    }
+
+    private static class ActivitiesByTask {
+        final ArrayMap<Task, OpenCloseActivities> mActivitiesMap = new ArrayMap<>();
+
+        void put(ActivityRecord ar) {
+            OpenCloseActivities activities = mActivitiesMap.get(ar.getTask());
+            if (activities == null) {
+                activities = new OpenCloseActivities();
+                mActivitiesMap.put(ar.getTask(), activities);
+            }
+            activities.add(ar);
+        }
+
+        void recordSnapshot(ActivitySnapshotController controller) {
+            for (int i = mActivitiesMap.size() - 1; i >= 0; i--) {
+                final OpenCloseActivities pair = mActivitiesMap.valueAt(i);
+                pair.recordSnapshot(controller);
+            }
+        }
+
+        static class OpenCloseActivities {
+            final ArrayList<ActivityRecord> mOpenActivities = new ArrayList<>();
+            final ArrayList<ActivityRecord> mCloseActivities = new ArrayList<>();
+
+            void add(ActivityRecord ar) {
+                if (ar.isVisibleRequested()) {
+                    mOpenActivities.add(ar);
+                } else {
+                    mCloseActivities.add(ar);
+                }
+            }
+
+            boolean allOpensOptInOnBackInvoked() {
+                if (mOpenActivities.isEmpty()) {
+                    return false;
+                }
+                for (int i = mOpenActivities.size() - 1; i >= 0; --i) {
+                    if (!mOpenActivities.get(i).mOptInOnBackInvoked) {
+                        return false;
+                    }
+                }
+                return true;
+            }
+
+            void recordSnapshot(ActivitySnapshotController controller) {
+                if (!allOpensOptInOnBackInvoked() || mCloseActivities.isEmpty()) {
+                    return;
+                }
+                for (int i = mCloseActivities.size() - 1; i >= 0; --i) {
+                    controller.recordSnapshot(mCloseActivities.get(i));
                 }
             }
         }