Updating predictions if hotseat items get removed/added as a result of model callback

Bug: 179284787
Test: Verified on device
Change-Id: Iecfc7f291f54aacccc6d2afe674ccd2f20e2035d
diff --git a/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java b/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
index cbe0eb0..edcd0a2 100644
--- a/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
@@ -355,10 +355,6 @@
         // populating workspace.
         // TODO: Find a better place for this
         WellbeingModel.INSTANCE.get(this);
-
-        if (mTaskbarController != null) {
-            mTaskbarController.onHotseatUpdated();
-        }
     }
 
     @Override
diff --git a/quickstep/src/com/android/launcher3/hybridhotseat/HotseatPredictionController.java b/quickstep/src/com/android/launcher3/hybridhotseat/HotseatPredictionController.java
index aa6601b..88cfacb 100644
--- a/quickstep/src/com/android/launcher3/hybridhotseat/HotseatPredictionController.java
+++ b/quickstep/src/com/android/launcher3/hybridhotseat/HotseatPredictionController.java
@@ -20,6 +20,7 @@
 import static com.android.launcher3.hybridhotseat.HotseatEduController.getSettingsIntent;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_HOTSEAT_PREDICTION_PINNED;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_HOTSEAT_RANKED;
+import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 
 import android.animation.Animator;
 import android.animation.AnimatorSet;
@@ -70,7 +71,11 @@
  */
 public class HotseatPredictionController implements DragController.DragListener,
         SystemShortcut.Factory<QuickstepLauncher>, InvariantDeviceProfile.OnIDPChangeListener,
-        DragSource {
+        DragSource, ViewGroup.OnHierarchyChangeListener {
+
+    private static final int FLAG_UPDATE_PAUSED = 1 << 0;
+    private static final int FLAG_DRAG_IN_PROGRESS = 1 << 1;
+    private static final int FLAG_FILL_IN_PROGRESS = 1 << 2;
 
     private int mHotSeatItemsCount;
 
@@ -80,8 +85,7 @@
     private List<ItemInfo> mPredictedItems = Collections.emptyList();
 
     private AnimatorSet mIconRemoveAnimators;
-    private boolean mUIUpdatePaused = false;
-    private boolean mDragInProgress = false;
+    private int mPauseFlags = 0;
 
     private List<PredictedAppIcon.PredictedIconOutlineDrawing> mOutlineDrawings = new ArrayList<>();
 
@@ -114,6 +118,30 @@
         mLauncher.getDragController().addDragListener(this);
 
         launcher.getDeviceProfile().inv.addOnChangeListener(this);
+        mHotseat.getShortcutsAndWidgets().setOnHierarchyChangeListener(this);
+    }
+
+    @Override
+    public void onChildViewAdded(View parent, View child) {
+        onHotseatHierarchyChanged();
+    }
+
+    @Override
+    public void onChildViewRemoved(View parent, View child) {
+        onHotseatHierarchyChanged();
+    }
+
+    private void onHotseatHierarchyChanged() {
+        if (mPauseFlags == 0 && !mLauncher.isWorkspaceLoading()) {
+            // Post update after a single frame to avoid layout within layout
+            MAIN_EXECUTOR.getHandler().post(this::updateFillIfNotLoading);
+        }
+    }
+
+    private void updateFillIfNotLoading() {
+        if (mPauseFlags == 0 && !mLauncher.isWorkspaceLoading()) {
+            fillGapsWithPrediction(true);
+        }
     }
 
     /**
@@ -160,11 +188,11 @@
     }
 
     private void fillGapsWithPrediction() {
-        fillGapsWithPrediction(false, null);
+        fillGapsWithPrediction(false);
     }
 
-    private void fillGapsWithPrediction(boolean animate, Runnable callback) {
-        if (mUIUpdatePaused || mDragInProgress) {
+    private void fillGapsWithPrediction(boolean animate) {
+        if (mPauseFlags != 0) {
             return;
         }
 
@@ -175,12 +203,14 @@
             mIconRemoveAnimators.addListener(new AnimationSuccessListener() {
                 @Override
                 public void onAnimationSuccess(Animator animator) {
-                    fillGapsWithPrediction(animate, callback);
+                    fillGapsWithPrediction(animate);
                     mIconRemoveAnimators.removeListener(this);
                 }
             });
             return;
         }
+
+        mPauseFlags |= FLAG_FILL_IN_PROGRESS;
         for (int rank = 0; rank < mHotSeatItemsCount; rank++) {
             View child = mHotseat.getChildAt(
                     mHotseat.getCellXFromOrder(rank),
@@ -207,10 +237,12 @@
             }
             preparePredictionInfo(predictedItem, rank);
         }
-        bindItems(newItems, animate, callback);
+        bindItems(newItems, animate);
+
+        mPauseFlags &= ~FLAG_FILL_IN_PROGRESS;
     }
 
-    private void bindItems(List<WorkspaceItemInfo> itemsToAdd, boolean animate, Runnable callback) {
+    private void bindItems(List<WorkspaceItemInfo> itemsToAdd, boolean animate) {
         AnimatorSet animationSet = new AnimatorSet();
         for (WorkspaceItemInfo item : itemsToAdd) {
             PredictedAppIcon icon = PredictedAppIcon.createIcon(mHotseat, item);
@@ -221,12 +253,11 @@
             }
         }
         if (animate) {
-            if (callback != null) {
-                animationSet.addListener(AnimationSuccessListener.forRunnable(callback));
-            }
+            animationSet.addListener(AnimationSuccessListener
+                    .forRunnable(this::removeOutlineDrawings));
             animationSet.start();
         } else {
-            if (callback != null) callback.run();
+            removeOutlineDrawings();
         }
 
         if (mLauncher.getTaskbarController() != null) {
@@ -234,6 +265,16 @@
         }
     }
 
+    private void removeOutlineDrawings() {
+        if (mOutlineDrawings.isEmpty()) return;
+        for (PredictedAppIcon.PredictedIconOutlineDrawing outlineDrawing : mOutlineDrawings) {
+            mHotseat.removeDelegatedCellDrawing(outlineDrawing);
+        }
+        mHotseat.invalidate();
+        mOutlineDrawings.clear();
+    }
+
+
     /**
      * Unregisters callbacks and frees resources
      */
@@ -245,11 +286,9 @@
      * start and pauses predicted apps update on the hotseat
      */
     public void setPauseUIUpdate(boolean paused) {
-        if (mLauncher.getTaskbarController() != null) {
-            // Taskbar is present, always allow updates since hotseat is still visible.
-            return;
-        }
-        mUIUpdatePaused = paused;
+        mPauseFlags = paused
+                ? (mPauseFlags | FLAG_UPDATE_PAUSED)
+                : (mPauseFlags & ~FLAG_UPDATE_PAUSED);
         if (!paused) {
             fillGapsWithPrediction();
         }
@@ -365,14 +404,14 @@
         for (PredictedAppIcon.PredictedIconOutlineDrawing outlineDrawing : mOutlineDrawings) {
             mHotseat.addDelegatedCellDrawing(outlineDrawing);
         }
-        mDragInProgress = true;
+        mPauseFlags |= FLAG_DRAG_IN_PROGRESS;
         mHotseat.invalidate();
     }
 
     @Override
     public void onDragEnd() {
-        mDragInProgress = false;
-        fillGapsWithPrediction(true, this::removeOutlineDrawings);
+        mPauseFlags &= ~FLAG_DRAG_IN_PROGRESS;
+        fillGapsWithPrediction(true);
     }
 
     @Nullable
@@ -393,15 +432,6 @@
         itemInfo.screenId = rank;
     }
 
-    private void removeOutlineDrawings() {
-        if (mOutlineDrawings.isEmpty()) return;
-        for (PredictedAppIcon.PredictedIconOutlineDrawing outlineDrawing : mOutlineDrawings) {
-            mHotseat.removeDelegatedCellDrawing(outlineDrawing);
-        }
-        mHotseat.invalidate();
-        mOutlineDrawings.clear();
-    }
-
     @Override
     public void onIdpChanged(int changeFlags, InvariantDeviceProfile profile) {
         this.mHotSeatItemsCount = profile.numHotseatIcons;
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
index a00ce56..3a26f25 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
@@ -17,6 +17,7 @@
 
 import static android.view.accessibility.AccessibilityEvent.TYPE_VIEW_FOCUSED;
 
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APPLICATION;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_SHORTCUT;
@@ -52,6 +53,7 @@
 import com.android.launcher3.logging.StatsLogManager.StatsLogger;
 import com.android.launcher3.model.BgDataModel.FixedContainerItems;
 import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.popup.SystemShortcut;
 import com.android.launcher3.search.DeviceSearchAdapterProvider;
 import com.android.launcher3.statemanager.StateManager.AtomicAnimationFactory;
@@ -80,6 +82,7 @@
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
 import java.util.ArrayList;
+import java.util.List;
 import java.util.Objects;
 import java.util.stream.Stream;
 
@@ -152,7 +155,8 @@
     @Override
     public boolean startActivitySafely(View v, Intent intent, ItemInfo item) {
         if (mHotseatPredictionController != null) {
-            mHotseatPredictionController.setPauseUIUpdate(true);
+            // Only pause is taskbar controller is not present
+            mHotseatPredictionController.setPauseUIUpdate(getTaskbarController() == null);
         }
         return super.startActivitySafely(v, intent, item);
     }
@@ -218,6 +222,15 @@
     }
 
     @Override
+    public void bindWorkspaceItemsChanged(List<WorkspaceItemInfo> updated) {
+        super.bindWorkspaceItemsChanged(updated);
+        if (getTaskbarController() != null && updated.stream()
+                .filter(w -> w.container == CONTAINER_HOTSEAT).findFirst().isPresent()) {
+            getTaskbarController().onHotseatUpdated();
+        }
+    }
+
+    @Override
     public void onDestroy() {
         super.onDestroy();
         getAppsView().getSearchUiManager().destroy();