Fix back presses dismissing bubbles when IME is up

If bubbles is expanded & shade is not use the virtual display id of the
expanded activity view (if it's valid, otherwise do whats normal).

Test: manual  1) have expanded bubble with IME up
              2) hit back button
              => note that the IME goes away but the bubble remains

              1) have expanded bubble up
              2) pull down shade
              3) hit back button
              => shade goes away

              Do above with gesture nav turned on & do back gesture

Test: atest KeyButtonViewTest
Fixes: 122535136
Change-Id: Iabaace66cfb4d3d80b1a3a3c7b0773deb48fbcaa
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/EdgeBackGestureHandler.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/EdgeBackGestureHandler.java
index 79bf6b3..8aa50a7 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/EdgeBackGestureHandler.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/EdgeBackGestureHandler.java
@@ -15,6 +15,8 @@
  */
 package com.android.systemui.statusbar.phone;
 
+import static android.view.Display.INVALID_DISPLAY;
+
 import android.content.Context;
 import android.content.pm.ParceledListSlice;
 import android.content.res.Resources;
@@ -46,7 +48,9 @@
 import android.view.WindowManager;
 import android.view.WindowManagerGlobal;
 
+import com.android.systemui.Dependency;
 import com.android.systemui.R;
+import com.android.systemui.bubbles.BubbleController;
 import com.android.systemui.recents.OverviewProxyService;
 import com.android.systemui.shared.system.InputChannelCompat.InputEventReceiver;
 import com.android.systemui.shared.system.QuickStepContract;
@@ -358,6 +362,13 @@
                 0 /* metaState */, KeyCharacterMap.VIRTUAL_KEYBOARD, 0 /* scancode */,
                 KeyEvent.FLAG_FROM_SYSTEM | KeyEvent.FLAG_VIRTUAL_HARD_KEY,
                 InputDevice.SOURCE_KEYBOARD);
+
+        // Bubble controller will give us a valid display id if it should get the back event
+        BubbleController bubbleController = Dependency.get(BubbleController.class);
+        int bubbleDisplayId = bubbleController.getExpandedDisplayId(mContext);
+        if (code == KeyEvent.KEYCODE_BACK && bubbleDisplayId != INVALID_DISPLAY) {
+            ev.setDisplayId(bubbleDisplayId);
+        }
         InputManager.getInstance().injectInputEvent(ev, InputManager.INJECT_INPUT_EVENT_MODE_ASYNC);
     }
 }