diff --git a/quickstep/res/layout/taskbar_overflow_button.xml b/quickstep/res/layout/taskbar_overflow_view.xml
similarity index 86%
rename from quickstep/res/layout/taskbar_overflow_button.xml
rename to quickstep/res/layout/taskbar_overflow_view.xml
index 20104f2..7444e59 100644
--- a/quickstep/res/layout/taskbar_overflow_button.xml
+++ b/quickstep/res/layout/taskbar_overflow_view.xml
@@ -15,8 +15,7 @@
 -->
 
 <!-- Note: The actual size will match the taskbar icon sizes in TaskbarView#onLayout(). -->
-<com.android.launcher3.views.IconButtonView xmlns:android="http://schemas.android.com/apk/res/android"
-    style="@style/BaseIcon.Workspace.Taskbar"
+<com.android.launcher3.taskbar.TaskbarOverflowView xmlns:android="http://schemas.android.com/apk/res/android"
     android:layout_width="@dimen/taskbar_icon_min_touch_size"
     android:layout_height="@dimen/taskbar_icon_min_touch_size"
     android:backgroundTint="@android:color/transparent"
diff --git a/quickstep/res/values/dimens.xml b/quickstep/res/values/dimens.xml
index 5f35007..451ba55 100644
--- a/quickstep/res/values/dimens.xml
+++ b/quickstep/res/values/dimens.xml
@@ -360,6 +360,7 @@
     <dimen name="taskbar_running_app_indicator_width">12dp</dimen>
     <dimen name="taskbar_running_app_indicator_top_margin">4dp</dimen>
     <dimen name="taskbar_minimized_app_indicator_width">6dp</dimen>
+    <dimen name="taskbar_overflow_button_preview_stroke">2dp</dimen>
 
     <!-- Transient taskbar -->
     <dimen name="transient_taskbar_padding">12dp</dimen>
@@ -502,6 +503,7 @@
     <dimen name="keyboard_quick_switch_recents_icon_size">20dp</dimen>
     <dimen name="keyboard_quick_switch_desktop_icon_size">32dp</dimen>
     <dimen name="keyboard_quick_switch_margin_top">56dp</dimen>
+    <dimen name="keyboard_quick_switch_margin_bottom">24dp</dimen>
     <dimen name="keyboard_quick_switch_margin_ends">16dp</dimen>
     <dimen name="keyboard_quick_switch_view_spacing">16dp</dimen>
     <dimen name="keyboard_quick_switch_view_small_spacing">4dp</dimen>
diff --git a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchViewController.java b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchViewController.java
index a80c11c..2902d55 100644
--- a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchViewController.java
@@ -20,6 +20,7 @@
 
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
+import android.content.res.Resources;
 import android.view.Gravity;
 import android.view.KeyEvent;
 import android.view.MotionEvent;
@@ -30,6 +31,7 @@
 import androidx.annotation.Nullable;
 
 import com.android.internal.jank.Cuj;
+import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.anim.AnimatorListeners;
 import com.android.launcher3.desktop.DesktopAppLaunchTransition;
@@ -115,6 +117,12 @@
 
         BaseDragLayer.LayoutParams lp = new BaseDragLayer.LayoutParams(
                 mKeyboardQuickSwitchView.getLayoutParams());
+        final Resources resources = mKeyboardQuickSwitchView.getResources();
+        final int marginHorizontal = resources.getDimensionPixelSize(
+                R.dimen.keyboard_quick_switch_margin_ends);
+        final int marginBottom = resources.getDimensionPixelSize(
+                R.dimen.keyboard_quick_switch_margin_bottom);
+        lp.setMargins(marginHorizontal, 0, marginHorizontal, marginBottom);
         lp.width = BaseDragLayer.LayoutParams.WRAP_CONTENT;
         lp.gravity = Gravity.BOTTOM | Gravity.CENTER_HORIZONTAL;
         mKeyboardQuickSwitchView.setLayoutParams(lp);
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index f3741b2..6f1e96f 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -297,6 +297,7 @@
             BubbleStashController bubbleStashController = isTransientTaskbar
                     ? new TransientBubbleStashController(dimensionsProvider, this)
                     : new PersistentBubbleStashController(dimensionsProvider);
+            bubbleStashController.setHotseatVerticalCenter(launcherDp.getHotseatVerticalCenter());
             bubbleControllersOptional = Optional.of(new BubbleControllers(
                     new BubbleBarController(this, bubbleBarView),
                     new BubbleBarViewController(this, bubbleBarView, bubbleBarContainer),
@@ -362,8 +363,11 @@
     /** Updates {@link DeviceProfile} instances for any Taskbar windows. */
     public void updateDeviceProfile(DeviceProfile launcherDp) {
         applyDeviceProfile(launcherDp);
-
         mControllers.taskbarOverlayController.updateLauncherDeviceProfile(launcherDp);
+        mControllers.bubbleControllers.ifPresent(bubbleControllers -> {
+            int hotseatVertCenter = launcherDp.getHotseatVerticalCenter();
+            bubbleControllers.bubbleStashController.setHotseatVerticalCenter(hotseatVertCenter);
+        });
         AbstractFloatingView.closeAllOpenViewsExcept(this, false, TYPE_REBIND_SAFE);
         // Reapply fullscreen to take potential new screen size into account.
         setTaskbarWindowFullscreen(mIsFullscreen);
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarEduTooltipController.kt b/quickstep/src/com/android/launcher3/taskbar/TaskbarEduTooltipController.kt
index ade8f8c..a89bc3a 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarEduTooltipController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarEduTooltipController.kt
@@ -265,7 +265,8 @@
                 !DisplayController.isPinnedTaskbar(activityContext) ||
                 !isTooltipEnabled ||
                 !shouldShowSearchEdu ||
-                userHasSeenSearchEdu
+                userHasSeenSearchEdu ||
+                !controllers.taskbarStashController.isTaskbarVisibleAndNotStashing
         ) {
             return
         }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarInsetsController.kt b/quickstep/src/com/android/launcher3/taskbar/TaskbarInsetsController.kt
index 058dd07..4a6b6d4 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarInsetsController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarInsetsController.kt
@@ -145,6 +145,7 @@
             // if bubble bar is visible or animating new bubble, add bar bounds to the touch region
             if (isBubbleBarVisible || isAnimatingNewBubble) {
                 defaultTouchableRegion.addBoundsToRegion(bubbleBarViewController.bubbleBarBounds)
+                defaultTouchableRegion.addBoundsToRegion(bubbleBarViewController.flyoutBounds)
             }
         }
         if (
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarOverflowView.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarOverflowView.java
new file mode 100644
index 0000000..fad5ca3
--- /dev/null
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarOverflowView.java
@@ -0,0 +1,212 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.taskbar;
+
+import android.content.Context;
+import android.graphics.Canvas;
+import android.graphics.Paint;
+import android.graphics.drawable.Drawable;
+import android.util.AttributeSet;
+import android.view.LayoutInflater;
+import android.view.ViewGroup;
+import android.widget.FrameLayout;
+
+import com.android.launcher3.R;
+import com.android.launcher3.Reorderable;
+import com.android.launcher3.Utilities;
+import com.android.launcher3.util.MultiTranslateDelegate;
+import com.android.systemui.shared.recents.model.Task;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * View used as overflow icon within task bar, when the list of recent/running apps overflows the
+ * available display bounds - if display is not wide enough to show all running apps in the taskbar,
+ * this icon is added to the taskbar as an entry point to open UI that surfaces all running apps.
+ * The icon contains icon representations of up to 4 more recent tasks in overflow, stacked on top
+ * each other in counter clockwise manner (icons of tasks partially overlapping with each other).
+ */
+public class TaskbarOverflowView extends FrameLayout implements Reorderable {
+    private final List<Task> mItems = new ArrayList<Task>();
+    private int mIconSize;
+    private int mPadding;
+    private Paint mItemBackgroundPaint;
+    private final MultiTranslateDelegate mTranslateDelegate = new MultiTranslateDelegate(this);
+    private float mScaleForReorderBounce = 1f;
+
+    public TaskbarOverflowView(Context context, AttributeSet attrs) {
+        super(context, attrs);
+        init();
+    }
+
+    public TaskbarOverflowView(Context context) {
+        super(context);
+        init();
+    }
+
+    /**
+     * Inflates the taskbar overflow button view.
+     * @param resId The resource to inflate the view from.
+     * @param group The parent view.
+     * @param iconSize The size of the overflow button icon.
+     * @param padding The internal padding of the overflow view.
+     * @return A taskbar overflow button.
+     */
+    public static TaskbarOverflowView inflateIcon(int resId, ViewGroup group, int iconSize,
+            int padding) {
+        LayoutInflater inflater = LayoutInflater.from(group.getContext());
+        TaskbarOverflowView icon = (TaskbarOverflowView) inflater.inflate(resId, group, false);
+
+        icon.mIconSize = iconSize;
+        icon.mPadding = padding;
+        return icon;
+    }
+
+    private void init() {
+        mItemBackgroundPaint = new Paint(Paint.ANTI_ALIAS_FLAG);
+        mItemBackgroundPaint.setColor(getContext().getColor(R.color.taskbar_background));
+
+        setWillNotDraw(false);
+    }
+
+    @Override
+    protected void onDraw(Canvas canvas) {
+        super.onDraw(canvas);
+
+        boolean isRtlLayout = Utilities.isRtl(getResources());
+        float radius = mIconSize / 2.0f - mPadding;
+        float itemPreviewStrokeWidth =
+                getResources().getDimension(R.dimen.taskbar_overflow_button_preview_stroke);
+
+        int itemsToShow = Math.min(mItems.size(), 4);
+        for (int i = itemsToShow - 1; i >= 0; --i) {
+            Drawable icon = mItems.get(mItems.size() - i - 1).icon;
+            if (icon == null) {
+                continue;
+            }
+
+            // Set the item icon size so two items fit within the overflow icon with stroke width
+            // included, and overlap of 4 stroke width sizes between base item preview items.
+            // 2 * strokeWidth + 2 * itemIconSize - 4 * strokeWidth = iconSize = 2 * radius.
+            float itemIconSize = radius + itemPreviewStrokeWidth;
+            // Offset item icon from center so item icon stroke edge mateches the parent icon edge.
+            float itemCenterOffset = radius - itemIconSize / 2 - itemPreviewStrokeWidth;
+
+            float itemCenterX = getItemXOffset(itemCenterOffset, isRtlLayout, i, itemsToShow);
+            float itemCenterY = getItemYOffset(itemCenterOffset, i, itemsToShow);
+
+            Drawable iconCopy = icon.getConstantState().newDrawable().mutate();
+            iconCopy.setBounds(0, 0, (int) itemIconSize, (int) itemIconSize);
+
+            canvas.save();
+            float itemIconRadius = itemIconSize / 2;
+            canvas.translate(
+                    mPadding + itemCenterX + radius - itemIconRadius,
+                    mPadding + itemCenterY + radius - itemIconRadius);
+            canvas.drawCircle(itemIconRadius, itemIconRadius,
+                    itemIconRadius + itemPreviewStrokeWidth, mItemBackgroundPaint);
+            iconCopy.draw(canvas);
+            canvas.restore();
+        }
+    }
+
+    /**
+     * Clears the list of tasks tracked by the view.
+     */
+    public void clearItems() {
+        mItems.clear();
+        invalidate();
+    }
+
+    /**
+     * Update the view to represent a new list of recent tasks.
+     * @param items Items to be shown in the view.
+     */
+    public void setItems(List<Task> items) {
+        mItems.clear();
+        mItems.addAll(items);
+        invalidate();
+    }
+
+    /**
+     * Called when a task is updated. If the task is contained within the view, it's cached value
+     * gets updated. If the task is shown within the icon, invalidates the view, so the task icon
+     * gets updated.
+     * @param task The updated task.
+     */
+    public void updateTaskIsShown(Task task) {
+        for (int i = 0; i < mItems.size(); ++i) {
+            if (mItems.get(i).key.id == task.key.id) {
+                mItems.set(i, task);
+                if (i >= mItems.size() - 4) {
+                    invalidate();
+                }
+                break;
+            }
+        }
+    }
+
+    @Override
+    public MultiTranslateDelegate getTranslateDelegate() {
+        return mTranslateDelegate;
+    }
+
+    @Override
+    public float getReorderBounceScale() {
+        return mScaleForReorderBounce;
+    }
+
+    @Override
+    public void setReorderBounceScale(float scale) {
+        mScaleForReorderBounce = scale;
+        super.setScaleX(scale);
+        super.setScaleY(scale);
+    }
+
+    private float getItemXOffset(float baseOffset, boolean isRtl, int itemIndex, int itemCount) {
+        // Item with index 1 is on the left in all cases.
+        if (itemIndex == 1) {
+            return (isRtl ? 1 : -1) * baseOffset;
+        }
+
+        // First item is centered if total number of items shown is 3, on the right otherwise.
+        if (itemIndex == 0) {
+            if (itemCount == 3) {
+                return 0;
+            }
+            return (isRtl ? -1 : 1) * baseOffset;
+        }
+
+        // Last item is on the right when there are more than 2 items (case which is already handled
+        // as `itemIndex == 1`).
+        if (itemIndex == itemCount - 1) {
+            return (isRtl ? -1 : 1) * baseOffset;
+        }
+
+        return (isRtl ? 1 : -1) * baseOffset;
+    }
+
+    private float getItemYOffset(float baseOffset, int itemIndex, int itemCount) {
+        // If icon contains two items, they are both centered vertically.
+        if (itemCount == 2) {
+            return 0;
+        }
+        // First half of items is on top, later half is on bottom.
+        return (itemIndex + 1 <= itemCount / 2 ? -1 : 1) * baseOffset;
+    }
+}
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
index fcb583a..6cc03ab 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
@@ -64,12 +64,12 @@
 import com.android.launcher3.util.LauncherBindableItemsContainer;
 import com.android.launcher3.util.Themes;
 import com.android.launcher3.views.ActivityContext;
-import com.android.launcher3.views.IconButtonView;
 import com.android.quickstep.util.GroupTask;
 import com.android.quickstep.views.TaskViewType;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.wm.shell.shared.bubbles.BubbleBarLocation;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.function.Predicate;
 
@@ -106,7 +106,7 @@
     @Nullable private TaskbarDividerContainer mTaskbarDividerContainer;
 
     // Only non-null when device supports having a Taskbar Overflow button.
-    @Nullable private IconButtonView mTaskbarOverflowView;
+    @Nullable private TaskbarOverflowView mTaskbarOverflowView;
 
     /**
      * Whether the divider is between Hotseat icons and Recents,
@@ -122,6 +122,8 @@
 
     private final int mMaxNumIcons;
 
+    private final int mAllAppsButtonTranslationOffset;
+
     public TaskbarView(@NonNull Context context) {
         this(context, null);
     }
@@ -141,8 +143,6 @@
         mActivityContext = ActivityContext.lookupContext(context);
         mIconLayoutBounds = mActivityContext.getTransientTaskbarBounds();
         Resources resources = getResources();
-        boolean isTransientTaskbar = DisplayController.isTransientTaskbar(mActivityContext)
-                && !mActivityContext.isPhoneMode();
         mIsRtl = Utilities.isRtl(resources);
         mTransientTaskbarMinWidth = resources.getDimension(R.dimen.transient_taskbar_min_width);
 
@@ -173,18 +173,19 @@
         setWillNotDraw(false);
 
         mAllAppsButtonContainer = new TaskbarAllAppsButtonContainer(context);
+        mAllAppsButtonTranslationOffset =  (int) getResources().getDimension(
+                mAllAppsButtonContainer.getAllAppsButtonTranslationXOffset(isTransientTaskbar()));
 
         if (enableTaskbarPinning() || enableRecentsInTaskbar()) {
             mTaskbarDividerContainer = new TaskbarDividerContainer(context);
         }
 
         if (Flags.taskbarOverflow()) {
-            mTaskbarOverflowView = (IconButtonView) LayoutInflater.from(context)
-                    .inflate(R.layout.taskbar_overflow_button, this, false);
-            mTaskbarOverflowView.setIconDrawable(
-                    resources.getDrawable(R.drawable.taskbar_overflow_icon));
-            mTaskbarOverflowView.setPadding(mItemPadding, mItemPadding, mItemPadding, mItemPadding);
+            mTaskbarOverflowView = TaskbarOverflowView.inflateIcon(
+                    R.layout.taskbar_overflow_view, this,
+                    mIconTouchSize, mItemPadding);
         }
+
         // TODO: Disable touch events on QSB otherwise it can crash.
         mQsb = LayoutInflater.from(context).inflate(R.layout.search_container_hotseat, this, false);
 
@@ -222,9 +223,7 @@
                 enableTaskbarPinning() && !mActivityContext.isThreeButtonNav();
         availableWidth -= iconSize - (int) getResources().getDimension(
                 mAllAppsButtonContainer.getAllAppsButtonTranslationXOffset(
-                        forceTransientTaskbarSize || (
-                                DisplayController.isTransientTaskbar(mActivityContext)
-                                        && !mActivityContext.isPhoneMode())));
+                        forceTransientTaskbarSize || isTransientTaskbar()));
         ++additionalIcons;
 
         return Math.floorDiv(availableWidth, iconSize) + additionalIcons;
@@ -449,24 +448,47 @@
         int nonTaskIconsToBeAdded = 1;
 
         boolean supportsOverflow = Flags.taskbarOverflow();
+        int overflowSize = 0;
         if (supportsOverflow) {
             int numberOfSupportedRecents = 0;
             for (GroupTask task : recentTasks) {
                 // TODO(b/343289567 and b/316004172): support app pairs and desktop mode.
-                if (!task.hasMultipleTasks()) {
+                if (!task.supportsMultipleTasks()) {
                     ++numberOfSupportedRecents;
                 }
             }
-            if (nextViewIndex + numberOfSupportedRecents + nonTaskIconsToBeAdded > mMaxNumIcons
-                    && mTaskbarOverflowView != null) {
+
+            overflowSize =
+                    nextViewIndex + numberOfSupportedRecents + nonTaskIconsToBeAdded - mMaxNumIcons;
+            if (overflowSize > 0 && mTaskbarOverflowView != null) {
                 addView(mTaskbarOverflowView, nextViewIndex++);
+            } else if (mTaskbarOverflowView != null) {
+                mTaskbarOverflowView.clearItems();
             }
         }
 
+        List<Task> overflownTasks = null;
+        // An extra item needs to be added to overflow button to account for the space taken up by
+        // the overflow button.
+        final int itemsToAddToOverflow = overflowSize > 0 ? overflowSize + 1 : 0;
+        if (overflowSize > 0) {
+            overflownTasks = new ArrayList<Task>(itemsToAddToOverflow);
+        }
+
         // Add Recent/Running icons.
         for (GroupTask task : recentTasks) {
-            if (supportsOverflow && nextViewIndex + nonTaskIconsToBeAdded >= mMaxNumIcons) {
-                break;
+            if (mTaskbarOverflowView != null && overflownTasks != null
+                    && overflownTasks.size() < itemsToAddToOverflow) {
+                // TODO(b/343289567 and b/316004172): support app pairs and desktop mode.
+                if (task.supportsMultipleTasks()) {
+                    continue;
+                }
+
+                overflownTasks.add(task.task1);
+                if (overflownTasks.size() == itemsToAddToOverflow) {
+                    mTaskbarOverflowView.setItems(overflownTasks);
+                }
+                continue;
             }
 
             // Replace any Recent views with the appropriate type if it's not already that type.
@@ -668,6 +690,15 @@
         mIconLayoutBounds.right = iconEnd;
         mIconLayoutBounds.top = (bottom - top - mIconTouchSize) / 2;
         mIconLayoutBounds.bottom = mIconLayoutBounds.top + mIconTouchSize;
+
+        // With rtl layout, the all apps button will be translated by `allAppsButtonOffset` after
+        // layout completion (by `TaskbarViewController`). Offset the icon end by the same amount
+        // when laying out icons, so the taskbar content remains centered after all apps button
+        // translation.
+        if (layoutRtl) {
+            iconEnd += mAllAppsButtonTranslationOffset;
+        }
+
         int count = getChildCount();
         for (int i = count; i > 0; i--) {
             View child = getChildAt(i - 1);
@@ -699,6 +730,15 @@
 
         mIconLayoutBounds.left = iconEnd;
 
+        // Adjust the icon layout bounds by the amount by which all apps button will be translated
+        // post layout to maintain margin between all apps button and the edge of the transient
+        // taskbar background. Done for ltr layout only - for rtl layout, the offset needs to be
+        // adjusted on the right, which is done by offsetting `iconEnd` after setting
+        // `mIconLayoutBounds.right`.
+        if (!layoutRtl) {
+            mIconLayoutBounds.left += mAllAppsButtonTranslationOffset;
+        }
+
         if (mIconLayoutBounds.right - mIconLayoutBounds.left < mTransientTaskbarMinWidth) {
             int center = mIconLayoutBounds.centerX();
             int distanceFromCenter = (int) mTransientTaskbarMinWidth / 2;
@@ -741,12 +781,13 @@
     /**
      * Returns the space used by the icons
      */
-    public int getIconLayoutWidth() {
+    private int getIconLayoutWidth() {
         int countExcludingQsb = getChildCount();
         DeviceProfile deviceProfile = mActivityContext.getDeviceProfile();
         if (deviceProfile.isQsbInline) {
             countExcludingQsb--;
         }
+
         int iconLayoutBoundsWidth =
                 countExcludingQsb * (mItemMarginLeftRight * 2 + mIconTouchSize);
 
@@ -755,17 +796,28 @@
             // All Apps icon, divider icon, and first app icon in taskbar
             iconLayoutBoundsWidth -= mItemMarginLeftRight * 4;
         }
+
+        // The all apps button container gets offset horizontally, reducing the overall taskbar
+        // view size.
+        iconLayoutBoundsWidth -= mAllAppsButtonTranslationOffset;
+
         return iconLayoutBoundsWidth;
     }
 
     /**
-     * Returns the app icons currently shown in the taskbar.
+     * Returns the app icons currently shown in the taskbar. The returned list does not include qsb,
+     * but it includes all apps button and icon divider views.
      */
     public View[] getIconViews() {
         final int count = getChildCount();
-        View[] icons = new View[count];
+        if (count == 0) {
+            return new View[0];
+        }
+        View[] icons = new View[count - (mActivityContext.getDeviceProfile().isQsbInline ? 1 : 0)];
+        int insertionPoint = 0;
         for (int i = 0; i < count; i++) {
-            icons[i] = getChildAt(i);
+            if (getChildAt(i)  == mQsb) continue;
+            icons[insertionPoint++] = getChildAt(i);
         }
         return icons;
     }
@@ -789,7 +841,7 @@
      * Returns the taskbar overflow view in the taskbar.
      */
     @Nullable
-    public IconButtonView getTaskbarOverflowView() {
+    public TaskbarOverflowView getTaskbarOverflowView() {
         return mTaskbarOverflowView;
     }
 
@@ -847,12 +899,25 @@
         // Ignore, we just implement Insettable to draw behind system insets.
     }
 
+    private boolean isTransientTaskbar() {
+        return DisplayController.isTransientTaskbar(mActivityContext)
+                && !mActivityContext.isPhoneMode();
+    }
+
     public boolean areIconsVisible() {
         // Consider the overall visibility
         return getVisibility() == VISIBLE;
     }
 
     /**
+     * @return The all apps button horizontal offset used to calculate the taskbar contents width
+     * during layout.
+     */
+    public int getAllAppsButtonTranslationXOffsetUsedForLayout() {
+        return mAllAppsButtonTranslationOffset;
+    }
+
+    /**
      * Maps {@code op} over all the child views.
      */
     public void mapOverItems(LauncherBindableItemsContainer.ItemOperator op) {
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java
index 253d025..6f2d459 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java
@@ -359,10 +359,6 @@
         return mTaskbarView.getIconLayoutBounds();
     }
 
-    public int getIconLayoutWidth() {
-        return mTaskbarView.getIconLayoutWidth();
-    }
-
     public View[] getIconViews() {
         return mTaskbarView.getIconViews();
     }
@@ -442,6 +438,19 @@
         float allAppIconTranslateRange = mapRange(scale, transientTaskbarAllAppsOffset,
                 persistentTaskbarAllAppsOffset);
 
+        // Task icons are laid out so the taskbar content is centered. The taskbar width (used for
+        // centering taskbar icons) depends on the all apps button X translation, and is different
+        // for persistent and transient taskbar. If the offset used for current taskbar layout is
+        // different than the offset used in final taskbar state, the icons may jump when the
+        // animation completes, and the taskbar is replaced. Adjust item transform to account for
+        // this mismatch.
+        float sizeDiffTranslationRange =
+                mapRange(scale,
+                        (mTaskbarView.getAllAppsButtonTranslationXOffsetUsedForLayout()
+                                - transientTaskbarAllAppsOffset) / 2,
+                        (mTaskbarView.getAllAppsButtonTranslationXOffsetUsedForLayout()
+                                - persistentTaskbarAllAppsOffset) / 2);
+
         // no x translation required when all apps button is the only icon in taskbar.
         if (iconViews.length <= 1) {
             allAppIconTranslateRange = 0f;
@@ -449,6 +458,7 @@
 
         if (mIsRtl) {
             allAppIconTranslateRange *= -1;
+            sizeDiffTranslationRange *= -1;
         }
 
         if (mActivity.isThreeButtonNav()) {
@@ -457,25 +467,18 @@
             return;
         }
 
-        float taskbarCenterX =
-                mTaskbarView.getLeft() + (mTaskbarView.getRight() - mTaskbarView.getLeft()) / 2.0f;
-
         float finalMarginScale = mapRange(scale, 0f, mTransientIconSize - mPersistentIconSize);
 
-        float halfIconCount = iconViews.length / 2.0f;
+        // The index of the "middle" icon which will be used as a index from which the icon margins
+        // will be scaled. If number of icons is even, using the middle point between indices of two
+        // central icons.
+        float middleIndex = (iconViews.length - 1) / 2.0f;
         for (int iconIndex = 0; iconIndex < iconViews.length; iconIndex++) {
             View iconView = iconViews[iconIndex];
             MultiTranslateDelegate translateDelegate =
                     ((Reorderable) iconView).getTranslateDelegate();
-            float iconCenterX =
-                    iconView.getLeft() + (iconView.getRight() - iconView.getLeft()) / 2.0f;
-            if (iconCenterX <= taskbarCenterX) {
-                translateDelegate.getTranslationX(INDEX_TASKBAR_PINNING_ANIM).setValue(
-                        finalMarginScale * (halfIconCount - iconIndex));
-            } else {
-                translateDelegate.getTranslationX(INDEX_TASKBAR_PINNING_ANIM).setValue(
-                        -finalMarginScale * (iconIndex - halfIconCount));
-            }
+            translateDelegate.getTranslationX(INDEX_TASKBAR_PINNING_ANIM).setValue(
+                    finalMarginScale * (middleIndex - iconIndex) + sizeDiffTranslationRange);
 
             if (iconView.equals(mTaskbarView.getAllAppsButtonContainer())) {
                 mTaskbarView.getAllAppsButtonContainer().setTranslationXForTaskbarAllAppsIcon(
@@ -488,18 +491,14 @@
      * Calculates visual taskbar view width.
      */
     public float getCurrentVisualTaskbarWidth() {
-        if (mTaskbarView.getIconViews().length == 0) {
+        View[] iconViews = mTaskbarView.getIconViews();
+        if (iconViews.length == 0) {
             return 0;
         }
 
-        View[] iconViews = mTaskbarView.getIconViews();
+        float left = iconViews[0].getX();
 
-        int leftIndex = mActivity.getDeviceProfile().isQsbInline && !mIsRtl ? 1 : 0;
-        int rightIndex = mActivity.getDeviceProfile().isQsbInline && mIsRtl
-                ? iconViews.length - 2
-                : iconViews.length - 1;
-
-        float left = iconViews[leftIndex].getX();
+        int rightIndex = iconViews.length - 1;
         float right = iconViews[rightIndex].getRight() + iconViews[rightIndex].getTranslationX();
 
         return right - left + (2 * mTaskbarLeftRightMargin);
@@ -1090,6 +1089,8 @@
                 if (groupTask.containsTask(task.key.id)) {
                     mTaskbarView.applyGroupTaskToBubbleTextView(btv, groupTask);
                 }
+            } else if (view instanceof TaskbarOverflowView overflowButton) {
+                overflowButton.updateTaskIsShown(task);
             }
         }
     }
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
index 76d3606..717ee95 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
@@ -45,6 +45,7 @@
 import com.android.launcher3.taskbar.bubbles.animation.BubbleBarViewAnimator;
 import com.android.launcher3.taskbar.bubbles.flyout.BubbleBarFlyoutController;
 import com.android.launcher3.taskbar.bubbles.flyout.BubbleBarFlyoutPositioner;
+import com.android.launcher3.taskbar.bubbles.flyout.FlyoutCallbacks;
 import com.android.launcher3.taskbar.bubbles.stashing.BubbleStashController;
 import com.android.launcher3.util.MultiPropertyFactory;
 import com.android.launcher3.util.MultiValueAlpha;
@@ -146,7 +147,7 @@
         mTaskbarStashController = controllers.taskbarStashController;
         mTaskbarInsetsController = controllers.taskbarInsetsController;
         mBubbleBarFlyoutController = new BubbleBarFlyoutController(
-                mBubbleBarContainer, createFlyoutPositioner(), createFlyoutTopBoundaryListener());
+                mBubbleBarContainer, createFlyoutPositioner(), createFlyoutCallbacks());
         mBubbleBarViewAnimator = new BubbleBarViewAnimator(
                 mBarView, mBubbleStashController, mBubbleBarFlyoutController,
                 mBubbleBarController::showExpandedView);
@@ -175,7 +176,9 @@
 
             @Override
             public void onBubbleBarTouched() {
-                BubbleBarViewController.this.onBubbleBarTouched();
+                if (isAnimatingNewBubble()) {
+                    interruptAnimationForTouch();
+                }
             }
 
             @Override
@@ -273,8 +276,8 @@
         };
     }
 
-    private BubbleBarFlyoutController.TopBoundaryListener createFlyoutTopBoundaryListener() {
-        return new BubbleBarFlyoutController.TopBoundaryListener() {
+    private FlyoutCallbacks createFlyoutCallbacks() {
+        return new FlyoutCallbacks() {
             @Override
             public void extendTopBoundary(int space) {
                 int defaultSize = mActivity.getDefaultTaskbarWindowSize();
@@ -285,6 +288,12 @@
             public void resetTopBoundary() {
                 mActivity.setTaskbarWindowSize(mActivity.getDefaultTaskbarWindowSize());
             }
+
+            @Override
+            public void flyoutClicked() {
+                interruptAnimationForTouch();
+                expandBubbleBar();
+            }
         };
     }
 
@@ -304,12 +313,10 @@
         }
     }
 
-    private void onBubbleBarTouched() {
-        if (isAnimatingNewBubble()) {
-            mBubbleBarViewAnimator.onBubbleBarTouchedWhileAnimating();
-            mBubbleStashController.onNewBubbleAnimationInterrupted(false,
-                    mBarView.getTranslationY());
-        }
+    /** Interrupts the running animation for a touch event on the bubble bar or flyout. */
+    private void interruptAnimationForTouch() {
+        mBubbleBarViewAnimator.interruptForTouch();
+        mBubbleStashController.onNewBubbleAnimationInterrupted(false, mBarView.getTranslationY());
     }
 
     private void expandBubbleBar() {
@@ -463,6 +470,12 @@
         return mBarView.getBubbleBarBounds();
     }
 
+    /** Returns the bounds of the flyout view if it exists, or {@code null} otherwise. */
+    @Nullable
+    public Rect getFlyoutBounds() {
+        return mBubbleBarFlyoutController.getFlyoutBounds();
+    }
+
     /** Checks that bubble bar is visible and that the motion event is within bounds. */
     public boolean isEventOverBubbleBar(MotionEvent event) {
         if (!isBubbleBarVisible()) return false;
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/animation/BubbleBarViewAnimator.kt b/quickstep/src/com/android/launcher3/taskbar/bubbles/animation/BubbleBarViewAnimator.kt
index 8a52ca9..78e5dbd 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/animation/BubbleBarViewAnimator.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/animation/BubbleBarViewAnimator.kt
@@ -471,8 +471,8 @@
         bubbleStashController.updateTaskbarTouchRegion()
     }
 
-    /** Handles touching the animating bubble bar. */
-    fun onBubbleBarTouchedWhileAnimating() {
+    /** Interrupts the animation due to touching the bubble bar or flyout. */
+    fun interruptForTouch() {
         PhysicsAnimator.getInstance(bubbleBarView).cancelIfRunning()
         bubbleStashController.getStashedHandlePhysicsAnimator().cancelIfRunning()
         cancelFlyout()
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutController.kt b/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutController.kt
index d6400bb..fdbbbb0 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutController.kt
@@ -16,13 +16,13 @@
 
 package com.android.launcher3.taskbar.bubbles.flyout
 
+import android.graphics.Rect
 import android.view.Gravity
 import android.view.ViewGroup
 import android.widget.FrameLayout
 import androidx.core.animation.ValueAnimator
 import com.android.launcher3.R
 import com.android.systemui.util.addListener
-import com.android.systemui.util.doOnEnd
 
 /** Creates and manages the visibility of the [BubbleBarFlyoutView]. */
 class BubbleBarFlyoutController
@@ -30,7 +30,7 @@
 constructor(
     private val container: FrameLayout,
     private val positioner: BubbleBarFlyoutPositioner,
-    private val topBoundaryListener: TopBoundaryListener,
+    private val callbacks: FlyoutCallbacks,
     private val flyoutScheduler: FlyoutScheduler = HandlerScheduler(container),
 ) {
 
@@ -47,6 +47,15 @@
         FADE,
     }
 
+    /** The bounds of the flyout. */
+    val flyoutBounds: Rect?
+        get() {
+            val flyout = this.flyout ?: return null
+            val rect = Rect(flyout.bounds)
+            rect.offset(0, flyout.translationY.toInt())
+            return rect
+        }
+
     fun setUpAndShowFlyout(message: BubbleBarFlyoutMessage, onEnd: () -> Unit) {
         flyout?.let(container::removeView)
         val flyout = BubbleBarFlyoutView(container.context, positioner, flyoutScheduler)
@@ -72,9 +81,12 @@
                 val flyoutTop = flyout.top + flyout.translationY
                 // If the top position of the flyout is negative, then it's bleeding over the
                 // top boundary of its parent view
-                if (flyoutTop < 0) topBoundaryListener.extendTopBoundary(space = -flyoutTop.toInt())
+                if (flyoutTop < 0) callbacks.extendTopBoundary(space = -flyoutTop.toInt())
             },
-            onEnd = { onEnd() },
+            onEnd = {
+                onEnd()
+                flyout.setOnClickListener { callbacks.flyoutClicked() }
+            },
         )
         flyout.showFromCollapsed(message) { animator.start() }
         this.flyout = flyout
@@ -100,21 +112,15 @@
                     flyout.updateExpansionProgress(animator.animatedValue as Float)
                 }
         }
-        animator.doOnEnd {
-            container.removeView(flyout)
-            this@BubbleBarFlyoutController.flyout = null
-            topBoundaryListener.resetTopBoundary()
-            endAction()
-        }
+        animator.addListener(
+            onStart = { flyout.setOnClickListener(null) },
+            onEnd = {
+                container.removeView(flyout)
+                this@BubbleBarFlyoutController.flyout = null
+                callbacks.resetTopBoundary()
+                endAction()
+            },
+        )
         animator.start()
     }
-
-    /** Notifies when the top boundary of the flyout view changes. */
-    interface TopBoundaryListener {
-        /** Requests to extend the top boundary of the parent to fully include the flyout. */
-        fun extendTopBoundary(space: Int)
-
-        /** Resets the top boundary of the parent. */
-        fun resetTopBoundary()
-    }
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutView.kt b/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutView.kt
index 6903c87..bb8a392 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutView.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutView.kt
@@ -24,6 +24,7 @@
 import android.graphics.Paint
 import android.graphics.Path
 import android.graphics.PointF
+import android.graphics.Rect
 import android.graphics.RectF
 import android.view.LayoutInflater
 import android.view.View
@@ -138,6 +139,9 @@
      */
     private val backgroundPaint = Paint(Paint.ANTI_ALIAS_FLAG or Paint.FILTER_BITMAP_FLAG)
 
+    /** The bounds of the flyout relative to the parent view. */
+    val bounds = Rect()
+
     init {
         LayoutInflater.from(context).inflate(R.layout.bubblebar_flyout, this, true)
         id = R.id.bubble_bar_flyout_view
@@ -174,6 +178,14 @@
         applyConfigurationColors(resources.configuration)
     }
 
+    override fun onLayout(changed: Boolean, left: Int, top: Int, right: Int, bottom: Int) {
+        super.onLayout(changed, left, top, right, bottom)
+        bounds.left = left
+        bounds.top = top
+        bounds.right = right
+        bounds.bottom = bottom
+    }
+
     /** Sets the data for the flyout and starts playing the expand animation. */
     fun showFromCollapsed(flyoutMessage: BubbleBarFlyoutMessage, expandAnimation: () -> Unit) {
         icon.alpha = 0f
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/FlyoutCallbacks.kt b/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/FlyoutCallbacks.kt
new file mode 100644
index 0000000..e2f010a
--- /dev/null
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/FlyoutCallbacks.kt
@@ -0,0 +1,29 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.taskbar.bubbles.flyout
+
+/** Callbacks that the flyout uses to notify of events. */
+interface FlyoutCallbacks {
+    /** Requests to extend the top boundary of the parent to fully include the flyout. */
+    fun extendTopBoundary(space: Int)
+
+    /** Resets the top boundary of the parent. */
+    fun resetTopBoundary()
+
+    /** The flyout was clicked. */
+    fun flyoutClicked()
+}
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/BubbleStashController.kt b/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/BubbleStashController.kt
index a78890b..831faa1 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/BubbleStashController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/BubbleStashController.kt
@@ -42,12 +42,6 @@
 
         /** Provides taskbar height in pixels. */
         fun getTaskbarHeight(): Int
-
-        /** Provides hotseat bottom space in pixels. */
-        fun getHotseatBottomSpace(): Int
-
-        /** Provides hotseat height in pixels. */
-        fun getHotseatHeight(): Int
     }
 
     /** Execute passed action only after controllers are initiated. */
@@ -94,7 +88,7 @@
         taskbarInsetsController: TaskbarInsetsController,
         bubbleBarViewController: BubbleBarViewController,
         bubbleStashedHandleViewController: BubbleStashedHandleViewController?,
-        controllersAfterInitAction: ControllersAfterInitAction
+        controllersAfterInitAction: ControllersAfterInitAction,
     )
 
     /** Shows the bubble bar at [bubbleBarTranslationY] position immediately without animation. */
@@ -127,6 +121,9 @@
     /** Set a bubble bar location */
     fun setBubbleBarLocation(bubbleBarLocation: BubbleBarLocation)
 
+    /** Set the hotseat vertical center that bubble bar will align with. */
+    fun setHotseatVerticalCenter(hotseatVerticalCenter: Int)
+
     /**
      * Stashes the bubble bar (transform to the handle view), or just shrink width of the expanded
      * bubble bar based on the controller implementation.
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/DeviceProfileDimensionsProviderAdapter.kt b/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/DeviceProfileDimensionsProviderAdapter.kt
index a55763b..886b9f0 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/DeviceProfileDimensionsProviderAdapter.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/DeviceProfileDimensionsProviderAdapter.kt
@@ -27,13 +27,9 @@
 class DeviceProfileDimensionsProviderAdapter(
     private val taskbarActivityContext: TaskbarActivityContext
 ) : TaskbarHotseatDimensionsProvider {
-    override fun getTaskbarBottomSpace(): Int = deviceProfile().taskbarBottomMargin
+    override fun getTaskbarBottomSpace(): Int = taskbarDp().taskbarBottomMargin
 
-    override fun getTaskbarHeight(): Int = deviceProfile().taskbarHeight
+    override fun getTaskbarHeight(): Int = taskbarDp().taskbarHeight
 
-    override fun getHotseatBottomSpace(): Int = deviceProfile().hotseatBarBottomSpacePx
-
-    override fun getHotseatHeight(): Int = deviceProfile().hotseatCellHeightPx
-
-    private fun deviceProfile(): DeviceProfile = taskbarActivityContext.deviceProfile
+    private fun taskbarDp(): DeviceProfile = taskbarActivityContext.deviceProfile
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/PersistentBubbleStashController.kt b/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/PersistentBubbleStashController.kt
index 722dfe7..c117ad4 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/PersistentBubbleStashController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/PersistentBubbleStashController.kt
@@ -36,7 +36,7 @@
 import com.android.wm.shell.shared.bubbles.BubbleBarLocation
 
 class PersistentBubbleStashController(
-    private val taskbarHotseatDimensionsProvider: TaskbarHotseatDimensionsProvider,
+    private val taskbarHotseatDimensionsProvider: TaskbarHotseatDimensionsProvider
 ) : BubbleStashController {
 
     private lateinit var taskbarInsetsController: TaskbarInsetsController
@@ -45,6 +45,7 @@
     private lateinit var bubbleBarAlphaAnimator: MultiPropertyFactory<View>.MultiProperty
     private lateinit var bubbleBarScaleAnimator: AnimatedFloat
     private lateinit var controllersAfterInitAction: ControllersAfterInitAction
+    private var hotseatVerticalCenter: Int = 0
 
     override var launcherState: BubbleLauncherState = BubbleLauncherState.IN_APP
         set(state) {
@@ -92,17 +93,15 @@
 
     override val bubbleBarTranslationYForHotseat: Float
         get() {
-            val hotseatBottomSpace = taskbarHotseatDimensionsProvider.getHotseatBottomSpace()
-            val hotseatCellHeight = taskbarHotseatDimensionsProvider.getHotseatHeight()
-            val bubbleBarHeight: Float = bubbleBarViewController.bubbleBarCollapsedHeight
-            return -hotseatBottomSpace - (hotseatCellHeight - bubbleBarHeight) / 2
+            val bubbleBarHeight = bubbleBarViewController.bubbleBarCollapsedHeight
+            return -hotseatVerticalCenter + bubbleBarHeight / 2
         }
 
     override fun init(
         taskbarInsetsController: TaskbarInsetsController,
         bubbleBarViewController: BubbleBarViewController,
         bubbleStashedHandleViewController: BubbleStashedHandleViewController?,
-        controllersAfterInitAction: ControllersAfterInitAction
+        controllersAfterInitAction: ControllersAfterInitAction,
     ) {
         this.taskbarInsetsController = taskbarInsetsController
         this.bubbleBarViewController = bubbleBarViewController
@@ -119,13 +118,17 @@
             animatorSet.playTogether(
                 bubbleBarScaleAnimator.animateToValue(1f),
                 bubbleBarTranslationYAnimator.animateToValue(bubbleBarTranslationY),
-                bubbleBarAlphaAnimator.animateToValue(1f)
+                bubbleBarAlphaAnimator.animateToValue(1f),
             )
         }
         updateTouchRegionOnAnimationEnd(animatorSet)
         animatorSet.setDuration(BAR_STASH_DURATION).start()
     }
 
+    override fun setHotseatVerticalCenter(hotseatVerticalCenter: Int) {
+        this.hotseatVerticalCenter = hotseatVerticalCenter
+    }
+
     override fun showBubbleBarImmediate() = showBubbleBarImmediate(bubbleBarTranslationY)
 
     override fun showBubbleBarImmediate(bubbleBarTranslationY: Float) {
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/TransientBubbleStashController.kt b/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/TransientBubbleStashController.kt
index 9e7d1c4..fbeecaa 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/TransientBubbleStashController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/TransientBubbleStashController.kt
@@ -78,6 +78,7 @@
         context.resources.getDimensionPixelSize(R.dimen.bubblebar_stashed_size) / 2f
 
     private var animator: AnimatorSet? = null
+    private var hotseatVerticalCenter: Int = 0
 
     override var isStashed: Boolean = false
         @VisibleForTesting set
@@ -118,10 +119,8 @@
 
     override val bubbleBarTranslationYForHotseat: Float
         get() {
-            val hotseatBottomSpace = taskbarHotseatDimensionsProvider.getHotseatBottomSpace()
-            val hotseatCellHeight = taskbarHotseatDimensionsProvider.getHotseatHeight()
-            val bubbleBarHeight: Float = bubbleBarViewController.bubbleBarCollapsedHeight
-            return -hotseatBottomSpace - (hotseatCellHeight - bubbleBarHeight) / 2
+            val bubbleBarHeight = bubbleBarViewController.bubbleBarCollapsedHeight
+            return -hotseatVerticalCenter + bubbleBarHeight / 2
         }
 
     override val bubbleBarTranslationYForTaskbar: Float =
@@ -176,6 +175,10 @@
             .start()
     }
 
+    override fun setHotseatVerticalCenter(hotseatVerticalCenter: Int) {
+        this.hotseatVerticalCenter = hotseatVerticalCenter
+    }
+
     override fun showBubbleBarImmediate() {
         showBubbleBarImmediate(bubbleBarTranslationY)
     }
diff --git a/quickstep/src/com/android/quickstep/util/ScalingWorkspaceRevealAnim.kt b/quickstep/src/com/android/quickstep/util/ScalingWorkspaceRevealAnim.kt
index 2f0a6df..db02f55 100644
--- a/quickstep/src/com/android/quickstep/util/ScalingWorkspaceRevealAnim.kt
+++ b/quickstep/src/com/android/quickstep/util/ScalingWorkspaceRevealAnim.kt
@@ -16,10 +16,13 @@
 
 package com.android.quickstep.util
 
+import android.animation.Animator
+import android.animation.AnimatorListenerAdapter
 import android.animation.AnimatorSet
 import android.graphics.Matrix
 import android.graphics.Path
 import android.graphics.RectF
+import android.util.Log
 import android.view.View
 import android.view.animation.PathInterpolator
 import androidx.core.graphics.transform
@@ -41,6 +44,8 @@
 import com.android.launcher3.uioverrides.QuickstepLauncher
 import com.android.quickstep.views.RecentsView
 
+const val TAG = "ScalingWorkspaceRevealAnim"
+
 /**
  * Creates an animation where the workspace and hotseat fade in while revealing from the center of
  * the screen outwards radially. This is used in conjunction with the swipe up to home animation.
@@ -197,6 +202,19 @@
         workspace.setLayerType(View.LAYER_TYPE_HARDWARE, null)
         hotseat.setLayerType(View.LAYER_TYPE_HARDWARE, null)
         animation.addListener(
+            object : AnimatorListenerAdapter() {
+                override fun onAnimationCancel(animation: Animator) {
+                    super.onAnimationCancel(animation)
+                    Log.d(TAG, "onAnimationCancel")
+                }
+
+                override fun onAnimationPause(animation: Animator) {
+                    super.onAnimationPause(animation)
+                    Log.d(TAG, "onAnimationPause")
+                }
+            }
+        )
+        animation.addListener(
             AnimatorListeners.forEndCallback(
                 Runnable {
                     workspace.setLayerType(View.LAYER_TYPE_NONE, null)
@@ -207,6 +225,8 @@
                         Animations.setOngoingAnimation(workspace, animation = null)
                         Animations.setOngoingAnimation(hotseat, animation = null)
                     }
+
+                    Log.d(TAG, "alpha of workspace at the end of animation: ${workspace.alpha}")
                 }
             )
         )
diff --git a/quickstep/src/com/android/quickstep/util/WorkspaceRevealAnim.java b/quickstep/src/com/android/quickstep/util/WorkspaceRevealAnim.java
index 0a97793..32e0e13 100644
--- a/quickstep/src/com/android/quickstep/util/WorkspaceRevealAnim.java
+++ b/quickstep/src/com/android/quickstep/util/WorkspaceRevealAnim.java
@@ -30,6 +30,7 @@
 import android.animation.AnimatorSet;
 import android.animation.ObjectAnimator;
 import android.util.FloatProperty;
+import android.util.Log;
 import android.view.View;
 
 import com.android.app.animation.Interpolators;
@@ -51,6 +52,8 @@
  */
 public class WorkspaceRevealAnim {
 
+    private static final String TAG = "WorkspaceRevealAnim";
+
     // Should be used for animations running alongside this WorkspaceRevealAnim.
     public static final int DURATION_MS = 350;
     private static final FloatProperty<Workspace<?>> WORKSPACE_SCALE_PROPERTY =
@@ -97,6 +100,19 @@
 
         mAnimators.setDuration(DURATION_MS);
         mAnimators.setInterpolator(Interpolators.DECELERATED_EASE);
+        mAnimators.addListener(new AnimatorListenerAdapter() {
+            @Override
+            public void onAnimationCancel(Animator animation) {
+                super.onAnimationCancel(animation);
+                Log.d(TAG, "onAnimationCancel");
+            }
+
+            @Override
+            public void onAnimationEnd(Animator animation) {
+                super.onAnimationEnd(animation);
+                Log.d(TAG, "onAnimationEnd: workspace alpha = " + workspace.getAlpha());
+            }
+        });
     }
 
     private <T extends View>  void addRevealAnimatorsForView(T v, FloatProperty<T> scaleProperty) {
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/animation/BubbleBarViewAnimatorTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/animation/BubbleBarViewAnimatorTest.kt
index b37048a..b0d01d3 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/animation/BubbleBarViewAnimatorTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/animation/BubbleBarViewAnimatorTest.kt
@@ -40,6 +40,7 @@
 import com.android.launcher3.taskbar.bubbles.flyout.BubbleBarFlyoutController
 import com.android.launcher3.taskbar.bubbles.flyout.BubbleBarFlyoutMessage
 import com.android.launcher3.taskbar.bubbles.flyout.BubbleBarFlyoutPositioner
+import com.android.launcher3.taskbar.bubbles.flyout.FlyoutCallbacks
 import com.android.launcher3.taskbar.bubbles.flyout.FlyoutScheduler
 import com.android.launcher3.taskbar.bubbles.stashing.BubbleStashController
 import com.android.wm.shell.shared.animation.PhysicsAnimator
@@ -179,9 +180,7 @@
         // verify the hide bubble animation is pending
         assertThat(animatorScheduler.delayedBlock).isNotNull()
 
-        InstrumentationRegistry.getInstrumentation().runOnMainSync {
-            animator.onBubbleBarTouchedWhileAnimating()
-        }
+        InstrumentationRegistry.getInstrumentation().runOnMainSync { animator.interruptForTouch() }
 
         waitForFlyoutToHide()
 
@@ -992,18 +991,20 @@
                 override val collapsedElevation = 1f
                 override val distanceToRevealTriangle = 10f
             }
-        val topBoundaryListener =
-            object : BubbleBarFlyoutController.TopBoundaryListener {
+        val flyoutCallbacks =
+            object : FlyoutCallbacks {
                 override fun extendTopBoundary(space: Int) {}
 
                 override fun resetTopBoundary() {}
+
+                override fun flyoutClicked() {}
             }
         val flyoutScheduler = FlyoutScheduler { block -> block.invoke() }
         flyoutController =
             BubbleBarFlyoutController(
                 flyoutContainer,
                 flyoutPositioner,
-                topBoundaryListener,
+                flyoutCallbacks,
                 flyoutScheduler,
             )
     }
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutControllerTest.kt
index 527bdaa..0eea741 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutControllerTest.kt
@@ -44,7 +44,7 @@
 
     private lateinit var flyoutController: BubbleBarFlyoutController
     private lateinit var flyoutContainer: FrameLayout
-    private lateinit var topBoundaryListener: FakeTopBoundaryListener
+    private lateinit var flyoutCallbacks: FakeFlyoutCallbacks
     private val context = ApplicationProvider.getApplicationContext<Context>()
     private val flyoutMessage = BubbleBarFlyoutMessage(icon = null, "sender name", "message")
     private var onLeft = true
@@ -67,15 +67,10 @@
                 override val collapsedElevation = 1f
                 override val distanceToRevealTriangle = 50f
             }
-        topBoundaryListener = FakeTopBoundaryListener()
+        flyoutCallbacks = FakeFlyoutCallbacks()
         val flyoutScheduler = FlyoutScheduler { block -> block.invoke() }
         flyoutController =
-            BubbleBarFlyoutController(
-                flyoutContainer,
-                positioner,
-                topBoundaryListener,
-                flyoutScheduler,
-            )
+            BubbleBarFlyoutController(flyoutContainer, positioner, flyoutCallbacks, flyoutScheduler)
     }
 
     @Test
@@ -140,7 +135,7 @@
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
             animatorTestRule.advanceTimeBy(300)
         }
-        assertThat(topBoundaryListener.topBoundaryExtendedSpace).isEqualTo(50)
+        assertThat(flyoutCallbacks.topBoundaryExtendedSpace).isEqualTo(50)
     }
 
     @Test
@@ -153,7 +148,7 @@
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
             animatorTestRule.advanceTimeBy(300)
         }
-        assertThat(topBoundaryListener.topBoundaryExtendedSpace).isEqualTo(0)
+        assertThat(flyoutCallbacks.topBoundaryExtendedSpace).isEqualTo(0)
     }
 
     @Test
@@ -164,7 +159,7 @@
             flyoutController.collapseFlyout {}
             animatorTestRule.advanceTimeBy(300)
         }
-        assertThat(topBoundaryListener.topBoundaryReset).isTrue()
+        assertThat(flyoutCallbacks.topBoundaryReset).isTrue()
     }
 
     @Test
@@ -178,13 +173,27 @@
             animatorTestRule.advanceTimeBy(300)
             assertThat(flyoutView.alpha).isEqualTo(0f)
         }
-        assertThat(topBoundaryListener.topBoundaryReset).isTrue()
+        assertThat(flyoutCallbacks.topBoundaryReset).isTrue()
     }
 
-    class FakeTopBoundaryListener : BubbleBarFlyoutController.TopBoundaryListener {
+    @Test
+    fun clickFlyout_notifiesCallback() {
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {
+            flyoutController.setUpAndShowFlyout(flyoutMessage) {}
+            assertThat(flyoutContainer.childCount).isEqualTo(1)
+            val flyoutView = flyoutContainer.findViewById<View>(R.id.bubble_bar_flyout_view)
+            assertThat(flyoutView.alpha).isEqualTo(1f)
+            animatorTestRule.advanceTimeBy(300)
+            flyoutView.performClick()
+        }
+        assertThat(flyoutCallbacks.flyoutClicked).isTrue()
+    }
+
+    class FakeFlyoutCallbacks : FlyoutCallbacks {
 
         var topBoundaryExtendedSpace = 0
         var topBoundaryReset = false
+        var flyoutClicked = false
 
         override fun extendTopBoundary(space: Int) {
             topBoundaryExtendedSpace = space
@@ -193,5 +202,9 @@
         override fun resetTopBoundary() {
             topBoundaryReset = true
         }
+
+        override fun flyoutClicked() {
+            flyoutClicked = true
+        }
     }
 }
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/stashing/PersistentBubbleStashControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/stashing/PersistentBubbleStashControllerTest.kt
index 5dc78a9..00b42bc 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/stashing/PersistentBubbleStashControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/stashing/PersistentBubbleStashControllerTest.kt
@@ -48,6 +48,7 @@
 
     companion object {
         const val BUBBLE_BAR_HEIGHT = 100f
+        const val HOTSEAT_VERTICAL_CENTER = 95
         const val HOTSEAT_TRANSLATION_Y = -45f
         const val TASK_BAR_TRANSLATION_Y = -5f
     }
@@ -74,11 +75,12 @@
             PersistentBubbleStashController(DefaultDimensionsProvider())
         setUpBubbleBarView()
         setUpBubbleBarController()
+        persistentTaskBarStashController.setHotseatVerticalCenter(HOTSEAT_VERTICAL_CENTER)
         persistentTaskBarStashController.init(
             taskbarInsetsController,
             bubbleBarViewController,
             null,
-            ImmediateAction()
+            ImmediateAction(),
         )
     }
 
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/stashing/StashingTestUtils.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/stashing/StashingTestUtils.kt
index 0f8a2c3..96c2f45 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/stashing/StashingTestUtils.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/stashing/StashingTestUtils.kt
@@ -23,21 +23,13 @@
 class DefaultDimensionsProvider(
     private val taskBarBottomSpace: Int = TASKBAR_BOTTOM_SPACE,
     private val taskBarHeight: Int = TASKBAR_HEIGHT,
-    private val hotseatBottomSpace: Int = HOTSEAT_BOTTOM_SPACE,
-    private val hotseatHeight: Int = HOTSEAT_HEIGHT
 ) : BubbleStashController.TaskbarHotseatDimensionsProvider {
     override fun getTaskbarBottomSpace(): Int = taskBarBottomSpace
 
     override fun getTaskbarHeight(): Int = taskBarHeight
 
-    override fun getHotseatBottomSpace(): Int = hotseatBottomSpace
-
-    override fun getHotseatHeight(): Int = hotseatHeight
-
     companion object {
         const val TASKBAR_BOTTOM_SPACE = 0
         const val TASKBAR_HEIGHT = 110
-        const val HOTSEAT_BOTTOM_SPACE = 20
-        const val HOTSEAT_HEIGHT = 150
     }
 }
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/stashing/TransientBubbleStashControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/stashing/TransientBubbleStashControllerTest.kt
index 8b277e7..64416dd 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/stashing/TransientBubbleStashControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/stashing/TransientBubbleStashControllerTest.kt
@@ -60,6 +60,7 @@
 
     companion object {
         const val TASKBAR_BOTTOM_SPACE = 5
+        const val HOTSEAT_VERTICAL_CENTER = 95
         const val BUBBLE_BAR_WIDTH = 200
         const val BUBBLE_BAR_HEIGHT = 100
         const val HOTSEAT_TRANSLATION_Y = -45f
@@ -108,6 +109,7 @@
         setUpStashedHandleView()
         setUpBubbleStashedHandleViewController()
         PhysicsAnimatorTestUtils.prepareForTest()
+        mTransientBubbleStashController.setHotseatVerticalCenter(HOTSEAT_VERTICAL_CENTER)
         mTransientBubbleStashController.init(
             taskbarInsetsController,
             bubbleBarViewController,
diff --git a/src/com/android/launcher3/BaseDraggingActivity.java b/src/com/android/launcher3/BaseDraggingActivity.java
index 177b28c..5f5fb72 100644
--- a/src/com/android/launcher3/BaseDraggingActivity.java
+++ b/src/com/android/launcher3/BaseDraggingActivity.java
@@ -50,6 +50,8 @@
     // automatically when user interacts with the launcher.
     public static final Object AUTO_CANCEL_ACTION_MODE = new Object();
 
+    private boolean mIsThemeUpdatedBeforeRecreate;
+
     private ActionMode mCurrentActionMode;
 
     private int mThemeRes = R.style.AppTheme;
@@ -80,8 +82,13 @@
         updateTheme();
     }
 
-    private void updateTheme() {
+    public boolean isThemeUpdatedBeforeRecreate() {
+        return mIsThemeUpdatedBeforeRecreate;
+    }
+
+    protected void updateTheme() {
         if (mThemeRes != Themes.getActivityThemeRes(this)) {
+            mIsThemeUpdatedBeforeRecreate = true;
             recreate();
         }
     }
diff --git a/src/com/android/launcher3/DeviceProfile.java b/src/com/android/launcher3/DeviceProfile.java
index 8862550..4305703 100644
--- a/src/com/android/launcher3/DeviceProfile.java
+++ b/src/com/android/launcher3/DeviceProfile.java
@@ -24,6 +24,7 @@
 import static com.android.launcher3.InvariantDeviceProfile.INDEX_TWO_PANEL_LANDSCAPE;
 import static com.android.launcher3.InvariantDeviceProfile.INDEX_TWO_PANEL_PORTRAIT;
 import static com.android.launcher3.Utilities.dpiFromPx;
+import static com.android.launcher3.Utilities.isEnglishLanguage;
 import static com.android.launcher3.Utilities.pxFromSp;
 import static com.android.launcher3.folder.ClippedFolderIconLayoutRule.ICON_OVERLAP_FACTOR;
 import static com.android.launcher3.icons.GraphicsUtils.getShapePath;
@@ -1344,8 +1345,14 @@
         }
         if ((Flags.enableTwolineToggle()
                 && LauncherPrefs.ENABLE_TWOLINE_ALLAPPS_TOGGLE.get(context))) {
-            // Add extra textHeight to the existing allAppsCellHeight.
-            allAppsCellHeightPx += Utilities.calculateTextHeight(allAppsIconTextSizePx);
+            if (!isEnglishLanguage(context)) {
+                // Set toggle preference value to false if not english here as it's possible the
+                // preference is stale after language change.
+                LauncherPrefs.get(context).put(LauncherPrefs.ENABLE_TWOLINE_ALLAPPS_TOGGLE, false);
+            } else {
+                // Add extra textHeight to the existing allAppsCellHeight.
+                allAppsCellHeightPx += Utilities.calculateTextHeight(allAppsIconTextSizePx);
+            }
         }
 
         updateHotseatSizes(iconSizePx);
@@ -2019,6 +2026,18 @@
     }
 
     /**
+     * Returns the number of pixels the hotseat icons vertical center is translated from the bottom
+     * of the screen.
+     */
+    public int getHotseatVerticalCenter() {
+        return hotseatBarSizePx
+                - (isQsbInline ? 0 : hotseatQsbVisualHeight)
+                - hotseatQsbSpace
+                - (hotseatCellHeightPx / 2)
+                + ((hotseatCellHeightPx - iconSizePx) / 2);
+    }
+
+    /**
      * Returns the number of pixels the taskbar is translated from the bottom of the screen.
      */
     public int getTaskbarOffsetY() {
diff --git a/src/com/android/launcher3/LauncherBackupAgent.java b/src/com/android/launcher3/LauncherBackupAgent.java
index 2617b93..a96495d 100644
--- a/src/com/android/launcher3/LauncherBackupAgent.java
+++ b/src/com/android/launcher3/LauncherBackupAgent.java
@@ -1,5 +1,7 @@
 package com.android.launcher3;
 
+import static com.android.launcher3.LauncherPrefs.NO_DB_FILES_RESTORED;
+
 import android.app.backup.BackupAgent;
 import android.app.backup.BackupDataInput;
 import android.app.backup.BackupDataOutput;
@@ -10,10 +12,13 @@
 
 import java.io.File;
 import java.io.IOException;
+import java.util.Arrays;
+import java.util.stream.Collectors;
 
 public class LauncherBackupAgent extends BackupAgent {
-
     private static final String TAG = "LauncherBackupAgent";
+    private static final String DB_FILE_PREFIX = "launcher";
+    private static final String DB_FILE_SUFFIX = ".db";
 
     @Override
     public void onCreate() {
@@ -47,7 +52,34 @@
 
     @Override
     public void onRestoreFinished() {
-        FileLog.d(TAG, "onRestoreFinished: set pending for RestoreDbTask");
         RestoreDbTask.setPending(this);
+        FileLog.d(TAG, "onRestoreFinished: set pending for RestoreDbTask");
+        markIfFilesWereNotActuallyRestored();
+    }
+
+    /**
+     * When restore is finished, we check to see if any db files were successfully restored. If not,
+     * our restore will fail later, but will report a different cause. This is important to split
+     * out the metric failures that are launcher's fault, and those that are due to bugs in the
+     * backup/restore code itself.
+     */
+    private void markIfFilesWereNotActuallyRestored() {
+        File directory = new File(getDatabasePath(InvariantDeviceProfile.INSTANCE.get(this).dbFile)
+                .getParent());
+        if (!directory.exists()) {
+            FileLog.e(TAG, "restore failed as target database directory doesn't exist");
+        } else {
+            // Check for any db file that was restored, and collect as list
+            String fileNames = Arrays.stream(directory.listFiles())
+                    .map(File::getName)
+                    .filter(n -> n.startsWith(DB_FILE_PREFIX) && n.endsWith(DB_FILE_SUFFIX))
+                    .collect(Collectors.joining(", "));
+            if (fileNames.isBlank()) {
+                FileLog.e(TAG, "no database files were successfully restored");
+                LauncherPrefs.get(this).putSync(NO_DB_FILES_RESTORED.to(true));
+            } else {
+                FileLog.d(TAG, "database files successfully restored: " + fileNames);
+            }
+        }
     }
 }
diff --git a/src/com/android/launcher3/LauncherPrefs.kt b/src/com/android/launcher3/LauncherPrefs.kt
index 7ebfc18..5c03644 100644
--- a/src/com/android/launcher3/LauncherPrefs.kt
+++ b/src/com/android/launcher3/LauncherPrefs.kt
@@ -126,6 +126,9 @@
                 EncryptionType.ENCRYPTED,
             )
         @JvmField
+        val NO_DB_FILES_RESTORED =
+            nonRestorableItem("no_db_files_restored", false, EncryptionType.DEVICE_PROTECTED)
+        @JvmField
         val IS_FIRST_LOAD_AFTER_RESTORE =
             nonRestorableItem(FIRST_LOAD_AFTER_RESTORE_KEY, false, EncryptionType.ENCRYPTED)
         @JvmField val APP_WIDGET_IDS = backedUpItem(RestoreDbTask.APPWIDGET_IDS, "")
diff --git a/src/com/android/launcher3/Utilities.java b/src/com/android/launcher3/Utilities.java
index 9192e13..71a2589 100644
--- a/src/com/android/launcher3/Utilities.java
+++ b/src/com/android/launcher3/Utilities.java
@@ -616,6 +616,14 @@
     }
 
     /**
+     * Utility method to know if a device's primary language is English.
+     */
+    public static boolean isEnglishLanguage(Context context) {
+        return context.getResources().getConfiguration().locale.getLanguage()
+                .equals(Locale.ENGLISH.getLanguage());
+    }
+
+    /**
      * Returns the full drawable for info as multiple layers of AdaptiveIconDrawable. The second
      * drawable in the Pair is the badge used with the icon.
      *
diff --git a/src/com/android/launcher3/backuprestore/LauncherRestoreEventLogger.kt b/src/com/android/launcher3/backuprestore/LauncherRestoreEventLogger.kt
index e6654b1..b05539a 100644
--- a/src/com/android/launcher3/backuprestore/LauncherRestoreEventLogger.kt
+++ b/src/com/android/launcher3/backuprestore/LauncherRestoreEventLogger.kt
@@ -24,9 +24,10 @@
         RestoreError.WIDGETS_DISABLED,
         RestoreError.PROFILE_NOT_RESTORED,
         RestoreError.WIDGET_REMOVED,
+        RestoreError.DATABASE_FILE_NOT_RESTORED,
         RestoreError.GRID_MIGRATION_FAILURE,
         RestoreError.NO_SEARCH_WIDGET,
-        RestoreError.INVALID_WIDGET_ID
+        RestoreError.INVALID_WIDGET_ID,
     )
     annotation class RestoreError {
         companion object {
@@ -38,6 +39,7 @@
             const val APP_NOT_INSTALLED = "app_not_installed"
             const val WIDGETS_DISABLED = "widgets_disabled"
             const val PROFILE_NOT_RESTORED = "profile_not_restored"
+            const val DATABASE_FILE_NOT_RESTORED = "db_file_not_restored"
             const val WIDGET_REMOVED = "widget_not_found"
             const val GRID_MIGRATION_FAILURE = "grid_migration_failed"
             const val NO_SEARCH_WIDGET = "no_search_widget"
@@ -52,7 +54,7 @@
             return ResourceBasedOverride.Overrides.getObject(
                 LauncherRestoreEventLogger::class.java,
                 context,
-                R.string.launcher_restore_event_logger_class
+                R.string.launcher_restore_event_logger_class,
             )
         }
     }
@@ -117,7 +119,7 @@
     open fun logFavoritesItemsRestoreFailed(
         favoritesId: Int,
         count: Int,
-        @RestoreError error: String?
+        @RestoreError error: String?,
     ) {
         // no-op
     }
diff --git a/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java b/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
index e3c2d36..1dd7d45 100644
--- a/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
+++ b/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
@@ -55,7 +55,7 @@
 import com.android.launcher3.model.BaseLauncherBinder;
 import com.android.launcher3.model.BgDataModel;
 import com.android.launcher3.model.BgDataModel.Callbacks;
-import com.android.launcher3.model.GridSizeMigrationUtil;
+import com.android.launcher3.model.GridSizeMigrationDBController;
 import com.android.launcher3.model.LoaderTask;
 import com.android.launcher3.model.ModelDbController;
 import com.android.launcher3.provider.LauncherDbUtils;
@@ -284,7 +284,7 @@
     private void loadModelData() {
         final Context inflationContext = getPreviewContext();
         final InvariantDeviceProfile idp = new InvariantDeviceProfile(inflationContext, mGridName);
-        if (GridSizeMigrationUtil.needsToMigrate(inflationContext, idp)) {
+        if (GridSizeMigrationDBController.needsToMigrate(inflationContext, idp)) {
             // Start the migration
             PreviewContext previewContext = new PreviewContext(inflationContext, idp);
             // Copy existing data to preview DB
diff --git a/src/com/android/launcher3/model/DbEntry.java b/src/com/android/launcher3/model/DbEntry.java
new file mode 100644
index 0000000..c0c51da
--- /dev/null
+++ b/src/com/android/launcher3/model/DbEntry.java
@@ -0,0 +1,155 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.model;
+
+import android.content.ContentValues;
+import android.content.Intent;
+import android.util.Log;
+
+import androidx.annotation.NonNull;
+
+import com.android.launcher3.LauncherSettings;
+import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.util.ContentWriter;
+
+import java.net.URISyntaxException;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+public class DbEntry extends ItemInfo implements Comparable<DbEntry> {
+
+    private static final String TAG = "DbEntry";
+
+    String mIntent;
+    String mProvider;
+    Map<String, Set<Integer>> mFolderItems = new HashMap<>();
+
+    /**
+     * Id of the specific widget.
+     */
+    public int appWidgetId = NO_ID;
+
+    /** Comparator according to the reading order */
+    @Override
+    public int compareTo(DbEntry another) {
+        if (screenId != another.screenId) {
+            return Integer.compare(screenId, another.screenId);
+        }
+        if (cellY != another.cellY) {
+            return Integer.compare(cellY, another.cellY);
+        }
+        return Integer.compare(cellX, another.cellX);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (!(o instanceof DbEntry)) return false;
+        DbEntry entry = (DbEntry) o;
+        return Objects.equals(getEntryMigrationId(), entry.getEntryMigrationId());
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(getEntryMigrationId());
+    }
+
+    /**
+     *  Puts the updated DbEntry values into ContentValues which we then use to insert the
+     *  entry to the DB.
+     */
+    public void updateContentValues(ContentValues values) {
+        values.put(LauncherSettings.Favorites.SCREEN, screenId);
+        values.put(LauncherSettings.Favorites.CELLX, cellX);
+        values.put(LauncherSettings.Favorites.CELLY, cellY);
+        values.put(LauncherSettings.Favorites.SPANX, spanX);
+        values.put(LauncherSettings.Favorites.SPANY, spanY);
+    }
+
+    @Override
+    public void writeToValues(@NonNull ContentWriter writer) {
+        super.writeToValues(writer);
+        writer.put(LauncherSettings.Favorites.APPWIDGET_ID, appWidgetId);
+    }
+
+    @Override
+    public void readFromValues(@NonNull ContentValues values) {
+        super.readFromValues(values);
+        appWidgetId = values.getAsInteger(LauncherSettings.Favorites.APPWIDGET_ID);
+    }
+
+    /**
+     * This id is not used in the DB is only used while doing the migration and it identifies
+     * an entry on each workspace. For example two calculator icons would have the same
+     * migration id even thought they have different database ids.
+     */
+    public String getEntryMigrationId() {
+        switch (itemType) {
+            case LauncherSettings.Favorites.ITEM_TYPE_FOLDER:
+            case LauncherSettings.Favorites.ITEM_TYPE_APP_PAIR:
+                return getFolderMigrationId();
+            case LauncherSettings.Favorites.ITEM_TYPE_APPWIDGET:
+                // mProvider is the app the widget belongs to and appWidgetId it's the unique
+                // is of the widget, we need both because if you remove a widget and then add it
+                // again, then it can change and the WidgetProvider would not know the widget.
+                return mProvider + appWidgetId;
+            case LauncherSettings.Favorites.ITEM_TYPE_APPLICATION:
+                final String intentStr = cleanIntentString(mIntent);
+                try {
+                    Intent i = Intent.parseUri(intentStr, 0);
+                    return Objects.requireNonNull(i.getComponent()).toString();
+                } catch (Exception e) {
+                    return intentStr;
+                }
+            default:
+                return cleanIntentString(mIntent);
+        }
+    }
+
+    /**
+     * This method should return an id that should be the same for two folders containing the
+     * same elements.
+     */
+    @NonNull
+    private String getFolderMigrationId() {
+        return mFolderItems.keySet().stream()
+                .map(intentString -> mFolderItems.get(intentString).size()
+                        + cleanIntentString(intentString))
+                .sorted()
+                .collect(Collectors.joining(","));
+    }
+
+    /**
+     * This is needed because sourceBounds can change and make the id of two equal items
+     * different.
+     */
+    @NonNull
+    private String cleanIntentString(@NonNull String intentStr) {
+        try {
+            Intent i = Intent.parseUri(intentStr, 0);
+            i.setSourceBounds(null);
+            return i.toURI();
+        } catch (URISyntaxException e) {
+            Log.e(TAG, "Unable to parse Intent string", e);
+            return intentStr;
+        }
+
+    }
+}
diff --git a/src/com/android/launcher3/model/GridSizeMigrationUtil.java b/src/com/android/launcher3/model/GridSizeMigrationDBController.java
similarity index 83%
rename from src/com/android/launcher3/model/GridSizeMigrationUtil.java
rename to src/com/android/launcher3/model/GridSizeMigrationDBController.java
index 4c017e9..9531d5b 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationUtil.java
+++ b/src/com/android/launcher3/model/GridSizeMigrationDBController.java
@@ -28,7 +28,6 @@
 import android.content.ComponentName;
 import android.content.ContentValues;
 import android.content.Context;
-import android.content.Intent;
 import android.database.Cursor;
 import android.database.DatabaseUtils;
 import android.database.sqlite.SQLiteDatabase;
@@ -45,15 +44,12 @@
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.config.FeatureFlags;
-import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.provider.LauncherDbUtils.SQLiteTransaction;
-import com.android.launcher3.util.ContentWriter;
 import com.android.launcher3.util.GridOccupancy;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.WidgetManagerHelper;
 
-import java.net.URISyntaxException;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
@@ -61,7 +57,6 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
-import java.util.Objects;
 import java.util.Set;
 import java.util.stream.Collectors;
 
@@ -69,12 +64,12 @@
  * This class takes care of shrinking the workspace (by maximum of one row and one column), as a
  * result of restoring from a larger device or device density change.
  */
-public class GridSizeMigrationUtil {
+public class GridSizeMigrationDBController {
 
-    private static final String TAG = "GridSizeMigrationUtil";
+    private static final String TAG = "GridSizeMigrationDBController";
     private static final boolean DEBUG = true;
 
-    private GridSizeMigrationUtil() {
+    private GridSizeMigrationDBController() {
         // Util class should not be instantiated
     }
 
@@ -85,7 +80,7 @@
         return needsToMigrate(new DeviceGridState(context), new DeviceGridState(idp));
     }
 
-    private static boolean needsToMigrate(
+    static boolean needsToMigrate(
             DeviceGridState srcDeviceState, DeviceGridState destDeviceState) {
         boolean needsToMigrate = !destDeviceState.isCompatible(srcDeviceState);
         if (needsToMigrate) {
@@ -95,6 +90,9 @@
         return needsToMigrate;
     }
 
+    /**
+     * @return all the workspace and hotseat entries in the db.
+     */
     @VisibleForTesting
     public static List<DbEntry> readAllEntries(SQLiteDatabase db, String tableName,
             Context context) {
@@ -198,7 +196,7 @@
                     Collectors.joining(",\n", "[", "]"))
                     + "\n Removing Items:"
                     + dstWorkspaceItems.stream().filter(entry ->
-                            toBeRemoved.contains(entry.id)).map(DbEntry::toString).collect(
+                    toBeRemoved.contains(entry.id)).map(DbEntry::toString).collect(
                     Collectors.joining(",\n", "[", "]"))
                     + "\n Adding Workspace Items:"
                     + workspaceToBeAdded.stream().map(DbEntry::toString).collect(
@@ -291,7 +289,7 @@
         });
     }
 
-    private static void insertEntryInDb(DatabaseHelper helper, DbEntry entry,
+    static void insertEntryInDb(DatabaseHelper helper, DbEntry entry,
             String srcTableName, String destTableName, List<Integer> idsInUse) {
         int id = copyEntryAndUpdate(helper, entry, srcTableName, destTableName, idsInUse);
         if (entry.itemType == LauncherSettings.Favorites.ITEM_TYPE_FOLDER
@@ -341,7 +339,7 @@
         return newId;
     }
 
-    private static void removeEntryFromDb(SQLiteDatabase db, String tableName, IntArray entryIds) {
+    static void removeEntryFromDb(SQLiteDatabase db, String tableName, IntArray entryIds) {
         db.delete(tableName,
                 Utilities.createDbSelectionQuery(LauncherSettings.Favorites._ID, entryIds), null);
     }
@@ -387,7 +385,7 @@
     private static boolean findPlacementForEntry(@NonNull final DbEntry entry,
             @NonNull final Point next, @NonNull final Point trg,
             @NonNull final GridOccupancy occupied, final int screenId) {
-        for (int y = next.y; y <  trg.y; y++) {
+        for (int y = next.y; y < trg.y; y++) {
             for (int x = next.x; x < trg.x; x++) {
                 boolean fits = occupied.isRegionVacant(x, y, entry.spanX, entry.spanY);
                 boolean minFits = occupied.isRegionVacant(x, y, entry.minSpanX,
@@ -413,7 +411,7 @@
     private static void solveHotseatPlacement(
             @NonNull final DatabaseHelper helper, final int hotseatSize,
             @NonNull final DbReader srcReader, @NonNull final DbReader destReader,
-            @NonNull final  List<DbEntry> placedHotseatItems,
+            @NonNull final List<DbEntry> placedHotseatItems,
             @NonNull final List<DbEntry> itemsToPlace, List<Integer> idsInUse) {
 
         final boolean[] occupied = new boolean[hotseatSize];
@@ -436,15 +434,26 @@
         }
     }
 
-    @VisibleForTesting
+    static void copyCurrentGridToNewGrid(
+            @NonNull Context context,
+            @NonNull DeviceGridState destDeviceState,
+            @NonNull DatabaseHelper target,
+            @NonNull SQLiteDatabase source) {
+        // Only use this strategy when comparing the previous grid to the new grid and the
+        // columns are the same and the destination has more rows
+        copyTable(source, TABLE_NAME, target.getWritableDatabase(), TABLE_NAME, context);
+        destDeviceState.writeToPrefs(context);
+    }
+
+    @VisibleForTesting(otherwise = VisibleForTesting.PACKAGE_PRIVATE)
     public static class DbReader {
 
-        private final SQLiteDatabase mDb;
-        private final String mTableName;
-        private final Context mContext;
-        private int mLastScreenId = -1;
+        final SQLiteDatabase mDb;
+        final String mTableName;
+        final Context mContext;
+        int mLastScreenId = -1;
 
-        private final Map<Integer, ArrayList<DbEntry>> mWorkspaceEntriesByScreenId =
+        final Map<Integer, ArrayList<DbEntry>> mWorkspaceEntriesByScreenId =
                 new ArrayMap<>();
 
         public DbReader(SQLiteDatabase db, String tableName, Context context) {
@@ -529,7 +538,7 @@
                             LauncherSettings.Favorites.INTENT,               // 7
                             LauncherSettings.Favorites.APPWIDGET_PROVIDER,   // 8
                             LauncherSettings.Favorites.APPWIDGET_ID},        // 9
-                        LauncherSettings.Favorites.CONTAINER + " = "
+                    LauncherSettings.Favorites.CONTAINER + " = "
                             + LauncherSettings.Favorites.CONTAINER_DESKTOP);
             final int indexId = c.getColumnIndexOrThrow(LauncherSettings.Favorites._ID);
             final int indexItemType = c.getColumnIndexOrThrow(LauncherSettings.Favorites.ITEM_TYPE);
@@ -648,118 +657,4 @@
             return mDb.query(mTableName, columns, where, null, null, null, null);
         }
     }
-
-    public static class DbEntry extends ItemInfo implements Comparable<DbEntry> {
-
-        private String mIntent;
-        private String mProvider;
-        private Map<String, Set<Integer>> mFolderItems = new HashMap<>();
-
-        /**
-         * Id of the specific widget.
-         */
-        public int appWidgetId = NO_ID;
-
-        /** Comparator according to the reading order */
-        @Override
-        public int compareTo(DbEntry another) {
-            if (screenId != another.screenId) {
-                return Integer.compare(screenId, another.screenId);
-            }
-            if (cellY != another.cellY) {
-                return Integer.compare(cellY, another.cellY);
-            }
-            return Integer.compare(cellX, another.cellX);
-        }
-
-        @Override
-        public boolean equals(Object o) {
-            if (this == o) return true;
-            if (o == null || getClass() != o.getClass()) return false;
-            DbEntry entry = (DbEntry) o;
-            return Objects.equals(getEntryMigrationId(), entry.getEntryMigrationId());
-        }
-
-        @Override
-        public int hashCode() {
-            return Objects.hash(getEntryMigrationId());
-        }
-
-        public void updateContentValues(ContentValues values) {
-            values.put(LauncherSettings.Favorites.SCREEN, screenId);
-            values.put(LauncherSettings.Favorites.CELLX, cellX);
-            values.put(LauncherSettings.Favorites.CELLY, cellY);
-            values.put(LauncherSettings.Favorites.SPANX, spanX);
-            values.put(LauncherSettings.Favorites.SPANY, spanY);
-        }
-
-        @Override
-        public void writeToValues(@NonNull ContentWriter writer) {
-            super.writeToValues(writer);
-            writer.put(LauncherSettings.Favorites.APPWIDGET_ID, appWidgetId);
-        }
-
-        @Override
-        public void readFromValues(@NonNull ContentValues values) {
-            super.readFromValues(values);
-            appWidgetId = values.getAsInteger(LauncherSettings.Favorites.APPWIDGET_ID);
-        }
-
-        /** This id is not used in the DB is only used while doing the migration and it identifies
-         * an entry on each workspace. For example two calculator icons would have the same
-         * migration id even thought they have different database ids.
-         */
-        public String getEntryMigrationId() {
-            switch (itemType) {
-                case LauncherSettings.Favorites.ITEM_TYPE_FOLDER:
-                case LauncherSettings.Favorites.ITEM_TYPE_APP_PAIR:
-                    return getFolderMigrationId();
-                case LauncherSettings.Favorites.ITEM_TYPE_APPWIDGET:
-                    // mProvider is the app the widget belongs to and appWidgetId it's the unique
-                    // is of the widget, we need both because if you remove a widget and then add it
-                    // again, then it can change and the WidgetProvider would not know the widget.
-                    return mProvider + appWidgetId;
-                case LauncherSettings.Favorites.ITEM_TYPE_APPLICATION:
-                    final String intentStr = cleanIntentString(mIntent);
-                    try {
-                        Intent i = Intent.parseUri(intentStr, 0);
-                        return Objects.requireNonNull(i.getComponent()).toString();
-                    } catch (Exception e) {
-                        return intentStr;
-                    }
-                default:
-                    return cleanIntentString(mIntent);
-            }
-        }
-
-        /**
-         * This method should return an id that should be the same for two folders containing the
-         * same elements.
-         */
-        @NonNull
-        private String getFolderMigrationId() {
-            return mFolderItems.keySet().stream()
-                    .map(intentString -> mFolderItems.get(intentString).size()
-                            + cleanIntentString(intentString))
-                    .sorted()
-                    .collect(Collectors.joining(","));
-        }
-
-        /**
-         * This is needed because sourceBounds can change and make the id of two equal items
-         * different.
-         */
-        @NonNull
-        private String cleanIntentString(@NonNull String intentStr) {
-            try {
-                Intent i = Intent.parseUri(intentStr, 0);
-                i.setSourceBounds(null);
-                return i.toURI();
-            } catch (URISyntaxException e) {
-                Log.e(TAG, "Unable to parse Intent string", e);
-                return intentStr;
-            }
-
-        }
-    }
 }
diff --git a/src/com/android/launcher3/model/GridSizeMigrationLogic.java b/src/com/android/launcher3/model/GridSizeMigrationLogic.java
new file mode 100644
index 0000000..12a14b2
--- /dev/null
+++ b/src/com/android/launcher3/model/GridSizeMigrationLogic.java
@@ -0,0 +1,466 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.model;
+
+import static com.android.launcher3.Flags.enableSmartspaceRemovalToggle;
+import static com.android.launcher3.LauncherPrefs.IS_FIRST_LOAD_AFTER_RESTORE;
+import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
+import static com.android.launcher3.LauncherSettings.Favorites.TMP_TABLE;
+import static com.android.launcher3.Utilities.SHOULD_SHOW_FIRST_PAGE_WIDGET;
+import static com.android.launcher3.model.GridSizeMigrationDBController.copyCurrentGridToNewGrid;
+import static com.android.launcher3.model.GridSizeMigrationDBController.insertEntryInDb;
+import static com.android.launcher3.model.GridSizeMigrationDBController.needsToMigrate;
+import static com.android.launcher3.model.GridSizeMigrationDBController.removeEntryFromDb;
+import static com.android.launcher3.model.LoaderTask.SMARTSPACE_ON_HOME_SCREEN;
+import static com.android.launcher3.provider.LauncherDbUtils.copyTable;
+import static com.android.launcher3.provider.LauncherDbUtils.dropTable;
+
+import android.content.Context;
+import android.database.sqlite.SQLiteDatabase;
+import android.graphics.Point;
+import android.util.Log;
+
+import androidx.annotation.NonNull;
+
+import com.android.launcher3.Flags;
+import com.android.launcher3.LauncherPrefs;
+import com.android.launcher3.LauncherSettings;
+import com.android.launcher3.config.FeatureFlags;
+import com.android.launcher3.provider.LauncherDbUtils;
+import com.android.launcher3.util.CellAndSpan;
+import com.android.launcher3.util.GridOccupancy;
+import com.android.launcher3.util.IntArray;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+public class GridSizeMigrationLogic {
+
+    private static final String TAG = "GridSizeMigrationLogic";
+    private static final boolean DEBUG = true;
+
+    /**
+     * Migrates the grid size from srcDeviceState to destDeviceState and make those changes
+     * in the target DB, using the source DB to determine what to add/remove/move/resize
+     * in the destination DB.
+     */
+    public void migrateGrid(
+            @NonNull Context context,
+            @NonNull DeviceGridState srcDeviceState,
+            @NonNull DeviceGridState destDeviceState,
+            @NonNull DatabaseHelper target,
+            @NonNull SQLiteDatabase source) {
+        if (!needsToMigrate(srcDeviceState, destDeviceState)) {
+            return;
+        }
+
+        boolean isFirstLoad = LauncherPrefs.get(context).get(IS_FIRST_LOAD_AFTER_RESTORE);
+        Log.d(TAG, "Begin grid migration. First load: " + isFirstLoad);
+
+        // This is a special case where if the grid is the same amount of columns but a larger
+        // amount of rows we simply copy over the source grid to the destination grid, rather
+        // than undergoing the general grid migration.
+        if (shouldMigrateToStrictlyTallerGrid(isFirstLoad, srcDeviceState, destDeviceState)) {
+            copyCurrentGridToNewGrid(context, destDeviceState, target, source);
+            return;
+        }
+        copyTable(source, TABLE_NAME, target.getWritableDatabase(), TMP_TABLE, context);
+
+        long migrationStartTime = System.currentTimeMillis();
+        try (LauncherDbUtils.SQLiteTransaction t =
+                     new LauncherDbUtils.SQLiteTransaction(target.getWritableDatabase())) {
+            GridSizeMigrationDBController.DbReader srcReader = new GridSizeMigrationDBController
+                    .DbReader(t.getDb(), TMP_TABLE, context);
+            GridSizeMigrationDBController.DbReader destReader =
+                    new GridSizeMigrationDBController.DbReader(
+                            t.getDb(), TABLE_NAME, context);
+
+            Point targetSize = new Point(destDeviceState.getColumns(), destDeviceState.getRows());
+            // Migrate hotseat.
+            migrateHotseat(destDeviceState.getNumHotseat(), srcReader, destReader, target);
+            // Migrate workspace.
+            migrateWorkspace(srcReader, destReader, target, targetSize);
+
+            dropTable(t.getDb(), TMP_TABLE);
+            t.commit();
+        } catch (Exception e) {
+            Log.e(TAG, "Error during grid migration", e);
+        } finally {
+            Log.v(TAG, "Workspace migration completed in "
+                    + (System.currentTimeMillis() - migrationStartTime));
+
+            // Save current configuration, so that the migration does not run again.
+            destDeviceState.writeToPrefs(context);
+        }
+    }
+
+    private void migrateHotseat(int destHotseatSize,
+            GridSizeMigrationDBController.DbReader srcReader,
+            GridSizeMigrationDBController.DbReader destReader, DatabaseHelper helper) {
+        final List<DbEntry> srcHotseatItems =
+                srcReader.loadHotseatEntries();
+        final List<DbEntry> dstHotseatItems =
+                destReader.loadHotseatEntries();
+
+
+        final List<DbEntry> hotseatToBeAdded =
+                getItemsToBeAdded(srcHotseatItems, dstHotseatItems);
+        final IntArray toBeRemoved = new IntArray();
+        toBeRemoved.addAll(getItemsToBeRemoved(srcHotseatItems, dstHotseatItems));
+
+        if (DEBUG) {
+            Log.d(TAG, "Start hotseat migration:"
+                    + "\n Removing Hotseat Items:"
+                    + dstHotseatItems.stream().filter(entry -> toBeRemoved
+                    .contains(entry.id)).map(DbEntry::toString)
+                    .collect(Collectors.joining(",\n", "[", "]"))
+                    + "\n Adding Hotseat Items:"
+                    + hotseatToBeAdded.stream().map(DbEntry::toString)
+                    .collect(Collectors.joining(",\n", "[", "]"))
+            );
+        }
+
+        // Removes the items that we need to remove from the destination DB.
+        if (!toBeRemoved.isEmpty()) {
+            removeEntryFromDb(destReader.mDb, destReader.mTableName, toBeRemoved);
+        }
+
+        placeHotseatItems(
+                hotseatToBeAdded, dstHotseatItems, destHotseatSize, helper, srcReader, destReader);
+    }
+
+    private void placeHotseatItems(List<DbEntry> hotseatToBeAdded,
+            List<DbEntry> dstHotseatItems, int destHotseatSize,
+            DatabaseHelper helper, GridSizeMigrationDBController.DbReader srcReader,
+            GridSizeMigrationDBController.DbReader destReader) {
+        if (hotseatToBeAdded.isEmpty()) {
+            return;
+        }
+
+        List<Integer> idsInUse = dstHotseatItems.stream().map(entry -> entry.id).toList();
+
+        Collections.sort(hotseatToBeAdded);
+
+        List<DbEntry> placementSolutionHotseat =
+                solveHotseatPlacement(destHotseatSize, dstHotseatItems, hotseatToBeAdded);
+        for (DbEntry entryToPlace: placementSolutionHotseat) {
+            insertEntryInDb(helper, entryToPlace, srcReader.mTableName, destReader.mTableName,
+                    idsInUse);
+        }
+    }
+
+    private void migrateWorkspace(GridSizeMigrationDBController.DbReader srcReader,
+            GridSizeMigrationDBController.DbReader destReader, DatabaseHelper helper,
+            Point targetSize) {
+
+
+        final List<DbEntry> srcWorkspaceItems =
+                srcReader.loadAllWorkspaceEntries();
+
+        final List<DbEntry> dstWorkspaceItems =
+                destReader.loadAllWorkspaceEntries();
+
+        final IntArray toBeRemoved = new IntArray();
+
+        List<DbEntry> workspaceToBeAdded =
+                getItemsToBeAdded(srcWorkspaceItems, dstWorkspaceItems);
+        toBeRemoved.addAll(getItemsToBeRemoved(srcWorkspaceItems, dstWorkspaceItems));
+
+        if (DEBUG) {
+            Log.d(TAG, "Start workspace migration:"
+                    + "\n Source Device:"
+                    + srcWorkspaceItems.stream().map(
+                            DbEntry::toString)
+                    .collect(Collectors.joining(",\n", "[", "]"))
+                    + "\n Target Device:"
+                    + dstWorkspaceItems.stream().map(
+                            DbEntry::toString)
+                    .collect(Collectors.joining(",\n", "[", "]"))
+                    + "\n Removing Workspace Items:"
+                    + dstWorkspaceItems.stream().filter(entry -> toBeRemoved
+                            .contains(entry.id)).map(
+                            DbEntry::toString)
+                    .collect(Collectors.joining(",\n", "[", "]"))
+                    + "\n Adding Workspace Items:"
+                    + workspaceToBeAdded.stream().map(
+                            DbEntry::toString)
+                    .collect(Collectors.joining(",\n", "[", "]"))
+            );
+        }
+
+        // Removes the items that we need to remove from the destination DB.
+        if (!toBeRemoved.isEmpty()) {
+            removeEntryFromDb(destReader.mDb, destReader.mTableName, toBeRemoved);
+        }
+
+        placeWorkspaceItems(workspaceToBeAdded, dstWorkspaceItems, targetSize.x, targetSize.y,
+                helper, srcReader, destReader);
+    }
+
+    private void placeWorkspaceItems(
+            List<DbEntry> workspaceToBeAdded,
+            List<DbEntry> dstWorkspaceItems,
+            int trgX, int trgY, DatabaseHelper helper,
+            GridSizeMigrationDBController.DbReader srcReader,
+            GridSizeMigrationDBController.DbReader destReader) {
+        if (workspaceToBeAdded.isEmpty()) {
+            return;
+        }
+
+        List<Integer> idsInUse = dstWorkspaceItems.stream().map(entry -> entry.id).collect(
+                Collectors.toList());
+
+        Collections.sort(workspaceToBeAdded);
+
+
+        // First we create a collection of the screens
+        List<Integer> screens = new ArrayList<>();
+        for (int screenId = 0; screenId <= destReader.mLastScreenId; screenId++) {
+            screens.add(screenId);
+        }
+
+        // Then we place the items on the screens
+        WorkspaceItemsToPlace itemsToPlace =
+                new WorkspaceItemsToPlace(workspaceToBeAdded);
+        for (int screenId : screens) {
+            if (DEBUG) {
+                Log.d(TAG, "Migrating " + screenId);
+            }
+            itemsToPlace = solveGridPlacement(
+                    destReader.mContext, screenId, trgX, trgY, itemsToPlace.mRemainingItemsToPlace,
+                    destReader.mWorkspaceEntriesByScreenId.get(screenId));
+            placeItems(itemsToPlace, helper, srcReader, destReader, idsInUse);
+            while (!itemsToPlace.mPlacementSolution.isEmpty()) {
+                insertEntryInDb(helper, itemsToPlace.mPlacementSolution.remove(0),
+                        srcReader.mTableName, destReader.mTableName, idsInUse);
+            }
+            if (itemsToPlace.mRemainingItemsToPlace.isEmpty()) {
+                break;
+            }
+        }
+
+        // In case the new grid is smaller, there might be some leftover items that don't fit on
+        // any of the screens, in this case we add them to new screens until all of them are placed.
+        int screenId = destReader.mLastScreenId + 1;
+        while (!itemsToPlace.mRemainingItemsToPlace.isEmpty()) {
+            itemsToPlace = solveGridPlacement(destReader.mContext, screenId,
+                    trgX, trgY, itemsToPlace.mRemainingItemsToPlace,
+                    destReader.mWorkspaceEntriesByScreenId.get(screenId));
+            placeItems(itemsToPlace, helper, srcReader, destReader, idsInUse);
+            screenId++;
+        }
+    }
+
+    private void placeItems(WorkspaceItemsToPlace itemsToPlace, DatabaseHelper helper,
+            GridSizeMigrationDBController.DbReader srcReader,
+            GridSizeMigrationDBController.DbReader destReader, List<Integer> idsInUse) {
+        while (!itemsToPlace.mPlacementSolution.isEmpty()) {
+            insertEntryInDb(helper, itemsToPlace.mPlacementSolution.remove(0),
+                    srcReader.mTableName, destReader.mTableName, idsInUse);
+        }
+    }
+
+
+    /**
+     * Only migrate the grid in this manner if the target grid is taller and not wider.
+     */
+    private boolean shouldMigrateToStrictlyTallerGrid(boolean isFirstLoad,
+            @NonNull DeviceGridState srcDeviceState,
+            @NonNull DeviceGridState destDeviceState) {
+        if (isFirstLoad
+                && Flags.enableGridMigrationFix()
+                && srcDeviceState.getColumns().equals(destDeviceState.getColumns())
+                && srcDeviceState.getRows() < destDeviceState.getRows()) {
+            return true;
+        }
+        return false;
+    }
+
+    /**
+     * Finds all the items that are in the old grid which aren't in the new grid, meaning they
+     * need to be added to the new grid.
+     *
+     * @return a list of DbEntry's which we need to add.
+     */
+    private List<DbEntry> getItemsToBeAdded(
+            @NonNull final List<DbEntry> src,
+            @NonNull final List<DbEntry> dest) {
+        Map<DbEntry, Integer> entryCountDiff =
+                calcDiff(src, dest);
+        List<DbEntry> toBeAdded = new ArrayList<>();
+        src.forEach(entry -> {
+            if (entryCountDiff.get(entry) > 0) {
+                toBeAdded.add(entry);
+                entryCountDiff.put(entry, entryCountDiff.get(entry) - 1);
+            }
+        });
+        return toBeAdded;
+    }
+
+    /**
+     * Finds all the items that are in the new grid which aren't in the old grid, meaning they
+     * need to be removed from the new grid.
+     *
+     * @return an IntArray of item id's which we need to remove.
+     */
+    private IntArray getItemsToBeRemoved(
+            @NonNull final List<DbEntry> src,
+            @NonNull final List<DbEntry> dest) {
+        Map<DbEntry, Integer> entryCountDiff =
+                calcDiff(src, dest);
+        IntArray toBeRemoved = new IntArray();
+        dest.forEach(entry -> {
+            if (entryCountDiff.get(entry) < 0) {
+                toBeRemoved.add(entry.id);
+                if (entry.itemType == LauncherSettings.Favorites.ITEM_TYPE_FOLDER) {
+                    entry.mFolderItems.values().forEach(ids -> ids.forEach(toBeRemoved::add));
+                }
+                entryCountDiff.put(entry, entryCountDiff.get(entry) + 1);
+            }
+        });
+        return toBeRemoved;
+    }
+
+    /**
+     * Calculates the difference between the old and new grid items in terms of how many of each
+     * item there are. E.g. if the old grid had 2 Calculator icons but the new grid has 0, then the
+     * difference there would be 2. While if the old grid has 0 Calculator icons and the
+     * new grid has 1, then the difference would be -1.
+     *
+     * @return a Map with each DbEntry as a key and the count of said entry as the value.
+     */
+    private Map<DbEntry, Integer> calcDiff(
+            @NonNull final List<DbEntry> src,
+            @NonNull final List<DbEntry> dest) {
+        Map<DbEntry, Integer> entryCountDiff = new HashMap<>();
+        src.forEach(entry ->
+                entryCountDiff.put(entry, entryCountDiff.getOrDefault(entry, 0) + 1));
+        dest.forEach(entry ->
+                entryCountDiff.put(entry, entryCountDiff.getOrDefault(entry, 0) - 1));
+        return entryCountDiff;
+    }
+
+    private List<DbEntry> solveHotseatPlacement(final int hotseatSize,
+            @NonNull final List<DbEntry> placedHotseatItems,
+            @NonNull final List<DbEntry> itemsToPlace) {
+        List<DbEntry> placementSolution = new ArrayList<>();
+        List<DbEntry> remainingItemsToPlace =
+                new ArrayList<>(itemsToPlace);
+        final boolean[] occupied = new boolean[hotseatSize];
+        for (DbEntry entry : placedHotseatItems) {
+            occupied[entry.screenId] = true;
+        }
+
+        for (int i = 0; i < occupied.length; i++) {
+            if (!occupied[i] && !remainingItemsToPlace.isEmpty()) {
+                DbEntry entry = remainingItemsToPlace.remove(0);
+                entry.screenId = i;
+                // These values does not affect the item position, but we should set them
+                // to something other than -1.
+                entry.cellX = i;
+                entry.cellY = 0;
+
+                placementSolution.add(entry);
+                occupied[entry.screenId] = true;
+            }
+        }
+        return placementSolution;
+    }
+
+    private WorkspaceItemsToPlace solveGridPlacement(
+            Context context,
+            final int screenId, final int trgX, final int trgY,
+            @NonNull final List<DbEntry> sortedItemsToPlace,
+            List<DbEntry> existedEntries) {
+        WorkspaceItemsToPlace itemsToPlace = new WorkspaceItemsToPlace(sortedItemsToPlace);
+        final GridOccupancy occupied = new GridOccupancy(trgX, trgY);
+        final Point trg = new Point(trgX, trgY);
+        final Point next = new Point(0, screenId == 0
+                && (FeatureFlags.QSB_ON_FIRST_SCREEN
+                && (!enableSmartspaceRemovalToggle() || LauncherPrefs.getPrefs(context)
+                .getBoolean(SMARTSPACE_ON_HOME_SCREEN, true))
+                && !SHOULD_SHOW_FIRST_PAGE_WIDGET)
+                ? 1 /* smartspace */ : 0);
+        if (existedEntries != null) {
+            for (DbEntry entry : existedEntries) {
+                occupied.markCells(entry, true);
+            }
+        }
+        Iterator<DbEntry> iterator =
+                itemsToPlace.mRemainingItemsToPlace.iterator();
+        while (iterator.hasNext()) {
+            final DbEntry entry = iterator.next();
+            if (entry.minSpanX > trgX || entry.minSpanY > trgY) {
+                iterator.remove();
+                continue;
+            }
+            CellAndSpan placement = findPlacementForEntry(
+                    entry, next.x, next.y, trg, occupied);
+            if (placement != null) {
+                entry.screenId = screenId;
+                entry.cellX = placement.cellX;
+                entry.cellY = placement.cellY;
+                entry.spanX = placement.spanX;
+                entry.spanY = placement.spanY;
+                occupied.markCells(entry, true);
+                next.set(entry.cellX + entry.spanX, entry.cellY);
+                itemsToPlace.mPlacementSolution.add(entry);
+                iterator.remove();
+            }
+        }
+        return itemsToPlace;
+    }
+
+    /**
+     * Search for the next possible placement of an item. (mNextStartX, mNextStartY) serves as
+     * a memoization of last placement, we can start our search for next placement from there
+     * to speed up the search.
+     *
+     * @return NewEntryPlacement object if we found a valid placement, null if we didn't.
+     */
+    private CellAndSpan findPlacementForEntry(
+            @NonNull final DbEntry entry,
+            int startPosX, int startPosY, @NonNull final Point trg,
+            @NonNull final GridOccupancy occupied) {
+        for (int y = startPosY; y <  trg.y; y++) {
+            for (int x = startPosX; x < trg.x; x++) {
+                boolean minFits = occupied.isRegionVacant(x, y, entry.minSpanX, entry.minSpanY);
+                if (minFits) {
+                    return (new CellAndSpan(x, y, entry.minSpanX, entry.minSpanY));
+                }
+            }
+            startPosX = 0;
+        }
+        return null;
+    }
+
+    private static class WorkspaceItemsToPlace {
+        List<DbEntry> mRemainingItemsToPlace;
+        List<DbEntry> mPlacementSolution;
+
+        WorkspaceItemsToPlace(List<DbEntry> sortedItemsToPlace) {
+            mRemainingItemsToPlace = new ArrayList<>(sortedItemsToPlace);
+            mPlacementSolution = new ArrayList<>();
+        }
+
+    }
+}
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index 09d1146..b0108c2 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -435,7 +435,15 @@
         final WidgetInflater widgetInflater = new WidgetInflater(context);
 
         ModelDbController dbController = mApp.getModel().getModelDbController();
-        dbController.tryMigrateDB(restoreEventLogger);
+        if (Flags.gridMigrationRefactor()) {
+            try {
+                dbController.attemptMigrateDb(restoreEventLogger);
+            } catch (Exception e) {
+                FileLog.e(TAG, "Failed to migrate grid", e);
+            }
+        } else {
+            dbController.tryMigrateDB(restoreEventLogger);
+        }
         Log.d(TAG, "loadWorkspace: loading default favorites");
         dbController.loadDefaultFavoritesIfNecessary();
 
diff --git a/src/com/android/launcher3/model/ModelDbController.java b/src/com/android/launcher3/model/ModelDbController.java
index 5d66d16..4f0f162 100644
--- a/src/com/android/launcher3/model/ModelDbController.java
+++ b/src/com/android/launcher3/model/ModelDbController.java
@@ -20,15 +20,16 @@
 import static android.util.Base64.NO_WRAP;
 
 import static com.android.launcher3.DefaultLayoutParser.RES_PARTNER_DEFAULT_LAYOUT;
+import static com.android.launcher3.LauncherPrefs.NO_DB_FILES_RESTORED;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APP_PAIR;
 import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
 import static com.android.launcher3.LauncherSettings.Favorites.addTableToDb;
 import static com.android.launcher3.LauncherSettings.Settings.BLOB_KEY_PREFIX;
-import static com.android.launcher3.LauncherSettings.Settings.LAYOUT_PROVIDER_KEY;
 import static com.android.launcher3.LauncherSettings.Settings.LAYOUT_DIGEST_LABEL;
 import static com.android.launcher3.LauncherSettings.Settings.LAYOUT_DIGEST_TAG;
+import static com.android.launcher3.LauncherSettings.Settings.LAYOUT_PROVIDER_KEY;
 import static com.android.launcher3.provider.LauncherDbUtils.tableExists;
 
 import android.app.blob.BlobHandle;
@@ -97,6 +98,7 @@
 
     private static final String EMPTY_DATABASE_CREATED = "EMPTY_DATABASE_CREATED";
     public static final String EXTRA_DB_NAME = "db_name";
+    public static final String DATA_TYPE_DB_FILE = "database_file";
 
     protected DatabaseHelper mOpenHelper;
 
@@ -289,13 +291,115 @@
 
 
     /**
+     * Resets the launcher DB if we should reset it.
+     */
+    public void resetLauncherDb(@Nullable LauncherRestoreEventLogger restoreEventLogger) {
+        if (restoreEventLogger != null) {
+            sendMetricsForFailedMigration(restoreEventLogger, getDb());
+        }
+        FileLog.d(TAG, "Migration failed: resetting launcher database");
+        createEmptyDB();
+        LauncherPrefs.get(mContext).putSync(
+                getEmptyDbCreatedKey(mOpenHelper.getDatabaseName()).to(true));
+
+        // Write the grid state to avoid another migration
+        new DeviceGridState(LauncherAppState.getIDP(mContext)).writeToPrefs(mContext);
+    }
+
+    /**
+     * Determines if we should reset the DB.
+     */
+    private boolean shouldResetDb() {
+        if (isThereExistingDb()) {
+            return true;
+        }
+        if (!isGridMigrationNecessary()) {
+            return false;
+        }
+        if (isCurrentDbSameAsTarget()) {
+            return true;
+        }
+        return false;
+    }
+
+    private boolean isThereExistingDb() {
+        if (LauncherPrefs.get(mContext).get(getEmptyDbCreatedKey())) {
+            // If we already have a new DB, ignore migration
+            Log.d(TAG, "migrateGridIfNeeded: new DB already created, skipping migration");
+            return true;
+        }
+        return false;
+    }
+
+    private boolean isGridMigrationNecessary() {
+        InvariantDeviceProfile idp = LauncherAppState.getIDP(mContext);
+        if (GridSizeMigrationDBController.needsToMigrate(mContext, idp)) {
+            return true;
+        }
+        Log.d(TAG, "migrateGridIfNeeded: no grid migration needed");
+        return false;
+    }
+
+    private boolean isCurrentDbSameAsTarget() {
+        InvariantDeviceProfile idp = LauncherAppState.getIDP(mContext);
+        String targetDbName = new DeviceGridState(idp).getDbFile();
+        if (TextUtils.equals(targetDbName, mOpenHelper.getDatabaseName())) {
+            Log.e(TAG, "migrateGridIfNeeded: target db is same as current: " + targetDbName);
+            return true;
+        }
+        return false;
+    }
+
+    /**
+     * Migrates the DB. If the migration failed, it clears the DB.
+     */
+    public void attemptMigrateDb(LauncherRestoreEventLogger restoreEventLogger) throws Exception {
+        createDbIfNotExists();
+
+        if (shouldResetDb()) {
+            resetLauncherDb(restoreEventLogger);
+            return;
+        }
+
+        InvariantDeviceProfile idp = LauncherAppState.getIDP(mContext);
+        DatabaseHelper oldHelper = mOpenHelper;
+        mOpenHelper = (mContext instanceof SandboxContext) ? oldHelper
+                : createDatabaseHelper(true /* forMigration */);
+        try {
+            // This is the current grid we have, given by the mContext
+            DeviceGridState srcDeviceState = new DeviceGridState(mContext);
+            // This is the state we want to migrate to that is given by the idp
+            DeviceGridState destDeviceState = new DeviceGridState(idp);
+
+            GridSizeMigrationLogic gridSizeMigrationLogic = new GridSizeMigrationLogic();
+            gridSizeMigrationLogic.migrateGrid(mContext, srcDeviceState, destDeviceState,
+                    mOpenHelper, oldHelper.getWritableDatabase());
+        } catch (Exception e) {
+            resetLauncherDb(restoreEventLogger);
+            throw new Exception("Failed to migrate grid", e);
+        } finally {
+            if (mOpenHelper != oldHelper) {
+                oldHelper.close();
+            }
+        }
+    }
+
+    /**
      * Migrates the DB if needed. If the migration failed, it clears the DB.
      */
     public void tryMigrateDB(@Nullable LauncherRestoreEventLogger restoreEventLogger) {
 
         if (!migrateGridIfNeeded()) {
             if (restoreEventLogger != null) {
-                sendMetricsForFailedMigration(restoreEventLogger, getDb());
+                if (LauncherPrefs.get(mContext).get(NO_DB_FILES_RESTORED)) {
+                    restoreEventLogger.logLauncherItemsRestoreFailed(DATA_TYPE_DB_FILE, 1,
+                            RestoreError.DATABASE_FILE_NOT_RESTORED);
+                    LauncherPrefs.get(mContext).put(NO_DB_FILES_RESTORED, false);
+                    FileLog.d(TAG, "There is no data to migrate: resetting launcher database");
+                } else {
+                    restoreEventLogger.logLauncherItemsRestored(DATA_TYPE_DB_FILE, 1);
+                    sendMetricsForFailedMigration(restoreEventLogger, getDb());
+                }
             }
             FileLog.d(TAG, "Migration failed: resetting launcher database");
             createEmptyDB();
@@ -304,6 +408,8 @@
 
             // Write the grid state to avoid another migration
             new DeviceGridState(LauncherAppState.getIDP(mContext)).writeToPrefs(mContext);
+        } else if (restoreEventLogger != null) {
+            restoreEventLogger.logLauncherItemsRestored(DATA_TYPE_DB_FILE, 1);
         }
     }
 
@@ -317,17 +423,17 @@
         createDbIfNotExists();
         if (LauncherPrefs.get(mContext).get(getEmptyDbCreatedKey())) {
             // If we have already create a new DB, ignore migration
-            Log.d(TAG, "migrateGridIfNeeded: new DB already created, skipping migration");
+            FileLog.d(TAG, "migrateGridIfNeeded: new DB already created, skipping migration");
             return false;
         }
         InvariantDeviceProfile idp = LauncherAppState.getIDP(mContext);
-        if (!GridSizeMigrationUtil.needsToMigrate(mContext, idp)) {
+        if (!GridSizeMigrationDBController.needsToMigrate(mContext, idp)) {
             Log.d(TAG, "migrateGridIfNeeded: no grid migration needed");
             return true;
         }
         String targetDbName = new DeviceGridState(idp).getDbFile();
         if (TextUtils.equals(targetDbName, mOpenHelper.getDatabaseName())) {
-            Log.e(TAG, "migrateGridIfNeeded: target db is same as current: " + targetDbName);
+            FileLog.e(TAG, "migrateGridIfNeeded: target db is same as current: " + targetDbName);
             return false;
         }
         DatabaseHelper oldHelper = mOpenHelper;
@@ -338,7 +444,7 @@
             DeviceGridState srcDeviceState = new DeviceGridState(mContext);
             // This is the state we want to migrate to that is given by the idp
             DeviceGridState destDeviceState = new DeviceGridState(idp);
-            return GridSizeMigrationUtil.migrateGridIfNeeded(mContext, srcDeviceState,
+            return GridSizeMigrationDBController.migrateGridIfNeeded(mContext, srcDeviceState,
                     destDeviceState, mOpenHelper, oldHelper.getWritableDatabase());
         } catch (Exception e) {
             FileLog.e(TAG, "Failed to migrate grid", e);
diff --git a/tests/multivalentTests/src/com/android/launcher3/celllayout/FavoriteItemsTransaction.java b/tests/multivalentTests/src/com/android/launcher3/celllayout/FavoriteItemsTransaction.java
index 0c3081f..a9082e2 100644
--- a/tests/multivalentTests/src/com/android/launcher3/celllayout/FavoriteItemsTransaction.java
+++ b/tests/multivalentTests/src/com/android/launcher3/celllayout/FavoriteItemsTransaction.java
@@ -22,6 +22,7 @@
 
 import android.content.Context;
 
+import com.android.launcher3.Flags;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherModel;
 import com.android.launcher3.LauncherSettings;
@@ -59,7 +60,11 @@
         runOnExecutorSync(MODEL_EXECUTOR, () -> {
             ModelDbController controller = model.getModelDbController();
             // Migrate any previous data so that the DB state is correct
-            controller.tryMigrateDB(null /* restoreEventLogger */);
+            if (Flags.gridMigrationRefactor()) {
+                controller.attemptMigrateDb(null /* restoreEventLogger */);
+            } else {
+                controller.tryMigrateDB(null /* restoreEventLogger */);
+            }
 
             // Create DB again to load fresh data
             controller.createEmptyDB();
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/GridSizeMigrationUtilTest.kt b/tests/multivalentTests/src/com/android/launcher3/model/GridSizeMigrationDBControllerTest.kt
similarity index 83%
rename from tests/multivalentTests/src/com/android/launcher3/model/GridSizeMigrationUtilTest.kt
rename to tests/multivalentTests/src/com/android/launcher3/model/GridSizeMigrationDBControllerTest.kt
index f57e8a1..c6f291d 100644
--- a/tests/multivalentTests/src/com/android/launcher3/model/GridSizeMigrationUtilTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/model/GridSizeMigrationDBControllerTest.kt
@@ -22,13 +22,16 @@
 import android.database.sqlite.SQLiteDatabase
 import android.graphics.Point
 import android.os.Process
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.launcher3.Flags
 import com.android.launcher3.InvariantDeviceProfile
 import com.android.launcher3.LauncherPrefs
 import com.android.launcher3.LauncherPrefs.Companion.WORKSPACE_SIZE
 import com.android.launcher3.LauncherSettings.Favorites.*
-import com.android.launcher3.model.GridSizeMigrationUtil.DbReader
+import com.android.launcher3.model.GridSizeMigrationDBController.DbReader
 import com.android.launcher3.pm.UserCache
 import com.android.launcher3.provider.LauncherDbUtils
 import com.android.launcher3.util.LauncherModelHelper
@@ -82,9 +85,22 @@
         modelHelper.destroy()
     }
 
-    /** Old migration logic, should be modified once is not needed anymore */
     @Test
     @Throws(Exception::class)
+    @EnableFlags(Flags.FLAG_GRID_MIGRATION_REFACTOR)
+    fun testMigrationRefactorFlagOn() {
+        testMigration()
+    }
+
+    @Test
+    @Throws(Exception::class)
+    @DisableFlags(Flags.FLAG_GRID_MIGRATION_REFACTOR)
+    fun testMigrationRefactorFlagOff() {
+        testMigration()
+    }
+
+    /** Old migration logic, should be modified once is not needed anymore */
+    @Throws(Exception::class)
     fun testMigration() {
         // Src Hotseat icons
         addItem(ITEM_TYPE_APPLICATION, 0, CONTAINER_HOTSEAT, 0, 0, testPackage1, 1, TMP_TABLE)
@@ -113,15 +129,26 @@
         idp.numRows = 4
         val srcReader = DbReader(db, TMP_TABLE, context)
         val destReader = DbReader(db, TABLE_NAME, context)
-        GridSizeMigrationUtil.migrate(
-            dbHelper,
-            srcReader,
-            destReader,
-            idp.numDatabaseHotseatIcons,
-            Point(idp.numColumns, idp.numRows),
-            DeviceGridState(context),
-            DeviceGridState(idp),
-        )
+        if (Flags.gridMigrationRefactor()) {
+            val gridSizeMigrationLogic = GridSizeMigrationLogic()
+            gridSizeMigrationLogic.migrateGrid(
+                context,
+                DeviceGridState(context),
+                DeviceGridState(idp),
+                dbHelper,
+                db,
+            )
+        } else {
+            GridSizeMigrationDBController.migrate(
+                dbHelper,
+                srcReader,
+                destReader,
+                idp.numDatabaseHotseatIcons,
+                Point(idp.numColumns, idp.numRows),
+                DeviceGridState(context),
+                DeviceGridState(idp),
+            )
+        }
 
         // Check hotseat items
         var c =
@@ -187,9 +214,22 @@
         assertThat(locMap[testPackage9]).isEqualTo(Point(0, 2))
     }
 
-    /** Old migration logic, should be modified once is not needed anymore */
     @Test
     @Throws(Exception::class)
+    @EnableFlags(Flags.FLAG_GRID_MIGRATION_REFACTOR)
+    fun testMigrationBackAndForthRefactorFlagOn() {
+        testMigrationBackAndForth()
+    }
+
+    @Test
+    @Throws(Exception::class)
+    @DisableFlags(Flags.FLAG_GRID_MIGRATION_REFACTOR)
+    fun testMigrationBackAndForthRefactorFlagOff() {
+        testMigrationBackAndForth()
+    }
+
+    /** Old migration logic, should be modified once is not needed anymore */
+    @Throws(Exception::class)
     fun testMigrationBackAndForth() {
         // Hotseat items in grid A
         // 1 2 _ 3 4
@@ -224,15 +264,26 @@
         val readerGridA = DbReader(db, TMP_TABLE, context)
         val readerGridB = DbReader(db, TABLE_NAME, context)
         // migrate from A -> B
-        GridSizeMigrationUtil.migrate(
-            dbHelper,
-            readerGridA,
-            readerGridB,
-            idp.numDatabaseHotseatIcons,
-            Point(idp.numColumns, idp.numRows),
-            DeviceGridState(context),
-            DeviceGridState(idp),
-        )
+        if (Flags.gridMigrationRefactor()) {
+            var gridSizeMigrationLogic = GridSizeMigrationLogic()
+            gridSizeMigrationLogic.migrateGrid(
+                context,
+                DeviceGridState(context),
+                DeviceGridState(idp),
+                dbHelper,
+                db,
+            )
+        } else {
+            GridSizeMigrationDBController.migrate(
+                dbHelper,
+                readerGridA,
+                readerGridB,
+                idp.numDatabaseHotseatIcons,
+                Point(idp.numColumns, idp.numRows),
+                DeviceGridState(context),
+                DeviceGridState(idp),
+            )
+        }
 
         // Check hotseat items in grid B
         var c =
@@ -280,15 +331,8 @@
         addItem(ITEM_TYPE_APPLICATION, 0, CONTAINER_DESKTOP, 0, 2, testPackage9)
 
         // migrate from B -> A
-        GridSizeMigrationUtil.migrate(
-            dbHelper,
-            readerGridB,
-            readerGridA,
-            5,
-            Point(5, 5),
-            DeviceGridState(idp),
-            DeviceGridState(context),
-        )
+        migrateGrid(dbHelper, readerGridB, readerGridA, 5, 5, 5)
+
         // Check hotseat items in grid A
         c =
             db.query(
@@ -339,14 +383,13 @@
         db.delete(TMP_TABLE, "$_ID=7", null)
 
         // migrate from A -> B
-        GridSizeMigrationUtil.migrate(
+        migrateGrid(
             dbHelper,
             readerGridA,
             readerGridB,
             idp.numDatabaseHotseatIcons,
-            Point(idp.numColumns, idp.numRows),
-            DeviceGridState(context),
-            DeviceGridState(idp),
+            idp.numColumns,
+            idp.numRows,
         )
 
         // Check hotseat items in grid B
@@ -392,6 +435,36 @@
         assertThat(locMap[testPackage9]).isEqualTo(Triple(0, 0, 2))
     }
 
+    private fun migrateGrid(
+        dbHelper: DatabaseHelper,
+        srcReader: DbReader,
+        destReader: DbReader,
+        destHotseatSize: Int,
+        pointX: Int,
+        pointY: Int,
+    ) {
+        if (Flags.gridMigrationRefactor()) {
+            var gridSizeMigrationLogic = GridSizeMigrationLogic()
+            gridSizeMigrationLogic.migrateGrid(
+                context,
+                DeviceGridState(context),
+                DeviceGridState(idp),
+                dbHelper,
+                db,
+            )
+        } else {
+            GridSizeMigrationDBController.migrate(
+                dbHelper,
+                srcReader,
+                destReader,
+                destHotseatSize,
+                Point(pointX, pointY),
+                DeviceGridState(idp),
+                DeviceGridState(context),
+            )
+        }
+    }
+
     private fun verifyHotseat(c: Cursor, idp: InvariantDeviceProfile, expected: List<String?>) {
         assertThat(c.count).isEqualTo(idp.numDatabaseHotseatIcons)
         val screenIndex = c.getColumnIndex(SCREEN)
@@ -421,6 +494,17 @@
     }
 
     @Test
+    @EnableFlags(Flags.FLAG_GRID_MIGRATION_REFACTOR)
+    fun migrateToLargerHotseatRefactorFlagOn() {
+        migrateToLargerHotseat()
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_GRID_MIGRATION_REFACTOR)
+    fun migrateToLargerHotseatRefactorFlagOff() {
+        migrateToLargerHotseat()
+    }
+
     fun migrateToLargerHotseat() {
         val srcHotseatItems =
             intArrayOf(
@@ -471,14 +555,13 @@
         idp.numRows = 4
         val srcReader = DbReader(db, TMP_TABLE, context)
         val destReader = DbReader(db, TABLE_NAME, context)
-        GridSizeMigrationUtil.migrate(
+        migrateGrid(
             dbHelper,
             srcReader,
             destReader,
             idp.numDatabaseHotseatIcons,
-            Point(idp.numColumns, idp.numRows),
-            DeviceGridState(context),
-            DeviceGridState(idp),
+            idp.numColumns,
+            idp.numRows,
         )
 
         // Check hotseat items
@@ -516,6 +599,17 @@
     }
 
     @Test
+    @EnableFlags(Flags.FLAG_GRID_MIGRATION_REFACTOR)
+    fun migrateFromLargerHotseatRefactorFlagOn() {
+        migrateFromLargerHotseat()
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_GRID_MIGRATION_REFACTOR)
+    fun migrateFromLargerHotseatRefactorFlagOff() {
+        migrateFromLargerHotseat()
+    }
+
     fun migrateFromLargerHotseat() {
         addItem(ITEM_TYPE_APPLICATION, 0, CONTAINER_HOTSEAT, 0, 0, testPackage1, 1, TMP_TABLE)
         addItem(ITEM_TYPE_DEEP_SHORTCUT, 2, CONTAINER_HOTSEAT, 0, 0, testPackage2, 2, TMP_TABLE)
@@ -528,14 +622,13 @@
         idp.numRows = 4
         val srcReader = DbReader(db, TMP_TABLE, context)
         val destReader = DbReader(db, TABLE_NAME, context)
-        GridSizeMigrationUtil.migrate(
+        migrateGrid(
             dbHelper,
             srcReader,
             destReader,
             idp.numDatabaseHotseatIcons,
-            Point(idp.numColumns, idp.numRows),
-            DeviceGridState(context),
-            DeviceGridState(idp),
+            idp.numColumns,
+            idp.numRows,
         )
 
         // Check hotseat items
@@ -573,11 +666,24 @@
         c.close()
     }
 
+    @Test
+    @Throws(Exception::class)
+    @EnableFlags(Flags.FLAG_GRID_MIGRATION_REFACTOR)
+    fun migrateFromSmallerGridBigDifferenceRefactorFlagOn() {
+        migrateFromSmallerGridBigDifference()
+    }
+
+    @Test
+    @Throws(Exception::class)
+    @DisableFlags(Flags.FLAG_GRID_MIGRATION_REFACTOR)
+    fun migrateFromSmallerGridBigDifferenceRefactorFlagOff() {
+        migrateFromSmallerGridBigDifference()
+    }
+
     /**
      * Migrating from a smaller grid to a large one should reflow the pages if the column difference
      * is more than 2
      */
-    @Test
     @Throws(Exception::class)
     fun migrateFromSmallerGridBigDifference() {
         enableNewMigrationLogic("2,2")
@@ -594,14 +700,13 @@
         idp.numRows = 5
         val srcReader = DbReader(db, TMP_TABLE, context)
         val destReader = DbReader(db, TABLE_NAME, context)
-        GridSizeMigrationUtil.migrate(
+        migrateGrid(
             dbHelper,
             srcReader,
             destReader,
             idp.numDatabaseHotseatIcons,
-            Point(idp.numColumns, idp.numRows),
-            DeviceGridState(context),
-            DeviceGridState(idp),
+            idp.numColumns,
+            idp.numRows,
         )
 
         // Get workspace items
@@ -636,9 +741,22 @@
         assertThat(locMap[testPackage5]).isEqualTo(0)
     }
 
-    /** Migrating from a larger grid to a smaller, we reflow from page 0 */
     @Test
     @Throws(Exception::class)
+    @EnableFlags(Flags.FLAG_GRID_MIGRATION_REFACTOR)
+    fun migrateFromLargerGridRefactorFlagOn() {
+        migrateFromLargerGrid()
+    }
+
+    @Test
+    @Throws(Exception::class)
+    @DisableFlags(Flags.FLAG_GRID_MIGRATION_REFACTOR)
+    fun migrateFromLargerGridRefactorFlagOff() {
+        migrateFromLargerGrid()
+    }
+
+    /** Migrating from a larger grid to a smaller, we reflow from page 0 */
+    @Throws(Exception::class)
     fun migrateFromLargerGrid() {
         enableNewMigrationLogic("5,5")
 
@@ -654,14 +772,13 @@
         idp.numRows = 4
         val srcReader = DbReader(db, TMP_TABLE, context)
         val destReader = DbReader(db, TABLE_NAME, context)
-        GridSizeMigrationUtil.migrate(
+        migrateGrid(
             dbHelper,
             srcReader,
             destReader,
             idp.numDatabaseHotseatIcons,
-            Point(idp.numColumns, idp.numRows),
-            DeviceGridState(context),
-            DeviceGridState(idp),
+            idp.numColumns,
+            idp.numRows,
         )
 
         // Get workspace items
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/ModelTestExtensions.kt b/tests/multivalentTests/src/com/android/launcher3/util/ModelTestExtensions.kt
index 6bd182b..8d072d8 100644
--- a/tests/multivalentTests/src/com/android/launcher3/util/ModelTestExtensions.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/util/ModelTestExtensions.kt
@@ -1,6 +1,7 @@
 package com.android.launcher3.util
 
 import android.content.ContentValues
+import com.android.launcher3.Flags
 import com.android.launcher3.LauncherModel
 import com.android.launcher3.LauncherSettings.Favorites
 import com.android.launcher3.LauncherSettings.Favorites.APPWIDGET_ID
@@ -30,7 +31,8 @@
         loadModelSync()
         TestUtil.runOnExecutorSync(Executors.MODEL_EXECUTOR) {
             modelDbController.run {
-                tryMigrateDB(null /* restoreEventLogger */)
+                if (Flags.gridMigrationRefactor()) attemptMigrateDb(null /* restoreEventLogger */)
+                else tryMigrateDB(null /* restoreEventLogger */)
                 createEmptyDB()
                 clearEmptyDbFlag()
             }
@@ -67,12 +69,12 @@
         tableName: String = Favorites.TABLE_NAME,
         appWidgetId: Int = -1,
         appWidgetSource: Int = -1,
-        appWidgetProvider: String? = null
+        appWidgetProvider: String? = null,
     ) {
         loadModelSync()
         TestUtil.runOnExecutorSync(Executors.MODEL_EXECUTOR) {
             val controller: ModelDbController = modelDbController
-            controller.tryMigrateDB(null /* restoreEventLogger */)
+            controller.attemptMigrateDb(null /* restoreEventLogger */)
             modelDbController.newTransaction().use { transaction ->
                 val values =
                     ContentValues().apply {
diff --git a/tests/src/com/android/launcher3/backuprestore/BackupAndRestoreDBSelectionTest.kt b/tests/src/com/android/launcher3/backuprestore/BackupAndRestoreDBSelectionTest.kt
index 35ac0a1..b4ee090 100644
--- a/tests/src/com/android/launcher3/backuprestore/BackupAndRestoreDBSelectionTest.kt
+++ b/tests/src/com/android/launcher3/backuprestore/BackupAndRestoreDBSelectionTest.kt
@@ -16,6 +16,8 @@
 
 package com.android.launcher3.backuprestore
 
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
 import android.platform.test.flag.junit.SetFlagsRule
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.MediumTest
@@ -52,10 +54,24 @@
         setFlagsRule.setFlags(true, Flags.FLAG_ENABLE_NARROW_GRID_RESTORE)
     }
 
+    @EnableFlags(Flags.FLAG_GRID_MIGRATION_REFACTOR)
+    fun oldDatabasesNotPresentAfterRestoreRefactorFlagEnabled() {
+        oldDatabasesNotPresentAfterRestore()
+    }
+
+    @DisableFlags(Flags.FLAG_GRID_MIGRATION_REFACTOR)
+    fun oldDatabasesNotPresentAfterRestoreRefactorFlagDisabled() {
+        oldDatabasesNotPresentAfterRestore()
+    }
+
     @Test
     fun oldDatabasesNotPresentAfterRestore() {
         val dbController = ModelDbController(getInstrumentation().targetContext)
-        dbController.tryMigrateDB(null)
+        if (Flags.gridMigrationRefactor()) {
+            dbController.attemptMigrateDb(null)
+        } else {
+            dbController.tryMigrateDB(null)
+        }
         TestUtil.runOnExecutorSync(MODEL_EXECUTOR) {
             assert(backAndRestoreRule.getDatabaseFiles().size == 1) {
                 "There should only be one database after restoring, the last one used. Actual databases ${backAndRestoreRule.getDatabaseFiles()}"
diff --git a/tests/src/com/android/launcher3/model/GridMigrationTest.kt b/tests/src/com/android/launcher3/model/GridMigrationTest.kt
index 15222a4..666ec16 100644
--- a/tests/src/com/android/launcher3/model/GridMigrationTest.kt
+++ b/tests/src/com/android/launcher3/model/GridMigrationTest.kt
@@ -52,11 +52,15 @@
             phoneContext,
             dbFileName,
             { UserCache.INSTANCE.get(phoneContext).getSerialNumberForUser(it) },
-            {}
+            {},
         )
 
-    fun readEntries(): List<GridSizeMigrationUtil.DbEntry> =
-        GridSizeMigrationUtil.readAllEntries(dbHelper.readableDatabase, TABLE_NAME, phoneContext)
+    fun readEntries(): List<DbEntry> =
+        GridSizeMigrationDBController.readAllEntries(
+            dbHelper.readableDatabase,
+            TABLE_NAME,
+            phoneContext,
+        )
 }
 
 /**
@@ -80,7 +84,7 @@
         TestToPhoneFileCopier(
             src = "databases/GridMigrationTest/$DB_FILE",
             dest = "databases/$DB_FILE",
-            removeOnFinish = true
+            removeOnFinish = true,
         )
 
     @Before
@@ -89,13 +93,24 @@
     }
 
     private fun migrate(src: GridMigrationData, dst: GridMigrationData) {
-        GridSizeMigrationUtil.migrateGridIfNeeded(
-            phoneContext,
-            src.gridState,
-            dst.gridState,
-            dst.dbHelper,
-            src.dbHelper.readableDatabase
-        )
+        if (Flags.gridMigrationRefactor()) {
+            val gridSizeMigrationLogic = GridSizeMigrationLogic()
+            gridSizeMigrationLogic.migrateGrid(
+                phoneContext,
+                src.gridState,
+                dst.gridState,
+                dst.dbHelper,
+                src.dbHelper.readableDatabase,
+            )
+        } else {
+            GridSizeMigrationDBController.migrateGridIfNeeded(
+                phoneContext,
+                src.gridState,
+                dst.gridState,
+                dst.dbHelper,
+                src.dbHelper.readableDatabase,
+            )
+        }
     }
 
     /**
@@ -115,10 +130,8 @@
     }
 
     private fun compare(dst: GridMigrationData, target: GridMigrationData) {
-        val sort = compareBy<GridSizeMigrationUtil.DbEntry>({ it.cellX }, { it.cellY })
-        val mapF = { it: GridSizeMigrationUtil.DbEntry ->
-            EntryData(it.cellX, it.cellY, it.spanX, it.spanY, it.rank)
-        }
+        val sort = compareBy<DbEntry>({ it.cellX }, { it.cellY })
+        val mapF = { it: DbEntry -> EntryData(it.cellX, it.cellY, it.spanX, it.spanY, it.rank) }
         val entriesDst = dst.readEntries().sortedWith(sort).map(mapF)
         val entriesTarget = target.readEntries().sortedWith(sort).map(mapF)
 
@@ -149,7 +162,7 @@
         TestToPhoneFileCopier(
             src = "databases/GridMigrationTest/result5x5to3x3.db",
             dest = "databases/result5x5to3x3.db",
-            removeOnFinish = true
+            removeOnFinish = true,
         )
 
     @Test
@@ -160,10 +173,10 @@
                 GridMigrationData(
                     null, // in memory db, to download a new db change null for the filename of the
                     // db name to store it. Do not use existing names.
-                    DeviceGridState(3, 3, 3, TYPE_PHONE, "")
+                    DeviceGridState(3, 3, 3, TYPE_PHONE, ""),
                 ),
             target =
-                GridMigrationData("result5x5to3x3.db", DeviceGridState(3, 3, 3, TYPE_PHONE, ""))
+                GridMigrationData("result5x5to3x3.db", DeviceGridState(3, 3, 3, TYPE_PHONE, "")),
         )
 
     @JvmField
@@ -172,7 +185,7 @@
         TestToPhoneFileCopier(
             src = "databases/GridMigrationTest/result5x5to4x7.db",
             dest = "databases/result5x5to4x7.db",
-            removeOnFinish = true
+            removeOnFinish = true,
         )
 
     @Test
@@ -183,10 +196,10 @@
                 GridMigrationData(
                     null, // in memory db, to download a new db change null for the filename of the
                     // db name to store it. Do not use existing names.
-                    DeviceGridState(4, 7, 4, TYPE_PHONE, "")
+                    DeviceGridState(4, 7, 4, TYPE_PHONE, ""),
                 ),
             target =
-                GridMigrationData("result5x5to4x7.db", DeviceGridState(4, 7, 4, TYPE_PHONE, ""))
+                GridMigrationData("result5x5to4x7.db", DeviceGridState(4, 7, 4, TYPE_PHONE, "")),
         )
 
     @JvmField
@@ -195,7 +208,7 @@
         TestToPhoneFileCopier(
             src = "databases/GridMigrationTest/result5x5to5x8.db",
             dest = "databases/result5x5to5x8.db",
-            removeOnFinish = true
+            removeOnFinish = true,
         )
 
     @Test
@@ -206,10 +219,10 @@
                 GridMigrationData(
                     null, // in memory db, to download a new db change null for the filename of the
                     // db name to store it. Do not use existing names.
-                    DeviceGridState(5, 8, 5, TYPE_PHONE, "")
+                    DeviceGridState(5, 8, 5, TYPE_PHONE, ""),
                 ),
             target =
-                GridMigrationData("result5x5to5x8.db", DeviceGridState(5, 8, 5, TYPE_PHONE, ""))
+                GridMigrationData("result5x5to5x8.db", DeviceGridState(5, 8, 5, TYPE_PHONE, "")),
         )
 
     @JvmField
@@ -218,7 +231,7 @@
         TestToPhoneFileCopier(
             src = "databases/GridMigrationTest/flagged_result5x5to5x8.db",
             dest = "databases/flagged_result5x5to5x8.db",
-            removeOnFinish = true
+            removeOnFinish = true,
         )
 
     @Test
@@ -230,13 +243,13 @@
                 GridMigrationData(
                     null, // in memory db, to download a new db change null for the filename of the
                     // db name to store it. Do not use existing names.
-                    DeviceGridState(5, 8, 5, TYPE_PHONE, "")
+                    DeviceGridState(5, 8, 5, TYPE_PHONE, ""),
                 ),
             target =
                 GridMigrationData(
                     "flagged_result5x5to5x8.db",
-                    DeviceGridState(5, 8, 5, TYPE_PHONE, "")
-                )
+                    DeviceGridState(5, 8, 5, TYPE_PHONE, ""),
+                ),
         )
     }
 }
diff --git a/tests/src/com/android/launcher3/model/gridmigration/ValidGridMigrationUnitTest.kt b/tests/src/com/android/launcher3/model/gridmigration/ValidGridMigrationUnitTest.kt
index 03d0195..c08237c 100644
--- a/tests/src/com/android/launcher3/model/gridmigration/ValidGridMigrationUnitTest.kt
+++ b/tests/src/com/android/launcher3/model/gridmigration/ValidGridMigrationUnitTest.kt
@@ -20,17 +20,21 @@
 import android.database.sqlite.SQLiteDatabase
 import android.graphics.Point
 import android.os.Process
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
 import android.util.Log
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import androidx.test.platform.app.InstrumentationRegistry
+import com.android.launcher3.Flags
 import com.android.launcher3.InvariantDeviceProfile
 import com.android.launcher3.LauncherSettings.Favorites
 import com.android.launcher3.celllayout.testgenerator.ValidGridMigrationTestCaseGenerator
 import com.android.launcher3.celllayout.testgenerator.generateItemsForTest
 import com.android.launcher3.model.DatabaseHelper
 import com.android.launcher3.model.DeviceGridState
-import com.android.launcher3.model.GridSizeMigrationUtil
+import com.android.launcher3.model.GridSizeMigrationDBController
+import com.android.launcher3.model.GridSizeMigrationLogic
 import com.android.launcher3.pm.UserCache
 import com.android.launcher3.provider.LauncherDbUtils
 import com.android.launcher3.util.rule.TestStabilityRule
@@ -130,22 +134,44 @@
         addItemsToDb(dbHelper.writableDatabase, dstGrid)
 
         LauncherDbUtils.SQLiteTransaction(dbHelper.writableDatabase).use {
-            GridSizeMigrationUtil.migrate(
-                dbHelper,
-                GridSizeMigrationUtil.DbReader(it.db, srcGrid.tableName, context),
-                GridSizeMigrationUtil.DbReader(it.db, dstGrid.tableName, context),
-                dstGrid.size.x,
-                dstGrid.size,
-                srcGrid.toGridState(),
-                dstGrid.toGridState(),
-            )
+            if (Flags.gridMigrationRefactor()) {
+                val gridSizeMigrationLogic = GridSizeMigrationLogic()
+                gridSizeMigrationLogic.migrateGrid(
+                    context,
+                    srcGrid.toGridState(),
+                    dstGrid.toGridState(),
+                    dbHelper,
+                    it.db,
+                )
+            } else {
+                GridSizeMigrationDBController.migrate(
+                    dbHelper,
+                    GridSizeMigrationDBController.DbReader(it.db, srcGrid.tableName, context),
+                    GridSizeMigrationDBController.DbReader(it.db, dstGrid.tableName, context),
+                    dstGrid.size.x,
+                    dstGrid.size,
+                    srcGrid.toGridState(),
+                    dstGrid.toGridState(),
+                )
+            }
             it.commit()
         }
         return readDb(dstGrid.tableName, dbHelper.readableDatabase)
     }
 
     @Test
-    fun runTestCase() {
+    @EnableFlags(Flags.FLAG_GRID_MIGRATION_REFACTOR)
+    fun runTestCaseRefactorFlagEnabled() {
+        runTestCase()
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_GRID_MIGRATION_REFACTOR)
+    fun runTestCaseRefactorFlagDisabled() {
+        runTestCase()
+    }
+
+    private fun runTestCase() {
         val caseGenerator = ValidGridMigrationTestCaseGenerator(Random(SEED.toLong()))
         for (i in 0..SMALL_TEST_SIZE) {
             val testCase = caseGenerator.generateTestCase(isDestEmpty = true)
@@ -163,7 +189,18 @@
     }
 
     @Test
-    fun mergeBoards() {
+    @EnableFlags(Flags.FLAG_GRID_MIGRATION_REFACTOR)
+    fun mergeBoardsRefactorFlagEnabled() {
+        mergeBoards()
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_GRID_MIGRATION_REFACTOR)
+    fun mergeBoardsRefactorFlagDisabled() {
+        mergeBoards()
+    }
+
+    private fun mergeBoards() {
         val caseGenerator = ValidGridMigrationTestCaseGenerator(Random(SEED.toLong()))
         for (i in 0..SMALL_TEST_SIZE) {
             val testCase = caseGenerator.generateTestCase(isDestEmpty = false)
@@ -187,7 +224,20 @@
     // This test takes about 4 minutes, there is no need to run it in presubmit.
     @Stability(flavors = TestStabilityRule.LOCAL or TestStabilityRule.PLATFORM_POSTSUBMIT)
     @Test
-    fun runExtensiveTestCases() {
+    @EnableFlags(Flags.FLAG_GRID_MIGRATION_REFACTOR)
+    fun runExtensiveTestCasesRefactorFlagEnabled() {
+        runExtensiveTestCases()
+    }
+
+    // This test takes about 4 minutes, there is no need to run it in presubmit.
+    @Stability(flavors = TestStabilityRule.LOCAL or TestStabilityRule.PLATFORM_POSTSUBMIT)
+    @Test
+    @DisableFlags(Flags.FLAG_GRID_MIGRATION_REFACTOR)
+    fun runExtensiveTestCasesRefactorFlagDisabled() {
+        runExtensiveTestCases()
+    }
+
+    private fun runExtensiveTestCases() {
         val caseGenerator = ValidGridMigrationTestCaseGenerator(Random(SEED.toLong()))
         for (i in 0..LARGE_TEST_SIZE) {
             val testCase = caseGenerator.generateTestCase(isDestEmpty = true)