Merge "Transfer IME callbacks when app's focused root view changes." into tm-dev

core/java/android/view/inputmethod/InputMethodManager.java

@@ -782,6 +782,7 @@ public final class InputMethodManager {
         @Override
         public void setCurrentRootView(ViewRootImpl rootView) {
             synchronized (mH) {
+                mImeDispatcher.switchRootView(mCurRootView, rootView);
                 mCurRootView = rootView;
             }
         }

core/java/android/window/ImeOnBackInvokedDispatcher.java

@@ -25,8 +25,9 @@ import android.os.Parcelable;
 import android.os.RemoteException;
 import android.os.ResultReceiver;
 import android.util.Log;
+import android.view.ViewRootImpl;
 
-import java.util.HashMap;
+import java.util.ArrayList;
 
 /**
  * A {@link OnBackInvokedDispatcher} for IME that forwards {@link OnBackInvokedCallback}
@@ -117,7 +118,7 @@ public class ImeOnBackInvokedDispatcher implements OnBackInvokedDispatcher, Parc
                 }
             };
 
-    private final HashMap<Integer, OnBackInvokedCallback> mImeCallbackMap = new HashMap<>();
+    private final ArrayList<ImeOnBackInvokedCallback> mImeCallbacks = new ArrayList<>();
 
     private void receive(
             int resultCode, Bundle resultData,
@@ -140,39 +141,55 @@ public class ImeOnBackInvokedDispatcher implements OnBackInvokedDispatcher, Parc
             int callbackId,
             @NonNull OnBackInvokedDispatcher receivingDispatcher) {
         final ImeOnBackInvokedCallback imeCallback =
-                new ImeOnBackInvokedCallback(iCallback);
+                new ImeOnBackInvokedCallback(iCallback, callbackId, priority);
-        mImeCallbackMap.put(callbackId, imeCallback);
+        mImeCallbacks.add(imeCallback);
         receivingDispatcher.registerOnBackInvokedCallback(priority, imeCallback);
     }
 
     private void unregisterReceivedCallback(
             int callbackId, OnBackInvokedDispatcher receivingDispatcher) {
-        final OnBackInvokedCallback callback = mImeCallbackMap.get(callbackId);
+        ImeOnBackInvokedCallback callback = null;
+        for (ImeOnBackInvokedCallback imeCallback : mImeCallbacks) {
+            if (imeCallback.getId() == callbackId) {
+                callback = imeCallback;
+                break;
+            }
+        }
         if (callback == null) {
             Log.e(TAG, "Ime callback not found. Ignoring unregisterReceivedCallback. "
                     + "callbackId: " + callbackId);
             return;
         }
         receivingDispatcher.unregisterOnBackInvokedCallback(callback);
+        mImeCallbacks.remove(callback);
     }
 
     /** Clears all registered callbacks on the instance. */
     public void clear() {
         // Unregister previously registered callbacks if there's any.
         if (getReceivingDispatcher() != null) {
-            for (OnBackInvokedCallback callback : mImeCallbackMap.values()) {
+            for (ImeOnBackInvokedCallback callback : mImeCallbacks) {
                 getReceivingDispatcher().unregisterOnBackInvokedCallback(callback);
             }
         }
-        mImeCallbackMap.clear();
+        mImeCallbacks.clear();
     }
 
     static class ImeOnBackInvokedCallback implements OnBackInvokedCallback {
         @NonNull
         private final IOnBackInvokedCallback mIOnBackInvokedCallback;
+        /**
+         * The hashcode of the callback instance in the IME process, used as a unique id to
+         * identify the callback when it's passed between processes.
+         */
+        private final int mId;
+        private final int mPriority;
 
-        ImeOnBackInvokedCallback(@NonNull IOnBackInvokedCallback iCallback) {
+        ImeOnBackInvokedCallback(@NonNull IOnBackInvokedCallback iCallback, int id,
+                @Priority int priority) {
             mIOnBackInvokedCallback = iCallback;
+            mId = id;
+            mPriority = priority;
         }
 
         @Override
@@ -186,8 +203,33 @@ public class ImeOnBackInvokedDispatcher implements OnBackInvokedDispatcher, Parc
             }
         }
 
+        private int getId() {
+            return mId;
+        }
+
         IOnBackInvokedCallback getIOnBackInvokedCallback() {
             return mIOnBackInvokedCallback;
         }
     }
+
+    /**
+     * Transfers {@link ImeOnBackInvokedCallback}s registered on one {@link ViewRootImpl} to
+     * another {@link ViewRootImpl} on focus change.
+     *
+     * @param previous the previously focused {@link ViewRootImpl}.
+     * @param current the currently focused {@link ViewRootImpl}.
+     */
+    // TODO(b/232845902): Add CTS to test IME back behavior when there's root view change while
+    // IME is up.
+    public void switchRootView(ViewRootImpl previous, ViewRootImpl current) {
+        for (ImeOnBackInvokedCallback imeCallback : mImeCallbacks) {
+            if (previous != null) {
+                previous.getOnBackInvokedDispatcher().unregisterOnBackInvokedCallback(imeCallback);
+            }
+            if (current != null) {
+                current.getOnBackInvokedDispatcher().registerOnBackInvokedCallback(
+                        imeCallback.mPriority, imeCallback);
+            }
+        }
+    }
 }

core/java/android/window/WindowOnBackInvokedDispatcher.java

@@ -157,15 +157,16 @@ public class WindowOnBackInvokedDispatcher implements OnBackInvokedDispatcher {
 
     /** Clears all registered callbacks on the instance. */
     public void clear() {
+        if (mImeDispatcher != null) {
+            mImeDispatcher.clear();
+            mImeDispatcher = null;
+        }
         if (!mAllCallbacks.isEmpty()) {
             // Clear binder references in WM.
             setTopOnBackInvokedCallback(null);
         }
         mAllCallbacks.clear();
         mOnBackInvokedCallbacks.clear();
-        if (mImeDispatcher != null) {
-            mImeDispatcher = null;
-        }
     }
 
     private void setTopOnBackInvokedCallback(@Nullable OnBackInvokedCallback callback) {