Show last selected bubble when stack expands

Fixes: 152888262

Test: manual
        1. Have 5 bubbles, select -middle- bubble, select overflow
        2. Collapse stack, reopen stack
        -> See that middle bubble is now first in stack and selected

Test: manual
        1. Have 5 bubbles, select -first- bubble, select overflow
        2. Collapse stack, reopen stack
        -> See that first bubble is selected again

Test: manual (overflow disabled)
        -> Previously selected bubble shows first in stack on expand

Change-Id: I8749306be247be3136d2b8a36852e796a5195ef7
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleData.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleData.java
index 4c149dd..a1393cd 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleData.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleData.java
@@ -120,6 +120,7 @@
     /** Bubbles that are being loaded but haven't been added to the stack just yet. */
     private final List<Bubble> mPendingBubbles;
     private Bubble mSelectedBubble;
+    private boolean mShowingOverflow;
     private boolean mExpanded;
     private final int mMaxBubbles;
     private final int mMaxOverflowBubbles;
@@ -215,6 +216,10 @@
         dispatchPendingChanges();
     }
 
+    void setShowingOverflow(boolean showingOverflow) {
+        mShowingOverflow = showingOverflow;
+    }
+
     private void moveOverflowBubbleToPending(Bubble b) {
         // Preserve new order for next repack, which sorts by last updated time.
         b.markUpdatedAt(mTimeSource.currentTimeMillis());
@@ -513,9 +518,11 @@
         if (DEBUG_BUBBLE_DATA) {
             Log.d(TAG, "setSelectedBubbleInternal: " + bubble);
         }
-        if (Objects.equals(bubble, mSelectedBubble)) {
+        if (!mShowingOverflow && Objects.equals(bubble, mSelectedBubble)) {
             return;
         }
+        // Otherwise, if we are showing the overflow menu, return to the previously selected bubble.
+
         if (bubble != null && !mBubbles.contains(bubble) && !mOverflowBubbles.contains(bubble)) {
             Log.e(TAG, "Cannot select bubble which doesn't exist!"
                     + " (" + bubble + ") bubbles=" + mBubbles);
@@ -559,6 +566,10 @@
             mStateChange.orderChanged |= repackAll();
             // Save the state which should be returned to when expanded (with no other changes)
 
+            if (mShowingOverflow) {
+                // Show previously selected bubble instead of overflow menu on next expansion.
+                setSelectedBubbleInternal(mSelectedBubble);
+            }
             if (mBubbles.indexOf(mSelectedBubble) > 0) {
                 // Move the selected bubble to the top while collapsed.
                 if (!mSelectedBubble.isOngoing() && mBubbles.get(0).isOngoing()) {
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java
index 644e54f..6e66139 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleStackView.java
@@ -439,8 +439,8 @@
                     // that means overflow was previously expanded. Set the selected bubble
                     // internally without going through BubbleData (which would ignore it since it's
                     // already selected).
+                    mBubbleData.setShowingOverflow(true);
                     setSelectedBubble(clickedBubble);
-
                 }
             } else {
                 // Otherwise, we either tapped the stack (which means we're collapsed
@@ -1227,8 +1227,12 @@
         if (mExpandedBubble != null && mExpandedBubble.equals(bubbleToSelect)) {
             return;
         }
+        if (bubbleToSelect == null || bubbleToSelect.getKey() != BubbleOverflow.KEY) {
+            mBubbleData.setShowingOverflow(false);
+        }
         final BubbleViewProvider previouslySelected = mExpandedBubble;
         mExpandedBubble = bubbleToSelect;
+        updatePointerPosition();
 
         if (mIsExpanded) {
             // Make the container of the expanded view transparent before removing the expanded view
@@ -1238,7 +1242,6 @@
             mSurfaceSynchronizer.syncSurfaceAndRun(() -> {
                 previouslySelected.setContentVisibility(false);
                 updateExpandedBubble();
-                updatePointerPosition();
                 requestUpdate();
 
                 logBubbleEvent(previouslySelected,