From 13048f123bfe3e128320a6bd3086123dc14ca45d Mon Sep 17 00:00:00 2001
From: Vadim Caen <caen@google.com>
Date: Wed, 11 May 2022 19:50:48 +0200
Subject: [PATCH] Transfer the actual dispatcher when the window is preserved

When recreating the Activity, if the window is preserved,
onAttachedToWindow won't be called and PhoneWindow.onViewRootImplSet
won't be called either, where we set the actual dispatcher to the
Activiy's proxy dispatcher.

So when the window is preserved, we directly assign the actual Dispatcher
from the preserved Window to the new PhoneWindow

Fixes: 228571903
Test: In progress
Test: Open an app with OnBackInvokedCallback enabled, split to top,
      close the other application, do a back gesture
Change-Id: I050247501df8f6cf6cd0024a6233f88a035f6279
---
 .../window/ProxyOnBackInvokedDispatcher.java    | 17 +++++++++++++++--
 .../android/internal/policy/PhoneWindow.java    |  4 +++-
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/core/java/android/window/ProxyOnBackInvokedDispatcher.java b/core/java/android/window/ProxyOnBackInvokedDispatcher.java
index bedf5038f6695..10d43e89a09d3 100644
--- a/core/java/android/window/ProxyOnBackInvokedDispatcher.java
+++ b/core/java/android/window/ProxyOnBackInvokedDispatcher.java
@@ -90,7 +90,11 @@ public class ProxyOnBackInvokedDispatcher implements OnBackInvokedDispatcher {
         synchronized (mLock) {
             mCallbacks.add(Pair.create(callback, priority));
             if (mActualDispatcher != null) {
-                mActualDispatcher.registerOnBackInvokedCallback(priority, callback);
+                if (priority <= PRIORITY_SYSTEM) {
+                    mActualDispatcher.registerSystemOnBackInvokedCallback(callback);
+                } else {
+                    mActualDispatcher.registerOnBackInvokedCallback(priority, callback);
+                }
             }
         }
     }
@@ -171,7 +175,16 @@ public class ProxyOnBackInvokedDispatcher implements OnBackInvokedDispatcher {
                 return;
             }
             clearCallbacksOnDispatcher();
-            mActualDispatcher = actualDispatcher;
+            if (actualDispatcher instanceof ProxyOnBackInvokedDispatcher) {
+                // We don't want to nest ProxyDispatchers, so if we are given on, we unwrap its
+                // actual dispatcher.
+                // This can happen when an Activity is recreated but the Window is preserved (e.g.
+                // when going from split-screen back to single screen)
+                mActualDispatcher =
+                        ((ProxyOnBackInvokedDispatcher) actualDispatcher).mActualDispatcher;
+            } else {
+                mActualDispatcher = actualDispatcher;
+            }
             transferCallbacksToDispatcher();
         }
     }
diff --git a/core/java/com/android/internal/policy/PhoneWindow.java b/core/java/com/android/internal/policy/PhoneWindow.java
index 488fb180723e7..f727d80bba120 100644
--- a/core/java/com/android/internal/policy/PhoneWindow.java
+++ b/core/java/com/android/internal/policy/PhoneWindow.java
@@ -342,7 +342,7 @@ public class PhoneWindow extends Window implements MenuBuilder.Callback {
 
     boolean mDecorFitsSystemWindows = true;
 
-    private ProxyOnBackInvokedDispatcher mProxyOnBackInvokedDispatcher =
+    private final ProxyOnBackInvokedDispatcher mProxyOnBackInvokedDispatcher =
             new ProxyOnBackInvokedDispatcher();
 
     static class WindowManagerHolder {
@@ -378,6 +378,8 @@ public class PhoneWindow extends Window implements MenuBuilder.Callback {
             // window, as we'll be skipping the addView in handleResumeActivity(), and
             // the token will not be updated as for a new window.
             getAttributes().token = preservedWindow.getAttributes().token;
+            mProxyOnBackInvokedDispatcher.setActualDispatcher(
+                    preservedWindow.getOnBackInvokedDispatcher());
         }
         // Even though the device doesn't support picture-in-picture mode,
         // an user can force using it through developer options.