Preserve bubble bar state.

Added logic to save and then restore existing bubbles in the
TaskbarSharedState object.

Bug: 375482371
Test: Manual. Switch between pinned and unpinned taskbar. Video:
http://recall/-/gx8ASgewUeUS3QYohfrd1J/d8hVHxksAokx85J8P7etoC
Flag: com.android.wm.shell.enable_bubble_bar
Change-Id: I34f580f1c652859f427385edc0ccadd447e1400f
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
index 4a85acc..5a63ca6 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
@@ -169,7 +169,7 @@
         taskbarOverlayController.init(this);
         taskbarAllAppsController.init(this, sharedState.allAppsVisible);
         navButtonController.init(this);
-        bubbleControllers.ifPresent(controllers -> controllers.init(this));
+        bubbleControllers.ifPresent(controllers -> controllers.init(sharedState, this));
         taskbarInsetsController.init(this);
         voiceInteractionWindowController.init(this);
         taskbarRecentAppsController.init(this);
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarSharedState.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarSharedState.java
index 729cbe9..a64dab1 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarSharedState.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarSharedState.java
@@ -30,6 +30,10 @@
 import android.view.InsetsFrameProvider;
 
 import com.android.systemui.shared.system.QuickStepContract.SystemUiStateFlags;
+import com.android.wm.shell.shared.bubbles.BubbleBarLocation;
+import com.android.wm.shell.shared.bubbles.BubbleInfo;
+
+import java.util.List;
 
 /**
  * State shared across different taskbar instance
@@ -69,6 +73,15 @@
 
     public boolean allAppsVisible = false;
 
+    public BubbleBarLocation bubbleBarLocation;
+
+    public List<BubbleInfo> bubbleInfoItems;
+
+    /** Returns whether there are a saved bubbles. */
+    public boolean hasSavedBubbles() {
+        return bubbleInfoItems != null && !bubbleInfoItems.isEmpty();
+    }
+
     // LauncherTaskbarUIController#mTaskbarInAppDisplayProgressMultiProp
     public float[] inAppDisplayProgressMultiPropValues = new float[DISPLAY_PROGRESS_COUNT];
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
index b22fd6f..30e4e47 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
@@ -35,6 +35,7 @@
 import android.util.ArrayMap;
 import android.util.Log;
 
+import com.android.launcher3.taskbar.TaskbarSharedState;
 import com.android.launcher3.taskbar.bubbles.stashing.BubbleStashController;
 import com.android.launcher3.util.Executors.SimpleThreadFactory;
 import com.android.quickstep.SystemUiProxy;
@@ -47,6 +48,7 @@
 import com.android.wm.shell.shared.bubbles.RemovedBubble;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
@@ -112,6 +114,7 @@
 
     private BubbleBarItem mSelectedBubble;
 
+    private TaskbarSharedState mSharedState;
     private ImeVisibilityChecker mImeVisibilityChecker;
     private BubbleBarViewController mBubbleBarViewController;
     private BubbleStashController mBubbleStashController;
@@ -173,12 +176,25 @@
 
     public void onDestroy() {
         mSystemUiProxy.setBubblesListener(null);
+        // Saves bubble bar state
+        BubbleInfo[] bubbleInfoItems = new BubbleInfo[mBubbles.size()];
+        mBubbles.values().forEach(bubbleBarBubble -> {
+            int index = mBubbleBarViewController.bubbleViewIndex(bubbleBarBubble.getView());
+            if (index < 0 || index >= bubbleInfoItems.length) {
+                Log.e(TAG, "Found improper index: " + index + " for " + bubbleBarBubble);
+            } else {
+                bubbleInfoItems[index] = bubbleBarBubble.getInfo();
+            }
+        });
+        mSharedState.bubbleInfoItems = Arrays.asList(bubbleInfoItems);
     }
 
     /** Initializes controllers. */
     public void init(BubbleControllers bubbleControllers,
             BubbleBarLocationListener bubbleBarLocationListener,
-            ImeVisibilityChecker imeVisibilityChecker) {
+            ImeVisibilityChecker imeVisibilityChecker,
+            TaskbarSharedState sharedState) {
+        mSharedState = sharedState;
         mImeVisibilityChecker = imeVisibilityChecker;
         mBubbleBarViewController = bubbleControllers.bubbleBarViewController;
         mBubbleStashController = bubbleControllers.bubbleStashController;
@@ -188,6 +204,7 @@
         mBubbleBarLocationListener = bubbleBarLocationListener;
 
         bubbleControllers.runAfterInit(() -> {
+            restoreSavedState(sharedState);
             mBubbleBarViewController.setHiddenForBubbles(
                     !sBubbleBarEnabled || mBubbles.isEmpty());
             mBubbleStashedHandleViewController.ifPresent(
@@ -266,6 +283,26 @@
         }
     }
 
+    private void restoreSavedState(TaskbarSharedState sharedState) {
+        if (sharedState.bubbleBarLocation != null) {
+            updateBubbleBarLocationInternal(sharedState.bubbleBarLocation);
+        }
+        List<BubbleInfo> bubbleInfos = sharedState.bubbleInfoItems;
+        if (bubbleInfos == null || bubbleInfos.isEmpty()) return;
+        // Iterate in reverse because new bubbles are added in front and the list is in order.
+        for (int i = bubbleInfos.size() - 1; i >= 0; i--) {
+            BubbleBarBubble bubble = mBubbleCreator.populateBubble(mContext,
+                    bubbleInfos.get(i), mBarView, /* existingBubble = */ null);
+            if (bubble == null) {
+                Log.e(TAG, "Could not instantiate BubbleBarBubble for " + bubbleInfos.get(i));
+                continue;
+            }
+            addBubbleInternally(bubble,  /* showAppBadge = */
+                    mBubbleBarViewController.isExpanded() || i == 0,
+                    /* isExpanding = */ false,  /* suppressAnimation = */ true);
+        }
+    }
+
     private void applyViewChanges(BubbleBarViewUpdate update) {
         final boolean isCollapsed = (update.expandedChanged && !update.expanded)
                 || (!update.expandedChanged && !mBubbleBarViewController.isExpanded());
@@ -277,6 +314,12 @@
                 update.initialState || mBubbleBarViewController.isHiddenForSysui()
                         || mImeVisibilityChecker.isImeVisible();
 
+        if (update.initialState && mSharedState.hasSavedBubbles()) {
+            // clear restored state
+            mBubbleBarViewController.removeAllBubbles();
+            mBubbles.clear();
+        }
+
         BubbleBarBubble bubbleToSelect = null;
 
         if (Flags.enableOptionalBubbleOverflow()
@@ -347,8 +390,8 @@
             for (int i = update.currentBubbles.size() - 1; i >= 0; i--) {
                 BubbleBarBubble bubble = update.currentBubbles.get(i);
                 if (bubble != null) {
-                    mBubbles.put(bubble.getKey(), bubble);
-                    mBubbleBarViewController.addBubble(bubble, isExpanding, suppressAnimation);
+                    addBubbleInternally(bubble, /* showAppBadge = */ !isCollapsed || i == 0,
+                            isExpanding, suppressAnimation);
                     if (isCollapsed) {
                         // If we're collapsed, the most recently added bubble will be selected.
                         bubbleToSelect = bubble;
@@ -420,6 +463,7 @@
             }
         }
         if (update.bubbleBarLocation != null) {
+            mSharedState.bubbleBarLocation = update.bubbleBarLocation;
             if (update.bubbleBarLocation != mBubbleBarViewController.getBubbleBarLocation()) {
                 updateBubbleBarLocationInternal(update.bubbleBarLocation);
             }
@@ -519,6 +563,14 @@
         }
     }
 
+    private void addBubbleInternally(BubbleBarBubble bubble, boolean showAppBadge,
+            boolean isExpanding, boolean suppressAnimation) {
+        //TODO(b/360652359): remove setting scale to the app badge once issue is fixed
+        bubble.getView().setBadgeScale(showAppBadge ? 1 : 0);
+        mBubbles.put(bubble.getKey(), bubble);
+        mBubbleBarViewController.addBubble(bubble, isExpanding, suppressAnimation);
+    }
+
     /** Interface for checking whether the IME is visible. */
     public interface ImeVisibilityChecker {
         /** Whether the IME is visible. */
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
index 31b1ea0..bbb0b8b 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
@@ -324,6 +324,7 @@
     }
 
     private void onBubbleClicked(BubbleView bubbleView) {
+        if (mBubbleBarPinning.isAnimating()) return;
         bubbleView.markSeen();
         BubbleBarItem bubble = bubbleView.getBubble();
         if (bubble == null) {
@@ -918,7 +919,7 @@
      * from Launcher.
      */
     public void setExpanded(boolean isExpanded) {
-        if (isExpanded != mBarView.isExpanded()) {
+        if (!mBubbleBarPinning.isAnimating() && isExpanded != mBarView.isExpanded()) {
             mBarView.setExpanded(isExpanded);
             adjustTaskbarAndHotseatToBubbleBarState(isExpanded);
             if (!isExpanded) {
@@ -1060,6 +1061,16 @@
         mSystemUiProxy.removeAllBubbles();
     }
 
+    /** Removes all existing bubble views */
+    public void removeAllBubbles() {
+        mBarView.removeAllViews();
+    }
+
+    /** Returns the view index of the existing bubble */
+    public int bubbleViewIndex(View bubbleView) {
+        return mBarView.indexOfChild(bubbleView);
+    }
+
     /**
      * Set listener to be notified when bubble bar bounds have changed
      */
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleControllers.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleControllers.java
index b5d94bd..d993685 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleControllers.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleControllers.java
@@ -21,6 +21,7 @@
 import android.view.View;
 
 import com.android.launcher3.taskbar.TaskbarControllers;
+import com.android.launcher3.taskbar.TaskbarSharedState;
 import com.android.launcher3.taskbar.bubbles.BubbleBarViewController.TaskbarViewPropertiesProvider;
 import com.android.launcher3.taskbar.bubbles.stashing.BubbleBarLocationOnDemandListener;
 import com.android.launcher3.taskbar.bubbles.stashing.BubbleStashController;
@@ -79,7 +80,7 @@
      * BubbleControllers instance, but should be careful to only access things that were created
      * in constructors for now, as some controllers may still be waiting for init().
      */
-    public void init(TaskbarControllers taskbarControllers) {
+    public void init(TaskbarSharedState taskbarSharedState, TaskbarControllers taskbarControllers) {
         BubbleBarLocationCompositeListener bubbleBarLocationListeners =
                 new BubbleBarLocationCompositeListener(
                         taskbarControllers.navbarButtonsViewController,
@@ -88,7 +89,8 @@
                 );
         bubbleBarController.init(this,
                 bubbleBarLocationListeners,
-                taskbarControllers.navbarButtonsViewController::isImeVisible);
+                taskbarControllers.navbarButtonsViewController::isImeVisible,
+                taskbarSharedState);
         bubbleStashedHandleViewController.ifPresent(
                 controller -> controller.init(/* bubbleControllers = */ this));
         bubbleStashController.init(