diff --git a/aconfig/launcher.aconfig b/aconfig/launcher.aconfig
index d334cf4..949acc1 100644
--- a/aconfig/launcher.aconfig
+++ b/aconfig/launcher.aconfig
@@ -499,3 +499,10 @@
    description: "Enforce system radius for widget corners instead of a separate 16.dp value"
    bug: "370950552"
 }
+
+flag {
+    name: "enable_contrast_tiles"
+    namespace: "launcher"
+    description: "Enable launcher app contrast tiles."
+    bug: "341217082"
+}
diff --git a/quickstep/res/layout/task_desktop.xml b/quickstep/res/layout/task_desktop.xml
index 1564653..0472007 100644
--- a/quickstep/res/layout/task_desktop.xml
+++ b/quickstep/res/layout/task_desktop.xml
@@ -19,16 +19,11 @@
     android:id="@+id/task_view_desktop"
     android:layout_width="match_parent"
     android:layout_height="match_parent"
-    android:clipChildren="true"
-    android:clipToPadding="true"
     android:contentDescription="@string/recent_task_desktop"
     android:defaultFocusHighlightEnabled="false"
     android:focusable="true"
-    android:padding="0.1dp"
     launcher:focusBorderColor="?attr/materialColorOutline"
     launcher:hoverBorderColor="?attr/materialColorPrimary">
-    <!-- Setting a padding of 0.1 dp since android:clipToPadding needs a non-zero value for
-    padding to work-->
     <View
         android:id="@+id/background"
         android:layout_width="match_parent"
@@ -40,4 +35,9 @@
         android:layout_height="wrap_content"
         android:inflatedId="@id/icon" />
 
+    <com.android.quickstep.views.DesktopTaskContentView
+        android:id="@+id/desktop_content"
+        android:layout_width="match_parent"
+        android:layout_height="match_parent" />
+
 </com.android.quickstep.views.DesktopTaskView>
diff --git a/quickstep/src/com/android/launcher3/taskbar/FallbackTaskbarUIController.java b/quickstep/src/com/android/launcher3/taskbar/FallbackTaskbarUIController.java
index 929e793..e3bcb0d 100644
--- a/quickstep/src/com/android/launcher3/taskbar/FallbackTaskbarUIController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/FallbackTaskbarUIController.java
@@ -25,20 +25,24 @@
 
 import com.android.launcher3.popup.SystemShortcut;
 import com.android.launcher3.statemanager.StateManager;
-import com.android.quickstep.RecentsActivity;
+import com.android.launcher3.statemanager.StatefulContainer;
 import com.android.quickstep.TopTaskTracker;
 import com.android.quickstep.fallback.RecentsState;
 import com.android.quickstep.util.TISBindHelper;
 import com.android.quickstep.views.RecentsView;
+import com.android.quickstep.views.RecentsViewContainer;
 
 import java.util.stream.Stream;
 
 /**
  * A data source which integrates with the fallback RecentsActivity instance (for 3P launchers).
+ * @param <T> The type of the RecentsViewContainer that will handle Recents state changes.
  */
-public class FallbackTaskbarUIController extends TaskbarUIController {
+public class FallbackTaskbarUIController
+        <T extends RecentsViewContainer & StatefulContainer<RecentsState>>
+        extends TaskbarUIController {
 
-    private final RecentsActivity mRecentsActivity;
+    private final T mRecentsContainer;
 
     private final StateManager.StateListener<RecentsState> mStateListener =
             new StateManager.StateListener<RecentsState>() {
@@ -63,23 +67,23 @@
                 }
             };
 
-    public FallbackTaskbarUIController(RecentsActivity recentsActivity) {
-        mRecentsActivity = recentsActivity;
+    public FallbackTaskbarUIController(T recentsContainer) {
+        mRecentsContainer = recentsContainer;
     }
 
     @Override
     protected void init(TaskbarControllers taskbarControllers) {
         super.init(taskbarControllers);
-        mRecentsActivity.setTaskbarUIController(this);
-        mRecentsActivity.getStateManager().addStateListener(mStateListener);
+        mRecentsContainer.setTaskbarUIController(this);
+        mRecentsContainer.getStateManager().addStateListener(mStateListener);
     }
 
     @Override
     protected void onDestroy() {
         super.onDestroy();
         getRecentsView().setTaskLaunchListener(null);
-        mRecentsActivity.setTaskbarUIController(null);
-        mRecentsActivity.getStateManager().removeStateListener(mStateListener);
+        mRecentsContainer.setTaskbarUIController(null);
+        mRecentsContainer.getStateManager().removeStateListener(mStateListener);
     }
 
     /**
@@ -109,7 +113,7 @@
 
     @Override
     public RecentsView getRecentsView() {
-        return mRecentsActivity.getOverviewPanel();
+        return mRecentsContainer.getOverviewPanel();
     }
 
     @Override
@@ -131,11 +135,11 @@
     @Nullable
     @Override
     protected TISBindHelper getTISBindHelper() {
-        return mRecentsActivity.getTISBindHelper();
+        return mRecentsContainer.getTISBindHelper();
     }
 
     @Override
     protected String getTaskbarUIControllerName() {
-        return "FallbackTaskbarUIController";
+        return "FallbackTaskbarUIController<" + mRecentsContainer.getClass().getSimpleName() + ">";
     }
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchViewController.java b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchViewController.java
index 3f71870..7a63f74 100644
--- a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchViewController.java
@@ -31,12 +31,14 @@
 import androidx.annotation.Nullable;
 
 import com.android.internal.jank.Cuj;
+import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.anim.AnimatorListeners;
 import com.android.launcher3.desktop.DesktopAppLaunchTransition;
 import com.android.launcher3.taskbar.overlay.TaskbarOverlayContext;
 import com.android.launcher3.taskbar.overlay.TaskbarOverlayDragLayer;
+import com.android.launcher3.util.DisplayController;
 import com.android.launcher3.views.BaseDragLayer;
 import com.android.quickstep.SystemUiProxy;
 import com.android.quickstep.util.GroupTask;
@@ -101,7 +103,16 @@
             boolean hasDesktopTask,
             boolean wasDesktopTaskFilteredOut,
             boolean wasOpenedFromTaskbar) {
-        positionView(wasOpenedFromTaskbar);
+        final boolean isTransientTaskBar = DisplayController.isTransientTaskbar(
+                mControllers.taskbarActivityContext);
+        positionView(wasOpenedFromTaskbar, isTransientTaskBar);
+
+        // Keep the taskbar unstashed if the KQS is opened.
+        if (wasOpenedFromTaskbar && isTransientTaskBar) {
+            mControllers.taskbarStashController.updateTaskbarTimeout(/* isAutohideSuspended= */
+                    true);
+        }
+
         mOverlayContext.getDragLayer().addView(mKeyboardQuickSwitchView);
         mOnDesktop = onDesktop;
         mWasDesktopTaskFilteredOut = wasDesktopTaskFilteredOut;
@@ -117,7 +128,7 @@
                 /* useDesktopTaskView= */ !onDesktop && hasDesktopTask);
     }
 
-    protected void positionView(boolean wasOpenedFromTaskbar) {
+    protected void positionView(boolean wasOpenedFromTaskbar, boolean isTransientTaskbar) {
         if (!wasOpenedFromTaskbar) {
             // Keep the default positioning.
             return;
@@ -128,8 +139,16 @@
         final Resources resources = mKeyboardQuickSwitchView.getResources();
         final int marginHorizontal = resources.getDimensionPixelSize(
                 R.dimen.keyboard_quick_switch_margin_ends);
-        final int marginBottom = resources.getDimensionPixelSize(
+
+        final DeviceProfile dp = mControllers.taskbarActivityContext.getDeviceProfile();
+        // Calculate the additional margin space that the KQS should move up for the transient
+        // taskbar. The value of spaceForTaskbar is the distance between the bottom of the KQS
+        // view with 0 bottom margin to the top of the transient taskbar view.
+        final int spaceForTaskbar = isTransientTaskbar ? dp.taskbarHeight + dp.taskbarBottomMargin
+                - dp.stashedTaskbarHeight : 0;
+        final int marginBottom = spaceForTaskbar + 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;
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
index 4a85acc..5a63ca6 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
@@ -169,7 +169,7 @@
         taskbarOverlayController.init(this);
         taskbarAllAppsController.init(this, sharedState.allAppsVisible);
         navButtonController.init(this);
-        bubbleControllers.ifPresent(controllers -> controllers.init(this));
+        bubbleControllers.ifPresent(controllers -> controllers.init(sharedState, this));
         taskbarInsetsController.init(this);
         voiceInteractionWindowController.init(this);
         taskbarRecentAppsController.init(this);
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
index c18cf28..ab4b1b6 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
@@ -71,7 +71,9 @@
 import com.android.quickstep.AllAppsActionManager;
 import com.android.quickstep.RecentsActivity;
 import com.android.quickstep.SystemUiProxy;
+import com.android.quickstep.fallback.window.RecentsWindowManager;
 import com.android.quickstep.util.ContextualSearchInvoker;
+import com.android.quickstep.views.RecentsViewContainer;
 import com.android.systemui.shared.statusbar.phone.BarTransitions;
 import com.android.systemui.shared.system.QuickStepContract;
 import com.android.systemui.shared.system.QuickStepContract.SystemUiStateFlags;
@@ -115,7 +117,6 @@
     private WindowManager mWindowManager;
     private FrameLayout mTaskbarRootLayout;
     private boolean mAddedWindow;
-    private boolean mIsSuspended;
     private final TaskbarNavButtonController mNavButtonController;
     private final ComponentCallbacks mComponentCallbacks;
 
@@ -131,6 +132,8 @@
 
     private TaskbarActivityContext mTaskbarActivityContext;
     private StatefulActivity mActivity;
+    private RecentsViewContainer mRecentsViewContainer;
+
     /**
      * Cache a copy here so we can initialize state whenever taskbar is recreated, since
      * this class does not get re-initialized w/ new taskbars.
@@ -404,9 +407,28 @@
         }
         mUnfoldProgressProvider.setSourceProvider(unfoldTransitionProgressProvider);
 
+        if (activity instanceof RecentsViewContainer recentsViewContainer) {
+            setRecentsViewContainer(recentsViewContainer);
+        }
+    }
+
+    /**
+     * Sets the current RecentsViewContainer, from which we create a TaskbarUIController.
+     */
+    public void setRecentsViewContainer(@NonNull RecentsViewContainer recentsViewContainer) {
+        if (mRecentsViewContainer == recentsViewContainer) {
+            return;
+        }
+        if (mRecentsViewContainer == mActivity) {
+            // When switching to RecentsWindowManager (not an Activity), the old mActivity is not
+            // destroyed, nor is there a new Activity to replace it. Thus if we don't clear it here,
+            // it will not get re-set properly if we return to the Activity (e.g. NexusLauncher).
+            mActivityOnDestroyCallback.run();
+        }
+        mRecentsViewContainer = recentsViewContainer;
         if (mTaskbarActivityContext != null) {
             mTaskbarActivityContext.setUIController(
-                    createTaskbarUIControllerForActivity(mActivity));
+                    createTaskbarUIControllerForRecentsViewContainer(mRecentsViewContainer));
         }
     }
 
@@ -429,12 +451,18 @@
     /**
      * Creates a {@link TaskbarUIController} to use while the given StatefulActivity is active.
      */
-    private TaskbarUIController createTaskbarUIControllerForActivity(StatefulActivity activity) {
-        if (activity instanceof QuickstepLauncher) {
-            return new LauncherTaskbarUIController((QuickstepLauncher) activity);
+    private TaskbarUIController createTaskbarUIControllerForRecentsViewContainer(
+            RecentsViewContainer container) {
+        if (container instanceof QuickstepLauncher quickstepLauncher) {
+            return new LauncherTaskbarUIController(quickstepLauncher);
         }
-        if (activity instanceof RecentsActivity) {
-            return new FallbackTaskbarUIController((RecentsActivity) activity);
+        // If a 3P Launcher is default, always use FallbackTaskbarUIController regardless of
+        // whether the recents container is RecentsActivity or RecentsWindowManager.
+        if (container instanceof RecentsActivity recentsActivity) {
+            return new FallbackTaskbarUIController<>(recentsActivity);
+        }
+        if (container instanceof RecentsWindowManager recentsWindowManager) {
+            return new FallbackTaskbarUIController<>(recentsWindowManager);
         }
         return TaskbarUIController.DEFAULT;
     }
@@ -446,8 +474,6 @@
      */
     @VisibleForTesting
     public synchronized void recreateTaskbar() {
-        if (mIsSuspended) return;
-
         Trace.beginSection("recreateTaskbar");
         try {
             DeviceProfile dp = mUserUnlocked ?
@@ -484,9 +510,9 @@
             mSharedState.allAppsVisible = mSharedState.allAppsVisible && isLargeScreenTaskbar;
             mTaskbarActivityContext.init(mSharedState);
 
-            if (mActivity != null) {
+            if (mRecentsViewContainer != null) {
                 mTaskbarActivityContext.setUIController(
-                    createTaskbarUIControllerForActivity(mActivity));
+                        createTaskbarUIControllerForRecentsViewContainer(mRecentsViewContainer));
             }
 
             if (enableTaskbarNoRecreate()) {
@@ -663,16 +689,6 @@
         }
     }
 
-    /**
-     * Removes Taskbar from the window manager and prevents recreation if {@code true}.
-     * <p>
-     * Suspending is for testing purposes only; avoid calling this method in production.
-     */
-    @VisibleForTesting
-    public void setSuspended(boolean isSuspended) {
-        mIsSuspended = isSuspended;
-    }
-
     private void addTaskbarRootViewToWindow() {
         if (enableTaskbarNoRecreate() && !mAddedWindow && mTaskbarActivityContext != null) {
             mWindowManager.addView(mTaskbarRootLayout,
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarSharedState.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarSharedState.java
index 729cbe9..a64dab1 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarSharedState.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarSharedState.java
@@ -30,6 +30,10 @@
 import android.view.InsetsFrameProvider;
 
 import com.android.systemui.shared.system.QuickStepContract.SystemUiStateFlags;
+import com.android.wm.shell.shared.bubbles.BubbleBarLocation;
+import com.android.wm.shell.shared.bubbles.BubbleInfo;
+
+import java.util.List;
 
 /**
  * State shared across different taskbar instance
@@ -69,6 +73,15 @@
 
     public boolean allAppsVisible = false;
 
+    public BubbleBarLocation bubbleBarLocation;
+
+    public List<BubbleInfo> bubbleInfoItems;
+
+    /** Returns whether there are a saved bubbles. */
+    public boolean hasSavedBubbles() {
+        return bubbleInfoItems != null && !bubbleInfoItems.isEmpty();
+    }
+
     // LauncherTaskbarUIController#mTaskbarInAppDisplayProgressMultiProp
     public float[] inAppDisplayProgressMultiPropValues = new float[DISPLAY_PROGRESS_COUNT];
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
index b22fd6f..30e4e47 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
@@ -35,6 +35,7 @@
 import android.util.ArrayMap;
 import android.util.Log;
 
+import com.android.launcher3.taskbar.TaskbarSharedState;
 import com.android.launcher3.taskbar.bubbles.stashing.BubbleStashController;
 import com.android.launcher3.util.Executors.SimpleThreadFactory;
 import com.android.quickstep.SystemUiProxy;
@@ -47,6 +48,7 @@
 import com.android.wm.shell.shared.bubbles.RemovedBubble;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
@@ -112,6 +114,7 @@
 
     private BubbleBarItem mSelectedBubble;
 
+    private TaskbarSharedState mSharedState;
     private ImeVisibilityChecker mImeVisibilityChecker;
     private BubbleBarViewController mBubbleBarViewController;
     private BubbleStashController mBubbleStashController;
@@ -173,12 +176,25 @@
 
     public void onDestroy() {
         mSystemUiProxy.setBubblesListener(null);
+        // Saves bubble bar state
+        BubbleInfo[] bubbleInfoItems = new BubbleInfo[mBubbles.size()];
+        mBubbles.values().forEach(bubbleBarBubble -> {
+            int index = mBubbleBarViewController.bubbleViewIndex(bubbleBarBubble.getView());
+            if (index < 0 || index >= bubbleInfoItems.length) {
+                Log.e(TAG, "Found improper index: " + index + " for " + bubbleBarBubble);
+            } else {
+                bubbleInfoItems[index] = bubbleBarBubble.getInfo();
+            }
+        });
+        mSharedState.bubbleInfoItems = Arrays.asList(bubbleInfoItems);
     }
 
     /** Initializes controllers. */
     public void init(BubbleControllers bubbleControllers,
             BubbleBarLocationListener bubbleBarLocationListener,
-            ImeVisibilityChecker imeVisibilityChecker) {
+            ImeVisibilityChecker imeVisibilityChecker,
+            TaskbarSharedState sharedState) {
+        mSharedState = sharedState;
         mImeVisibilityChecker = imeVisibilityChecker;
         mBubbleBarViewController = bubbleControllers.bubbleBarViewController;
         mBubbleStashController = bubbleControllers.bubbleStashController;
@@ -188,6 +204,7 @@
         mBubbleBarLocationListener = bubbleBarLocationListener;
 
         bubbleControllers.runAfterInit(() -> {
+            restoreSavedState(sharedState);
             mBubbleBarViewController.setHiddenForBubbles(
                     !sBubbleBarEnabled || mBubbles.isEmpty());
             mBubbleStashedHandleViewController.ifPresent(
@@ -266,6 +283,26 @@
         }
     }
 
+    private void restoreSavedState(TaskbarSharedState sharedState) {
+        if (sharedState.bubbleBarLocation != null) {
+            updateBubbleBarLocationInternal(sharedState.bubbleBarLocation);
+        }
+        List<BubbleInfo> bubbleInfos = sharedState.bubbleInfoItems;
+        if (bubbleInfos == null || bubbleInfos.isEmpty()) return;
+        // Iterate in reverse because new bubbles are added in front and the list is in order.
+        for (int i = bubbleInfos.size() - 1; i >= 0; i--) {
+            BubbleBarBubble bubble = mBubbleCreator.populateBubble(mContext,
+                    bubbleInfos.get(i), mBarView, /* existingBubble = */ null);
+            if (bubble == null) {
+                Log.e(TAG, "Could not instantiate BubbleBarBubble for " + bubbleInfos.get(i));
+                continue;
+            }
+            addBubbleInternally(bubble,  /* showAppBadge = */
+                    mBubbleBarViewController.isExpanded() || i == 0,
+                    /* isExpanding = */ false,  /* suppressAnimation = */ true);
+        }
+    }
+
     private void applyViewChanges(BubbleBarViewUpdate update) {
         final boolean isCollapsed = (update.expandedChanged && !update.expanded)
                 || (!update.expandedChanged && !mBubbleBarViewController.isExpanded());
@@ -277,6 +314,12 @@
                 update.initialState || mBubbleBarViewController.isHiddenForSysui()
                         || mImeVisibilityChecker.isImeVisible();
 
+        if (update.initialState && mSharedState.hasSavedBubbles()) {
+            // clear restored state
+            mBubbleBarViewController.removeAllBubbles();
+            mBubbles.clear();
+        }
+
         BubbleBarBubble bubbleToSelect = null;
 
         if (Flags.enableOptionalBubbleOverflow()
@@ -347,8 +390,8 @@
             for (int i = update.currentBubbles.size() - 1; i >= 0; i--) {
                 BubbleBarBubble bubble = update.currentBubbles.get(i);
                 if (bubble != null) {
-                    mBubbles.put(bubble.getKey(), bubble);
-                    mBubbleBarViewController.addBubble(bubble, isExpanding, suppressAnimation);
+                    addBubbleInternally(bubble, /* showAppBadge = */ !isCollapsed || i == 0,
+                            isExpanding, suppressAnimation);
                     if (isCollapsed) {
                         // If we're collapsed, the most recently added bubble will be selected.
                         bubbleToSelect = bubble;
@@ -420,6 +463,7 @@
             }
         }
         if (update.bubbleBarLocation != null) {
+            mSharedState.bubbleBarLocation = update.bubbleBarLocation;
             if (update.bubbleBarLocation != mBubbleBarViewController.getBubbleBarLocation()) {
                 updateBubbleBarLocationInternal(update.bubbleBarLocation);
             }
@@ -519,6 +563,14 @@
         }
     }
 
+    private void addBubbleInternally(BubbleBarBubble bubble, boolean showAppBadge,
+            boolean isExpanding, boolean suppressAnimation) {
+        //TODO(b/360652359): remove setting scale to the app badge once issue is fixed
+        bubble.getView().setBadgeScale(showAppBadge ? 1 : 0);
+        mBubbles.put(bubble.getKey(), bubble);
+        mBubbleBarViewController.addBubble(bubble, isExpanding, suppressAnimation);
+    }
+
     /** Interface for checking whether the IME is visible. */
     public interface ImeVisibilityChecker {
         /** Whether the IME is visible. */
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
index fef5c7f..96fadf7 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
@@ -57,6 +57,7 @@
 import com.android.launcher3.util.MultiPropertyFactory;
 import com.android.launcher3.util.MultiValueAlpha;
 import com.android.quickstep.SystemUiProxy;
+import com.android.wm.shell.Flags;
 import com.android.wm.shell.shared.bubbles.BubbleBarLocation;
 
 import java.io.PrintWriter;
@@ -177,6 +178,9 @@
         mBubbleBarClickListener = v -> expandBubbleBar();
         mBubbleDragController.setupBubbleBarView(mBarView);
         mOverflowBubble = bubbleControllers.bubbleCreator.createOverflow(mBarView);
+        if (!Flags.enableOptionalBubbleOverflow()) {
+            showOverflow(true);
+        }
         mBarView.setOnClickListener(mBubbleBarClickListener);
         mBarView.addOnLayoutChangeListener(
                 (v, left, top, right, bottom, oldLeft, oldTop, oldRight, oldBottom) -> {
@@ -324,6 +328,7 @@
     }
 
     private void onBubbleClicked(BubbleView bubbleView) {
+        if (mBubbleBarPinning.isAnimating()) return;
         bubbleView.markSeen();
         BubbleBarItem bubble = bubbleView.getBubble();
         if (bubble == null) {
@@ -923,7 +928,7 @@
      * from Launcher.
      */
     public void setExpanded(boolean isExpanded) {
-        if (isExpanded != mBarView.isExpanded()) {
+        if (!mBubbleBarPinning.isAnimating() && isExpanded != mBarView.isExpanded()) {
             mBarView.setExpanded(isExpanded);
             adjustTaskbarAndHotseatToBubbleBarState(isExpanded);
             if (!isExpanded) {
@@ -1065,6 +1070,16 @@
         mSystemUiProxy.removeAllBubbles();
     }
 
+    /** Removes all existing bubble views */
+    public void removeAllBubbles() {
+        mBarView.removeAllViews();
+    }
+
+    /** Returns the view index of the existing bubble */
+    public int bubbleViewIndex(View bubbleView) {
+        return mBarView.indexOfChild(bubbleView);
+    }
+
     /**
      * Set listener to be notified when bubble bar bounds have changed
      */
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleControllers.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleControllers.java
index b5d94bd..d993685 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleControllers.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleControllers.java
@@ -21,6 +21,7 @@
 import android.view.View;
 
 import com.android.launcher3.taskbar.TaskbarControllers;
+import com.android.launcher3.taskbar.TaskbarSharedState;
 import com.android.launcher3.taskbar.bubbles.BubbleBarViewController.TaskbarViewPropertiesProvider;
 import com.android.launcher3.taskbar.bubbles.stashing.BubbleBarLocationOnDemandListener;
 import com.android.launcher3.taskbar.bubbles.stashing.BubbleStashController;
@@ -79,7 +80,7 @@
      * BubbleControllers instance, but should be careful to only access things that were created
      * in constructors for now, as some controllers may still be waiting for init().
      */
-    public void init(TaskbarControllers taskbarControllers) {
+    public void init(TaskbarSharedState taskbarSharedState, TaskbarControllers taskbarControllers) {
         BubbleBarLocationCompositeListener bubbleBarLocationListeners =
                 new BubbleBarLocationCompositeListener(
                         taskbarControllers.navbarButtonsViewController,
@@ -88,7 +89,8 @@
                 );
         bubbleBarController.init(this,
                 bubbleBarLocationListeners,
-                taskbarControllers.navbarButtonsViewController::isImeVisible);
+                taskbarControllers.navbarButtonsViewController::isImeVisible,
+                taskbarSharedState);
         bubbleStashedHandleViewController.ifPresent(
                 controller -> controller.init(/* bubbleControllers = */ this));
         bubbleStashController.init(
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
index 228dc91..fe68ebc 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
@@ -139,6 +139,7 @@
 import com.android.launcher3.statemanager.StateManager.StateHandler;
 import com.android.launcher3.taskbar.LauncherTaskbarUIController;
 import com.android.launcher3.taskbar.TaskbarManager;
+import com.android.launcher3.taskbar.TaskbarUIController;
 import com.android.launcher3.testing.TestLogging;
 import com.android.launcher3.testing.shared.TestProtocol;
 import com.android.launcher3.uioverrides.QuickstepWidgetHolder.QuickstepHolderFactory;
@@ -1089,10 +1090,12 @@
         );
     }
 
-    public void setTaskbarUIController(LauncherTaskbarUIController taskbarUIController) {
-        mTaskbarUIController = taskbarUIController;
+    @Override
+    public void setTaskbarUIController(@Nullable TaskbarUIController taskbarUIController) {
+        mTaskbarUIController = (LauncherTaskbarUIController) taskbarUIController;
     }
 
+    @Override
     public @Nullable LauncherTaskbarUIController getTaskbarUIController() {
         return mTaskbarUIController;
     }
@@ -1399,6 +1402,7 @@
     }
 
     @NonNull
+    @Override
     public TISBindHelper getTISBindHelper() {
         return mTISBindHelper;
     }
diff --git a/quickstep/src/com/android/quickstep/RecentTasksList.java b/quickstep/src/com/android/quickstep/RecentTasksList.java
index 05bef35..a5cc32a 100644
--- a/quickstep/src/com/android/quickstep/RecentTasksList.java
+++ b/quickstep/src/com/android/quickstep/RecentTasksList.java
@@ -119,6 +119,11 @@
                     topTaskTracker.onTaskMovedToFront(taskInfo);
                 });
             }
+
+            @Override
+            public void onTaskInfoChanged(ActivityManager.RunningTaskInfo taskInfo) {
+                mMainThreadExecutor.execute(() -> topTaskTracker.onTaskChanged(taskInfo));
+            }
         });
         // We may receive onRunningTaskAppeared events later for tasks which have already been
         // included in the list returned by mSysUiProxy.getRunningTasks(), or may receive
diff --git a/quickstep/src/com/android/quickstep/RecentsActivity.java b/quickstep/src/com/android/quickstep/RecentsActivity.java
index b19f651..6075294 100644
--- a/quickstep/src/com/android/quickstep/RecentsActivity.java
+++ b/quickstep/src/com/android/quickstep/RecentsActivity.java
@@ -70,6 +70,7 @@
 import com.android.launcher3.statemanager.StatefulActivity;
 import com.android.launcher3.taskbar.FallbackTaskbarUIController;
 import com.android.launcher3.taskbar.TaskbarManager;
+import com.android.launcher3.taskbar.TaskbarUIController;
 import com.android.launcher3.util.ActivityOptionsWrapper;
 import com.android.launcher3.util.ContextTracker;
 import com.android.launcher3.util.RunnableList;
@@ -115,7 +116,7 @@
     private FallbackRecentsView mFallbackRecentsView;
     private OverviewActionsView<?> mActionsView;
     private TISBindHelper mTISBindHelper;
-    private @Nullable FallbackTaskbarUIController mTaskbarUIController;
+    private @Nullable FallbackTaskbarUIController<RecentsActivity> mTaskbarUIController;
 
     private StateManager<RecentsState, RecentsActivity> mStateManager;
 
@@ -174,11 +175,14 @@
         mTISBindHelper.runOnBindToTouchInteractionService(r);
     }
 
-    public void setTaskbarUIController(FallbackTaskbarUIController taskbarUIController) {
-        mTaskbarUIController = taskbarUIController;
+    @Override
+    public void setTaskbarUIController(@Nullable TaskbarUIController taskbarUIController) {
+        mTaskbarUIController = (FallbackTaskbarUIController<RecentsActivity>) taskbarUIController;
     }
 
-    public FallbackTaskbarUIController getTaskbarUIController() {
+    @Nullable
+    @Override
+    public FallbackTaskbarUIController<RecentsActivity> getTaskbarUIController() {
         return mTaskbarUIController;
     }
 
@@ -515,6 +519,7 @@
     }
 
     @NonNull
+    @Override
     public TISBindHelper getTISBindHelper() {
         return mTISBindHelper;
     }
diff --git a/quickstep/src/com/android/quickstep/TopTaskTracker.java b/quickstep/src/com/android/quickstep/TopTaskTracker.java
index 23a1ec7..71b6573 100644
--- a/quickstep/src/com/android/quickstep/TopTaskTracker.java
+++ b/quickstep/src/com/android/quickstep/TopTaskTracker.java
@@ -135,6 +135,15 @@
         }
     }
 
+    public void onTaskChanged(RunningTaskInfo taskInfo) {
+        for (int i = 0; i < mOrderedTaskList.size(); i++) {
+            if (mOrderedTaskList.get(i).taskId == taskInfo.taskId) {
+                mOrderedTaskList.set(i, taskInfo);
+                break;
+            }
+        }
+    }
+
     @Override
     public void onTaskStageChanged(int taskId, @StageType int stage, boolean visible) {
         // If a task is not visible anymore or has been moved to undefined, stop tracking it.
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index 032e755..e8f38be 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -776,10 +776,13 @@
         mAllAppsActionManager.setHomeAndOverviewSame(isHomeAndOverviewSame);
         RecentsViewContainer newOverviewContainer =
                 mOverviewComponentObserver.getContainerInterface().getCreatedContainer();
-        if (newOverviewContainer != null
-                && newOverviewContainer instanceof StatefulActivity activity) {
-            //TODO(b/368030750) refactor taskbarManager to accept RecentsViewContainer
-            mTaskbarManager.setActivity(activity);
+        if (newOverviewContainer != null) {
+            if (newOverviewContainer instanceof StatefulActivity activity) {
+                // This will also call setRecentsViewContainer() internally.
+                mTaskbarManager.setActivity(activity);
+            } else {
+                mTaskbarManager.setRecentsViewContainer(newOverviewContainer);
+            }
         }
         mTISBinder.onOverviewTargetChange();
     }
diff --git a/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowManager.kt b/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowManager.kt
index 74f9901..78224ae 100644
--- a/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowManager.kt
+++ b/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowManager.kt
@@ -40,6 +40,7 @@
 import com.android.launcher3.statemanager.StateManager
 import com.android.launcher3.statemanager.StateManager.AtomicAnimationFactory
 import com.android.launcher3.statemanager.StatefulContainer
+import com.android.launcher3.taskbar.TaskbarUIController
 import com.android.launcher3.util.ContextTracker
 import com.android.launcher3.util.DisplayController
 import com.android.launcher3.util.RunnableList
@@ -117,6 +118,7 @@
 
     private var callbacks: RecentsAnimationCallbacks? = null
 
+    private var taskbarUIController: TaskbarUIController? = null
     private var tisBindHelper: TISBindHelper = TISBindHelper(this) {}
 
     // Callback array that corresponds to events defined in @ActivityEvent
@@ -290,6 +292,18 @@
         return tisBindHelper.desktopVisibilityController
     }
 
+    override fun setTaskbarUIController(taskbarUIController: TaskbarUIController?) {
+        this.taskbarUIController = taskbarUIController
+    }
+
+    override fun getTaskbarUIController(): TaskbarUIController? {
+        return taskbarUIController
+    }
+
+    override fun getTISBindHelper(): TISBindHelper {
+        return tisBindHelper
+    }
+
     fun registerInitListener(onInitListener: Predicate<Boolean>) {
         this.onInitListener = onInitListener
     }
diff --git a/quickstep/src/com/android/quickstep/views/DesktopTaskContentView.kt b/quickstep/src/com/android/quickstep/views/DesktopTaskContentView.kt
new file mode 100644
index 0000000..481acac
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/views/DesktopTaskContentView.kt
@@ -0,0 +1,52 @@
+/*
+ * 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.quickstep.views
+
+import android.content.Context
+import android.graphics.Outline
+import android.graphics.Rect
+import android.util.AttributeSet
+import android.view.View
+import android.view.ViewOutlineProvider
+import android.widget.FrameLayout
+import com.android.quickstep.views.TaskView.FullscreenDrawParams
+
+class DesktopTaskContentView
+@JvmOverloads
+constructor(context: Context, attrs: AttributeSet? = null) : FrameLayout(context, attrs) {
+    private val currentFullscreenParams = FullscreenDrawParams(context)
+    private val taskCornerRadius: Float
+        get() = currentFullscreenParams.cornerRadius
+
+    private val bounds = Rect()
+
+    init {
+        clipToOutline = true
+        outlineProvider =
+            object : ViewOutlineProvider() {
+                override fun getOutline(view: View, outline: Outline) {
+                    outline.setRoundRect(bounds, taskCornerRadius)
+                }
+            }
+    }
+
+    override fun onSizeChanged(w: Int, h: Int, oldw: Int, oldh: Int) {
+        super.onSizeChanged(w, h, oldw, oldh)
+        bounds.set(0, 0, w, h)
+        invalidateOutline()
+    }
+}
diff --git a/quickstep/src/com/android/quickstep/views/DesktopTaskView.kt b/quickstep/src/com/android/quickstep/views/DesktopTaskView.kt
index 15b0a6b..5e842aa 100644
--- a/quickstep/src/com/android/quickstep/views/DesktopTaskView.kt
+++ b/quickstep/src/com/android/quickstep/views/DesktopTaskView.kt
@@ -26,6 +26,7 @@
 import android.util.Log
 import android.view.Gravity
 import android.view.View
+import android.widget.FrameLayout
 import androidx.core.content.res.ResourcesCompat
 import androidx.core.view.updateLayoutParams
 import com.android.launcher3.Flags.enableRefactorTaskThumbnail
@@ -81,12 +82,12 @@
     private val tempRect = Rect()
     private lateinit var backgroundView: View
     private lateinit var iconView: TaskViewIcon
-    private var childCountAtInflation = 0
+    private lateinit var contentView: FrameLayout
 
     override fun onFinishInflate() {
         super.onFinishInflate()
         backgroundView =
-            findViewById<View>(R.id.background)!!.apply {
+            findViewById<View>(R.id.background).apply {
                 updateLayoutParams<LayoutParams> {
                     topMargin = container.deviceProfile.overviewTaskThumbnailTopMarginPx
                 }
@@ -113,7 +114,12 @@
                 )
                 setText(resources.getText(R.string.recent_task_desktop))
             }
-        childCountAtInflation = childCount
+        contentView =
+            findViewById<FrameLayout>(R.id.desktop_content).apply {
+                updateLayoutParams<LayoutParams> {
+                    topMargin = container.deviceProfile.overviewTaskThumbnailTopMarginPx
+                }
+            }
     }
 
     /** Updates this desktop task to the gives task list defined in `tasks` */
@@ -137,13 +143,8 @@
                     } else {
                         taskThumbnailViewDeprecatedPool!!.view
                     }
+                contentView.addView(snapshotView, 0)
 
-                addView(
-                    snapshotView,
-                    // Add snapshotView to the front after initial views e.g. icon and
-                    // background.
-                    childCountAtInflation,
-                )
                 TaskContainer(
                     this,
                     task,
@@ -164,7 +165,7 @@
         super.onRecycle()
         visibility = VISIBLE
         taskContainers.forEach {
-            removeView(it.snapshotView)
+            contentView.removeView(it.snapshotView)
             if (enableRefactorTaskThumbnail()) {
                 taskThumbnailViewPool!!.recycle(it.thumbnailView)
             } else {
@@ -227,9 +228,7 @@
                 width = (taskSize.width() * scaleWidth).toInt()
                 height = (taskSize.height() * scaleHeight).toInt()
                 leftMargin = (positionInParent.x * scaleWidth).toInt()
-                topMargin =
-                    (positionInParent.y * scaleHeight).toInt() +
-                        container.deviceProfile.overviewTaskThumbnailTopMarginPx
+                topMargin = (positionInParent.y * scaleHeight).toInt()
             }
             if (DEBUG) {
                 with(it.snapshotView.layoutParams as LayoutParams) {
diff --git a/quickstep/src/com/android/quickstep/views/RecentsViewContainer.java b/quickstep/src/com/android/quickstep/views/RecentsViewContainer.java
index d8036aa..b04753b 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsViewContainer.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsViewContainer.java
@@ -16,8 +16,6 @@
 
 package com.android.quickstep.views;
 
-import android.app.Activity;
-import android.content.ComponentName;
 import android.content.Context;
 import android.content.ContextWrapper;
 import android.content.LocusId;
@@ -27,13 +25,16 @@
 import android.view.View;
 import android.view.Window;
 
+import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
 import com.android.launcher3.BaseActivity;
 import com.android.launcher3.logger.LauncherAtom;
 import com.android.launcher3.statehandlers.DesktopVisibilityController;
+import com.android.launcher3.taskbar.TaskbarUIController;
 import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.views.ScrimView;
+import com.android.quickstep.util.TISBindHelper;
 
 /**
  * Interface to be implemented by the parent view of RecentsView
@@ -212,4 +213,10 @@
 
     @Nullable
     DesktopVisibilityController getDesktopVisibilityController();
+
+    void setTaskbarUIController(@Nullable TaskbarUIController taskbarUIController);
+
+    @Nullable TaskbarUIController getTaskbarUIController();
+
+    @NonNull TISBindHelper getTISBindHelper();
 }
diff --git a/quickstep/src/com/android/quickstep/views/TaskContainer.kt b/quickstep/src/com/android/quickstep/views/TaskContainer.kt
index 959516f..25aba39 100644
--- a/quickstep/src/com/android/quickstep/views/TaskContainer.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskContainer.kt
@@ -151,7 +151,7 @@
         if (enableRefactorTaskThumbnail()) {
             bindThumbnailView()
         } else {
-            thumbnailViewDeprecated.bind(task, overlay)
+            thumbnailViewDeprecated.bind(task, overlay, taskView)
         }
         overlay.init()
     }
diff --git a/quickstep/src/com/android/quickstep/views/TaskThumbnailViewDeprecated.java b/quickstep/src/com/android/quickstep/views/TaskThumbnailViewDeprecated.java
index 56ca043..5dbc2ef 100644
--- a/quickstep/src/com/android/quickstep/views/TaskThumbnailViewDeprecated.java
+++ b/quickstep/src/com/android/quickstep/views/TaskThumbnailViewDeprecated.java
@@ -110,6 +110,7 @@
     private TaskView.FullscreenDrawParams mFullscreenParams;
     private ImageView mSplashView;
     private Drawable mSplashViewDrawable;
+    private TaskView mTaskView;
 
     @Nullable
     private Task mTask;
@@ -153,10 +154,11 @@
     /**
      * Updates the thumbnail to draw the provided task
      */
-    public void bind(Task task, TaskOverlay<?> overlay) {
+    public void bind(Task task, TaskOverlay<?> overlay, TaskView taskView) {
         mOverlay = overlay;
         mOverlay.reset();
         mTask = task;
+        mTaskView = taskView;
         int color = task == null ? Color.BLACK : task.colorBackground | 0xFF000000;
         mPaint.setColor(color);
         mBackgroundPaint.setColor(color);
@@ -292,8 +294,8 @@
 
     public void drawOnCanvas(Canvas canvas, float x, float y, float width, float height,
             float cornerRadius) {
-        if (mTask != null && getTaskView().isRunningTask()
-                && !getTaskView().getShouldShowScreenshot()) {
+        if (mTask != null && mTaskView.isRunningTask()
+                && !mTaskView.getShouldShowScreenshot()) {
             canvas.drawRoundRect(x, y, width, height, cornerRadius, cornerRadius, mClearPaint);
             canvas.drawRoundRect(x, y, width, height, cornerRadius, cornerRadius,
                     mDimmingPaintAfterClearing);
@@ -334,10 +336,6 @@
         }
     }
 
-    public TaskView getTaskView() {
-        return (TaskView) getParent();
-    }
-
     public void setOverlayEnabled(boolean overlayEnabled) {
         if (mOverlayEnabled != overlayEnabled) {
             mOverlayEnabled = overlayEnabled;
@@ -390,9 +388,9 @@
         float viewCenterY = viewHeight / 2f;
         float centeredDrawableLeft = (viewWidth - drawableWidth) / 2f;
         float centeredDrawableTop = (viewHeight - drawableHeight) / 2f;
-        float nonGridScale = getTaskView() == null ? 1 : 1 / getTaskView().getNonGridScale();
-        float recentsMaxScale = getTaskView() == null || getTaskView().getRecentsView() == null
-                ? 1 : 1 / getTaskView().getRecentsView().getMaxScaleForFullScreen();
+        float nonGridScale = mTaskView == null ? 1 : 1 / mTaskView.getNonGridScale();
+        float recentsMaxScale = mTaskView == null || mTaskView.getRecentsView() == null
+                ? 1 : 1 / mTaskView.getRecentsView().getMaxScaleForFullScreen();
         float scaleX = nonGridScale * recentsMaxScale * (1 / getScaleX());
         float scaleY = nonGridScale * recentsMaxScale * (1 / getScaleY());
 
@@ -419,7 +417,7 @@
     }
 
     private boolean isThumbnailRotationDifferentFromTask() {
-        RecentsView recents = getTaskView().getRecentsView();
+        RecentsView recents = mTaskView.getRecentsView();
         if (recents == null || mThumbnailData == null) {
             return false;
         }
@@ -467,7 +465,7 @@
         if (mBitmapShader != null && mThumbnailData != null) {
             mPreviewRect.set(0, 0, mThumbnailData.getThumbnail().getWidth(),
                     mThumbnailData.getThumbnail().getHeight());
-            int currentRotation = getTaskView().getOrientedState().getRecentsActivityRotation();
+            int currentRotation = mTaskView.getOrientedState().getRecentsActivityRotation();
             boolean isRtl = getLayoutDirection() == LAYOUT_DIRECTION_RTL;
             mPreviewPositionHelper.updateThumbnailMatrix(mPreviewRect, mThumbnailData,
                     getMeasuredWidth(), getMeasuredHeight(), dp.isTablet, currentRotation, isRtl);
@@ -475,7 +473,7 @@
             mBitmapShader.setLocalMatrix(mPreviewPositionHelper.getMatrix());
             mPaint.setShader(mBitmapShader);
         }
-        getTaskView().updateCurrentFullscreenParams();
+        mTaskView.updateCurrentFullscreenParams();
         invalidate();
     }
 
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarAutohideSuspendControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarAutohideSuspendControllerTest.kt
index 46a7733..cfa12e2 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarAutohideSuspendControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarAutohideSuspendControllerTest.kt
@@ -24,7 +24,6 @@
 import com.android.launcher3.taskbar.rules.TaskbarModeRule
 import com.android.launcher3.taskbar.rules.TaskbarModeRule.Mode.TRANSIENT
 import com.android.launcher3.taskbar.rules.TaskbarModeRule.TaskbarMode
-import com.android.launcher3.taskbar.rules.TaskbarSandboxComponent
 import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule
 import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.InjectController
 import com.android.launcher3.taskbar.rules.TaskbarWindowSandboxContext
@@ -42,12 +41,11 @@
 
     @get:Rule(order = 0)
     val context =
-        TaskbarWindowSandboxContext.create {
-            builder: TaskbarSandboxComponent.Builder,
-            sandboxContext: TaskbarWindowSandboxContext ->
+        TaskbarWindowSandboxContext.create { builder ->
             builder.bindSystemUiProxy(
-                object : SystemUiProxy(sandboxContext) {
+                object : SystemUiProxy(this) {
                     override fun notifyTaskbarAutohideSuspend(suspend: Boolean) {
+                        super.notifyTaskbarAutohideSuspend(suspend)
                         latestSuspendNotification = suspend
                     }
                 }
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarScrimViewControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarScrimViewControllerTest.kt
index b1bb472..3912051 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarScrimViewControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarScrimViewControllerTest.kt
@@ -24,7 +24,6 @@
 import com.android.launcher3.taskbar.rules.TaskbarModeRule.Mode.PINNED
 import com.android.launcher3.taskbar.rules.TaskbarModeRule.Mode.TRANSIENT
 import com.android.launcher3.taskbar.rules.TaskbarModeRule.TaskbarMode
-import com.android.launcher3.taskbar.rules.TaskbarSandboxComponent
 import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule
 import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.InjectController
 import com.android.launcher3.taskbar.rules.TaskbarWindowSandboxContext
@@ -45,11 +44,9 @@
 class TaskbarScrimViewControllerTest {
     @get:Rule(order = 0)
     val context =
-        TaskbarWindowSandboxContext.create {
-            builder: TaskbarSandboxComponent.Builder,
-            sandboxContext: TaskbarWindowSandboxContext ->
+        TaskbarWindowSandboxContext.create { builder ->
             builder.bindSystemUiProxy(
-                object : SystemUiProxy(sandboxContext) {
+                object : SystemUiProxy(this) {
                     override fun onBackPressed() {
                         super.onBackPressed()
                         backPressed = true
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt
index b0d706f..096f879 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt
@@ -19,12 +19,10 @@
 import android.app.Instrumentation
 import android.app.PendingIntent
 import android.content.IIntentSender
-import android.content.Intent
-import android.provider.Settings
 import android.provider.Settings.Secure.NAV_BAR_KIDS_MODE
 import android.provider.Settings.Secure.USER_SETUP_COMPLETE
+import android.provider.Settings.Secure.getUriFor
 import androidx.test.platform.app.InstrumentationRegistry
-import androidx.test.rule.ServiceTestRule
 import com.android.launcher3.LauncherAppState
 import com.android.launcher3.statehandlers.DesktopVisibilityController
 import com.android.launcher3.taskbar.TaskbarActivityContext
@@ -35,16 +33,12 @@
 import com.android.launcher3.taskbar.bubbles.BubbleControllers
 import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.InjectController
 import com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR
-import com.android.launcher3.util.LauncherMultivalentJUnit.Companion.isRunningInRobolectric
 import com.android.launcher3.util.TestUtil
 import com.android.quickstep.AllAppsActionManager
-import com.android.quickstep.TouchInteractionService
-import com.android.quickstep.TouchInteractionService.TISBinder
 import java.lang.reflect.Field
 import java.lang.reflect.ParameterizedType
 import java.util.Optional
 import org.junit.Assume.assumeTrue
-import org.junit.rules.RuleChain
 import org.junit.rules.TestRule
 import org.junit.runner.Description
 import org.junit.runners.model.Statement
@@ -81,11 +75,6 @@
 ) : TestRule {
 
     private val instrumentation = InstrumentationRegistry.getInstrumentation()
-    private val serviceTestRule = ServiceTestRule()
-
-    private val userSetupCompleteRule = TaskbarSecureSettingRule(USER_SETUP_COMPLETE)
-    private val kidsModeRule = TaskbarSecureSettingRule(NAV_BAR_KIDS_MODE)
-    private val settingRules = RuleChain.outerRule(userSetupCompleteRule).around(kidsModeRule)
 
     private lateinit var taskbarManager: TaskbarManager
 
@@ -96,10 +85,6 @@
         }
 
     override fun apply(base: Statement, description: Description): Statement {
-        return settingRules.apply(createStatement(base, description), description)
-    }
-
-    private fun createStatement(base: Statement, description: Description): Statement {
         return object : Statement() {
             override fun evaluate() {
 
@@ -111,34 +96,10 @@
                 }
 
                 // Process secure setting annotations.
-                instrumentation.runOnMainSync {
-                    userSetupCompleteRule.putInt(
-                        if (description.getAnnotation(UserSetupMode::class.java) != null) {
-                            0
-                        } else {
-                            1
-                        }
-                    )
-                    kidsModeRule.putInt(
-                        if (description.getAnnotation(NavBarKidsMode::class.java) != null) 1 else 0
-                    )
-                }
-
-                // Check for existing Taskbar instance from Launcher process.
-                val launcherTaskbarManager: TaskbarManager? =
-                    if (!isRunningInRobolectric) {
-                        try {
-                            val tisBinder =
-                                serviceTestRule.bindService(
-                                    Intent(context, TouchInteractionService::class.java)
-                                ) as? TISBinder
-                            tisBinder?.taskbarManager
-                        } catch (_: Exception) {
-                            null
-                        }
-                    } else {
-                        null
-                    }
+                context.settingsCacheSandbox[getUriFor(USER_SETUP_COMPLETE)] =
+                    if (description.getAnnotation(UserSetupMode::class.java) != null) 0 else 1
+                context.settingsCacheSandbox[getUriFor(NAV_BAR_KIDS_MODE)] =
+                    if (description.getAnnotation(NavBarKidsMode::class.java) != null) 1 else 0
 
                 taskbarManager =
                     TestUtil.getOnUiThread {
@@ -161,20 +122,12 @@
                 try {
                     TaskbarViewController.enableModelLoadingForTests(false)
 
-                    // Replace Launcher Taskbar window with test instance.
-                    instrumentation.runOnMainSync {
-                        launcherTaskbarManager?.setSuspended(true)
-                        taskbarManager.onUserUnlocked() // Required to complete initialization.
-                    }
+                    // Required to complete initialization.
+                    instrumentation.runOnMainSync { taskbarManager.onUserUnlocked() }
 
                     base.evaluate()
                 } finally {
-                    // Revert Taskbar window.
-                    instrumentation.runOnMainSync {
-                        taskbarManager.destroy()
-                        launcherTaskbarManager?.setSuspended(false)
-                    }
-
+                    instrumentation.runOnMainSync { taskbarManager.destroy() }
                     TaskbarViewController.enableModelLoadingForTests(true)
                 }
             }
@@ -238,25 +191,4 @@
     @Retention(AnnotationRetention.RUNTIME)
     @Target(AnnotationTarget.CLASS, AnnotationTarget.FUNCTION)
     annotation class NavBarKidsMode
-
-    /** Rule for Taskbar integer-based secure settings. */
-    private inner class TaskbarSecureSettingRule(private val settingName: String) : TestRule {
-
-        override fun apply(base: Statement, description: Description): Statement {
-            return object : Statement() {
-                override fun evaluate() {
-                    val originalValue =
-                        Settings.Secure.getInt(context.contentResolver, settingName, /* def= */ 0)
-                    try {
-                        base.evaluate()
-                    } finally {
-                        instrumentation.runOnMainSync { putInt(originalValue) }
-                    }
-                }
-            }
-        }
-
-        /** Puts [value] into secure settings under [settingName]. */
-        fun putInt(value: Int) = Settings.Secure.putInt(context.contentResolver, settingName, value)
-    }
 }
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContext.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContext.kt
index a331a25..8c51216 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContext.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContext.kt
@@ -28,12 +28,20 @@
 import com.android.launcher3.dagger.LauncherAppSingleton
 import com.android.launcher3.util.MainThreadInitializedObject.ObjectSandbox
 import com.android.launcher3.util.SandboxApplication
+import com.android.launcher3.util.SettingsCache
+import com.android.launcher3.util.SettingsCacheSandbox
 import com.android.quickstep.SystemUiProxy
 import dagger.BindsInstance
 import dagger.Component
 import org.junit.rules.ExternalResource
 import org.junit.rules.RuleChain
 import org.junit.rules.TestRule
+import org.junit.runner.Description
+import org.junit.runners.model.Statement
+
+/** Include additional bindings when building a [TaskbarSandboxComponent]. */
+typealias TaskbarComponentBinder =
+    TaskbarWindowSandboxContext.(TaskbarSandboxComponent.Builder) -> Unit
 
 /**
  * [SandboxApplication] for running Taskbar tests.
@@ -42,27 +50,46 @@
  * [DEFAULT_DISPLAY] (i.e. test is executing on a device).
  */
 class TaskbarWindowSandboxContext
-private constructor(base: SandboxApplication, val virtualDisplay: VirtualDisplay) :
-    ContextWrapper(base),
-    ObjectSandbox by base,
-    TestRule by RuleChain.outerRule(virtualDisplayRule(virtualDisplay)).around(base) {
+private constructor(
+    private val base: SandboxApplication,
+    val virtualDisplay: VirtualDisplay,
+    private val componentBinder: TaskbarComponentBinder?,
+) : ContextWrapper(base), ObjectSandbox by base, TestRule {
 
-    init {
-        putObject(LauncherPrefs.INSTANCE, FakeLauncherPrefs(this))
+    val settingsCacheSandbox = SettingsCacheSandbox()
+
+    private val virtualDisplayRule =
+        object : ExternalResource() {
+            override fun after() = virtualDisplay.release()
+        }
+
+    private val singletonSetupRule =
+        object : ExternalResource() {
+            override fun before() {
+                val context = this@TaskbarWindowSandboxContext
+                val builder =
+                    DaggerTaskbarSandboxComponent.builder()
+                        .bindSystemUiProxy(SystemUiProxy(context))
+                        .bindSettingsCache(settingsCacheSandbox.cache)
+                componentBinder?.invoke(context, builder)
+                base.initDaggerComponent(builder)
+
+                putObject(LauncherPrefs.INSTANCE, FakeLauncherPrefs(context))
+            }
+        }
+
+    override fun apply(statement: Statement, description: Description): Statement {
+        return RuleChain.outerRule(virtualDisplayRule)
+            .around(base)
+            .around(singletonSetupRule)
+            .apply(statement, description)
     }
 
     companion object {
         private const val VIRTUAL_DISPLAY_NAME = "TaskbarSandboxDisplay"
 
         /** Creates a [SandboxApplication] for Taskbar tests. */
-        fun create(
-            daggerComponentBinder:
-                ((
-                    builder: TaskbarSandboxComponent.Builder,
-                    sandboxContext: TaskbarWindowSandboxContext,
-                ) -> TaskbarSandboxComponent.Builder)? =
-                null
-        ): TaskbarWindowSandboxContext {
+        fun create(componentBinder: TaskbarComponentBinder? = null): TaskbarWindowSandboxContext {
             val base = ApplicationProvider.getApplicationContext<Context>()
             val displayManager = checkNotNull(base.getSystemService(DisplayManager::class.java))
 
@@ -79,31 +106,15 @@
                     )
                 }
 
-            val sandboxApplication =
-                SandboxApplication(base.createDisplayContext(virtualDisplay.display))
-            val taskbarWindowSandboxContext =
-                TaskbarWindowSandboxContext(sandboxApplication, virtualDisplay)
-
-            if (daggerComponentBinder != null) {
-                sandboxApplication.initDaggerComponent(
-                    daggerComponentBinder(
-                        DaggerTaskbarSandboxComponent.builder(),
-                        taskbarWindowSandboxContext,
-                    )
-                )
-            }
-
-            return taskbarWindowSandboxContext
+            return TaskbarWindowSandboxContext(
+                SandboxApplication(base.createDisplayContext(virtualDisplay.display)),
+                virtualDisplay,
+                componentBinder,
+            )
         }
     }
 }
 
-private fun virtualDisplayRule(virtualDisplay: VirtualDisplay): TestRule {
-    return object : ExternalResource() {
-        override fun after() = virtualDisplay.release()
-    }
-}
-
 @LauncherAppSingleton
 @Component
 interface TaskbarSandboxComponent : LauncherAppComponent {
@@ -111,6 +122,8 @@
     interface Builder : LauncherAppComponent.Builder {
         @BindsInstance fun bindSystemUiProxy(proxy: SystemUiProxy): Builder
 
+        @BindsInstance fun bindSettingsCache(settingsCache: SettingsCache): Builder
+
         override fun build(): TaskbarSandboxComponent
     }
 }
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/util/SettingsCacheSandbox.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/util/SettingsCacheSandbox.kt
new file mode 100644
index 0000000..dcd5352
--- /dev/null
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/util/SettingsCacheSandbox.kt
@@ -0,0 +1,47 @@
+/*
+ * 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.util
+
+import android.net.Uri
+import org.mockito.kotlin.any
+import org.mockito.kotlin.doAnswer
+import org.mockito.kotlin.mock
+
+/**
+ * Provides a sandboxed [SettingsCache] for testing.
+ *
+ * Note that listeners registered to [cache] will never be invoked.
+ */
+class SettingsCacheSandbox {
+    private val values = mutableMapOf<Uri, Int>()
+
+    /** Fake cache that delegates [SettingsCache.getValue] to [values]. */
+    val cache =
+        mock<SettingsCache> {
+            on { getValue(any<Uri>()) } doAnswer { mock.getValue(it.getArgument(0), 1) }
+            on { getValue(any<Uri>(), any<Int>()) } doAnswer
+                {
+                    values.getOrDefault(it.getArgument(0), it.getArgument(1)) == 1
+                }
+        }
+
+    operator fun get(key: Uri): Int? = values[key]
+
+    operator fun set(key: Uri, value: Int) {
+        values[key] = value
+    }
+}
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/LauncherSwipeHandlerV2Test.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/LauncherSwipeHandlerV2Test.kt
index 6212b59..c3d865f 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/LauncherSwipeHandlerV2Test.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/LauncherSwipeHandlerV2Test.kt
@@ -20,10 +20,13 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.launcher3.R
-import com.android.launcher3.taskbar.rules.DaggerTaskbarSandboxComponent
+import com.android.launcher3.dagger.LauncherAppComponent
+import com.android.launcher3.dagger.LauncherAppSingleton
 import com.android.launcher3.util.LauncherModelHelper
 import com.android.systemui.contextualeducation.GestureType
 import com.android.systemui.shared.system.InputConsumerController
+import dagger.BindsInstance
+import dagger.Component
 import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
@@ -60,7 +63,7 @@
     @Before
     fun setup() {
         sandboxContext.initDaggerComponent(
-            DaggerTaskbarSandboxComponent.builder().bindSystemUiProxy(systemUiProxy)
+            DaggerTestComponent.builder().bindSystemUiProxy(systemUiProxy)
         )
         val deviceState = mock(RecentsAnimationDeviceState::class.java)
         whenever(deviceState.rotationTouchHelper).thenReturn(mock(RotationTouchHelper::class.java))
@@ -100,3 +103,14 @@
             )
     }
 }
+
+@LauncherAppSingleton
+@Component
+interface TestComponent : LauncherAppComponent {
+    @Component.Builder
+    interface Builder : LauncherAppComponent.Builder {
+        @BindsInstance fun bindSystemUiProxy(proxy: SystemUiProxy): Builder
+
+        override fun build(): TestComponent
+    }
+}
diff --git a/quickstep/tests/src/com/android/launcher3/taskbar/FallbackTaskbarUIControllerTest.kt b/quickstep/tests/src/com/android/launcher3/taskbar/FallbackTaskbarUIControllerTest.kt
index 04012c0..df98606 100644
--- a/quickstep/tests/src/com/android/launcher3/taskbar/FallbackTaskbarUIControllerTest.kt
+++ b/quickstep/tests/src/com/android/launcher3/taskbar/FallbackTaskbarUIControllerTest.kt
@@ -33,7 +33,7 @@
 @RunWith(AndroidJUnit4::class)
 class FallbackTaskbarUIControllerTest : TaskbarBaseTestCase() {
 
-    lateinit var fallbackTaskbarUIController: FallbackTaskbarUIController
+    lateinit var fallbackTaskbarUIController: FallbackTaskbarUIController<RecentsActivity>
     lateinit var stateListener: StateManager.StateListener<RecentsState>
 
     private val recentsActivity: RecentsActivity = mock()
diff --git a/res/values-fr-rCA/strings.xml b/res/values-fr-rCA/strings.xml
index cae77dc..a941d88 100644
--- a/res/values-fr-rCA/strings.xml
+++ b/res/values-fr-rCA/strings.xml
@@ -192,7 +192,7 @@
     <string name="private_space_label" msgid="2359721649407947001">"Espace privé"</string>
     <string name="private_space_secondary_label" msgid="9203933341714508907">"Touchez pour configurer ou ouvrir"</string>
     <string name="ps_container_title" msgid="4391796149519594205">"Privé"</string>
-    <string name="ps_container_settings" msgid="6059734123353320479">"Paramètres de l\'Espace privé"</string>
+    <string name="ps_container_settings" msgid="6059734123353320479">"Paramètres de l\'espace privé"</string>
     <string name="ps_container_unlock_button_content_description" msgid="9181551784092204234">"Privé, déverrouillé."</string>
     <string name="ps_container_lock_button_content_description" msgid="5961993384382649530">"Privé, verrouillé."</string>
     <string name="ps_container_lock_title" msgid="2640257399982364682">"Verrouiller"</string>
diff --git a/res/values-it/strings.xml b/res/values-it/strings.xml
index 731f839..9cdb5aa 100644
--- a/res/values-it/strings.xml
+++ b/res/values-it/strings.xml
@@ -192,7 +192,7 @@
     <string name="private_space_label" msgid="2359721649407947001">"Spazio privato"</string>
     <string name="private_space_secondary_label" msgid="9203933341714508907">"Tocca per configurare o aprire"</string>
     <string name="ps_container_title" msgid="4391796149519594205">"Privato"</string>
-    <string name="ps_container_settings" msgid="6059734123353320479">"Impostazioni dello Spazio privato"</string>
+    <string name="ps_container_settings" msgid="6059734123353320479">"Impostazioni dello spazio privato"</string>
     <string name="ps_container_unlock_button_content_description" msgid="9181551784092204234">"Privato, sbloccato."</string>
     <string name="ps_container_lock_button_content_description" msgid="5961993384382649530">"Privato, bloccato."</string>
     <string name="ps_container_lock_title" msgid="2640257399982364682">"Blocca"</string>
diff --git a/res/values/attrs.xml b/res/values/attrs.xml
index 57c9bc7..16ea0cd 100644
--- a/res/values/attrs.xml
+++ b/res/values/attrs.xml
@@ -208,8 +208,16 @@
         <attr name="layout_sticky" format="boolean" />
     </declare-styleable>
 
+    <declare-styleable name="NumRows">
+        <attr name="minDeviceWidthPx" format="float"/>
+        <attr name="minDeviceHeightPx" format="float"/>
+        <attr name="numRowsNew" format="integer"/>
+        <attr name="dbFile" />
+    </declare-styleable>
+
     <declare-styleable name="GridDisplayOption">
         <attr name="name" format="string" />
+        <attr name="title" />
 
         <attr name="numRows" format="integer" />
         <attr name="numColumns" format="integer" />
@@ -294,6 +302,7 @@
         <!-- File that contains the specs for all apps icon and text size.
         Needs FeatureFlags.ENABLE_RESPONSIVE_WORKSPACE enabled -->
         <attr name="allAppsCellSpecsId" format="reference" />
+        <attr name="rowCountSpecsId" format="reference" />
         <!-- defaults to allAppsCellSpecsId, if not specified -->
         <attr name="allAppsCellSpecsTwoPanelId" format="reference" />
 
diff --git a/res/values/strings.xml b/res/values/strings.xml
index 9d06021..d918698 100644
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -355,8 +355,9 @@
     <!-- Title for an app whose download has been started. -->
     <string name="app_waiting_download_title"><xliff:g id="name" example="Messenger">%1$s</xliff:g> waiting to install</string>
     <!-- Title for an app which is archived. -->
-    <string name="app_archived_title"><xliff:g id="name" example="Messenger">%1$s</xliff:g> is archived. Tap to download and restore.</string>
-
+    <string name="app_archived_title"><xliff:g id="name" example="Messenger">%1$s</xliff:g> is archived.</string>
+    <!-- Accessibility Action for an app which is archived. -->
+    <string name="app_unarchiving_action">download and restore</string>
 
     <!-- Title shown on the alert dialog prompting the user to update the application in market
      in order to re-enable the disabled shortcuts -->
diff --git a/res/xml/backupscheme.xml b/res/xml/backupscheme.xml
index 0f0dde2..27fddc8 100644
--- a/res/xml/backupscheme.xml
+++ b/res/xml/backupscheme.xml
@@ -3,6 +3,8 @@
 
     <include domain="database" path="launcher.db" />
     <include domain="database" path="launcher_6_by_5.db" />
+    <include domain="database" path="launcher_5_by_6.db" />
+    <include domain="database" path="launcher_4_by_6.db" />
     <include domain="database" path="launcher_4_by_5.db" />
     <include domain="database" path="launcher_4_by_4.db" />
     <include domain="database" path="launcher_3_by_3.db" />
diff --git a/src/com/android/launcher3/BubbleTextView.java b/src/com/android/launcher3/BubbleTextView.java
index 909272e..34cf56b 100644
--- a/src/com/android/launcher3/BubbleTextView.java
+++ b/src/com/android/launcher3/BubbleTextView.java
@@ -57,12 +57,14 @@
 import android.view.MotionEvent;
 import android.view.View;
 import android.view.ViewDebug;
+import android.view.accessibility.AccessibilityNodeInfo;
 import android.widget.TextView;
 
 import androidx.annotation.DrawableRes;
 import androidx.annotation.Nullable;
 import androidx.annotation.UiThread;
 import androidx.annotation.VisibleForTesting;
+import androidx.core.view.accessibility.AccessibilityNodeInfoCompat;
 
 import com.android.launcher3.accessibility.BaseAccessibilityDelegate;
 import com.android.launcher3.dot.DotInfo;
@@ -519,6 +521,16 @@
         }
     }
 
+    @Override
+    public void onInitializeAccessibilityNodeInfo(AccessibilityNodeInfo info) {
+        super.onInitializeAccessibilityNodeInfo(info);
+        if (getTag() instanceof ItemInfoWithIcon infoWithIcon && infoWithIcon.isInactiveArchive()) {
+            info.addAction(new AccessibilityNodeInfo.AccessibilityAction(
+                    AccessibilityNodeInfoCompat.ACTION_CLICK,
+                    getContext().getString(R.string.app_unarchiving_action)));
+        }
+    }
+
     /** This is used for testing to forcefully set the display. */
     @VisibleForTesting
     public void setDisplay(int display) {
diff --git a/src/com/android/launcher3/InvariantDeviceProfile.java b/src/com/android/launcher3/InvariantDeviceProfile.java
index ece6540..7acba75 100644
--- a/src/com/android/launcher3/InvariantDeviceProfile.java
+++ b/src/com/android/launcher3/InvariantDeviceProfile.java
@@ -58,9 +58,9 @@
 import com.android.launcher3.testing.shared.ResourceUtils;
 import com.android.launcher3.util.DisplayController;
 import com.android.launcher3.util.DisplayController.Info;
-import com.android.launcher3.util.LockedUserState;
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.util.Partner;
+import com.android.launcher3.util.ResourceHelper;
 import com.android.launcher3.util.SafeCloseable;
 import com.android.launcher3.util.WindowBounds;
 import com.android.launcher3.util.window.WindowManagerProxy;
@@ -186,6 +186,8 @@
     @XmlRes
     public int workspaceSpecsId = INVALID_RESOURCE_HANDLE;
     @XmlRes
+    public int rowCountSpecsId = INVALID_RESOURCE_HANDLE;;
+    @XmlRes
     public int workspaceSpecsTwoPanelId = INVALID_RESOURCE_HANDLE;
     @XmlRes
     public int allAppsSpecsId = INVALID_RESOURCE_HANDLE;
@@ -232,8 +234,6 @@
         if (!newGridName.equals(gridName)) {
             LauncherPrefs.get(context).put(GRID_NAME, newGridName);
         }
-        LockedUserState.get(context).runOnUserUnlocked(() ->
-            new DeviceGridState(this).writeToPrefs(context));
 
         DisplayController.INSTANCE.get(context).setPriorityListener(
                 (displayContext, info, flags) -> {
@@ -340,15 +340,29 @@
         Info displayInfo = DisplayController.INSTANCE.get(context).getInfo();
         @DeviceType int deviceType = displayInfo.getDeviceType();
 
-        ArrayList<DisplayOption> allOptions =
+        List<DisplayOption> allOptions =
                 getPredefinedDeviceProfiles(context, gridName, deviceType,
                         RestoreDbTask.isPending(context));
+
+        // Filter out options that don't have the same number of columns as the grid
+        DeviceGridState deviceGridState = new DeviceGridState(context);
+        List<DisplayOption> allOptionsFilteredByColCount =
+                filterByColumnCount(allOptions, deviceGridState.getColumns());
+
         DisplayOption displayOption =
-                invDistWeightedInterpolate(displayInfo, allOptions, deviceType);
+                invDistWeightedInterpolate(displayInfo, allOptionsFilteredByColCount.isEmpty()
+                        ? new ArrayList<>(allOptions)
+                        : new ArrayList<>(allOptionsFilteredByColCount), deviceType);
         initGrid(context, displayInfo, displayOption, deviceType);
         return displayOption.grid.name;
     }
 
+    private List<DisplayOption> filterByColumnCount(
+            List<DisplayOption> allOptions, int numColumns) {
+        return allOptions.stream().filter(
+                option -> option.grid.numColumns == numColumns).toList();
+    }
+
     /**
      * @deprecated This is a temporary solution because on the backup and restore case we modify the
      * IDP, this resets it. b/332974074
@@ -383,6 +397,7 @@
         isScalable = closestProfile.isScalable;
         devicePaddingId = closestProfile.devicePaddingId;
         workspaceSpecsId = closestProfile.mWorkspaceSpecsId;
+        rowCountSpecsId = closestProfile.mRowCountSpecsId;
         workspaceSpecsTwoPanelId = closestProfile.mWorkspaceSpecsTwoPanelId;
         allAppsSpecsId = closestProfile.mAllAppsSpecsId;
         allAppsSpecsTwoPanelId = closestProfile.mAllAppsSpecsTwoPanelId;
@@ -495,7 +510,6 @@
         mChangeListeners.remove(listener);
     }
 
-
     public void setCurrentGrid(Context context, String gridName) {
         LauncherPrefs.get(context).put(GRID_NAME, gridName);
         MAIN_EXECUTOR.execute(() -> {
@@ -526,7 +540,7 @@
         }
     }
 
-    private static ArrayList<DisplayOption> getPredefinedDeviceProfiles(Context context,
+    private List<DisplayOption> getPredefinedDeviceProfiles(Context context,
             String gridName, @DeviceType int deviceType, boolean allowDisabledGrid) {
         ArrayList<DisplayOption> profiles = new ArrayList<>();
 
@@ -539,7 +553,8 @@
                         && GridOption.TAG_NAME.equals(parser.getName())) {
 
                     GridOption gridOption = new GridOption(context, Xml.asAttributeSet(parser));
-                    if (gridOption.isEnabled(deviceType) || allowDisabledGrid) {
+                    if ((gridOption.isEnabled(deviceType) || allowDisabledGrid)
+                            && (Flags.oneGridSpecs() == gridOption.isNewGridOption())) {
                         final int displayDepth = parser.getDepth();
                         while (((type = parser.next()) != XmlPullParser.END_TAG
                                 || parser.getDepth() > displayDepth)
@@ -566,13 +581,16 @@
                 }
             }
         }
-        if (filteredProfiles.isEmpty()) {
-            // No grid found, use the default options
+        if (filteredProfiles.isEmpty() && TextUtils.isEmpty(gridName)) {
+            // Use the default options since gridName is empty and there's no valid grids.
             for (DisplayOption option : profiles) {
                 if (option.canBeDefault) {
                     filteredProfiles.add(option);
                 }
             }
+        } else if (filteredProfiles.isEmpty()) {
+            // In this case we had a grid selected but we couldn't find it.
+            filteredProfiles.addAll(profiles);
         }
         if (filteredProfiles.isEmpty()) {
             throw new RuntimeException("No display option with canBeDefault=true");
@@ -581,6 +599,72 @@
     }
 
     /**
+     * Parses through the xml to find NumRows specs. Then calls findBestRowCount to get the correct
+     * row count for this GridOption.
+     *
+     * @return the result of {@link #findBestRowCount(List, Context, int)}.
+     */
+    public static NumRows getRowCount(ResourceHelper resourceHelper, Context context,
+            int deviceType) {
+        ArrayList<NumRows> rowCounts = new ArrayList<>();
+
+        try (XmlResourceParser parser = resourceHelper.getXml()) {
+            final int depth = parser.getDepth();
+            int type;
+            while (((type = parser.next()) != XmlPullParser.END_TAG
+                    || parser.getDepth() > depth) && type != XmlPullParser.END_DOCUMENT) {
+                if ((type == XmlPullParser.START_TAG)
+                        && "NumRows".equals(parser.getName())) {
+                    rowCounts.add(new NumRows(context, Xml.asAttributeSet(parser)));
+                }
+            }
+        } catch (IOException | XmlPullParserException e) {
+            throw new RuntimeException(e);
+        }
+
+        return findBestRowCount(rowCounts, context, deviceType);
+    }
+
+    /**
+     * @return the biggest row count that fits the display dimensions spec using NumRows to
+     * determine that. If no best row count is found, return -1.
+     */
+    public static NumRows findBestRowCount(List<NumRows> list, Context context,
+            @DeviceType int deviceType) {
+        Info displayInfo = DisplayController.INSTANCE.get(context).getInfo();
+        int minWidthPx = Integer.MAX_VALUE;
+        int minHeightPx = Integer.MAX_VALUE;
+        for (WindowBounds bounds : displayInfo.supportedBounds) {
+            boolean isTablet = displayInfo.isTablet(bounds);
+            if (isTablet && deviceType == TYPE_MULTI_DISPLAY) {
+                // For split displays, take half width per page
+                minWidthPx = Math.min(minWidthPx, bounds.availableSize.x / 2);
+                minHeightPx = Math.min(minHeightPx, bounds.availableSize.y);
+            } else if (!isTablet && bounds.isLandscape()) {
+                // We will use transposed layout in this case
+                minWidthPx = Math.min(minWidthPx, bounds.availableSize.y);
+                minHeightPx = Math.min(minHeightPx, bounds.availableSize.x);
+            } else {
+                minWidthPx = Math.min(minWidthPx, bounds.availableSize.x);
+                minHeightPx = Math.min(minHeightPx, bounds.availableSize.y);
+            }
+        }
+
+        NumRows selectedRow = null;
+        for (NumRows item: list) {
+            if (minWidthPx >= item.mMinDeviceWidthPx && minHeightPx >= item.mMinDeviceHeightPx) {
+                if (selectedRow == null || selectedRow.mNumRowsNew < item.mNumRowsNew) {
+                    selectedRow = item;
+                }
+            }
+        }
+        if (selectedRow != null) {
+            return selectedRow;
+        }
+        return null;
+    }
+
+    /**
      * Returns the GridOption associated to the given file name or null if the fileName is not
      * supported.
      * Ej, launcher.db -> "normal grid", launcher_4_by_4.db -> "practical grid"
@@ -626,6 +710,7 @@
         return parseAllDefinedGridOptions(context)
                 .stream()
                 .filter(go -> go.isEnabled(deviceType))
+                .filter(go -> (Flags.oneGridSpecs() == go.isNewGridOption()))
                 .collect(Collectors.toList());
     }
 
@@ -709,7 +794,7 @@
     }
 
     private static DisplayOption invDistWeightedInterpolate(
-            Info displayInfo, ArrayList<DisplayOption> points, @DeviceType int deviceType) {
+            Info displayInfo, List<DisplayOption> points, @DeviceType int deviceType) {
         int minWidthPx = Integer.MAX_VALUE;
         int minHeightPx = Integer.MAX_VALUE;
         for (WindowBounds bounds : displayInfo.supportedBounds) {
@@ -733,7 +818,7 @@
         float height = dpiFromPx(minHeightPx, displayInfo.getDensityDpi());
 
         // Sort the profiles based on the closeness to the device size
-        Collections.sort(points, (a, b) ->
+        points.sort((a, b) ->
                 Float.compare(dist(width, height, a.minWidthDps, a.minHeightDps),
                         dist(width, height, b.minWidthDps, b.minHeightDps)));
 
@@ -855,6 +940,7 @@
         private static final int DONT_INLINE_QSB = 0;
 
         public final String name;
+        public final String title;
         public final int numRows;
         public final int numColumns;
         public final int numSearchContainerColumns;
@@ -894,17 +980,30 @@
         private final int mWorkspaceCellSpecsTwoPanelId;
         private final int mAllAppsCellSpecsId;
         private final int mAllAppsCellSpecsTwoPanelId;
+        private final int mRowCountSpecsId;
 
         public GridOption(Context context, AttributeSet attrs) {
             TypedArray a = context.obtainStyledAttributes(
                     attrs, R.styleable.GridDisplayOption);
             name = a.getString(R.styleable.GridDisplayOption_name);
-            numRows = a.getInt(R.styleable.GridDisplayOption_numRows, 0);
+            title = a.getString(R.styleable.GridDisplayOption_title);
+            deviceCategory = a.getInt(R.styleable.GridDisplayOption_deviceCategory,
+                    DEVICE_CATEGORY_ALL);
+            mRowCountSpecsId = a.getResourceId(
+                    R.styleable.GridDisplayOption_rowCountSpecsId, INVALID_RESOURCE_HANDLE);
+            if (mRowCountSpecsId != INVALID_RESOURCE_HANDLE) {
+                ResourceHelper resourceHelper = new ResourceHelper(context, mRowCountSpecsId);
+                NumRows numR = getRowCount(resourceHelper, context, deviceCategory);
+                numRows = numR.mNumRowsNew;
+                dbFile = numR.mDbFile;
+            } else {
+                numRows = a.getInt(R.styleable.GridDisplayOption_numRows, 0);
+                dbFile = a.getString(R.styleable.GridDisplayOption_dbFile);
+            }
+
             numColumns = a.getInt(R.styleable.GridDisplayOption_numColumns, 0);
             numSearchContainerColumns = a.getInt(
                     R.styleable.GridDisplayOption_numSearchContainerColumns, numColumns);
-
-            dbFile = a.getString(R.styleable.GridDisplayOption_dbFile);
             defaultLayoutId = a.getResourceId(
                     R.styleable.GridDisplayOption_defaultLayoutId, 0);
             demoModeLayoutId = a.getResourceId(
@@ -969,8 +1068,6 @@
                     R.styleable.GridDisplayOption_isScalable, false);
             devicePaddingId = a.getResourceId(
                     R.styleable.GridDisplayOption_devicePaddingId, INVALID_RESOURCE_HANDLE);
-            deviceCategory = a.getInt(R.styleable.GridDisplayOption_deviceCategory,
-                    DEVICE_CATEGORY_ALL);
 
             if (FeatureFlags.enableResponsiveWorkspace()) {
                 mWorkspaceSpecsId = a.getResourceId(
@@ -1053,6 +1150,28 @@
                     return false;
             }
         }
+
+        public boolean isNewGridOption() {
+            return mRowCountSpecsId != INVALID_RESOURCE_HANDLE;
+        }
+    }
+
+    public static final class NumRows {
+        final int mNumRowsNew;
+        final float mMinDeviceWidthPx;
+        final float mMinDeviceHeightPx;
+        final String mDbFile;
+
+        NumRows(Context context, AttributeSet attrs) {
+            TypedArray a = context.obtainStyledAttributes(attrs, R.styleable.NumRows);
+
+            mNumRowsNew = (int) a.getFloat(R.styleable.NumRows_numRowsNew, 0);
+            mMinDeviceWidthPx = a.getFloat(R.styleable.NumRows_minDeviceWidthPx, 0);
+            mMinDeviceHeightPx = a.getFloat(R.styleable.NumRows_minDeviceHeightPx, 0);
+            mDbFile = a.getString(R.styleable.NumRows_dbFile);
+
+            a.recycle();
+        }
     }
 
     @VisibleForTesting
diff --git a/src/com/android/launcher3/LauncherFiles.java b/src/com/android/launcher3/LauncherFiles.java
index d730cea..1148f79 100644
--- a/src/com/android/launcher3/LauncherFiles.java
+++ b/src/com/android/launcher3/LauncherFiles.java
@@ -18,6 +18,8 @@
     public static final String LAUNCHER_DB = "launcher.db";
     public static final String LAUNCHER_6_BY_5_DB = "launcher_6_by_5.db";
     public static final String LAUNCHER_4_BY_5_DB = "launcher_4_by_5.db";
+    public static final String LAUNCHER_4_BY_6_DB = "launcher_4_by_6.db";
+    public static final String LAUNCHER_5_BY_6_DB = "launcher_5_by_6.db";
     public static final String LAUNCHER_4_BY_4_DB = "launcher_4_by_4.db";
     public static final String LAUNCHER_3_BY_3_DB = "launcher_3_by_3.db";
     public static final String LAUNCHER_2_BY_2_DB = "launcher_2_by_2.db";
@@ -35,6 +37,8 @@
             LAUNCHER_DB,
             LAUNCHER_6_BY_5_DB,
             LAUNCHER_4_BY_5_DB,
+            LAUNCHER_4_BY_6_DB,
+            LAUNCHER_5_BY_6_DB,
             LAUNCHER_4_BY_4_DB,
             LAUNCHER_3_BY_3_DB,
             LAUNCHER_2_BY_2_DB));
diff --git a/src/com/android/launcher3/allapps/AllAppsRecyclerView.java b/src/com/android/launcher3/allapps/AllAppsRecyclerView.java
index e705d94..51d1c9f 100644
--- a/src/com/android/launcher3/allapps/AllAppsRecyclerView.java
+++ b/src/com/android/launcher3/allapps/AllAppsRecyclerView.java
@@ -331,6 +331,9 @@
 
     public void setLettersToScrollLayout(
             List<AlphabeticalAppsList.FastScrollSectionInfo> fastScrollSections) {
+        if (fastScrollSections.isEmpty()) {
+            return;
+        }
         if (mLetterList != null) {
             mLetterList.removeAllViews();
         }
@@ -364,6 +367,8 @@
         mLetterList.addView(lastLetterListTextView);
         constraintTextViewsVertically(mLetterList, textViews);
         mLetterList.setVisibility(VISIBLE);
+        // Set the alpha to 0 to avoid the letter list being shown when it shouldn't be.
+        mLetterList.setAlpha(0);
     }
 
     private void constraintTextViewsVertically(ConstraintLayout constraintLayout,
diff --git a/src/com/android/launcher3/dagger/LauncherBaseAppComponent.java b/src/com/android/launcher3/dagger/LauncherBaseAppComponent.java
index 5d0c44e..0fa275e 100644
--- a/src/com/android/launcher3/dagger/LauncherBaseAppComponent.java
+++ b/src/com/android/launcher3/dagger/LauncherBaseAppComponent.java
@@ -18,6 +18,7 @@
 
 import android.content.Context;
 
+import com.android.launcher3.graphics.IconShape;
 import com.android.launcher3.pm.InstallSessionHelper;
 import com.android.launcher3.util.ApiWrapper;
 import com.android.launcher3.util.DaggerSingletonTracker;
@@ -40,12 +41,13 @@
  */
 public interface LauncherBaseAppComponent {
     DaggerSingletonTracker getDaggerSingletonTracker();
-    RefreshRateTracker getRefreshRateTracker();
-    InstallSessionHelper getInstallSessionHelper();
     ApiWrapper getApiWrapper();
+    CustomWidgetManager getCustomWidgetManager();
+    IconShape getIconShape();
+    InstallSessionHelper getInstallSessionHelper();
+    RefreshRateTracker getRefreshRateTracker();
     ScreenOnTracker getScreenOnTracker();
     SettingsCache getSettingsCache();
-    CustomWidgetManager getCustomWidgetManager();
     PluginManagerWrapper getPluginManagerWrapper();
     PackageManagerHelper getPackageManagerHelper();
 
diff --git a/src/com/android/launcher3/graphics/GridCustomizationsProvider.java b/src/com/android/launcher3/graphics/GridCustomizationsProvider.java
index bc51a66..a5bcd0f 100644
--- a/src/com/android/launcher3/graphics/GridCustomizationsProvider.java
+++ b/src/com/android/launcher3/graphics/GridCustomizationsProvider.java
@@ -36,24 +36,18 @@
 import android.text.TextUtils;
 import android.util.Log;
 
-import androidx.annotation.NonNull;
-
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.InvariantDeviceProfile.GridOption;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherModel;
 import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.model.BgDataModel;
-import com.android.launcher3.shapes.AppShape;
-import com.android.launcher3.shapes.AppShapesProvider;
 import com.android.launcher3.util.Executors;
 import com.android.launcher3.util.Preconditions;
 import com.android.launcher3.util.RunnableList;
 import com.android.systemui.shared.Flags;
 
 import java.util.Collections;
-import java.util.List;
-import java.util.Optional;
 import java.util.Set;
 import java.util.WeakHashMap;
 import java.util.concurrent.ExecutionException;
@@ -61,44 +55,32 @@
 /**
  * Exposes various launcher grid options and allows the caller to change them.
  * APIs:
- *      /shape_options: List of various available shape options, where each has following fields
- *          shape_key: key of the shape option
- *          title: translated title of the shape option
- *          path: path of the shape, assuming drawn on 100x100 view port
- *          is_default: true if this shape option is currently set to the system
- *
- *      /grid_options: List the various available grid options, where each has following fields
- *          name: key of the grid option
+ *      /list_options: List the various available grip options, has following columns
+ *          name: name of the grid
  *          rows: number of rows in the grid
  *          cols: number of columns in the grid
  *          preview_count: number of previews available for this grid option. The preview uri
  *                         looks like /preview/<grid-name>/<preview index starting with 0>
- *          is_default: true if this grid option is currently set to the system
+ *          is_default: true if this grid is currently active
  *
- *     /get_preview: Open a file stream for the grid preview
+ *     /preview: Opens a file stream for the grid preview
  *
- *     /default_grid: Call update to set the current shape and grid, with values
- *          shape_key: key of the shape to apply
- *          name: key of the grid to apply
+ *     /default_grid: Call update to set the current grid, with values
+ *          name: name of the grid to apply
  */
 public class GridCustomizationsProvider extends ContentProvider {
 
     private static final String TAG = "GridCustomizationsProvider";
 
-    private static final String KEY_SHAPE_KEY = "shape_key";
-    private static final String KEY_TITLE = "title";
-    private static final String KEY_PATH = "path";
-    // is_default means if a certain option is currently set to the system
-    private static final String KEY_IS_DEFAULT = "is_default";
-    // Key of grid option. We do not change the name to grid_key for backward compatibility
-    private static final String KEY_GRID_KEY = "name";
+    private static final String KEY_NAME = "name";
+    private static final String KEY_GRID_TITLE = "grid_title";
     private static final String KEY_ROWS = "rows";
     private static final String KEY_COLS = "cols";
     private static final String KEY_PREVIEW_COUNT = "preview_count";
+    private static final String KEY_IS_DEFAULT = "is_default";
 
-    private static final String KEY_SHAPE_OPTIONS = "/shape_options";
-    private static final String KEY_GRID_OPTIONS = "/grid_options";
-    private static final String KEY_SHAPE_GRID = "/default_grid";
+    private static final String KEY_LIST_OPTIONS = "/list_options";
+    private static final String KEY_DEFAULT_GRID = "/default_grid";
 
     private static final String METHOD_GET_PREVIEW = "get_preview";
 
@@ -110,9 +92,9 @@
     private static final String KEY_SURFACE_PACKAGE = "surface_package";
     private static final String KEY_CALLBACK = "callback";
     public static final String KEY_HIDE_BOTTOM_ROW = "hide_bottom_row";
+    public static final String KEY_GRID_NAME = "grid_name";
 
     private static final int MESSAGE_ID_UPDATE_PREVIEW = 1337;
-    private static final int MESSAGE_ID_UPDATE_SHAPE = 2586;
     private static final int MESSAGE_ID_UPDATE_GRID = 7414;
     private static final int MESSAGE_ID_UPDATE_COLOR = 856;
 
@@ -128,39 +110,15 @@
     @Override
     public Cursor query(Uri uri, String[] projection, String selection,
             String[] selectionArgs, String sortOrder) {
-        Context context = getContext();
-        String path = uri.getPath();
-        if (context == null || path == null) {
-            return null;
-        }
-        switch (path) {
-            case KEY_SHAPE_OPTIONS: {
-                if (Flags.newCustomizationPickerUi()) {
-                    MatrixCursor cursor = new MatrixCursor(new String[]{
-                            KEY_SHAPE_KEY, KEY_TITLE, KEY_PATH, KEY_IS_DEFAULT});
-                    List<AppShape> shapes =  AppShapesProvider.INSTANCE.getShapes();
-                    for (int i = 0; i < shapes.size(); i++) {
-                        AppShape shape = shapes.get(i);
-                        cursor.newRow()
-                                .add(KEY_SHAPE_KEY, shape.getKey())
-                                .add(KEY_TITLE, shape.getTitle())
-                                .add(KEY_PATH, shape.getPath())
-                                // TODO (b/348664593): We should fetch the currently-set shape
-                                //  option from the preferences.
-                                .add(KEY_IS_DEFAULT, i == 0);
-                    }
-                    return cursor;
-                } else  {
-                    return null;
-                }
-            }
-            case KEY_GRID_OPTIONS: {
+        switch (uri.getPath()) {
+            case KEY_LIST_OPTIONS: {
                 MatrixCursor cursor = new MatrixCursor(new String[]{
-                        KEY_GRID_KEY, KEY_ROWS, KEY_COLS, KEY_PREVIEW_COUNT, KEY_IS_DEFAULT});
-                InvariantDeviceProfile idp = InvariantDeviceProfile.INSTANCE.get(context);
-                for (GridOption gridOption : idp.parseAllGridOptions(context)) {
+                        KEY_NAME, KEY_ROWS, KEY_COLS, KEY_PREVIEW_COUNT, KEY_IS_DEFAULT});
+                InvariantDeviceProfile idp = InvariantDeviceProfile.INSTANCE.get(getContext());
+                for (GridOption gridOption : idp.parseAllGridOptions(getContext())) {
                     cursor.newRow()
-                            .add(KEY_GRID_KEY, gridOption.name)
+                            .add(KEY_NAME, gridOption.name)
+                            .add(KEY_GRID_TITLE, gridOption.title)
                             .add(KEY_ROWS, gridOption.numRows)
                             .add(KEY_COLS, gridOption.numColumns)
                             .add(KEY_PREVIEW_COUNT, 1)
@@ -203,22 +161,14 @@
             return 0;
         }
         switch (path) {
-            case KEY_SHAPE_GRID: {
-                if (Flags.newCustomizationPickerUi()) {
-                    String shapeKey = values.getAsString(KEY_SHAPE_KEY);
-                    Optional<AppShape> optionalShape = AppShapesProvider.INSTANCE.getShapes()
-                            .stream().filter(shape -> shape.getKey().equals(shapeKey)).findFirst();
-                    String pathToSet = optionalShape.map(AppShape::getPath).orElse(null);
-                    // TODO (b/348664593): Apply shapeName to the system. This needs to be a
-                    //  synchronous call.
-                }
-                String gridKey = values.getAsString(KEY_GRID_KEY);
+            case KEY_DEFAULT_GRID: {
+                String gridName = values.getAsString(KEY_NAME);
                 InvariantDeviceProfile idp = InvariantDeviceProfile.INSTANCE.get(context);
                 // Verify that this is a valid grid option
                 GridOption match = null;
                 for (GridOption option : idp.parseAllGridOptions(context)) {
                     String name = option.name;
-                    if (name != null && name.equals(gridKey)) {
+                    if (name != null && name.equals(gridName)) {
                         match = option;
                         break;
                     }
@@ -227,7 +177,7 @@
                     return 0;
                 }
 
-                idp.setCurrentGrid(context, gridKey);
+                idp.setCurrentGrid(context, gridName);
                 if (Flags.newCustomizationPickerUi()) {
                     try {
                         // Wait for device profile to be fully reloaded and applied to the launcher
@@ -269,30 +219,20 @@
     }
 
     @Override
-    public Bundle call(@NonNull String method, String arg, Bundle extras) {
-        Context context = getContext();
-        if (context == null) {
-            return null;
-        }
-
-        if (context.checkPermission("android.permission.BIND_WALLPAPER",
+    public Bundle call(String method, String arg, Bundle extras) {
+        if (getContext().checkPermission("android.permission.BIND_WALLPAPER",
                 Binder.getCallingPid(), Binder.getCallingUid())
                 != PackageManager.PERMISSION_GRANTED) {
             return null;
         }
 
-        if (METHOD_GET_PREVIEW.equals(method)) {
-            return getPreview(extras);
-        } else {
+        if (!METHOD_GET_PREVIEW.equals(method)) {
             return null;
         }
+        return getPreview(extras);
     }
 
     private synchronized Bundle getPreview(Bundle request) {
-        Context context = getContext();
-        if (context == null) {
-            return null;
-        }
         RunnableList lifeCycleTracker = new RunnableList();
         try {
             PreviewSurfaceRenderer renderer = new PreviewSurfaceRenderer(
@@ -330,9 +270,7 @@
         public final PreviewSurfaceRenderer renderer;
         public boolean destroyed = false;
 
-        PreviewLifecycleObserver(
-                RunnableList lifeCycleTracker,
-                PreviewSurfaceRenderer renderer) {
+        PreviewLifecycleObserver(RunnableList lifeCycleTracker, PreviewSurfaceRenderer renderer) {
             this.lifeCycleTracker = lifeCycleTracker;
             this.renderer = renderer;
             lifeCycleTracker.add(() -> destroyed = true);
@@ -348,21 +286,10 @@
                 case MESSAGE_ID_UPDATE_PREVIEW:
                     renderer.hideBottomRow(message.getData().getBoolean(KEY_HIDE_BOTTOM_ROW));
                     break;
-                case MESSAGE_ID_UPDATE_SHAPE:
-                    if (Flags.newCustomizationPickerUi()) {
-                        String shapeKey = message.getData().getString(KEY_SHAPE_KEY);
-                        Optional<AppShape> optionalShape = AppShapesProvider.INSTANCE.getShapes()
-                                .stream()
-                                .filter(shape -> shape.getKey().equals(shapeKey))
-                                .findFirst();
-                        String pathToSet = optionalShape.map(AppShape::getPath).orElse(null);
-                        // TODO (b/348664593): Update launcher preview with the given shape
-                    }
-                    break;
                 case MESSAGE_ID_UPDATE_GRID:
-                    String gridKey = message.getData().getString(KEY_GRID_KEY);
-                    if (!TextUtils.isEmpty(gridKey)) {
-                        renderer.updateGrid(gridKey);
+                    String gridName = message.getData().getString(KEY_GRID_NAME);
+                    if (!TextUtils.isEmpty(gridName)) {
+                        renderer.updateGrid(gridName);
                     }
                     break;
                 case MESSAGE_ID_UPDATE_COLOR:
diff --git a/src/com/android/launcher3/graphics/IconShape.java b/src/com/android/launcher3/graphics/IconShape.java
index 5f8f2dc..cb14587 100644
--- a/src/com/android/launcher3/graphics/IconShape.java
+++ b/src/com/android/launcher3/graphics/IconShape.java
@@ -41,10 +41,12 @@
 
 import com.android.launcher3.R;
 import com.android.launcher3.anim.RoundedRectRevealOutlineProvider;
+import com.android.launcher3.dagger.ApplicationContext;
+import com.android.launcher3.dagger.LauncherAppSingleton;
+import com.android.launcher3.dagger.LauncherBaseAppComponent;
 import com.android.launcher3.icons.GraphicsUtils;
 import com.android.launcher3.icons.IconNormalizer;
-import com.android.launcher3.util.MainThreadInitializedObject;
-import com.android.launcher3.util.SafeCloseable;
+import com.android.launcher3.util.DaggerSingletonObject;
 import com.android.launcher3.views.ClipPathView;
 
 import org.xmlpull.v1.XmlPullParser;
@@ -54,19 +56,22 @@
 import java.util.ArrayList;
 import java.util.List;
 
+import javax.inject.Inject;
+
 /**
  * Abstract representation of the shape of an icon shape
  */
-public final class IconShape implements SafeCloseable {
+@LauncherAppSingleton
+public final class IconShape {
 
-    public static final MainThreadInitializedObject<IconShape> INSTANCE =
-            new MainThreadInitializedObject<>(IconShape::new);
-
+    public static DaggerSingletonObject<IconShape> INSTANCE =
+            new DaggerSingletonObject<>(LauncherBaseAppComponent::getIconShape);
 
     private ShapeDelegate mDelegate = new Circle();
     private float mNormalizationScale = ICON_VISIBLE_AREA_FACTOR;
 
-    private IconShape(Context context) {
+    @Inject
+    public IconShape(@ApplicationContext Context context) {
         pickBestShape(context);
     }
 
@@ -78,9 +83,6 @@
         return mNormalizationScale;
     }
 
-    @Override
-    public void close() { }
-
     /**
      * Initializes the shape which is closest to the {@link AdaptiveIconDrawable}
      */
diff --git a/src/com/android/launcher3/model/DbEntry.java b/src/com/android/launcher3/model/DbEntry.java
deleted file mode 100644
index c0c51da..0000000
--- a/src/com/android/launcher3/model/DbEntry.java
+++ /dev/null
@@ -1,155 +0,0 @@
-/*
- * 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/DbEntry.kt b/src/com/android/launcher3/model/DbEntry.kt
new file mode 100644
index 0000000..b79d312
--- /dev/null
+++ b/src/com/android/launcher3/model/DbEntry.kt
@@ -0,0 +1,141 @@
+/*
+ * 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 com.android.launcher3.LauncherSettings
+import com.android.launcher3.LauncherSettings.Favorites.CELLX
+import com.android.launcher3.LauncherSettings.Favorites.CELLY
+import com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APPLICATION
+import com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APPWIDGET
+import com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APP_PAIR
+import com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_FOLDER
+import com.android.launcher3.LauncherSettings.Favorites.SCREEN
+import com.android.launcher3.LauncherSettings.Favorites.SPANX
+import com.android.launcher3.LauncherSettings.Favorites.SPANY
+import com.android.launcher3.model.data.ItemInfo
+import com.android.launcher3.util.ContentWriter
+import java.net.URISyntaxException
+import java.util.Objects
+
+class DbEntry : ItemInfo(), Comparable<DbEntry> {
+    @JvmField var mIntent: String? = null
+    @JvmField var mProvider: String? = null
+    @JvmField var mFolderItems: MutableMap<String, Set<Int>> = HashMap()
+
+    /** Id of the specific widget. */
+    @JvmField var appWidgetId: Int = NO_ID
+
+    /** Comparator according to the reading order */
+    override fun compareTo(other: DbEntry): Int {
+        if (screenId != other.screenId) {
+            return screenId.compareTo(other.screenId)
+        }
+        if (cellY != other.cellY) {
+            return cellY.compareTo(other.cellY)
+        }
+        return cellX.compareTo(other.cellX)
+    }
+
+    override fun equals(other: Any?): Boolean {
+        if (this === other) return true
+        if (other !is DbEntry) return false
+        return getEntryMigrationId() == other.getEntryMigrationId()
+    }
+
+    override fun hashCode(): Int = Objects.hash(getEntryMigrationId())
+
+    /**
+     * Puts the updated DbEntry values into ContentValues which we then use to insert the entry to
+     * the DB.
+     */
+    fun updateContentValues(values: ContentValues) =
+        values.apply {
+            put(SCREEN, screenId)
+            put(CELLX, cellX)
+            put(CELLY, cellY)
+            put(SPANX, spanX)
+            put(SPANY, spanY)
+        }
+
+    override fun writeToValues(writer: ContentWriter) {
+        super.writeToValues(writer)
+        writer.put(LauncherSettings.Favorites.APPWIDGET_ID, appWidgetId)
+    }
+
+    override fun readFromValues(values: ContentValues) {
+        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.
+     */
+    private fun getEntryMigrationId(): String? {
+        when (itemType) {
+            ITEM_TYPE_FOLDER,
+            ITEM_TYPE_APP_PAIR -> return getFolderMigrationId()
+            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
+            ITEM_TYPE_APPLICATION -> {
+                val intentStr = mIntent?.let { cleanIntentString(it) }
+                try {
+                    val i = Intent.parseUri(intentStr, 0)
+                    return Objects.requireNonNull(i.component).toString()
+                } catch (e: Exception) {
+                    return intentStr
+                }
+            }
+
+            else -> return mIntent?.let { cleanIntentString(it) }
+        }
+    }
+
+    /**
+     * This method should return an id that should be the same for two folders containing the same
+     * elements.
+     */
+    private fun getFolderMigrationId(): String =
+        mFolderItems.keys
+            .map { intentString: String ->
+                mFolderItems[intentString]?.size.toString() + cleanIntentString(intentString)
+            }
+            .sorted()
+            .joinToString(",")
+
+    /**
+     * This is needed because sourceBounds can change and make the id of two equal items different.
+     */
+    private fun cleanIntentString(intentStr: String): String {
+        try {
+            return Intent.parseUri(intentStr, 0).apply { sourceBounds = null }.toURI()
+        } catch (e: URISyntaxException) {
+            Log.e(TAG, "Unable to parse Intent string", e)
+            return intentStr
+        }
+    }
+
+    companion object {
+        private const val TAG = "DbEntry"
+    }
+}
diff --git a/src/com/android/launcher3/model/DeviceGridState.java b/src/com/android/launcher3/model/DeviceGridState.java
index 729b381..90af215 100644
--- a/src/com/android/launcher3/model/DeviceGridState.java
+++ b/src/com/android/launcher3/model/DeviceGridState.java
@@ -156,10 +156,16 @@
     }
 
     public Integer getColumns() {
+        if (TextUtils.isEmpty(mGridSizeString)) {
+            return -1;
+        }
         return Integer.parseInt(String.valueOf(mGridSizeString.split(",")[0]));
     }
 
     public Integer getRows() {
+        if (TextUtils.isEmpty(mGridSizeString)) {
+            return -1;
+        }
         return Integer.parseInt(String.valueOf(mGridSizeString.split(",")[1]));
     }
 
diff --git a/src/com/android/launcher3/model/GridSizeMigrationDBController.java b/src/com/android/launcher3/model/GridSizeMigrationDBController.java
index d8ca095..bad7577 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationDBController.java
+++ b/src/com/android/launcher3/model/GridSizeMigrationDBController.java
@@ -452,7 +452,7 @@
         final Context mContext;
         int mLastScreenId = -1;
 
-        final Map<Integer, ArrayList<DbEntry>> mWorkspaceEntriesByScreenId =
+        final Map<Integer, List<DbEntry>> mWorkspaceEntriesByScreenId =
                 new ArrayMap<>();
 
         public DbReader(SQLiteDatabase db, String tableName, Context context) {
diff --git a/src/com/android/launcher3/model/GridSizeMigrationLogic.java b/src/com/android/launcher3/model/GridSizeMigrationLogic.java
deleted file mode 100644
index f8c8f77..0000000
--- a/src/com/android/launcher3/model/GridSizeMigrationLogic.java
+++ /dev/null
@@ -1,481 +0,0 @@
-/*
- * 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 androidx.annotation.VisibleForTesting;
-
-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());
-
-            // Here we keep all the DB ids we have in the destination DB such that we don't assign
-            // an item that we want to add to the destination DB the same id as an already existing
-            // item.
-            List<Integer> idsInUse = new ArrayList<>();
-
-            // Migrate hotseat.
-            migrateHotseat(destDeviceState.getNumHotseat(), srcReader, destReader, target,
-                    idsInUse);
-            // Migrate workspace.
-            migrateWorkspace(srcReader, destReader, target, targetSize, idsInUse);
-
-            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);
-        }
-    }
-
-    /**
-     * Handles hotseat migration.
-     */
-    @VisibleForTesting
-    public void migrateHotseat(int destHotseatSize,
-            GridSizeMigrationDBController.DbReader srcReader,
-            GridSizeMigrationDBController.DbReader destReader,
-            DatabaseHelper helper, List<Integer> idsInUse) {
-        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, idsInUse);
-    }
-
-    private void placeHotseatItems(List<DbEntry> hotseatToBeAdded,
-            List<DbEntry> dstHotseatItems, int destHotseatSize,
-            DatabaseHelper helper, GridSizeMigrationDBController.DbReader srcReader,
-            GridSizeMigrationDBController.DbReader destReader, List<Integer> idsInUse) {
-        if (hotseatToBeAdded.isEmpty()) {
-            return;
-        }
-
-        idsInUse.addAll(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);
-        }
-    }
-
-
-    /**
-     * Handles workspace migration.
-     */
-    @VisibleForTesting
-    public void migrateWorkspace(GridSizeMigrationDBController.DbReader srcReader,
-            GridSizeMigrationDBController.DbReader destReader, DatabaseHelper helper,
-            Point targetSize, List<Integer> idsInUse) {
-        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, idsInUse);
-    }
-
-    private void placeWorkspaceItems(
-            List<DbEntry> workspaceToBeAdded,
-            List<DbEntry> dstWorkspaceItems,
-            int trgX, int trgY, DatabaseHelper helper,
-            GridSizeMigrationDBController.DbReader srcReader,
-            GridSizeMigrationDBController.DbReader destReader, List<Integer> idsInUse) {
-        if (workspaceToBeAdded.isEmpty()) {
-            return;
-        }
-
-        idsInUse.addAll(dstWorkspaceItems.stream().map(entry -> entry.id).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/GridSizeMigrationLogic.kt b/src/com/android/launcher3/model/GridSizeMigrationLogic.kt
new file mode 100644
index 0000000..9470abf
--- /dev/null
+++ b/src/com/android/launcher3/model/GridSizeMigrationLogic.kt
@@ -0,0 +1,524 @@
+/*
+ * 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.Context
+import android.database.sqlite.SQLiteDatabase
+import android.graphics.Point
+import android.util.Log
+import androidx.annotation.VisibleForTesting
+import com.android.launcher3.Flags
+import com.android.launcher3.LauncherPrefs
+import com.android.launcher3.LauncherPrefs.Companion.get
+import com.android.launcher3.LauncherPrefs.Companion.getPrefs
+import com.android.launcher3.LauncherSettings
+import com.android.launcher3.Utilities
+import com.android.launcher3.config.FeatureFlags
+import com.android.launcher3.model.GridSizeMigrationDBController.DbReader
+import com.android.launcher3.provider.LauncherDbUtils
+import com.android.launcher3.provider.LauncherDbUtils.SQLiteTransaction
+import com.android.launcher3.util.CellAndSpan
+import com.android.launcher3.util.GridOccupancy
+import com.android.launcher3.util.IntArray
+
+class GridSizeMigrationLogic {
+    /**
+     * 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.
+     */
+    fun migrateGrid(
+        context: Context,
+        srcDeviceState: DeviceGridState,
+        destDeviceState: DeviceGridState,
+        target: DatabaseHelper,
+        source: SQLiteDatabase,
+    ) {
+        if (!GridSizeMigrationDBController.needsToMigrate(srcDeviceState, destDeviceState)) {
+            return
+        }
+
+        val isFirstLoad = get(context).get(LauncherPrefs.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)) {
+            GridSizeMigrationDBController.copyCurrentGridToNewGrid(
+                context,
+                destDeviceState,
+                target,
+                source,
+            )
+            return
+        }
+        LauncherDbUtils.copyTable(
+            source,
+            LauncherSettings.Favorites.TABLE_NAME,
+            target.writableDatabase,
+            LauncherSettings.Favorites.TMP_TABLE,
+            context,
+        )
+
+        val migrationStartTime = System.currentTimeMillis()
+        try {
+            SQLiteTransaction(target.writableDatabase).use { t ->
+                val srcReader = DbReader(t.db, LauncherSettings.Favorites.TMP_TABLE, context)
+                val destReader = DbReader(t.db, LauncherSettings.Favorites.TABLE_NAME, context)
+
+                val targetSize = Point(destDeviceState.columns, destDeviceState.rows)
+
+                // Here we keep all the DB ids we have in the destination DB such that we don't
+                // assign
+                // an item that we want to add to the destination DB the same id as an already
+                // existing
+                // item.
+                val idsInUse = mutableListOf<Int>()
+
+                // Migrate hotseat.
+                migrateHotseat(destDeviceState.numHotseat, srcReader, destReader, target, idsInUse)
+                // Migrate workspace.
+                migrateWorkspace(srcReader, destReader, target, targetSize, idsInUse)
+
+                LauncherDbUtils.dropTable(t.db, LauncherSettings.Favorites.TMP_TABLE)
+                t.commit()
+            }
+        } catch (e: Exception) {
+            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)
+        }
+    }
+
+    /** Handles hotseat migration. */
+    @VisibleForTesting
+    fun migrateHotseat(
+        destHotseatSize: Int,
+        srcReader: DbReader,
+        destReader: DbReader,
+        helper: DatabaseHelper,
+        idsInUse: MutableList<Int>,
+    ) {
+        val srcHotseatItems = srcReader.loadHotseatEntries()
+        val dstHotseatItems = destReader.loadHotseatEntries()
+
+        val hotseatToBeAdded = getItemsToBeAdded(srcHotseatItems, dstHotseatItems)
+        val toBeRemoved = IntArray()
+        toBeRemoved.addAll(getItemsToBeRemoved(srcHotseatItems, dstHotseatItems))
+
+        if (DEBUG) {
+            Log.d(
+                TAG,
+                """Start hotseat migration:
+            |Removing Hotseat Items: [${dstHotseatItems.filter { toBeRemoved.contains(it.id) }
+                .joinToString(",\n") { it.toString() }}]
+            |Adding Hotseat Items: [${hotseatToBeAdded
+                .joinToString(",\n") { it.toString() }}]
+            |"""
+                    .trimMargin(),
+            )
+        }
+
+        // Removes the items that we need to remove from the destination DB.
+        if (!toBeRemoved.isEmpty) {
+            GridSizeMigrationDBController.removeEntryFromDb(
+                destReader.mDb,
+                destReader.mTableName,
+                toBeRemoved,
+            )
+        }
+
+        placeHotseatItems(
+            hotseatToBeAdded,
+            dstHotseatItems,
+            destHotseatSize,
+            helper,
+            srcReader,
+            destReader,
+            idsInUse,
+        )
+    }
+
+    private fun placeHotseatItems(
+        hotseatToBeAdded: MutableList<DbEntry>,
+        dstHotseatItems: List<DbEntry>,
+        destHotseatSize: Int,
+        helper: DatabaseHelper,
+        srcReader: DbReader,
+        destReader: DbReader,
+        idsInUse: MutableList<Int>,
+    ) {
+        if (hotseatToBeAdded.isEmpty()) {
+            return
+        }
+
+        idsInUse.addAll(dstHotseatItems.map { entry: DbEntry -> entry.id })
+
+        hotseatToBeAdded.sort()
+
+        val placementSolutionHotseat =
+            solveHotseatPlacement(destHotseatSize, dstHotseatItems, hotseatToBeAdded)
+        for (entryToPlace in placementSolutionHotseat) {
+            GridSizeMigrationDBController.insertEntryInDb(
+                helper,
+                entryToPlace,
+                srcReader.mTableName,
+                destReader.mTableName,
+                idsInUse,
+            )
+        }
+    }
+
+    @VisibleForTesting
+    fun migrateWorkspace(
+        srcReader: DbReader,
+        destReader: DbReader,
+        helper: DatabaseHelper,
+        targetSize: Point,
+        idsInUse: MutableList<Int>,
+    ) {
+        val srcWorkspaceItems = srcReader.loadAllWorkspaceEntries()
+
+        val dstWorkspaceItems = destReader.loadAllWorkspaceEntries()
+
+        val toBeRemoved = IntArray()
+
+        val workspaceToBeAdded = getItemsToBeAdded(srcWorkspaceItems, dstWorkspaceItems)
+        toBeRemoved.addAll(getItemsToBeRemoved(srcWorkspaceItems, dstWorkspaceItems))
+
+        if (DEBUG) {
+            Log.d(
+                TAG,
+                """Start workspace migration:
+            |Source Device: [${srcWorkspaceItems.joinToString(",\n") { it.toString() }}]
+            |Target Device: [${dstWorkspaceItems.joinToString(",\n") { it.toString() }}]
+            |Removing Workspace Items: [${dstWorkspaceItems.filter { toBeRemoved.contains(it.id) }
+                .joinToString(",\n") { it.toString() }}]
+            |Adding Workspace Items: [${workspaceToBeAdded
+                .joinToString(",\n") { it.toString() }}]
+            |"""
+                    .trimMargin(),
+            )
+        }
+
+        // Removes the items that we need to remove from the destination DB.
+        if (!toBeRemoved.isEmpty) {
+            GridSizeMigrationDBController.removeEntryFromDb(
+                destReader.mDb,
+                destReader.mTableName,
+                toBeRemoved,
+            )
+        }
+
+        placeWorkspaceItems(
+            workspaceToBeAdded,
+            dstWorkspaceItems,
+            targetSize.x,
+            targetSize.y,
+            helper,
+            srcReader,
+            destReader,
+            idsInUse,
+        )
+    }
+
+    private fun placeWorkspaceItems(
+        workspaceToBeAdded: MutableList<DbEntry>,
+        dstWorkspaceItems: List<DbEntry>,
+        trgX: Int,
+        trgY: Int,
+        helper: DatabaseHelper,
+        srcReader: DbReader,
+        destReader: DbReader,
+        idsInUse: MutableList<Int>,
+    ) {
+        if (workspaceToBeAdded.isEmpty()) {
+            return
+        }
+
+        idsInUse.addAll(dstWorkspaceItems.map { entry: DbEntry -> entry.id })
+
+        workspaceToBeAdded.sort()
+
+        // First we create a collection of the screens
+        val screens: MutableList<Int> = ArrayList()
+        for (screenId in 0..destReader.mLastScreenId) {
+            screens.add(screenId)
+        }
+
+        // Then we place the items on the screens
+        var itemsToPlace = WorkspaceItemsToPlace(workspaceToBeAdded, mutableListOf())
+        for (screenId in screens) {
+            if (DEBUG) {
+                Log.d(TAG, "Migrating $screenId")
+            }
+            itemsToPlace =
+                solveGridPlacement(
+                    destReader.mContext,
+                    screenId,
+                    trgX,
+                    trgY,
+                    itemsToPlace.mRemainingItemsToPlace,
+                    destReader.mWorkspaceEntriesByScreenId[screenId],
+                )
+            placeItems(itemsToPlace, helper, srcReader, destReader, idsInUse)
+            while (itemsToPlace.mPlacementSolution.isNotEmpty()) {
+                GridSizeMigrationDBController.insertEntryInDb(
+                    helper,
+                    itemsToPlace.mPlacementSolution.removeAt(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.
+        var screenId = destReader.mLastScreenId + 1
+        while (itemsToPlace.mRemainingItemsToPlace.isNotEmpty()) {
+            itemsToPlace =
+                solveGridPlacement(
+                    destReader.mContext,
+                    screenId,
+                    trgX,
+                    trgY,
+                    itemsToPlace.mRemainingItemsToPlace,
+                    destReader.mWorkspaceEntriesByScreenId[screenId],
+                )
+            placeItems(itemsToPlace, helper, srcReader, destReader, idsInUse)
+            screenId++
+        }
+    }
+
+    private fun placeItems(
+        itemsToPlace: WorkspaceItemsToPlace,
+        helper: DatabaseHelper,
+        srcReader: DbReader,
+        destReader: DbReader,
+        idsInUse: List<Int>,
+    ) {
+        while (itemsToPlace.mPlacementSolution.isNotEmpty()) {
+            GridSizeMigrationDBController.insertEntryInDb(
+                helper,
+                itemsToPlace.mPlacementSolution.removeAt(0),
+                srcReader.mTableName,
+                destReader.mTableName,
+                idsInUse,
+            )
+        }
+    }
+
+    /** Only migrate the grid in this manner if the target grid is taller and not wider. */
+    private fun shouldMigrateToStrictlyTallerGrid(
+        isFirstLoad: Boolean,
+        srcDeviceState: DeviceGridState,
+        destDeviceState: DeviceGridState,
+    ): Boolean {
+        return (isFirstLoad && Flags.enableGridMigrationFix()) &&
+            srcDeviceState.columns == destDeviceState.columns &&
+            srcDeviceState.rows < destDeviceState.rows
+    }
+
+    /**
+     * 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 fun getItemsToBeAdded(src: List<DbEntry>, dest: List<DbEntry>): MutableList<DbEntry> {
+        val entryCountDiff = calcDiff(src, dest)
+        val toBeAdded: MutableList<DbEntry> = ArrayList()
+        src.forEach { entry ->
+            entryCountDiff[entry]?.let { entryDiff ->
+                if (entryDiff > 0) {
+                    toBeAdded.add(entry)
+                    entryCountDiff[entry] = entryDiff - 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 fun getItemsToBeRemoved(src: List<DbEntry>, dest: List<DbEntry>): IntArray {
+        val entryCountDiff = calcDiff(src, dest)
+        val toBeRemoved =
+            IntArray().apply {
+                dest.forEach { entry ->
+                    entryCountDiff[entry]?.let { entryDiff ->
+                        if (entryDiff < 0) {
+                            add(entry.id)
+                            if (entry.itemType == LauncherSettings.Favorites.ITEM_TYPE_FOLDER) {
+                                entry.mFolderItems.values.forEach { ids ->
+                                    ids.forEach { value -> add(value) }
+                                }
+                            }
+                        }
+                        entryCountDiff[entry] = entryDiff.plus(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 fun calcDiff(src: List<DbEntry>, dest: List<DbEntry>): MutableMap<DbEntry, Int> {
+        val entryCountDiff: MutableMap<DbEntry, Int> = HashMap()
+        src.forEach { entry -> entryCountDiff[entry] = entryCountDiff.getOrDefault(entry, 0) + 1 }
+        dest.forEach { entry -> entryCountDiff[entry] = entryCountDiff.getOrDefault(entry, 0) - 1 }
+        return entryCountDiff
+    }
+
+    private fun solveHotseatPlacement(
+        hotseatSize: Int,
+        placedHotseatItems: List<DbEntry>,
+        itemsToPlace: List<DbEntry>,
+    ): List<DbEntry> {
+        val placementSolution: MutableList<DbEntry> = ArrayList()
+        val remainingItemsToPlace: MutableList<DbEntry> = ArrayList(itemsToPlace)
+        val occupied = BooleanArray(hotseatSize)
+        for (entry in placedHotseatItems) {
+            occupied[entry.screenId] = true
+        }
+
+        for (i in occupied.indices) {
+            if (!occupied[i] && remainingItemsToPlace.isNotEmpty()) {
+                val entry: DbEntry =
+                    remainingItemsToPlace.removeAt(0).apply {
+                        screenId = i
+                        // These values does not affect the item position, but we should set them
+                        // to something other than -1.
+                        cellX = i
+                        cellY = 0
+                    }
+                placementSolution.add(entry)
+                occupied[entry.screenId] = true
+            }
+        }
+        return placementSolution
+    }
+
+    private fun solveGridPlacement(
+        context: Context,
+        screenId: Int,
+        trgX: Int,
+        trgY: Int,
+        sortedItemsToPlace: MutableList<DbEntry>,
+        existedEntries: MutableList<DbEntry>?,
+    ): WorkspaceItemsToPlace {
+        val itemsToPlace = WorkspaceItemsToPlace(sortedItemsToPlace, mutableListOf())
+        val occupied = GridOccupancy(trgX, trgY)
+        val trg = Point(trgX, trgY)
+        val next: Point =
+            if (
+                screenId == 0 &&
+                    (FeatureFlags.QSB_ON_FIRST_SCREEN &&
+                        (!Flags.enableSmartspaceRemovalToggle() ||
+                            getPrefs(context)
+                                .getBoolean(LoaderTask.SMARTSPACE_ON_HOME_SCREEN, true)) &&
+                        !Utilities.SHOULD_SHOW_FIRST_PAGE_WIDGET)
+            ) {
+                Point(0, 1 /* smartspace */)
+            } else {
+                Point(0, 0)
+            }
+        if (existedEntries != null) {
+            for (entry in existedEntries) {
+                occupied.markCells(entry, true)
+            }
+        }
+        val iterator = itemsToPlace.mRemainingItemsToPlace.iterator()
+        while (iterator.hasNext()) {
+            val entry = iterator.next()
+            if (entry.minSpanX > trgX || entry.minSpanY > trgY) {
+                iterator.remove()
+                continue
+            }
+            findPlacementForEntry(entry, next.x, next.y, trg, occupied)?.let {
+                entry.screenId = screenId
+                entry.cellX = it.cellX
+                entry.cellY = it.cellY
+                entry.spanX = it.spanX
+                entry.spanY = it.spanY
+                occupied.markCells(entry, true)
+                next[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 fun findPlacementForEntry(
+        entry: DbEntry,
+        startPosX: Int,
+        startPosY: Int,
+        trg: Point,
+        occupied: GridOccupancy,
+    ): CellAndSpan? {
+        var newStartPosX = startPosX
+        for (y in startPosY until trg.y) {
+            for (x in newStartPosX until trg.x) {
+                if (occupied.isRegionVacant(x, y, entry.minSpanX, entry.minSpanY)) {
+                    return (CellAndSpan(x, y, entry.minSpanX, entry.minSpanY))
+                }
+            }
+            newStartPosX = 0
+        }
+        return null
+    }
+
+    private data class WorkspaceItemsToPlace(
+        val mRemainingItemsToPlace: MutableList<DbEntry>,
+        val mPlacementSolution: MutableList<DbEntry>,
+    )
+
+    companion object {
+        private const val TAG = "GridSizeMigrationLogic"
+        private const val DEBUG = true
+    }
+}
diff --git a/src/com/android/launcher3/model/ModelDbController.java b/src/com/android/launcher3/model/ModelDbController.java
index 8c3e860..094798b 100644
--- a/src/com/android/launcher3/model/ModelDbController.java
+++ b/src/com/android/launcher3/model/ModelDbController.java
@@ -20,6 +20,7 @@
 import static android.util.Base64.NO_WRAP;
 
 import static com.android.launcher3.DefaultLayoutParser.RES_PARTNER_DEFAULT_LAYOUT;
+import static com.android.launcher3.LauncherPrefs.DB_FILE;
 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;
@@ -128,16 +129,20 @@
 
     private synchronized void createDbIfNotExists() {
         if (mOpenHelper == null) {
-            mOpenHelper = createDatabaseHelper(false /* forMigration */);
+            String dbFile = LauncherPrefs.get(mContext).get(DB_FILE);
+            if (dbFile.isEmpty()) {
+                dbFile = InvariantDeviceProfile.INSTANCE.get(mContext).dbFile;
+            }
+            mOpenHelper = createDatabaseHelper(false /* forMigration */, dbFile);
             printDBs("before: ");
             RestoreDbTask.restoreIfNeeded(mContext, this);
             printDBs("after: ");
         }
     }
 
-    protected DatabaseHelper createDatabaseHelper(boolean forMigration) {
+    protected DatabaseHelper createDatabaseHelper(boolean forMigration, String dbFile) {
         boolean isSandbox = mContext instanceof SandboxContext;
-        String dbName = isSandbox ? null : InvariantDeviceProfile.INSTANCE.get(mContext).dbFile;
+        String dbName = isSandbox ? null : dbFile;
 
         // Set the flag for empty DB
         Runnable onEmptyDbCreateCallback = forMigration ? () -> { }
@@ -364,7 +369,7 @@
         InvariantDeviceProfile idp = LauncherAppState.getIDP(mContext);
         DatabaseHelper oldHelper = mOpenHelper;
         mOpenHelper = (mContext instanceof SandboxContext) ? oldHelper
-                : createDatabaseHelper(true /* forMigration */);
+                : createDatabaseHelper(true, new DeviceGridState(idp).getDbFile());
         try {
             // This is the current grid we have, given by the mContext
             DeviceGridState srcDeviceState = new DeviceGridState(mContext);
@@ -388,7 +393,6 @@
      * Migrates the DB if needed. If the migration failed, it clears the DB.
      */
     public void tryMigrateDB(@Nullable LauncherRestoreEventLogger restoreEventLogger) {
-
         if (!migrateGridIfNeeded()) {
             if (restoreEventLogger != null) {
                 if (LauncherPrefs.get(mContext).get(NO_DB_FILES_RESTORED)) {
@@ -438,7 +442,7 @@
         }
         DatabaseHelper oldHelper = mOpenHelper;
         mOpenHelper = (mContext instanceof SandboxContext) ? oldHelper
-                : createDatabaseHelper(true /* forMigration */);
+                : createDatabaseHelper(true /* forMigration */, targetDbName);
         try {
             // This is the current grid we have, given by the mContext
             DeviceGridState srcDeviceState = new DeviceGridState(mContext);
diff --git a/src/com/android/launcher3/provider/RestoreDbTask.java b/src/com/android/launcher3/provider/RestoreDbTask.java
index 775d248..59c27af 100644
--- a/src/com/android/launcher3/provider/RestoreDbTask.java
+++ b/src/com/android/launcher3/provider/RestoreDbTask.java
@@ -125,7 +125,8 @@
         LauncherPrefs.get(context).removeSync(RESTORE_DEVICE);
 
         if (Flags.enableNarrowGridRestore()) {
-            String oldPhoneFileName = idp.dbFile;
+            DeviceGridState deviceGridState = new DeviceGridState(context);
+            String oldPhoneFileName = deviceGridState.getDbFile();
             List<String> previousDbs = existingDbs(context);
             removeOldDBs(context, oldPhoneFileName);
             // The idp before this contains data about the old phone, after this it becomes the idp
@@ -148,6 +149,7 @@
                 context, oldPhoneDbFileName);
         // The grid option could be null if current phone doesn't support the previous db.
         if (oldPhoneGridOption != null) {
+
             /* If the user only used the default db on the previous phone and the new default db is
              * bigger than or equal to the previous one, then keep the new default db */
             if (previousDbs.size() == 1 && oldPhoneGridOption.numColumns <= idp.numColumns
diff --git a/src/com/android/launcher3/shapes/AppShape.kt b/src/com/android/launcher3/shapes/AppShape.kt
deleted file mode 100644
index 68200a0..0000000
--- a/src/com/android/launcher3/shapes/AppShape.kt
+++ /dev/null
@@ -1,19 +0,0 @@
-/*
- * 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.shapes
-
-class AppShape(val key: String, val title: String, val path: String)
diff --git a/src/com/android/launcher3/shapes/AppShapesProvider.kt b/src/com/android/launcher3/shapes/AppShapesProvider.kt
deleted file mode 100644
index 41bac6a..0000000
--- a/src/com/android/launcher3/shapes/AppShapesProvider.kt
+++ /dev/null
@@ -1,54 +0,0 @@
-/*
- * 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.shapes
-
-object AppShapesProvider {
-
-    val shapes =
-        listOf(
-            AppShape(
-                "arch",
-                "arch",
-                "M100 83.46C100 85.471 100 86.476 99.9 87.321 99.116 93.916 93.916 99.116 87.321 99.9 86.476 100 85.471 100 83.46 100H16.54C14.529 100 13.524 100 12.679 99.9 6.084 99.116.884 93.916.1 87.321 0 86.476 0 85.471 0 83.46L0 50C0 22.386 22.386 0 50 0 77.614 0 100 22.386 100 50V83.46Z",
-            ),
-            AppShape(
-                "4_sided_cookie",
-                "4 sided cookie",
-                "M63.605 3C84.733-6.176 106.176 15.268 97 36.395L95.483 39.888C92.681 46.338 92.681 53.662 95.483 60.112L97 63.605C106.176 84.732 84.733 106.176 63.605 97L60.112 95.483C53.662 92.681 46.338 92.681 39.888 95.483L36.395 97C15.267 106.176-6.176 84.732 3 63.605L4.517 60.112C7.319 53.662 7.319 46.338 4.517 39.888L3 36.395C-6.176 15.268 15.267-6.176 36.395 3L39.888 4.517C46.338 7.319 53.662 7.319 60.112 4.517L63.605 3Z",
-            ),
-            AppShape(
-                "seven_sided_cookie",
-                "7 sided cookie",
-                "M35.209 4.878C36.326 3.895 36.884 3.404 37.397 3.006 44.82-2.742 55.18-2.742 62.603 3.006 63.116 3.404 63.674 3.895 64.791 4.878 65.164 5.207 65.351 5.371 65.539 5.529 68.167 7.734 71.303 9.248 74.663 9.932 74.902 9.981 75.147 10.025 75.637 10.113 77.1 10.375 77.831 10.506 78.461 10.66 87.573 12.893 94.032 21.011 94.176 30.412 94.186 31.062 94.151 31.805 94.08 33.293 94.057 33.791 94.045 34.04 94.039 34.285 93.958 37.72 94.732 41.121 96.293 44.18 96.404 44.399 96.522 44.618 96.759 45.056 97.467 46.366 97.821 47.021 98.093 47.611 102.032 56.143 99.727 66.266 92.484 72.24 91.983 72.653 91.381 73.089 90.177 73.961 89.774 74.254 89.572 74.4 89.377 74.548 86.647 76.626 84.477 79.353 83.063 82.483 82.962 82.707 82.865 82.936 82.671 83.395 82.091 84.766 81.8 85.451 81.51 86.033 77.31 94.44 67.977 98.945 58.801 96.994 58.166 96.859 57.451 96.659 56.019 96.259 55.54 96.125 55.3 96.058 55.063 95.998 51.74 95.154 48.26 95.154 44.937 95.998 44.699 96.058 44.46 96.125 43.981 96.259 42.549 96.659 41.834 96.859 41.199 96.994 32.023 98.945 22.69 94.44 18.49 86.033 18.2 85.451 17.909 84.766 17.329 83.395 17.135 82.936 17.038 82.707 16.937 82.483 15.523 79.353 13.353 76.626 10.623 74.548 10.428 74.4 10.226 74.254 9.823 73.961 8.619 73.089 8.017 72.653 7.516 72.24.273 66.266-2.032 56.143 1.907 47.611 2.179 47.021 2.533 46.366 3.241 45.056 3.478 44.618 3.596 44.399 3.707 44.18 5.268 41.121 6.042 37.72 5.961 34.285 5.955 34.04 5.943 33.791 5.92 33.293 5.849 31.805 5.814 31.062 5.824 30.412 5.968 21.011 12.427 12.893 21.539 10.66 22.169 10.506 22.9 10.375 24.363 10.113 24.853 10.025 25.098 9.981 25.337 9.932 28.697 9.248 31.833 7.734 34.461 5.529 34.649 5.371 34.836 5.207 35.209 4.878Z",
-            ),
-            AppShape(
-                "sunny",
-                "sunny",
-                "M42.846 4.873C46.084-.531 53.916-.531 57.154 4.873L60.796 10.951C62.685 14.103 66.414 15.647 69.978 14.754L76.851 13.032C82.962 11.5 88.5 17.038 86.968 23.149L85.246 30.022C84.353 33.586 85.897 37.315 89.049 39.204L95.127 42.846C100.531 46.084 100.531 53.916 95.127 57.154L89.049 60.796C85.897 62.685 84.353 66.414 85.246 69.978L86.968 76.851C88.5 82.962 82.962 88.5 76.851 86.968L69.978 85.246C66.414 84.353 62.685 85.898 60.796 89.049L57.154 95.127C53.916 100.531 46.084 100.531 42.846 95.127L39.204 89.049C37.315 85.898 33.586 84.353 30.022 85.246L23.149 86.968C17.038 88.5 11.5 82.962 13.032 76.851L14.754 69.978C15.647 66.414 14.103 62.685 10.951 60.796L4.873 57.154C-.531 53.916-.531 46.084 4.873 42.846L10.951 39.204C14.103 37.315 15.647 33.586 14.754 30.022L13.032 23.149C11.5 17.038 17.038 11.5 23.149 13.032L30.022 14.754C33.586 15.647 37.315 14.103 39.204 10.951L42.846 4.873Z",
-            ),
-            AppShape(
-                "circle",
-                "circle",
-                "M99.18 50C99.18 77.162 77.162 99.18 50 99.18 22.838 99.18.82 77.162.82 50 .82 22.839 22.838.82 50 .82 77.162.82 99.18 22.839 99.18 50Z",
-            ),
-            AppShape(
-                "square",
-                "square",
-                "M99.18 53.689C99.18 67.434 99.18 74.306 97.022 79.758 93.897 87.649 87.649 93.897 79.758 97.022 74.306 99.18 67.434 99.18 53.689 99.18H46.311C32.566 99.18 25.694 99.18 20.242 97.022 12.351 93.897 6.103 87.649 2.978 79.758.82 74.306.82 67.434.82 53.689L.82 46.311C.82 32.566.82 25.694 2.978 20.242 6.103 12.351 12.351 6.103 20.242 2.978 25.694.82 32.566.82 46.311.82L53.689.82C67.434.82 74.306.82 79.758 2.978 87.649 6.103 93.897 12.351 97.022 20.242 99.18 25.694 99.18 32.566 99.18 46.311V53.689Z\n",
-            ),
-        )
-}
diff --git a/src/com/android/launcher3/util/LogConfig.java b/src/com/android/launcher3/util/LogConfig.java
index f183f18..72e3e79 100644
--- a/src/com/android/launcher3/util/LogConfig.java
+++ b/src/com/android/launcher3/util/LogConfig.java
@@ -52,9 +52,9 @@
     public static final String WEB_APP_SEARCH_LOGGING = "WebAppSearchLogging";
 
     /**
-     * When turned on, we enable quick launch v2 related logging.
+     * When turned on, we enable quick launch related logging.
      */
-    public static final String QUICK_LAUNCH_V2 = "QuickLaunchV2";
+    public static final String QUICK_LAUNCH = "QuickLaunch";
 
     /**
      * When turned on, we enable Gms Play related logging.
diff --git a/tests/multivalentTests/src/com/android/launcher3/AbstractDeviceProfileTest.kt b/tests/multivalentTests/src/com/android/launcher3/AbstractDeviceProfileTest.kt
index 8770859..55a028b 100644
--- a/tests/multivalentTests/src/com/android/launcher3/AbstractDeviceProfileTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/AbstractDeviceProfileTest.kt
@@ -79,7 +79,7 @@
         val statusBarNaturalPx: Int,
         val statusBarRotatedPx: Int,
         val gesturePx: Int,
-        val cutoutPx: Int
+        val cutoutPx: Int,
     )
 
     open val deviceSpecs =
@@ -91,7 +91,7 @@
                     statusBarNaturalPx = 118,
                     statusBarRotatedPx = 74,
                     gesturePx = 63,
-                    cutoutPx = 118
+                    cutoutPx = 118,
                 ),
             "tablet" to
                 DeviceSpec(
@@ -100,7 +100,7 @@
                     statusBarNaturalPx = 104,
                     statusBarRotatedPx = 104,
                     gesturePx = 0,
-                    cutoutPx = 0
+                    cutoutPx = 0,
                 ),
             "twopanel-phone" to
                 DeviceSpec(
@@ -109,7 +109,7 @@
                     statusBarNaturalPx = 133,
                     statusBarRotatedPx = 110,
                     gesturePx = 63,
-                    cutoutPx = 133
+                    cutoutPx = 133,
                 ),
             "twopanel-tablet" to
                 DeviceSpec(
@@ -118,14 +118,14 @@
                     statusBarNaturalPx = 110,
                     statusBarRotatedPx = 133,
                     gesturePx = 0,
-                    cutoutPx = 0
-                )
+                    cutoutPx = 0,
+                ),
         )
 
     protected fun initializeVarsForPhone(
         deviceSpec: DeviceSpec,
         isGestureMode: Boolean = true,
-        isVerticalBar: Boolean = false
+        isVerticalBar: Boolean = false,
     ) {
         val (naturalX, naturalY) = deviceSpec.naturalSize
         val windowsBounds = phoneWindowsBounds(deviceSpec, isGestureMode, naturalX, naturalY)
@@ -137,14 +137,14 @@
             displayInfo,
             rotation = if (isVerticalBar) Surface.ROTATION_90 else Surface.ROTATION_0,
             isGestureMode,
-            densityDpi = deviceSpec.densityDpi
+            densityDpi = deviceSpec.densityDpi,
         )
     }
 
     protected fun initializeVarsForTablet(
         deviceSpec: DeviceSpec,
         isLandscape: Boolean = false,
-        isGestureMode: Boolean = true
+        isGestureMode: Boolean = true,
     ) {
         val (naturalX, naturalY) = deviceSpec.naturalSize
         val windowsBounds = tabletWindowsBounds(deviceSpec, naturalX, naturalY)
@@ -156,7 +156,7 @@
             displayInfo,
             rotation = if (isLandscape) Surface.ROTATION_0 else Surface.ROTATION_90,
             isGestureMode,
-            densityDpi = deviceSpec.densityDpi
+            densityDpi = deviceSpec.densityDpi,
         )
     }
 
@@ -165,7 +165,7 @@
         deviceSpecFolded: DeviceSpec,
         isLandscape: Boolean = false,
         isGestureMode: Boolean = true,
-        isFolded: Boolean = false
+        isFolded: Boolean = false,
     ) {
         val (unfoldedNaturalX, unfoldedNaturalY) = deviceSpecUnfolded.naturalSize
         val unfoldedWindowsBounds =
@@ -182,7 +182,7 @@
         val perDisplayBoundsCache =
             mapOf(
                 unfoldedDisplayInfo to unfoldedWindowsBounds,
-                foldedDisplayInfo to foldedWindowsBounds
+                foldedDisplayInfo to foldedWindowsBounds,
             )
 
         if (isFolded) {
@@ -191,7 +191,7 @@
                 displayInfo = foldedDisplayInfo,
                 rotation = if (isLandscape) Surface.ROTATION_90 else Surface.ROTATION_0,
                 isGestureMode = isGestureMode,
-                densityDpi = deviceSpecFolded.densityDpi
+                densityDpi = deviceSpecFolded.densityDpi,
             )
         } else {
             initializeCommonVars(
@@ -199,7 +199,7 @@
                 displayInfo = unfoldedDisplayInfo,
                 rotation = if (isLandscape) Surface.ROTATION_0 else Surface.ROTATION_90,
                 isGestureMode = isGestureMode,
-                densityDpi = deviceSpecUnfolded.densityDpi
+                densityDpi = deviceSpecUnfolded.densityDpi,
             )
         }
     }
@@ -208,7 +208,7 @@
         deviceSpec: DeviceSpec,
         isGestureMode: Boolean,
         naturalX: Int,
-        naturalY: Int
+        naturalY: Int,
     ): List<WindowBounds> {
         val buttonsNavHeight = Utilities.dpToPx(48f, deviceSpec.densityDpi)
 
@@ -217,14 +217,14 @@
                 0,
                 max(deviceSpec.statusBarNaturalPx, deviceSpec.cutoutPx),
                 0,
-                if (isGestureMode) deviceSpec.gesturePx else buttonsNavHeight
+                if (isGestureMode) deviceSpec.gesturePx else buttonsNavHeight,
             )
         val rotation90Insets =
             Rect(
                 deviceSpec.cutoutPx,
                 deviceSpec.statusBarRotatedPx,
                 if (isGestureMode) 0 else buttonsNavHeight,
-                if (isGestureMode) deviceSpec.gesturePx else 0
+                if (isGestureMode) deviceSpec.gesturePx else 0,
             )
         val rotation180Insets =
             Rect(
@@ -233,29 +233,29 @@
                 0,
                 max(
                     if (isGestureMode) deviceSpec.gesturePx else buttonsNavHeight,
-                    deviceSpec.cutoutPx
-                )
+                    deviceSpec.cutoutPx,
+                ),
             )
         val rotation270Insets =
             Rect(
                 if (isGestureMode) 0 else buttonsNavHeight,
                 deviceSpec.statusBarRotatedPx,
                 deviceSpec.cutoutPx,
-                if (isGestureMode) deviceSpec.gesturePx else 0
+                if (isGestureMode) deviceSpec.gesturePx else 0,
             )
 
         return listOf(
             WindowBounds(Rect(0, 0, naturalX, naturalY), rotation0Insets, Surface.ROTATION_0),
             WindowBounds(Rect(0, 0, naturalY, naturalX), rotation90Insets, Surface.ROTATION_90),
             WindowBounds(Rect(0, 0, naturalX, naturalY), rotation180Insets, Surface.ROTATION_180),
-            WindowBounds(Rect(0, 0, naturalY, naturalX), rotation270Insets, Surface.ROTATION_270)
+            WindowBounds(Rect(0, 0, naturalY, naturalX), rotation270Insets, Surface.ROTATION_270),
         )
     }
 
     private fun tabletWindowsBounds(
         deviceSpec: DeviceSpec,
         naturalX: Int,
-        naturalY: Int
+        naturalY: Int,
     ): List<WindowBounds> {
         val naturalInsets = Rect(0, deviceSpec.statusBarNaturalPx, 0, 0)
         val rotatedInsets = Rect(0, deviceSpec.statusBarRotatedPx, 0, 0)
@@ -264,7 +264,7 @@
             WindowBounds(Rect(0, 0, naturalX, naturalY), naturalInsets, Surface.ROTATION_0),
             WindowBounds(Rect(0, 0, naturalY, naturalX), rotatedInsets, Surface.ROTATION_90),
             WindowBounds(Rect(0, 0, naturalX, naturalY), naturalInsets, Surface.ROTATION_180),
-            WindowBounds(Rect(0, 0, naturalY, naturalX), rotatedInsets, Surface.ROTATION_270)
+            WindowBounds(Rect(0, 0, naturalY, naturalX), rotatedInsets, Surface.ROTATION_270),
         )
     }
 
@@ -273,7 +273,7 @@
         displayInfo: CachedDisplayInfo,
         rotation: Int,
         isGestureMode: Boolean = true,
-        densityDpi: Int
+        densityDpi: Int,
     ) {
         setFlagsRule.setFlags(true, Flags.FLAG_ENABLE_TWOLINE_TOGGLE)
         LauncherPrefs.get(testContext).put(LauncherPrefs.ENABLE_TWOLINE_ALLAPPS_TOGGLE, true)
@@ -307,6 +307,10 @@
 
         whenever(launcherPrefs.get(LauncherPrefs.TASKBAR_PINNING)).thenReturn(false)
         whenever(launcherPrefs.get(LauncherPrefs.TASKBAR_PINNING_IN_DESKTOP_MODE)).thenReturn(true)
+        whenever(launcherPrefs.get(LauncherPrefs.HOTSEAT_COUNT)).thenReturn(-1)
+        whenever(launcherPrefs.get(LauncherPrefs.DEVICE_TYPE)).thenReturn(-1)
+        whenever(launcherPrefs.get(LauncherPrefs.WORKSPACE_SIZE)).thenReturn("")
+        whenever(launcherPrefs.get(LauncherPrefs.DB_FILE)).thenReturn("")
         val info = spy(DisplayController.Info(context, windowManagerProxy, perDisplayBoundsCache))
         whenever(displayController.info).thenReturn(info)
         whenever(info.isTransientTaskbar).thenReturn(isGestureMode)