Ensure overflow isn't selected when bar is collapsed

After the bubble bar collapsing animation completes, and reordering
is finished, if the bubble overflow is selected, update the selected
bubble to be the first bubble. This ensures that the overflow
is not selected the next time the bubble bar is expanded.

This is done through a callback from the BubbleBarView to the
BubbleBarController with the key of the first bubble. Ideally the view
would just notify that the selected bubble needs to be updated, but
the controller doesn't hold ordering info.

Fixes: 289280431
Test: manual
       - Add 2 bubbles to the bubble bar
       - Expand the bubble bar. The first bubble A should be selected
       - Select the second bubble B
       - Select the overflow
       - Tap on the overflow again to collapse the bubble bar
       - Tap on the bubble bar to expand it
       - Observe that bubble B is first and selected
Change-Id: I0f9c2d628cfaf3a80da55ff634284687425e949a
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
index 6b5c962..b2d5940 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
@@ -184,6 +184,8 @@
         bubbleControllers.runAfterInit(() -> {
             mBubbleBarViewController.setHiddenForBubbles(!BUBBLE_BAR_ENABLED);
             mBubbleStashedHandleViewController.setHiddenForBubbles(!BUBBLE_BAR_ENABLED);
+            mBubbleBarViewController.setUpdateSelectedBubbleAfterCollapse(
+                    key -> setSelectedBubble(mBubbles.get(key)));
         });
     }
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
index 58c67e3..0b33645 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
@@ -32,6 +32,7 @@
 import com.android.launcher3.views.ActivityContext;
 
 import java.util.List;
+import java.util.function.Consumer;
 
 /**
  * The view that holds all the bubble views. Modifying this view should happen through
@@ -101,6 +102,9 @@
     @Nullable
     private Runnable mReorderRunnable;
 
+    @Nullable
+    private Consumer<String> mUpdateSelectedBubbleAfterCollapse;
+
     public BubbleBarView(Context context) {
         this(context, null);
     }
@@ -144,6 +148,13 @@
                     mReorderRunnable.run();
                     mReorderRunnable = null;
                 }
+                // If the bar was just collapsed and the overflow was the last bubble that was
+                // selected, set the first bubble as selected.
+                if (!mIsBarExpanded && mUpdateSelectedBubbleAfterCollapse != null
+                        && mSelectedBubbleView.getBubble() instanceof BubbleBarOverflow) {
+                    BubbleView firstBubble = (BubbleView) getChildAt(0);
+                    mUpdateSelectedBubbleAfterCollapse.accept(firstBubble.getBubble().getKey());
+                }
                 updateWidth();
             }
 
@@ -293,6 +304,11 @@
         }
     }
 
+    public void setUpdateSelectedBubbleAfterCollapse(
+            Consumer<String> updateSelectedBubbleAfterCollapse) {
+        mUpdateSelectedBubbleAfterCollapse = updateSelectedBubbleAfterCollapse;
+    }
+
     /**
      * Sets which bubble view should be shown as selected.
      */
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
index 52c144e..4e9f88a 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
@@ -36,6 +36,7 @@
 
 import java.util.List;
 import java.util.Objects;
+import java.util.function.Consumer;
 
 /**
  * Controller for {@link BubbleBarView}. Manages the visibility of the bubble bar as well as
@@ -184,6 +185,12 @@
         }
     }
 
+    /** Sets a callback that updates the selected bubble after the bubble bar collapses. */
+    public void setUpdateSelectedBubbleAfterCollapse(
+            Consumer<String> updateSelectedBubbleAfterCollapse) {
+        mBarView.setUpdateSelectedBubbleAfterCollapse(updateSelectedBubbleAfterCollapse);
+    }
+
     /**
      * Sets whether the bubble bar should be hidden due to SysUI state (e.g. on lockscreen).
      */