Fix bubble added from the overflow background arrow animation.

The new bubble animation now uses a "selected location" argument to
position the pointing arrow.

Test: BubbleAnimatorTest
Test: Manual. Create a bubble bar with overflow. Open the overflow menu
and click an overflow bubble icon. Observe that a new bubble is added
and immediately expands. Observe that the arrow animates to the new
bubble's position.
Test: Manual. Create a bubble bar with overflow. Open any bubble.
Trigger a new bubble notification. Observe that a new bubble is added
but does not expand, so the arrow remains at the currently open bubble.
Flag: com.android.wm.shell.enable_bubble_bar
Fixes: 359952121

Change-Id: I47f5b6c2aad775f2dd3e70f8c544a3711f192342
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
index db5e0d5..62dd748 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
@@ -342,26 +342,38 @@
         }
 
         BubbleBarBubble bubbleToSelect = null;
-
+        if (update.addedBubble != null) {
+            mBubbles.put(update.addedBubble.getKey(), update.addedBubble);
+        }
+        if (update.selectedBubbleKey != null) {
+            if (mSelectedBubble == null
+                    || !update.selectedBubbleKey.equals(mSelectedBubble.getKey())) {
+                BubbleBarBubble newlySelected = mBubbles.get(update.selectedBubbleKey);
+                if (newlySelected != null) {
+                    bubbleToSelect = newlySelected;
+                } else {
+                    Log.w(TAG, "trying to select bubble that doesn't exist:"
+                            + update.selectedBubbleKey);
+                }
+            }
+        }
         if (Flags.enableOptionalBubbleOverflow()
                 && update.showOverflowChanged && !update.showOverflow && update.addedBubble != null
                 && update.removedBubbles.isEmpty()
                 && !mBubbles.isEmpty()) {
             // A bubble was added from the overflow (& now it's empty / not showing)
-            mBubbles.put(update.addedBubble.getKey(), update.addedBubble);
             mBubbleBarViewController.removeOverflowAndAddBubble(update.addedBubble);
         } else if (update.addedBubble != null && update.removedBubbles.size() == 1) {
             // we're adding and removing a bubble at the same time. handle this as a single update.
             RemovedBubble removedBubble = update.removedBubbles.get(0);
             BubbleBarBubble bubbleToRemove = mBubbles.remove(removedBubble.getKey());
-            mBubbles.put(update.addedBubble.getKey(), update.addedBubble);
             boolean showOverflow = update.showOverflowChanged && update.showOverflow;
             if (bubbleToRemove != null) {
                 mBubbleBarViewController.addBubbleAndRemoveBubble(update.addedBubble,
                         bubbleToRemove, isExpanding, suppressAnimation, showOverflow);
             } else {
                 mBubbleBarViewController.addBubble(update.addedBubble, isExpanding,
-                        suppressAnimation);
+                        suppressAnimation, bubbleToSelect);
                 Log.w(TAG, "trying to remove bubble that doesn't exist: " + removedBubble.getKey());
             }
         } else {
@@ -384,9 +396,8 @@
                 }
             }
             if (update.addedBubble != null) {
-                mBubbles.put(update.addedBubble.getKey(), update.addedBubble);
                 mBubbleBarViewController.addBubble(update.addedBubble, isExpanding,
-                        suppressAnimation);
+                        suppressAnimation, bubbleToSelect);
             }
             if (Flags.enableOptionalBubbleOverflow()
                     && update.showOverflowChanged
@@ -400,7 +411,7 @@
             addBubbleInternally(update.updatedBubble, isExpanding, suppressAnimation);
         }
 
-        if (update.addedBubble != null && isCollapsed) {
+        if (update.addedBubble != null && isCollapsed && bubbleToSelect == null) {
             // If we're collapsed, the most recently added bubble will be selected.
             bubbleToSelect = update.addedBubble;
         }
@@ -411,7 +422,7 @@
                 BubbleBarBubble bubble = update.currentBubbles.get(i);
                 if (bubble != null) {
                     addBubbleInternally(bubble, isExpanding, suppressAnimation);
-                    if (isCollapsed) {
+                    if (isCollapsed && bubbleToSelect == null) {
                         // If we're collapsed, the most recently added bubble will be selected.
                         bubbleToSelect = bubble;
                     }
@@ -478,18 +489,6 @@
                 mBubbleBarViewController.reorderBubbles(newOrder);
             }
         }
-        if (update.selectedBubbleKey != null) {
-            if (mSelectedBubble == null
-                    || !update.selectedBubbleKey.equals(mSelectedBubble.getKey())) {
-                BubbleBarBubble newlySelected = mBubbles.get(update.selectedBubbleKey);
-                if (newlySelected != null) {
-                    bubbleToSelect = newlySelected;
-                } else {
-                    Log.w(TAG, "trying to select bubble that doesn't exist:"
-                            + update.selectedBubbleKey);
-                }
-            }
-        }
         if (bubbleToSelect != null) {
             setSelectedBubbleInternal(bubbleToSelect);
         }
@@ -608,7 +607,8 @@
     private void addBubbleInternally(BubbleBarBubble bubble, boolean isExpanding,
             boolean suppressAnimation) {
         mBubbles.put(bubble.getKey(), bubble);
-        mBubbleBarViewController.addBubble(bubble, isExpanding, suppressAnimation);
+        mBubbleBarViewController.addBubble(bubble, isExpanding,
+                suppressAnimation, /* bubbleToSelect = */ null);
     }
 
     /** Listener of {@link BubbleBarLocation} updates. */
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
index 0d0feff..aa6ad25 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
@@ -681,8 +681,18 @@
         return mRelativePivotY;
     }
 
-    /** Add a new bubble to the bubble bar. */
+    /** Add a new bubble to the bubble bar without updating the selected bubble. */
     public void addBubble(BubbleView bubble) {
+        addBubble(bubble, /* bubbleToSelect = */ null);
+    }
+
+    /**
+     * Add a new bubble to the bubble bar and selects the provided bubble.
+     *
+     * @param bubble         bubble to add
+     * @param bubbleToSelect if {@code null}, then selected bubble does not change
+     */
+    public void addBubble(BubbleView bubble, @Nullable BubbleView bubbleToSelect) {
         FrameLayout.LayoutParams lp = new FrameLayout.LayoutParams((int) mIconSize, (int) mIconSize,
                 Gravity.LEFT);
         final int index = bubble.isOverflow() ? getChildCount() : 0;
@@ -718,7 +728,12 @@
                     invalidate();
                 }
             };
-            mBubbleAnimator.animateNewBubble(indexOfChild(mSelectedBubbleView), listener);
+            if (bubbleToSelect != null) {
+                mBubbleAnimator.animateNewBubble(indexOfChild(mSelectedBubbleView),
+                        indexOfChild(bubbleToSelect), listener);
+            } else {
+                mBubbleAnimator.animateNewBubble(indexOfChild(mSelectedBubbleView), listener);
+            }
         } else {
             addView(bubble, index, lp);
         }
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
index d00959e..515c1fa 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
@@ -902,9 +902,15 @@
     /**
      * Adds the provided bubble to the bubble bar.
      */
-    public void addBubble(BubbleBarItem b, boolean isExpanding, boolean suppressAnimation) {
+    public void addBubble(BubbleBarItem b,
+            boolean isExpanding,
+            boolean suppressAnimation,
+            @Nullable BubbleBarBubble bubbleToSelect
+    ) {
         if (b != null) {
-            mBarView.addBubble(b.getView());
+            BubbleView bubbleToSelectView =
+                    bubbleToSelect == null ? null : bubbleToSelect.getView();
+            mBarView.addBubble(b.getView(), bubbleToSelectView);
             b.getView().setOnClickListener(mBubbleClickListener);
             mBubbleDragController.setupBubbleView(b.getView());
             b.getView().setController(mBubbleViewController);
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/animation/BubbleAnimator.kt b/quickstep/src/com/android/launcher3/taskbar/bubbles/animation/BubbleAnimator.kt
index 3604167..944e806 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/animation/BubbleAnimator.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/animation/BubbleAnimator.kt
@@ -39,9 +39,14 @@
     private var state: State = State.Idle
     private lateinit var animator: ValueAnimator
 
-    fun animateNewBubble(selectedBubbleIndex: Int, listener: Listener) {
+    @JvmOverloads
+    fun animateNewBubble(
+        selectedBubbleIndex: Int,
+        newlySelectedBubbleIndex: Int? = null,
+        listener: Listener,
+    ) {
         animator = createAnimator(listener)
-        state = State.AddingBubble(selectedBubbleIndex)
+        state = State.AddingBubble(selectedBubbleIndex, newlySelectedBubbleIndex)
         animator.start()
     }
 
@@ -180,16 +185,7 @@
     fun getArrowPosition(): Float {
         return when (val state = state) {
             State.Idle -> 0f
-            is State.AddingBubble -> {
-                val tx =
-                    getBubbleTranslationXWhileScalingBubble(
-                        bubbleIndex = state.selectedBubbleIndex,
-                        scalingBubbleIndex = 0,
-                        bubbleScale = animator.animatedFraction,
-                    )
-                tx + iconSize / 2f
-            }
-
+            is State.AddingBubble -> getArrowPositionWhenAddingBubble(state)
             is State.RemovingBubble -> getArrowPositionWhenRemovingBubble(state)
             is State.AddingAndRemoving -> {
                 // we never remove the selected bubble, so the arrow stays pointing to its center
@@ -205,6 +201,26 @@
         }
     }
 
+    private fun getArrowPositionWhenAddingBubble(state: State.AddingBubble): Float {
+        val scale = animator.animatedFraction
+        var tx = getBubbleTranslationXWhileScalingBubble(
+            bubbleIndex = state.selectedBubbleIndex,
+            scalingBubbleIndex = 0,
+            bubbleScale = scale
+        ) + iconSize / 2f
+        if (state.newlySelectedBubbleIndex != null) {
+            val selectedBubbleScale = if (state.newlySelectedBubbleIndex == 0) scale else 1f
+            val finalTx =
+                getBubbleTranslationXWhileScalingBubble(
+                    bubbleIndex = state.newlySelectedBubbleIndex,
+                    scalingBubbleIndex = 0,
+                    bubbleScale = 1f,
+                ) + iconSize * selectedBubbleScale / 2f
+            tx += (finalTx - tx) * animator.animatedFraction
+        }
+        return tx
+    }
+
     private fun getArrowPositionWhenRemovingBubble(state: State.RemovingBubble): Float =
         if (state.selectedBubbleIndex != state.bubbleIndex || state.removingLastRemainingBubble) {
             // if we're not removing the selected bubble or if we're removing the last remaining
@@ -378,7 +394,12 @@
         data object Idle : State
 
         /** A new bubble is being added to the bubble bar. */
-        data class AddingBubble(val selectedBubbleIndex: Int) : State
+        data class AddingBubble(
+            /** The index of the selected bubble. */
+            val selectedBubbleIndex: Int,
+            /** The index of the newly selected bubble. */
+            val newlySelectedBubbleIndex: Int?,
+        ) : State
 
         /** A bubble is being removed from the bubble bar. */
         data class RemovingBubble(
@@ -392,6 +413,7 @@
             val removingLastRemainingBubble: Boolean,
         ) : State
 
+        // TODO add index where bubble is being added, and index for newly selected bubble
         /** A new bubble is being added and an old bubble is being removed from the bubble bar. */
         data class AddingAndRemoving(val selectedBubbleIndex: Int, val removedBubbleIndex: Int) :
             State
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/animation/BubbleAnimatorTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/animation/BubbleAnimatorTest.kt
index eae181f..3ca36ec 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/animation/BubbleAnimatorTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/animation/BubbleAnimatorTest.kt
@@ -44,7 +44,7 @@
             )
         val listener = TestBubbleAnimatorListener()
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
-            bubbleAnimator.animateNewBubble(selectedBubbleIndex = 2, listener)
+            bubbleAnimator.animateNewBubble(selectedBubbleIndex = 2, listener = listener)
         }
 
         assertThat(bubbleAnimator.isRunning).isTrue()