Hit back button in overflow to collapse stack

getExpandedDisplayId(mContext) was returning -1 (INVALID_DISPLAY) for overflow

Add getDisplayId to BubbleViewProvider, override that in overflow and
return actual display id for overflow's expanded view.

Fixes: 150398525
Fixes: 149882551
Test: manual -> press back button in overflow -> stack collapses
Test: manual -> swipe right from left edge in gesture nav -> stack collapses
Test: adb shell am instrument -w com.android.systemui.tests
Change-Id: I9d43de3803a7c771eaabe38b8e8dc9aacd8bad94
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/Bubble.java b/packages/SystemUI/src/com/android/systemui/bubbles/Bubble.java
index 7ca2308..4240209c 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/Bubble.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/Bubble.java
@@ -279,7 +279,8 @@
     /**
      * @return the display id of the virtual display on which bubble contents is drawn.
      */
-    int getDisplayId() {
+    @Override
+    public int getDisplayId() {
         return mExpandedView != null ? mExpandedView.getVirtualDisplayId() : INVALID_DISPLAY;
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java
index f873f42..67299a6 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java
@@ -1175,23 +1175,17 @@
      * status bar, otherwise returns {@link Display#INVALID_DISPLAY}.
      */
     public int getExpandedDisplayId(Context context) {
-        final Bubble bubble = getExpandedBubble(context);
-        return bubble != null ? bubble.getDisplayId() : INVALID_DISPLAY;
-    }
-
-    @Nullable
-    private Bubble getExpandedBubble(Context context) {
         if (mStackView == null) {
-            return null;
+            return INVALID_DISPLAY;
         }
         final boolean defaultDisplay = context.getDisplay() != null
                 && context.getDisplay().getDisplayId() == DEFAULT_DISPLAY;
-        final Bubble expandedBubble = mStackView.getExpandedBubble();
-        if (defaultDisplay && expandedBubble != null && isStackExpanded()
+        final BubbleViewProvider expandedViewProvider = mStackView.getExpandedBubble();
+        if (defaultDisplay && expandedViewProvider != null && isStackExpanded()
                 && !mNotificationShadeWindowController.getPanelExpanded()) {
-            return expandedBubble;
+            return expandedViewProvider.getDisplayId();
         }
-        return null;
+        return INVALID_DISPLAY;
     }
 
     @VisibleForTesting
@@ -1256,7 +1250,7 @@
 
         @Override
         public void onSingleTaskDisplayEmpty(int displayId) {
-            final Bubble expandedBubble = mStackView != null
+            final BubbleViewProvider expandedBubble = mStackView != null
                     ? mStackView.getExpandedBubble()
                     : null;
             int expandedId = expandedBubble != null ? expandedBubble.getDisplayId() : -1;
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleDebugConfig.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleDebugConfig.java
index e800011..19733a5 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleDebugConfig.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleDebugConfig.java
@@ -59,13 +59,15 @@
         return FORCE_SHOW_USER_EDUCATION || forceShow;
     }
 
-    static String formatBubblesString(List<Bubble> bubbles, Bubble selected) {
+    static String formatBubblesString(List<Bubble> bubbles, BubbleViewProvider selected) {
         StringBuilder sb = new StringBuilder();
         for (Bubble bubble : bubbles) {
             if (bubble == null) {
                 sb.append("   <null> !!!!!\n");
             } else {
-                boolean isSelected = (selected != null && bubble == selected);
+                boolean isSelected = (selected != null
+                        && selected.getKey() != BubbleOverflow.KEY
+                        && bubble == selected);
                 String arrow = isSelected ? "=>" : "  ";
                 sb.append(String.format("%s Bubble{act=%12d, ongoing=%d, key=%s}\n",
                         arrow,
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleOverflow.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleOverflow.java
index 313bb42..4fb2d08 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleOverflow.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleOverflow.java
@@ -16,6 +16,7 @@
 
 package com.android.systemui.bubbles;
 
+import static android.view.Display.INVALID_DISPLAY;
 import static android.view.View.GONE;
 
 import static com.android.systemui.bubbles.BadgedImageView.DEFAULT_PATH_SIZE;
@@ -45,7 +46,7 @@
     public static final String KEY = "Overflow";
 
     private BadgedImageView mOverflowBtn;
-    private BubbleExpandedView mOverflowExpandedView;
+    private BubbleExpandedView mExpandedView;
     private LayoutInflater mInflater;
     private Context mContext;
     private Bitmap mIcon;
@@ -63,11 +64,11 @@
     }
 
     void setUpOverflow(ViewGroup parentViewGroup, BubbleStackView stackView) {
-        mOverflowExpandedView = (BubbleExpandedView) mInflater.inflate(
+        mExpandedView = (BubbleExpandedView) mInflater.inflate(
                 R.layout.bubble_expanded_view, parentViewGroup /* root */,
                 false /* attachToRoot */);
-        mOverflowExpandedView.setOverflow(true);
-        mOverflowExpandedView.setStackView(stackView);
+        mExpandedView.setOverflow(true);
+        mExpandedView.setStackView(stackView);
 
         updateIcon(mContext, parentViewGroup);
     }
@@ -124,7 +125,7 @@
 
     @Override
     public BubbleExpandedView getExpandedView() {
-        return mOverflowExpandedView;
+        return mExpandedView;
     }
 
     @Override
@@ -149,7 +150,7 @@
 
     @Override
     public void setContentVisibility(boolean visible) {
-        mOverflowExpandedView.setContentVisibility(visible);
+        mExpandedView.setContentVisibility(visible);
     }
 
     @Override
@@ -167,4 +168,9 @@
     public String getKey() {
         return BubbleOverflow.KEY;
     }
+
+    @Override
+    public int getDisplayId() {
+        return mExpandedView != null ? mExpandedView.getVirtualDisplayId() : INVALID_DISPLAY;
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java
index 0cf6d89..4ccc9f8 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java
@@ -801,14 +801,8 @@
      * The {@link Bubble} that is expanded, null if one does not exist.
      */
     @Nullable
-    Bubble getExpandedBubble() {
-        if (mExpandedBubble == null
-                || (BubbleExperimentConfig.allowBubbleOverflow(mContext)
-                    && mExpandedBubble.getIconView() == mBubbleOverflow.getBtn()
-                    && BubbleOverflow.KEY.equals(mExpandedBubble.getKey()))) {
-            return null;
-        }
-        return (Bubble) mExpandedBubble;
+    BubbleViewProvider getExpandedBubble() {
+        return mExpandedBubble;
     }
 
     // via BubbleData.Listener
@@ -1172,7 +1166,7 @@
         if (DEBUG_BUBBLE_STACK_VIEW) {
             Log.d(TAG, "animateCollapse");
             Log.d(TAG, BubbleDebugConfig.formatBubblesString(getBubblesOnScreen(),
-                    getExpandedBubble()));
+                    mExpandedBubble));
         }
         updateOverflowBtnVisibility(/* apply */ false);
         mBubbleContainer.cancelAllAnimations();
@@ -1811,17 +1805,16 @@
     }
 
     private void updatePointerPosition() {
-        Bubble expandedBubble = getExpandedBubble();
-        if (expandedBubble == null) {
+        if (mExpandedBubble == null) {
             return;
         }
-        int index = getBubbleIndex(expandedBubble);
+        int index = getBubbleIndex(mExpandedBubble);
         float bubbleLeftFromScreenLeft = mExpandedAnimationController.getBubbleLeft(index);
         float halfBubble = mBubbleSize / 2f;
         float bubbleCenter = bubbleLeftFromScreenLeft + halfBubble;
         // Padding might be adjusted for insets, so get it directly from the view
         bubbleCenter -= mExpandedViewContainer.getPaddingLeft();
-        expandedBubble.getExpandedView().setPointerPosition(bubbleCenter);
+        mExpandedBubble.getExpandedView().setPointerPosition(bubbleCenter);
     }
 
     /**
@@ -1843,11 +1836,10 @@
      * is between 0 and the bubble count minus 1.
      */
     int getBubbleIndex(@Nullable BubbleViewProvider provider) {
-        if (provider == null || provider.getKey() == BubbleOverflow.KEY) {
+        if (provider == null) {
             return 0;
         }
-        Bubble b = (Bubble) provider;
-        return mBubbleContainer.indexOfChild(b.getIconView());
+        return mBubbleContainer.indexOfChild(provider.getIconView());
     }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleViewProvider.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleViewProvider.java
index f04933a..ef84c73 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleViewProvider.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleViewProvider.java
@@ -41,4 +41,6 @@
     Path getDotPath();
 
     boolean showDot();
+
+    int getDisplayId();
 }