diff --git a/quickstep/res/layout/keyboard_quick_switch_desktop_taskview.xml b/quickstep/res/layout/keyboard_quick_switch_desktop_taskview.xml
index 71c782d..db47ff0 100644
--- a/quickstep/res/layout/keyboard_quick_switch_desktop_taskview.xml
+++ b/quickstep/res/layout/keyboard_quick_switch_desktop_taskview.xml
@@ -48,13 +48,13 @@
 
             app:layout_constraintVertical_chainStyle="packed"
             app:layout_constraintTop_toTopOf="parent"
-            app:layout_constraintBottom_toTopOf="@id/text"
+            app:layout_constraintBottom_toTopOf="@id/small_text"
             app:layout_constraintStart_toStartOf="parent"
             app:layout_constraintEnd_toEndOf="parent"/>
 
         <TextView
             style="@style/KeyboardQuickSwitchText.OnTaskView"
-            android:id="@+id/text"
+            android:id="@+id/small_text"
             android:layout_width="match_parent"
             android:layout_height="wrap_content"
             android:textAlignment="center"
diff --git a/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java b/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
index 4ba4e2b..c748385 100644
--- a/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
+++ b/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
@@ -411,7 +411,8 @@
             @NonNull RemoteAnimationTarget[] nonAppTargets, boolean launcherClosing) {
         TaskViewUtils.composeRecentsLaunchAnimator(anim, v, appTargets, wallpaperTargets,
                 nonAppTargets, launcherClosing, mLauncher.getStateManager(),
-                mLauncher.getOverviewPanel(), mLauncher.getDepthController());
+                mLauncher.getOverviewPanel(), mLauncher.getDepthController(),
+                /* transitionInfo= */ null);
     }
 
     private boolean areAllTargetsTranslucent(@NonNull RemoteAnimationTarget[] targets) {
diff --git a/quickstep/src/com/android/launcher3/model/WellbeingModel.java b/quickstep/src/com/android/launcher3/model/WellbeingModel.java
index 0f3aaa6..09433c5 100644
--- a/quickstep/src/com/android/launcher3/model/WellbeingModel.java
+++ b/quickstep/src/com/android/launcher3/model/WellbeingModel.java
@@ -109,9 +109,9 @@
                 ? Executors.UI_HELPER_EXECUTOR.getLooper()
                 : Executors.getPackageExecutor(mWellbeingProviderPkg).getLooper());
         mWellbeingAppChangeReceiver =
-                new SimpleBroadcastReceiver(mWorkerHandler, t -> restartObserver());
+                new SimpleBroadcastReceiver(context, mWorkerHandler, t -> restartObserver());
         mAppAddRemoveReceiver =
-                new SimpleBroadcastReceiver(mWorkerHandler, this::onAppPackageChanged);
+                new SimpleBroadcastReceiver(context, mWorkerHandler, this::onAppPackageChanged);
 
 
         mContentObserver = new ContentObserver(mWorkerHandler) {
@@ -148,8 +148,8 @@
     public void close() {
         if (!TextUtils.isEmpty(mWellbeingProviderPkg)) {
             mWorkerHandler.post(() -> {
-                mWellbeingAppChangeReceiver.unregisterReceiverSafely(mContext);
-                mAppAddRemoveReceiver.unregisterReceiverSafely(mContext);
+                mWellbeingAppChangeReceiver.unregisterReceiverSafely();
+                mAppAddRemoveReceiver.unregisterReceiverSafely();
                 mContext.getContentResolver().unregisterContentObserver(mContentObserver);
             });
         }
diff --git a/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.kt b/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.kt
index 03f5d96..3773d02 100644
--- a/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.kt
+++ b/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.kt
@@ -19,7 +19,9 @@
 import android.os.Debug
 import android.util.Log
 import android.util.SparseArray
+import android.view.Display.DEFAULT_DISPLAY
 import android.window.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY
+import androidx.core.util.forEach
 import com.android.launcher3.LauncherState
 import com.android.launcher3.dagger.ApplicationContext
 import com.android.launcher3.dagger.LauncherAppComponent
@@ -62,14 +64,14 @@
      * @property canCreateDesks true if it's possible to create new desks on the display represented
      *   by this object.
      * @property activeDeskId The ID of the active desk on the associated display (if any). It has a
-     *   value of `-1` if there are no active desks. Note that there can only be at most one active
-     *   desk on each display.
+     *   value of `INACTIVE_DESK_ID` (-1) if there are no active desks. Note that there can only be
+     *   at most one active desk on each display.
      * @property deskIds a set containing the IDs of the desks on the associated display.
      */
     private data class DisplayDeskConfig(
         val displayId: Int,
         var canCreateDesks: Boolean,
-        var activeDeskId: Int = -1,
+        var activeDeskId: Int = INACTIVE_DESK_ID,
         val deskIds: MutableSet<Int>,
     )
 
@@ -79,6 +81,8 @@
     private val desktopVisibilityListeners: MutableSet<DesktopVisibilityListener> = HashSet()
     private val taskbarDesktopModeListeners: MutableSet<TaskbarDesktopModeListener> = HashSet()
 
+    // TODO: b/394387739 - Deprecate this and replace it with something that tracks the count per
+    //  desk.
     /** Number of visible desktop windows in desktop mode. */
     var visibleDesktopTasksCount: Int = 0
         /**
@@ -103,7 +107,7 @@
                 field = visibleTasksCount
                 val areDesktopTasksVisibleNow = areDesktopTasksVisibleAndNotInOverview()
                 if (wereDesktopTasksVisibleBefore != areDesktopTasksVisibleNow) {
-                    notifyDesktopVisibilityListeners(areDesktopTasksVisibleNow)
+                    notifyIsInDesktopModeChanged(DEFAULT_DISPLAY, areDesktopTasksVisibleNow)
                 }
 
                 if (
@@ -142,8 +146,36 @@
         }
     }
 
+    /** Returns whether a desk is currently active on the display with the given [displayId]. */
+    fun isInDesktopMode(displayId: Int): Boolean {
+        if (!DesktopModeStatus.enableMultipleDesktops(context)) {
+            return areDesktopTasksVisible()
+        }
+
+        val isInDesktopMode = displaysDesksConfigsMap[displayId].activeDeskId != INACTIVE_DESK_ID
+        if (DEBUG) {
+            Log.d(TAG, "isInDesktopMode: $isInDesktopMode")
+        }
+        return isInDesktopMode
+    }
+
+    /**
+     * Returns whether a desk is currently active on the display with the given [displayId] and
+     * Overview is not active.
+     */
+    fun isInDesktopModeAndNotInOverview(displayId: Int): Boolean {
+        if (!DesktopModeStatus.enableMultipleDesktops(context)) {
+            return areDesktopTasksVisibleAndNotInOverview()
+        }
+
+        if (DEBUG) {
+            Log.d(TAG, "isInDesktopModeAndNotInOverview: overview=$inOverviewState")
+        }
+        return isInDesktopMode(displayId) && !inOverviewState
+    }
+
     /** Whether desktop tasks are visible in desktop mode. */
-    fun areDesktopTasksVisible(): Boolean {
+    private fun areDesktopTasksVisible(): Boolean {
         val desktopTasksVisible: Boolean = visibleDesktopTasksCount > 0
         if (DEBUG) {
             Log.d(TAG, "areDesktopTasksVisible: desktopVisible=$desktopTasksVisible")
@@ -152,7 +184,7 @@
     }
 
     /** Whether desktop tasks are visible in desktop mode. */
-    fun areDesktopTasksVisibleAndNotInOverview(): Boolean {
+    private fun areDesktopTasksVisibleAndNotInOverview(): Boolean {
         val desktopTasksVisible: Boolean = visibleDesktopTasksCount > 0
         if (DEBUG) {
             Log.d(
@@ -220,8 +252,23 @@
             val wereDesktopTasksVisibleBefore = areDesktopTasksVisibleAndNotInOverview()
             inOverviewState = overviewStateEnabled
             val areDesktopTasksVisibleNow = areDesktopTasksVisibleAndNotInOverview()
-            if (wereDesktopTasksVisibleBefore != areDesktopTasksVisibleNow) {
-                notifyDesktopVisibilityListeners(areDesktopTasksVisibleNow)
+
+            if (!DesktopModeStatus.enableMultipleDesktops(context)) {
+                if (wereDesktopTasksVisibleBefore != areDesktopTasksVisibleNow) {
+                    notifyIsInDesktopModeChanged(DEFAULT_DISPLAY, areDesktopTasksVisibleNow)
+                }
+            } else {
+                // When overview state changes, it changes together on all displays.
+                displaysDesksConfigsMap.forEach { displayId, deskConfig ->
+                    // Overview affects the state of desks only if desktop mode is active on this
+                    // display.
+                    if (isInDesktopMode(displayId)) {
+                        notifyIsInDesktopModeChanged(
+                            displayId,
+                            isInDesktopModeAndNotInOverview(displayId),
+                        )
+                    }
+                }
             }
 
             if (ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY.isTrue) {
@@ -249,12 +296,19 @@
         desktopVisibilityListeners.remove(listener)
     }
 
-    private fun notifyDesktopVisibilityListeners(areDesktopTasksVisible: Boolean) {
+    private fun notifyIsInDesktopModeChanged(
+        displayId: Int,
+        isInDesktopModeAndNotInOverview: Boolean,
+    ) {
         if (DEBUG) {
-            Log.d(TAG, "notifyDesktopVisibilityListeners: visible=$areDesktopTasksVisible")
+            Log.d(
+                TAG,
+                "notifyIsInDesktopModeChanged: displayId=$displayId, isInDesktopModeAndNotInOverview=$isInDesktopModeAndNotInOverview",
+            )
         }
+
         for (listener in desktopVisibilityListeners) {
-            listener.onDesktopVisibilityChanged(areDesktopTasksVisible)
+            listener.onIsInDesktopModeChanged(displayId, isInDesktopModeAndNotInOverview)
         }
     }
 
@@ -391,7 +445,7 @@
                 "Removing non-existing desk Id: $deskId on display: $displayId"
             }
             if (it.activeDeskId == deskId) {
-                it.activeDeskId = -1
+                it.activeDeskId = INACTIVE_DESK_ID
             }
         }
     }
@@ -401,6 +455,8 @@
             return
         }
 
+        val wasInDesktopMode = isInDesktopModeAndNotInOverview(displayId)
+
         getDisplayDeskConfig(displayId).also {
             check(oldActiveDesk == it.activeDeskId) {
                 "Mismatch between the Shell's oldActiveDesk: $oldActiveDesk, and Launcher's: ${it.activeDeskId}"
@@ -410,6 +466,10 @@
             }
             it.activeDeskId = newActiveDesk
         }
+
+        if (wasInDesktopMode != isInDesktopModeAndNotInOverview(displayId)) {
+            notifyIsInDesktopModeChanged(displayId, !wasInDesktopMode)
+        }
     }
 
     /** TODO: b/333533253 - Remove after flag rollout */
@@ -542,5 +602,7 @@
 
         private const val TAG = "DesktopVisController"
         private const val DEBUG = false
+
+        private const val INACTIVE_DESK_ID = -1
     }
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchTaskView.java b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchTaskView.java
index ce96556..bf5c0c8 100644
--- a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchTaskView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchTaskView.java
@@ -104,9 +104,18 @@
         mIcon2 = findViewById(R.id.icon_2);
         mContent = findViewById(R.id.content);
 
-        Resources resources = mContext.getResources();
-
         Preconditions.assertNotNull(mContent);
+
+        TypefaceUtils.setTypeface(
+                mContent.findViewById(R.id.large_text),
+                TypefaceUtils.FONT_FAMILY_HEADLINE_LARGE_EMPHASIZED
+        );
+        TypefaceUtils.setTypeface(
+                mContent.findViewById(R.id.small_text),
+                TypefaceUtils.FONT_FAMILY_LABEL_LARGE_BASELINE
+        );
+
+        Resources resources = mContext.getResources();
         mBorderAnimator = BorderAnimator.createScalingBorderAnimator(
                 /* borderRadiusPx= */ mBorderRadius != INVALID_BORDER_RADIUS
                         ? mBorderRadius
diff --git a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchView.java b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchView.java
index 4581119..4b4d68d 100644
--- a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchView.java
@@ -163,6 +163,10 @@
                 R.dimen.keyboard_quick_switch_view_small_spacing);
         mOutlineRadius = resources.getDimensionPixelSize(R.dimen.keyboard_quick_switch_view_radius);
         mIsRtl = Utilities.isRtl(resources);
+
+        TypefaceUtils.setTypeface(
+                mNoRecentItemsPane.findViewById(R.id.no_recent_items_text),
+                TypefaceUtils.FONT_FAMILY_LABEL_LARGE_BASELINE);
     }
 
     private void registerOnBackInvokedCallback() {
@@ -310,7 +314,7 @@
                     layoutInflater,
                     previousTaskView);
 
-            desktopButton.<TextView>findViewById(R.id.text).setText(
+            desktopButton.<TextView>findViewById(R.id.small_text).setText(
                     resources.getString(R.string.quick_switch_desktop));
         }
         mDisplayingRecentTasks = !groupTasks.isEmpty() || useDesktopTaskView;
diff --git a/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java b/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
index 4143157..3f3700b 100644
--- a/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
@@ -246,7 +246,7 @@
 
         if (!ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY.isTrue()
                 && mControllers.taskbarDesktopModeController
-                    .getAreDesktopTasksVisibleAndNotInOverview()) {
+                    .isInDesktopModeAndNotInOverview(mLauncher.getDisplayId())) {
             // TODO: b/333533253 - Remove after flag rollout
             isVisible = false;
         }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index 8880abd..b9d1275 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -1291,7 +1291,7 @@
 
     boolean areDesktopTasksVisible() {
         return mControllers != null
-                && mControllers.taskbarDesktopModeController.getAreDesktopTasksVisible();
+                && mControllers.taskbarDesktopModeController.isInDesktopMode(getDisplayId());
     }
 
     protected void onTaskbarIconClicked(View view) {
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
index af60f10..b244be9 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
@@ -194,7 +194,8 @@
                 voiceInteractionWindowController
         };
 
-        if (taskbarDesktopModeController.getAreDesktopTasksVisibleAndNotInOverview()) {
+        if (taskbarDesktopModeController.isInDesktopModeAndNotInOverview(
+                taskbarActivityContext.getDisplayId())) {
             mCornerRoundness.value = taskbarDesktopModeController.getTaskbarCornerRoundness(
                     mSharedState.showCornerRadiusInDesktopMode);
         } else {
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarDesktopModeController.kt b/quickstep/src/com/android/launcher3/taskbar/TaskbarDesktopModeController.kt
index cb399e8..f71dea9 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarDesktopModeController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarDesktopModeController.kt
@@ -30,18 +30,17 @@
     private lateinit var taskbarControllers: TaskbarControllers
     private lateinit var taskbarSharedState: TaskbarSharedState
 
-    val areDesktopTasksVisibleAndNotInOverview: Boolean
-        get() = desktopVisibilityController.areDesktopTasksVisibleAndNotInOverview()
-
-    val areDesktopTasksVisible: Boolean
-        get() = desktopVisibilityController.areDesktopTasksVisible()
-
     fun init(controllers: TaskbarControllers, sharedState: TaskbarSharedState) {
         taskbarControllers = controllers
         taskbarSharedState = sharedState
         desktopVisibilityController.registerTaskbarDesktopModeListener(this)
     }
 
+    fun isInDesktopMode(displayId: Int) = desktopVisibilityController.isInDesktopMode(displayId)
+
+    fun isInDesktopModeAndNotInOverview(displayId: Int) =
+        desktopVisibilityController.isInDesktopModeAndNotInOverview(displayId)
+
     override fun onTaskbarCornerRoundingUpdate(doesAnyTaskRequireTaskbarRounding: Boolean) {
         taskbarSharedState.showCornerRadiusInDesktopMode = doesAnyTaskRequireTaskbarRounding
         val cornerRadius = getTaskbarCornerRoundness(doesAnyTaskRequireTaskbarRounding)
@@ -49,7 +48,7 @@
     }
 
     fun shouldShowDesktopTasksInTaskbar(): Boolean {
-        return desktopVisibilityController.areDesktopTasksVisible() ||
+        return isInDesktopMode(context.displayId) ||
             DisplayController.showDesktopTaskbarForFreeformDisplay(context) ||
             (DisplayController.showLockedTaskbarOnHome(context) &&
                 taskbarControllers.taskbarStashController.isOnHome)
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarDragController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarDragController.java
index f36c481..b91f512 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarDragController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarDragController.java
@@ -347,7 +347,7 @@
         // Pre-drag has ended, start the global system drag.
         if (mDisallowGlobalDrag
                 || mControllers.taskbarDesktopModeController
-                    .getAreDesktopTasksVisibleAndNotInOverview()) {
+                    .isInDesktopModeAndNotInOverview(mActivity.getDisplayId())) {
             AbstractFloatingView.closeAllOpenViewsExcept(mActivity, TYPE_TASKBAR_ALL_APPS);
             return;
         }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java
index cada5a3..dd9f61e 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java
@@ -631,7 +631,8 @@
 
         float cornerRoundness = isInLauncher ? 0 : 1;
 
-        if (mControllers.taskbarDesktopModeController.getAreDesktopTasksVisibleAndNotInOverview()
+        if (mControllers.taskbarDesktopModeController.isInDesktopModeAndNotInOverview(
+                mControllers.taskbarActivityContext.getDisplayId())
                 && mControllers.getSharedState() != null) {
             cornerRoundness =
                     mControllers.taskbarDesktopModeController.getTaskbarCornerRoundness(
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
index f704254..96bcffd 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
@@ -124,8 +124,7 @@
     private final TaskbarNavButtonController mDefaultNavButtonController;
     private final ComponentCallbacks mDefaultComponentCallbacks;
 
-    private final SimpleBroadcastReceiver mShutdownReceiver =
-            new SimpleBroadcastReceiver(UI_HELPER_EXECUTOR, i -> destroyAllTaskbars());
+    private final SimpleBroadcastReceiver mShutdownReceiver;
 
     // The source for this provider is set when Launcher is available
     // We use 'non-destroyable' version here so the original provider won't be destroyed
@@ -183,8 +182,7 @@
 
     private boolean mUserUnlocked = false;
 
-    private final SimpleBroadcastReceiver mTaskbarBroadcastReceiver =
-            new SimpleBroadcastReceiver(UI_HELPER_EXECUTOR, this::showTaskbarFromBroadcast);
+    private final SimpleBroadcastReceiver mTaskbarBroadcastReceiver;
 
     private final AllAppsActionManager mAllAppsActionManager;
     private final RecentsDisplayModel mRecentsDisplayModel;
@@ -266,7 +264,13 @@
                 .register(NAV_BAR_KIDS_MODE, mOnSettingsChangeListener);
         Log.d(TASKBAR_NOT_DESTROYED_TAG, "registering component callbacks from constructor.");
         mPrimaryWindowContext.registerComponentCallbacks(mDefaultComponentCallbacks);
-        mShutdownReceiver.register(mPrimaryWindowContext, Intent.ACTION_SHUTDOWN);
+        mShutdownReceiver =
+                new SimpleBroadcastReceiver(
+                        mPrimaryWindowContext, UI_HELPER_EXECUTOR, i -> destroyAllTaskbars());
+        mTaskbarBroadcastReceiver =
+                new SimpleBroadcastReceiver(mPrimaryWindowContext,
+                        UI_HELPER_EXECUTOR, this::showTaskbarFromBroadcast);
+        mShutdownReceiver.register(Intent.ACTION_SHUTDOWN);
         UI_HELPER_EXECUTOR.execute(() -> {
             mSharedState.taskbarSystemActionPendingIntent = PendingIntent.getBroadcast(
                     mPrimaryWindowContext,
@@ -274,8 +278,7 @@
                     new Intent(ACTION_SHOW_TASKBAR).setPackage(
                             mPrimaryWindowContext.getPackageName()),
                     PendingIntent.FLAG_UPDATE_CURRENT | PendingIntent.FLAG_IMMUTABLE);
-            mTaskbarBroadcastReceiver.register(
-                    mPrimaryWindowContext, RECEIVER_NOT_EXPORTED, ACTION_SHOW_TASKBAR);
+            mTaskbarBroadcastReceiver.register(RECEIVER_NOT_EXPORTED, ACTION_SHOW_TASKBAR);
         });
 
         debugWhyTaskbarNotDestroyed("TaskbarManager created");
@@ -629,6 +632,7 @@
      */
     public void setSetupUIVisible(boolean isVisible) {
         mSharedState.setupUIVisible = isVisible;
+        mAllAppsActionManager.setSetupUiVisible(isVisible);
         TaskbarActivityContext taskbar = getTaskbarForDisplay(getDefaultDisplayId());
         if (taskbar != null) {
             taskbar.setSetupUIVisible(isVisible);
@@ -794,7 +798,7 @@
         mRecentsViewContainer = null;
         debugWhyTaskbarNotDestroyed("TaskbarManager#destroy()");
         removeActivityCallbacksAndListeners();
-        mTaskbarBroadcastReceiver.unregisterReceiverSafely(mPrimaryWindowContext);
+        mTaskbarBroadcastReceiver.unregisterReceiverSafely();
 
         if (mUserUnlocked) {
             DisplayController.INSTANCE.get(mPrimaryWindowContext).removeChangeListener(
@@ -806,7 +810,7 @@
                 .unregister(NAV_BAR_KIDS_MODE, mOnSettingsChangeListener);
         Log.d(TASKBAR_NOT_DESTROYED_TAG, "unregistering component callbacks from destroy().");
         mPrimaryWindowContext.unregisterComponentCallbacks(mDefaultComponentCallbacks);
-        mShutdownReceiver.unregisterReceiverSafely(mPrimaryWindowContext);
+        mShutdownReceiver.unregisterReceiverSafely();
         destroyAllTaskbars();
     }
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarPinningController.kt b/quickstep/src/com/android/launcher3/taskbar/TaskbarPinningController.kt
index 23c5070..7141bb8 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarPinningController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarPinningController.kt
@@ -58,8 +58,9 @@
                 }
                 val shouldPinTaskbar =
                     if (
-                        controllers.taskbarDesktopModeController
-                            .areDesktopTasksVisibleAndNotInOverview
+                        controllers.taskbarDesktopModeController.isInDesktopModeAndNotInOverview(
+                            context.displayId
+                        )
                     ) {
                         !launcherPrefs.get(TASKBAR_PINNING_IN_DESKTOP_MODE)
                     } else {
@@ -140,7 +141,11 @@
     @VisibleForTesting
     fun recreateTaskbarAndUpdatePinningValue() {
         updateIsAnimatingTaskbarPinningAndNotifyTaskbarDragLayer(false)
-        if (controllers.taskbarDesktopModeController.areDesktopTasksVisibleAndNotInOverview) {
+        if (
+            controllers.taskbarDesktopModeController.isInDesktopModeAndNotInOverview(
+                context.displayId
+            )
+        ) {
             launcherPrefs.put(
                 TASKBAR_PINNING_IN_DESKTOP_MODE,
                 !launcherPrefs.get(TASKBAR_PINNING_IN_DESKTOP_MODE),
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarPopupController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarPopupController.java
index feb9b33..5d8b821 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarPopupController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarPopupController.java
@@ -217,7 +217,7 @@
         }
         shortcuts.add(APP_INFO);
         if (!mControllers.taskbarDesktopModeController
-                .getAreDesktopTasksVisibleAndNotInOverview()) {
+                .isInDesktopModeAndNotInOverview(mContext.getDisplayId())) {
             shortcuts.addAll(mControllers.uiController.getSplitMenuOptions().toList());
         }
         if (BubbleAnythingFlagHelper.enableCreateAnyBubble()) {
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
index 2ded1bf..6016394 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
@@ -1189,7 +1189,7 @@
         if (mActivity.isHardwareKeyboard()
                 && mActivity.isThreeButtonNav()
                 && mControllers.taskbarDesktopModeController
-                    .getAreDesktopTasksVisibleAndNotInOverview()) {
+                    .isInDesktopModeAndNotInOverview(mActivity.getDisplayId())) {
             return false;
         }
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java
index e5d642d..89bcb41 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java
@@ -283,7 +283,7 @@
                                     foundTask,
                                     taskContainer.getIconView().getDrawable(),
                                     taskContainer.getSnapshotView(),
-                                    taskContainer.getSplitAnimationThumbnail(),
+                                    taskContainer.getThumbnail(),
                                     null /* intent */,
                                     null /* user */,
                                     info);
diff --git a/quickstep/src/com/android/launcher3/taskbar/TypefaceUtils.kt b/quickstep/src/com/android/launcher3/taskbar/TypefaceUtils.kt
new file mode 100644
index 0000000..fa551b8
--- /dev/null
+++ b/quickstep/src/com/android/launcher3/taskbar/TypefaceUtils.kt
@@ -0,0 +1,47 @@
+/*
+ * Copyright (C) 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.taskbar
+
+import android.graphics.Typeface
+import android.widget.TextView
+import com.android.launcher3.Flags
+
+/**
+ * Helper util class to set pre-defined typefaces to textviews
+ *
+ * If the typeface font family is already defined here, you can just reuse it directly. Otherwise,
+ * please define it here for future use. You do not need to define the font style. If you need
+ * anything other than [Typeface.NORMAL], pass it inline when calling [setTypeface]
+ */
+class TypefaceUtils {
+
+    companion object {
+        const val FONT_FAMILY_HEADLINE_LARGE_EMPHASIZED = "variable-headline-large-emphasized"
+        const val FONT_FAMILY_LABEL_LARGE_BASELINE = "variable-label-large"
+
+        @JvmStatic
+        @JvmOverloads
+        fun setTypeface(
+            textView: TextView?,
+            fontFamilyName: String,
+            fontStyle: Int = Typeface.NORMAL,
+        ) {
+            if (!Flags.expressiveThemeInTaskbarAndNavigation()) return
+            textView?.typeface = Typeface.create(fontFamilyName, fontStyle)
+        }
+    }
+}
diff --git a/quickstep/src/com/android/launcher3/uioverrides/PredictedAppIcon.java b/quickstep/src/com/android/launcher3/uioverrides/PredictedAppIcon.java
index b25f999..37b8dc7 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/PredictedAppIcon.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/PredictedAppIcon.java
@@ -51,7 +51,7 @@
 import com.android.launcher3.anim.AnimatorListeners;
 import com.android.launcher3.celllayout.CellLayoutLayoutParams;
 import com.android.launcher3.celllayout.DelegatedCellDrawing;
-import com.android.launcher3.graphics.IconShape;
+import com.android.launcher3.graphics.ThemeManager;
 import com.android.launcher3.icons.FastBitmapDrawable;
 import com.android.launcher3.icons.IconNormalizer;
 import com.android.launcher3.icons.LauncherIcons;
@@ -142,7 +142,7 @@
         int shadowSize = context.getResources().getDimensionPixelSize(
                 R.dimen.blur_size_thin_outline);
         mShadowFilter = new BlurMaskFilter(shadowSize, BlurMaskFilter.Blur.OUTER);
-        mShapePath = IconShape.INSTANCE.get(context).getShape().getPath(mNormalizedIconSize);
+        mShapePath = ThemeManager.INSTANCE.get(context).getIconShape().getPath(mNormalizedIconSize);
     }
 
     @Override
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
index 2111a80..bb57d6e 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
@@ -86,7 +86,6 @@
 import android.os.Trace;
 import android.os.UserHandle;
 import android.util.AttributeSet;
-import android.util.Log;
 import android.view.Display;
 import android.view.HapticFeedbackConstants;
 import android.view.KeyEvent;
@@ -247,7 +246,6 @@
     private SplitWithKeyboardShortcutController mSplitWithKeyboardShortcutController;
     private SplitToWorkspaceController mSplitToWorkspaceController;
     private BubbleBarLocation mBubbleBarLocation;
-    private static final String TRACKING_BUG = "b/395214062";
 
     /**
      * If Launcher restarted while in the middle of an Overview split select, it needs this data to
@@ -563,7 +561,6 @@
 
     @Override
     public void onDestroy() {
-        Log.d(TRACKING_BUG, "onDestroy: " + this.hashCode());
         if (mAppTransitionManager != null) {
             mAppTransitionManager.onActivityDestroyed();
         }
@@ -589,10 +586,7 @@
 
         RecentsView recentsView = getOverviewPanel();
         if (recentsView != null) {
-            Log.d(TRACKING_BUG, "onDestroy - recentsView.destroy(): " + this.hashCode());
             recentsView.destroy();
-        } else {
-            Log.d(TRACKING_BUG, "onDestroy - recentsView is null: " + this.hashCode());
         }
 
         super.onDestroy();
@@ -719,7 +713,6 @@
     @Override
     protected void onCreate(Bundle savedInstanceState) {
         super.onCreate(savedInstanceState);
-        Log.d(TRACKING_BUG, "onCreate: " + this.hashCode());
         if (savedInstanceState != null) {
             mPendingSplitSelectInfo = ObjectWrapper.unwrap(
                     savedInstanceState.getIBinder(PENDING_SPLIT_SELECT_INFO));
@@ -832,7 +825,7 @@
     @Override
     protected void onResume() {
         super.onResume();
-        Log.d(TRACKING_BUG, "onResume: " + this.hashCode());
+
         if (mLauncherUnfoldAnimationController != null) {
             mLauncherUnfoldAnimationController.onResume();
         }
@@ -867,7 +860,6 @@
     @Override
     protected void onStop() {
         super.onStop();
-        Log.d(TRACKING_BUG, "onStop: " + this.hashCode());
         if (mTaskbarUIController != null && FeatureFlags.enableHomeTransitionListener()) {
             mTaskbarUIController.onLauncherStop();
         }
@@ -1055,7 +1047,7 @@
         DesktopVisibilityController desktopVisibilityController =
                 DesktopVisibilityController.INSTANCE.get(this);
         if (!ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY.isTrue()
-                && desktopVisibilityController.areDesktopTasksVisibleAndNotInOverview()
+                && desktopVisibilityController.isInDesktopModeAndNotInOverview(getDisplayId())
                 && !desktopVisibilityController.isRecentsGestureInProgress()) {
             // Return early to skip setting activity to appear as resumed
             // TODO: b/333533253 - Remove after flag rollout
@@ -1377,7 +1369,7 @@
     @Override
     public boolean areDesktopTasksVisible() {
         return DesktopVisibilityController.INSTANCE.get(this)
-                .areDesktopTasksVisibleAndNotInOverview();
+                .isInDesktopModeAndNotInOverview(getDisplayId());
     }
 
     @Override
diff --git a/quickstep/src/com/android/launcher3/uioverrides/touchcontrollers/TaskViewDismissTouchController.kt b/quickstep/src/com/android/launcher3/uioverrides/touchcontrollers/TaskViewDismissTouchController.kt
index 77a05c1..98737a5 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/touchcontrollers/TaskViewDismissTouchController.kt
+++ b/quickstep/src/com/android/launcher3/uioverrides/touchcontrollers/TaskViewDismissTouchController.kt
@@ -26,10 +26,12 @@
 import com.android.launcher3.Utilities.isRtl
 import com.android.launcher3.Utilities.mapToRange
 import com.android.launcher3.touch.SingleAxisSwipeDetector
+import com.android.launcher3.util.MSDLPlayerWrapper
 import com.android.launcher3.util.TouchController
 import com.android.quickstep.views.RecentsView
 import com.android.quickstep.views.RecentsViewContainer
 import com.android.quickstep.views.TaskView
+import com.google.android.msdl.data.model.MSDLToken
 import kotlin.math.abs
 import kotlin.math.sign
 
@@ -53,6 +55,7 @@
     private var springAnimation: SpringAnimation? = null
     private var dismissLength: Int = 0
     private var verticalFactor: Int = 0
+    private var hasDismissThresholdHapticRun = false
     private var initialDisplacement: Float = 0f
 
     private fun canInterceptTouch(ev: MotionEvent): Boolean =
@@ -159,9 +162,30 @@
             }
             recentsView.redrawLiveTile()
         }
+        playDismissThresholdHaptic(displacement)
         return true
     }
 
+    /**
+     * Play a haptic to alert the user they have passed the dismiss threshold.
+     *
+     * <p>Check within a range of the threshold value, as the drag event does not necessarily happen
+     * at the exact threshold's displacement.
+     */
+    private fun playDismissThresholdHaptic(displacement: Float) {
+        val dismissThreshold = (DISMISS_THRESHOLD_FRACTION * dismissLength * verticalFactor)
+        val inHapticRange =
+            displacement >= (dismissThreshold - DISMISS_THRESHOLD_HAPTIC_RANGE) &&
+                displacement <= (dismissThreshold + DISMISS_THRESHOLD_HAPTIC_RANGE)
+        if (!inHapticRange) {
+            hasDismissThresholdHapticRun = false
+        } else if (!hasDismissThresholdHapticRun) {
+            MSDLPlayerWrapper.INSTANCE.get(recentsView.context)
+                .playToken(MSDLToken.SWIPE_THRESHOLD_INDICATOR)
+            hasDismissThresholdHapticRun = true
+        }
+    }
+
     override fun onDragEnd(velocity: Float) {
         val taskBeingDragged = taskBeingDragged ?: return
 
@@ -208,5 +232,6 @@
 
     companion object {
         private const val DISMISS_THRESHOLD_FRACTION = 0.5f
+        private const val DISMISS_THRESHOLD_HAPTIC_RANGE = 10f
     }
 }
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index f46f9ae..b43c3ac 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -2453,7 +2453,8 @@
     }
 
     @Override
-    public void onTasksAppeared(@NonNull RemoteAnimationTarget[] appearedTaskTargets) {
+    public void onTasksAppeared(@NonNull RemoteAnimationTarget[] appearedTaskTargets,
+            @Nullable TransitionInfo transitionInfo) {
         if (mRecentsAnimationController == null) {
             return;
         }
diff --git a/quickstep/src/com/android/quickstep/AllAppsActionManager.kt b/quickstep/src/com/android/quickstep/AllAppsActionManager.kt
index 6fd68d5..b807a4b 100644
--- a/quickstep/src/com/android/quickstep/AllAppsActionManager.kt
+++ b/quickstep/src/com/android/quickstep/AllAppsActionManager.kt
@@ -21,10 +21,16 @@
 import android.app.RemoteAction
 import android.content.Context
 import android.graphics.drawable.Icon
+import android.provider.Settings
+import android.provider.Settings.Secure.USER_SETUP_COMPLETE
 import android.view.accessibility.AccessibilityManager
 import com.android.launcher3.R
+import com.android.launcher3.util.SettingsCache
+import com.android.launcher3.util.SettingsCache.OnChangeListener
 import java.util.concurrent.Executor
 
+private val USER_SETUP_COMPLETE_URI = Settings.Secure.getUriFor(USER_SETUP_COMPLETE)
+
 /**
  * Registers a [RemoteAction] for toggling All Apps if needed.
  *
@@ -38,6 +44,12 @@
     private val createAllAppsPendingIntent: () -> PendingIntent,
 ) {
 
+    private val onSettingsChangeListener = OnChangeListener { v -> isUserSetupComplete = v }
+
+    init {
+        SettingsCache.INSTANCE[context].register(USER_SETUP_COMPLETE_URI, onSettingsChangeListener)
+    }
+
     /** `true` if home and overview are the same Activity. */
     var isHomeAndOverviewSame = false
         set(value) {
@@ -52,12 +64,27 @@
             updateSystemAction()
         }
 
+    /** `true` if the setup UI is visible. */
+    var isSetupUiVisible = false
+        set(value) {
+            field = value
+            updateSystemAction()
+        }
+
+    private var isUserSetupComplete =
+        SettingsCache.INSTANCE[context].getValue(USER_SETUP_COMPLETE_URI, 0)
+        set(value) {
+            field = value
+            updateSystemAction()
+        }
+
     /** `true` if the action should be registered. */
     var isActionRegistered = false
         private set
 
     private fun updateSystemAction() {
-        val shouldRegisterAction = isHomeAndOverviewSame || isTaskbarPresent
+        val isInSetupFlow = isSetupUiVisible || !isUserSetupComplete
+        val shouldRegisterAction = (isHomeAndOverviewSame || isTaskbarPresent) && !isInSetupFlow
         if (isActionRegistered == shouldRegisterAction) return
         isActionRegistered = shouldRegisterAction
 
@@ -84,8 +111,10 @@
         isActionRegistered = false
         context
             .getSystemService(AccessibilityManager::class.java)
-            ?.unregisterSystemAction(
-                GLOBAL_ACTION_ACCESSIBILITY_ALL_APPS,
-            )
+            ?.unregisterSystemAction(GLOBAL_ACTION_ACCESSIBILITY_ALL_APPS)
+        SettingsCache.INSTANCE[context].unregister(
+            USER_SETUP_COMPLETE_URI,
+            onSettingsChangeListener,
+        )
     }
 }
diff --git a/quickstep/src/com/android/quickstep/BaseActivityInterface.java b/quickstep/src/com/android/quickstep/BaseActivityInterface.java
index 7cab751..549c2f8 100644
--- a/quickstep/src/com/android/quickstep/BaseActivityInterface.java
+++ b/quickstep/src/com/android/quickstep/BaseActivityInterface.java
@@ -87,7 +87,7 @@
             // We were on our way to this state when we got canceled, end there instead.
             startState = stateFromGestureEndTarget(endTarget);
             if (DesktopVisibilityController.INSTANCE.get(activity)
-                    .areDesktopTasksVisibleAndNotInOverview()
+                    .isInDesktopModeAndNotInOverview(activity.getDisplayId())
                     && endTarget == LAST_TASK) {
                 // When we are cancelling the transition and going back to last task, move to
                 // rest state instead when desktop tasks are visible.
diff --git a/quickstep/src/com/android/quickstep/BaseContainerInterface.java b/quickstep/src/com/android/quickstep/BaseContainerInterface.java
index 6d588d9..c64067a 100644
--- a/quickstep/src/com/android/quickstep/BaseContainerInterface.java
+++ b/quickstep/src/com/android/quickstep/BaseContainerInterface.java
@@ -233,8 +233,10 @@
         if (endTarget != null) {
             // We were on our way to this state when we got canceled, end there instead.
             startState = stateFromGestureEndTarget(endTarget);
-            if (DesktopVisibilityController.INSTANCE.get(recentsView.getContext())
-                    .areDesktopTasksVisibleAndNotInOverview() && endTarget == LAST_TASK) {
+            final var context = recentsView.getContext();
+            if (DesktopVisibilityController.INSTANCE.get(context)
+                    .isInDesktopModeAndNotInOverview(context.getDisplayId())
+                    && endTarget == LAST_TASK) {
                 // When we are cancelling the transition and going back to last task, move to
                 // rest state instead when desktop tasks are visible.
                 // If a fullscreen task is visible, launcher goes to normal state when the
diff --git a/quickstep/src/com/android/quickstep/FallbackSwipeHandler.java b/quickstep/src/com/android/quickstep/FallbackSwipeHandler.java
index 331580c..7d8a53d 100644
--- a/quickstep/src/com/android/quickstep/FallbackSwipeHandler.java
+++ b/quickstep/src/com/android/quickstep/FallbackSwipeHandler.java
@@ -50,6 +50,7 @@
 import android.view.Surface;
 import android.view.SurfaceControl;
 import android.view.SurfaceControl.Transaction;
+import android.window.TransitionInfo;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
@@ -172,14 +173,15 @@
     }
 
     @Override
-    public void onTasksAppeared(@NonNull RemoteAnimationTarget[] appearedTaskTargets) {
+    public void onTasksAppeared(@NonNull RemoteAnimationTarget[] appearedTaskTargets,
+            @Nullable TransitionInfo transitionInfo) {
         if (mActiveAnimationFactory != null && mActiveAnimationFactory.handleHomeTaskAppeared(
                 appearedTaskTargets)) {
             mActiveAnimationFactory = null;
             return;
         }
 
-        super.onTasksAppeared(appearedTaskTargets);
+        super.onTasksAppeared(appearedTaskTargets, transitionInfo);
     }
 
     @Override
diff --git a/quickstep/src/com/android/quickstep/LauncherSwipeHandlerV2.java b/quickstep/src/com/android/quickstep/LauncherSwipeHandlerV2.java
index e1e9c99..0a77688 100644
--- a/quickstep/src/com/android/quickstep/LauncherSwipeHandlerV2.java
+++ b/quickstep/src/com/android/quickstep/LauncherSwipeHandlerV2.java
@@ -107,7 +107,7 @@
                 && workspaceView.isAttachedToWindow()
                 && workspaceView.getHeight() > 0
                 && !DesktopVisibilityController.INSTANCE.get(mContainer)
-                        .areDesktopTasksVisibleAndNotInOverview();
+                        .isInDesktopModeAndNotInOverview(mContainer.getDisplayId());
 
         mContainer.getRootView().setForceHideBackArrow(true);
 
diff --git a/quickstep/src/com/android/quickstep/OverviewComponentObserver.java b/quickstep/src/com/android/quickstep/OverviewComponentObserver.java
index 1f95c41..bc3de41 100644
--- a/quickstep/src/com/android/quickstep/OverviewComponentObserver.java
+++ b/quickstep/src/com/android/quickstep/OverviewComponentObserver.java
@@ -72,12 +72,9 @@
             new DaggerSingletonObject<>(LauncherAppComponent::getOverviewComponentObserver);
 
     // We register broadcast receivers on main thread to avoid missing updates.
-    private final SimpleBroadcastReceiver mUserPreferenceChangeReceiver =
-            new SimpleBroadcastReceiver(MAIN_EXECUTOR, this::updateOverviewTargets);
-    private final SimpleBroadcastReceiver mOtherHomeAppUpdateReceiver =
-            new SimpleBroadcastReceiver(MAIN_EXECUTOR, this::updateOverviewTargets);
+    private final SimpleBroadcastReceiver mUserPreferenceChangeReceiver;
+    private final SimpleBroadcastReceiver mOtherHomeAppUpdateReceiver;
 
-    private final Context mContext;
     private final RecentsDisplayModel mRecentsDisplayModel;
 
     private final Intent mCurrentHomeIntent;
@@ -101,10 +98,13 @@
             @ApplicationContext Context context,
             RecentsDisplayModel recentsDisplayModel,
             DaggerSingletonTracker lifecycleTracker) {
-        mContext = context;
+        mUserPreferenceChangeReceiver =
+                new SimpleBroadcastReceiver(context, MAIN_EXECUTOR, this::updateOverviewTargets);
+        mOtherHomeAppUpdateReceiver =
+                new SimpleBroadcastReceiver(context, MAIN_EXECUTOR, this::updateOverviewTargets);
         mRecentsDisplayModel = recentsDisplayModel;
         mCurrentHomeIntent = createHomeIntent();
-        mMyHomeIntent = new Intent(mCurrentHomeIntent).setPackage(mContext.getPackageName());
+        mMyHomeIntent = new Intent(mCurrentHomeIntent).setPackage(context.getPackageName());
         ResolveInfo info = context.getPackageManager().resolveActivity(mMyHomeIntent, 0);
         ComponentName myHomeComponent =
                 new ComponentName(context.getPackageName(), info.activityInfo.name);
@@ -112,7 +112,7 @@
         mConfigChangesMap.append(myHomeComponent.hashCode(), info.activityInfo.configChanges);
         mSetupWizardPkg = context.getString(R.string.setup_wizard_pkg);
 
-        ComponentName fallbackComponent = new ComponentName(mContext, RecentsActivity.class);
+        ComponentName fallbackComponent = new ComponentName(context, RecentsActivity.class);
         mFallbackIntent = new Intent(Intent.ACTION_MAIN)
                 .addCategory(Intent.CATEGORY_DEFAULT)
                 .setComponent(fallbackComponent)
@@ -124,7 +124,7 @@
             mConfigChangesMap.append(fallbackComponent.hashCode(), fallbackInfo.configChanges);
         } catch (PackageManager.NameNotFoundException ignored) { /* Impossible */ }
 
-        mUserPreferenceChangeReceiver.register(mContext, ACTION_PREFERRED_ACTIVITY_CHANGED);
+        mUserPreferenceChangeReceiver.register(ACTION_PREFERRED_ACTIVITY_CHANGED);
         updateOverviewTargets();
 
         lifecycleTracker.addCloseable(this::onDestroy);
@@ -224,7 +224,7 @@
 
                 mUpdateRegisteredPackage = defaultHome.getPackageName();
                 mOtherHomeAppUpdateReceiver.registerPkgActions(
-                        mContext, mUpdateRegisteredPackage, ACTION_PACKAGE_ADDED,
+                        mUpdateRegisteredPackage, ACTION_PACKAGE_ADDED,
                         ACTION_PACKAGE_CHANGED, ACTION_PACKAGE_REMOVED);
             }
         }
@@ -235,13 +235,13 @@
      * Clean up any registered receivers.
      */
     private void onDestroy() {
-        mUserPreferenceChangeReceiver.unregisterReceiverSafely(mContext);
+        mUserPreferenceChangeReceiver.unregisterReceiverSafely();
         unregisterOtherHomeAppUpdateReceiver();
     }
 
     private void unregisterOtherHomeAppUpdateReceiver() {
         if (mUpdateRegisteredPackage != null) {
-            mOtherHomeAppUpdateReceiver.unregisterReceiverSafely(mContext);
+            mOtherHomeAppUpdateReceiver.unregisterReceiverSafely();
             mUpdateRegisteredPackage = null;
         }
     }
diff --git a/quickstep/src/com/android/quickstep/RecentsActivity.java b/quickstep/src/com/android/quickstep/RecentsActivity.java
index fca67c3..3d12fdf 100644
--- a/quickstep/src/com/android/quickstep/RecentsActivity.java
+++ b/quickstep/src/com/android/quickstep/RecentsActivity.java
@@ -319,7 +319,7 @@
     /**
      * Composes the animations for a launch from the recents list if possible.
      */
-    private AnimatorSet  composeRecentsLaunchAnimator(
+    private AnimatorSet composeRecentsLaunchAnimator(
             @NonNull RecentsView recentsView,
             @NonNull TaskView taskView,
             RemoteAnimationTarget[] appTargets,
@@ -329,7 +329,8 @@
         boolean activityClosing = taskIsATargetWithMode(appTargets, getTaskId(), MODE_CLOSING);
         PendingAnimation pa = new PendingAnimation(RECENTS_LAUNCH_DURATION);
         createRecentsWindowAnimator(recentsView, taskView, !activityClosing, appTargets,
-                wallpaperTargets, nonAppTargets, null /* depthController */, pa);
+                wallpaperTargets, nonAppTargets, /* depthController= */ null ,
+                /* transitionInfo= */ null, pa);
         target.play(pa.buildAnim());
 
         // Found a visible recents task that matches the opening app, lets launch the app from there
diff --git a/quickstep/src/com/android/quickstep/RecentsAnimationCallbacks.java b/quickstep/src/com/android/quickstep/RecentsAnimationCallbacks.java
index c6b858b..d7152b5 100644
--- a/quickstep/src/com/android/quickstep/RecentsAnimationCallbacks.java
+++ b/quickstep/src/com/android/quickstep/RecentsAnimationCallbacks.java
@@ -21,6 +21,7 @@
 import static android.view.WindowManager.LayoutParams.TYPE_DOCK_DIVIDER;
 
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
+import static com.android.wm.shell.shared.TransitionUtil.TYPE_SPLIT_SCREEN_DIM_LAYER;
 
 import android.annotation.Nullable;
 import android.graphics.Rect;
@@ -163,11 +164,12 @@
 
     @BinderThread
     @Override
-    public void onTasksAppeared(RemoteAnimationTarget[] apps) {
+    public void onTasksAppeared(
+            RemoteAnimationTarget[] apps, @Nullable TransitionInfo transitionInfo) {
         Utilities.postAsyncCallback(MAIN_EXECUTOR.getHandler(), () -> {
             ActiveGestureProtoLogProxy.logRecentsAnimationCallbacksOnTasksAppeared();
             for (RecentsAnimationListener listener : getListeners()) {
-                listener.onTasksAppeared(apps);
+                listener.onTasksAppeared(apps, transitionInfo);
             }
         });
     }
@@ -189,7 +191,8 @@
             ArrayList<RemoteAnimationTarget> apps, ArrayList<RemoteAnimationTarget> nonApps) {
         for (int i = 0; i < appTargets.length; i++) {
             RemoteAnimationTarget target = appTargets[i];
-            if (target.windowType == TYPE_DOCK_DIVIDER) {
+            if (target.windowType == TYPE_DOCK_DIVIDER
+                    || target.windowType == TYPE_SPLIT_SCREEN_DIM_LAYER) {
                 nonApps.add(target);
             } else {
                 apps.add(target);
@@ -225,6 +228,7 @@
         /**
          * Callback made when a task started from the recents is ready for an app transition.
          */
-        default void onTasksAppeared(@NonNull RemoteAnimationTarget[] appearedTaskTarget) {}
+        default void onTasksAppeared(@NonNull RemoteAnimationTarget[] appearedTaskTarget,
+                @Nullable TransitionInfo transitionInfo) {}
     }
 }
diff --git a/quickstep/src/com/android/quickstep/SystemUiProxy.kt b/quickstep/src/com/android/quickstep/SystemUiProxy.kt
index a1ac39e..1f3eb2a 100644
--- a/quickstep/src/com/android/quickstep/SystemUiProxy.kt
+++ b/quickstep/src/com/android/quickstep/SystemUiProxy.kt
@@ -1218,8 +1218,12 @@
         override fun onAnimationCanceled(taskIds: IntArray?, taskSnapshots: Array<TaskSnapshot>?) =
             listener.onAnimationCanceled(wrap(taskIds, taskSnapshots))
 
-        override fun onTasksAppeared(apps: Array<RemoteAnimationTarget>?) =
-            listener.onTasksAppeared(apps)
+        override fun onTasksAppeared(
+            apps: Array<RemoteAnimationTarget>?,
+            transitionInfo: TransitionInfo?,
+        ) {
+            listener.onTasksAppeared(apps, transitionInfo)
+        }
     }
 
     //
diff --git a/quickstep/src/com/android/quickstep/TaskAnimationManager.java b/quickstep/src/com/android/quickstep/TaskAnimationManager.java
index 63b8aaf..9810308 100644
--- a/quickstep/src/com/android/quickstep/TaskAnimationManager.java
+++ b/quickstep/src/com/android/quickstep/TaskAnimationManager.java
@@ -231,7 +231,8 @@
             }
 
             @Override
-            public void onTasksAppeared(RemoteAnimationTarget[] appearedTaskTargets) {
+            public void onTasksAppeared(RemoteAnimationTarget[] appearedTaskTargets,
+                    @Nullable TransitionInfo transitionInfo) {
                 RemoteAnimationTarget appearedTaskTarget = appearedTaskTargets[0];
                 BaseContainerInterface containerInterface =
                         mLastGestureState.getContainerInterface();
@@ -264,7 +265,8 @@
                         recentsView.launchSideTaskInLiveTileMode(appearedTaskTarget.taskId,
                                 appearedTaskTargets,
                                 new RemoteAnimationTarget[0] /* wallpaper */,
-                                nonAppTargets /* nonApps */);
+                                nonAppTargets /* nonApps */,
+                                transitionInfo);
                         return;
                     } else {
                         ActiveGestureProtoLogProxy.logLaunchingSideTaskFailed();
diff --git a/quickstep/src/com/android/quickstep/TaskShortcutFactory.java b/quickstep/src/com/android/quickstep/TaskShortcutFactory.java
index 7990aae..a69d472 100644
--- a/quickstep/src/com/android/quickstep/TaskShortcutFactory.java
+++ b/quickstep/src/com/android/quickstep/TaskShortcutFactory.java
@@ -236,15 +236,14 @@
                         position[0] + width, position[1] + height);
 
                 // Take the thumbnail of the task without a scrim and apply it back after
-                // TODO(b/348643341) add ability to get override the scrim for this Bitmap retrieval
-                float alpha = 0f;
-                if (!enableRefactorTaskThumbnail()) {
-                    alpha = mTaskContainer.getThumbnailViewDeprecated().getDimAlpha();
+                Bitmap thumbnail;
+                if (enableRefactorTaskThumbnail()) {
+                    thumbnail = mTaskContainer.getThumbnail();
+                } else {
+                    float alpha = mTaskContainer.getThumbnailViewDeprecated().getDimAlpha();
                     mTaskContainer.getThumbnailViewDeprecated().setDimAlpha(0);
-                }
-                Bitmap thumbnail = RecentsTransition.drawViewIntoHardwareBitmap(
-                        taskBounds.width(), taskBounds.height(), snapShotView, 1f, Color.BLACK);
-                if (!enableRefactorTaskThumbnail()) {
+                    thumbnail = RecentsTransition.drawViewIntoHardwareBitmap(
+                            taskBounds.width(), taskBounds.height(), snapShotView, 1f, Color.BLACK);
                     mTaskContainer.getThumbnailViewDeprecated().setDimAlpha(alpha);
                 }
 
diff --git a/quickstep/src/com/android/quickstep/TaskViewUtils.java b/quickstep/src/com/android/quickstep/TaskViewUtils.java
index e47223b..35a6e72 100644
--- a/quickstep/src/com/android/quickstep/TaskViewUtils.java
+++ b/quickstep/src/com/android/quickstep/TaskViewUtils.java
@@ -39,6 +39,7 @@
 import static com.android.launcher3.Utilities.getDescendantCoordRelativeToAncestor;
 import static com.android.launcher3.util.MultiPropertyFactory.MULTI_PROPERTY_VALUE;
 import static com.android.quickstep.util.AnimUtils.clampToDuration;
+import static com.android.wm.shell.shared.TransitionUtil.TYPE_SPLIT_SCREEN_DIM_LAYER;
 
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
@@ -169,6 +170,7 @@
             @NonNull RemoteAnimationTarget[] wallpaperTargets,
             @NonNull RemoteAnimationTarget[] nonAppTargets,
             @Nullable DepthController depthController,
+            @Nullable TransitionInfo transitionInfo,
             PendingAnimation out) {
         boolean isQuickSwitch = v.isEndQuickSwitchCuj();
         v.setEndQuickSwitchCuj(false);
@@ -191,8 +193,7 @@
             RemoteTargetGluer gluer = new RemoteTargetGluer(v.getContext(),
                     recentsView.getSizeStrategy(), targets, forDesktop);
             if (forDesktop) {
-                remoteTargetHandles =
-                        gluer.assignTargetsForDesktop(targets, /* transitionInfo=*/ null);
+                remoteTargetHandles = gluer.assignTargetsForDesktop(targets, transitionInfo);
             } else if (v.containsMultipleTasks()) {
                 remoteTargetHandles = gluer.assignTargetsForSplitScreen(targets,
                         ((GroupedTaskView) v).getSplitBoundsConfig());
@@ -462,7 +463,7 @@
         final RecentsView recentsView = launchingTaskView.getRecentsView();
         composeRecentsLaunchAnimator(animatorSet, launchingTaskView, appTargets, wallpaperTargets,
                 nonAppTargets, /* launcherClosing */ true, stateManager, recentsView,
-                depthController);
+                depthController, /* transitionInfo= */ null);
 
         t.apply();
         animatorSet.start();
@@ -501,7 +502,7 @@
             composeRecentsLaunchAnimator(animatorSet, launchingTaskView,
                     appTargets, wallpaperTargets, nonAppTargets,
                     true, stateManager,
-                    recentsView, depthController);
+                    recentsView, depthController, /* transitionInfo= */ null);
             animatorSet.start();
             return;
         }
@@ -593,7 +594,7 @@
 
         composeRecentsLaunchAnimator(animatorSet, launchingTaskView, apps, wallpaper, nonApps,
                 true /* launcherClosing */, stateManager, launchingTaskView.getRecentsView(),
-                depthController);
+                depthController, transitionInfo);
 
         return animatorSet;
     }
@@ -603,13 +604,13 @@
             @NonNull RemoteAnimationTarget[] wallpaperTargets,
             @NonNull RemoteAnimationTarget[] nonAppTargets, boolean launcherClosing,
             @NonNull StateManager stateManager, @NonNull RecentsView recentsView,
-            @Nullable DepthController depthController) {
+            @Nullable DepthController depthController, @Nullable TransitionInfo transitionInfo) {
         boolean skipLauncherChanges = !launcherClosing;
 
         TaskView taskView = findTaskViewToLaunch(recentsView, v, appTargets);
         PendingAnimation pa = new PendingAnimation(RECENTS_LAUNCH_DURATION);
         createRecentsWindowAnimator(recentsView, taskView, skipLauncherChanges, appTargets,
-                wallpaperTargets, nonAppTargets, depthController, pa);
+                wallpaperTargets, nonAppTargets, depthController, transitionInfo, pa);
         if (launcherClosing) {
             // TODO(b/182592057): differentiate between "restore split" vs "launch fullscreen app"
             TaskViewUtils.createSplitAuxiliarySurfacesAnimator(nonAppTargets, true /*shown*/,
@@ -732,7 +733,9 @@
         List<SurfaceControl> auxiliarySurfaces = new ArrayList<>();
         for (RemoteAnimationTarget target : nonApps) {
             final SurfaceControl leash = target.leash;
-            if (target.windowType == TYPE_DOCK_DIVIDER && leash != null && leash.isValid()) {
+            if ((target.windowType == TYPE_DOCK_DIVIDER
+                    || target.windowType == TYPE_SPLIT_SCREEN_DIM_LAYER)
+                    && leash != null && leash.isValid()) {
                 auxiliarySurfaces.add(leash);
             }
         }
diff --git a/quickstep/src/com/android/quickstep/recents/di/RecentsDependencies.kt b/quickstep/src/com/android/quickstep/recents/di/RecentsDependencies.kt
index 977629f..4b3bef3 100644
--- a/quickstep/src/com/android/quickstep/recents/di/RecentsDependencies.kt
+++ b/quickstep/src/com/android/quickstep/recents/di/RecentsDependencies.kt
@@ -28,15 +28,14 @@
 import com.android.quickstep.recents.data.TasksRepository
 import com.android.quickstep.recents.domain.usecase.GetSysUiStatusNavFlagsUseCase
 import com.android.quickstep.recents.domain.usecase.GetTaskUseCase
+import com.android.quickstep.recents.domain.usecase.GetThumbnailPositionUseCase
 import com.android.quickstep.recents.domain.usecase.IsThumbnailValidUseCase
 import com.android.quickstep.recents.domain.usecase.OrganizeDesktopTasksUseCase
-import com.android.quickstep.recents.usecase.GetThumbnailPositionUseCase
 import com.android.quickstep.recents.usecase.GetThumbnailUseCase
 import com.android.quickstep.recents.viewmodel.RecentsViewData
 import com.android.quickstep.task.viewmodel.TaskOverlayViewModel
-import com.android.quickstep.task.viewmodel.TaskThumbnailViewModel
-import com.android.quickstep.task.viewmodel.TaskThumbnailViewModelImpl
 import com.android.systemui.shared.recents.model.Task
+import com.android.systemui.shared.recents.utilities.PreviewPositionHelper
 import kotlinx.coroutines.CoroutineName
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.SupervisorJob
@@ -173,8 +172,6 @@
         val instance: Any =
             when (modelClass) {
                 RecentsViewData::class.java -> RecentsViewData()
-                TaskThumbnailViewModel::class.java ->
-                    TaskThumbnailViewModelImpl(getThumbnailPositionUseCase = inject())
                 TaskOverlayViewModel::class.java -> {
                     val task = extras["Task"] as Task
                     TaskOverlayViewModel(
@@ -194,7 +191,7 @@
                     GetThumbnailPositionUseCase(
                         deviceProfileRepository = inject(),
                         rotationStateRepository = inject(),
-                        tasksRepository = inject(),
+                        previewPositionHelper = PreviewPositionHelper(),
                     )
                 OrganizeDesktopTasksUseCase::class.java -> OrganizeDesktopTasksUseCase()
                 else -> {
@@ -237,7 +234,7 @@
         fun initialize(view: View): RecentsDependencies = initialize(view.context)
 
         fun initialize(context: Context): RecentsDependencies {
-            Log.d(TAG, "initializing: $activeRecentsCount + 1 more")
+            Log.d(TAG, "initializing")
             synchronized(this) {
                 activeRecentsCount++
                 instance = RecentsDependencies(context.applicationContext)
@@ -277,7 +274,7 @@
                 Log.d(
                     TAG,
                     "RecentsDependencies was not destroyed. " +
-                        "There is still an active RecentsView instance: $activeRecentsCount",
+                        "There is still an active RecentsView instance.",
                 )
             }
         }
diff --git a/quickstep/src/com/android/quickstep/recents/usecase/GetThumbnailPositionUseCase.kt b/quickstep/src/com/android/quickstep/recents/domain/usecase/GetThumbnailPositionUseCase.kt
similarity index 62%
rename from quickstep/src/com/android/quickstep/recents/usecase/GetThumbnailPositionUseCase.kt
rename to quickstep/src/com/android/quickstep/recents/domain/usecase/GetThumbnailPositionUseCase.kt
index bea1d07..8501382 100644
--- a/quickstep/src/com/android/quickstep/recents/usecase/GetThumbnailPositionUseCase.kt
+++ b/quickstep/src/com/android/quickstep/recents/domain/usecase/GetThumbnailPositionUseCase.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2024 The Android Open Source Project
+ * Copyright (C) 2025 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.
@@ -14,28 +14,30 @@
  * limitations under the License.
  */
 
-package com.android.quickstep.recents.usecase
+package com.android.quickstep.recents.domain.usecase
 
 import android.graphics.Matrix
 import android.graphics.Rect
-import com.android.quickstep.recents.data.RecentTasksRepository
 import com.android.quickstep.recents.data.RecentsDeviceProfileRepository
 import com.android.quickstep.recents.data.RecentsRotationStateRepository
-import com.android.quickstep.recents.usecase.ThumbnailPositionState.MatrixScaling
-import com.android.quickstep.recents.usecase.ThumbnailPositionState.MissingThumbnail
+import com.android.systemui.shared.recents.model.ThumbnailData
 import com.android.systemui.shared.recents.utilities.PreviewPositionHelper
 
 /** Use case for retrieving [Matrix] for positioning Thumbnail in a View */
 class GetThumbnailPositionUseCase(
     private val deviceProfileRepository: RecentsDeviceProfileRepository,
     private val rotationStateRepository: RecentsRotationStateRepository,
-    private val tasksRepository: RecentTasksRepository,
-    private val previewPositionHelper: PreviewPositionHelper = PreviewPositionHelper(),
+    private val previewPositionHelper: PreviewPositionHelper,
 ) {
-    fun run(taskId: Int, width: Int, height: Int, isRtl: Boolean): ThumbnailPositionState {
-        val thumbnailData =
-            tasksRepository.getCurrentThumbnailById(taskId) ?: return MissingThumbnail
-        val thumbnail = thumbnailData.thumbnail ?: return MissingThumbnail
+    operator fun invoke(
+        thumbnailData: ThumbnailData?,
+        width: Int,
+        height: Int,
+        isRtl: Boolean,
+    ): ThumbnailPosition {
+        val thumbnail =
+            thumbnailData?.thumbnail ?: return ThumbnailPosition(Matrix.IDENTITY_MATRIX, false)
+
         previewPositionHelper.updateThumbnailMatrix(
             Rect(0, 0, thumbnail.width, thumbnail.height),
             thumbnailData,
@@ -45,9 +47,11 @@
             rotationStateRepository.getRecentsRotationState().activityRotation,
             isRtl,
         )
-        return MatrixScaling(
-            previewPositionHelper.matrix,
-            previewPositionHelper.isOrientationChanged,
+        return ThumbnailPosition(
+            matrix = previewPositionHelper.matrix,
+            isRotated = previewPositionHelper.isOrientationChanged,
         )
     }
 }
+
+data class ThumbnailPosition(val matrix: Matrix, val isRotated: Boolean)
diff --git a/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModel.kt b/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModel.kt
index 0b299ee..3c4a384 100644
--- a/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModel.kt
+++ b/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModel.kt
@@ -24,7 +24,9 @@
 import com.android.quickstep.recents.domain.model.TaskModel
 import com.android.quickstep.recents.domain.usecase.GetSysUiStatusNavFlagsUseCase
 import com.android.quickstep.recents.domain.usecase.GetTaskUseCase
+import com.android.quickstep.recents.domain.usecase.GetThumbnailPositionUseCase
 import com.android.quickstep.recents.domain.usecase.IsThumbnailValidUseCase
+import com.android.quickstep.recents.domain.usecase.ThumbnailPosition
 import com.android.quickstep.recents.viewmodel.RecentsViewData
 import com.android.quickstep.views.TaskViewType
 import com.android.systemui.shared.recents.model.ThumbnailData
@@ -48,6 +50,7 @@
     private val getTaskUseCase: GetTaskUseCase,
     private val getSysUiStatusNavFlagsUseCase: GetSysUiStatusNavFlagsUseCase,
     private val isThumbnailValidUseCase: IsThumbnailValidUseCase,
+    private val getThumbnailPositionUseCase: GetThumbnailPositionUseCase,
     dispatcherProvider: DispatcherProvider,
 ) {
     private var taskIds = MutableStateFlow(emptySet<Int>())
@@ -83,6 +86,19 @@
     fun isThumbnailValid(thumbnail: ThumbnailData?, width: Int, height: Int): Boolean =
         isThumbnailValidUseCase(thumbnail, width, height)
 
+    fun getThumbnailPosition(
+        thumbnail: ThumbnailData?,
+        width: Int,
+        height: Int,
+        isRtl: Boolean,
+    ): ThumbnailPosition =
+        getThumbnailPositionUseCase(
+            thumbnailData = thumbnail,
+            width = width,
+            height = height,
+            isRtl = isRtl,
+        )
+
     private fun mapToTaskTile(tasks: List<TaskData>, isLiveTile: Boolean): TaskTileUiState {
         val firstThumbnailData = (tasks.firstOrNull() as? TaskData.Data)?.thumbnailData
         return TaskTileUiState(
diff --git a/quickstep/src/com/android/quickstep/recents/usecase/ThumbnailPositionState.kt b/quickstep/src/com/android/quickstep/recents/usecase/ThumbnailPositionState.kt
deleted file mode 100644
index 1a1bef7..0000000
--- a/quickstep/src/com/android/quickstep/recents/usecase/ThumbnailPositionState.kt
+++ /dev/null
@@ -1,26 +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.quickstep.recents.usecase
-
-import android.graphics.Matrix
-
-/** State on how a task Thumbnail can fit on given canvas */
-sealed class ThumbnailPositionState {
-    data object MissingThumbnail : ThumbnailPositionState()
-
-    data class MatrixScaling(val matrix: Matrix, val isRotated: Boolean) : ThumbnailPositionState()
-}
diff --git a/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt b/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt
index 63e93ba..e672ec4 100644
--- a/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt
+++ b/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt
@@ -18,6 +18,7 @@
 
 import android.content.Context
 import android.graphics.Color
+import android.graphics.Matrix
 import android.graphics.Outline
 import android.graphics.Rect
 import android.graphics.drawable.ShapeDrawable
@@ -34,34 +35,15 @@
 import com.android.launcher3.R
 import com.android.launcher3.util.MultiPropertyFactory
 import com.android.launcher3.util.ViewPool
-import com.android.launcher3.util.coroutines.DispatcherProvider
-import com.android.quickstep.recents.di.RecentsDependencies
-import com.android.quickstep.recents.di.get
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.BackgroundOnly
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.LiveTile
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.Snapshot
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.SnapshotSplash
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.Uninitialized
-import com.android.quickstep.task.viewmodel.TaskThumbnailViewModel
 import com.android.quickstep.views.FixedSizeImageView
 import com.android.quickstep.views.TaskThumbnailViewHeader
-import kotlinx.coroutines.CoroutineName
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.SupervisorJob
-import kotlinx.coroutines.cancel
-import kotlinx.coroutines.launch
 
 class TaskThumbnailView : FrameLayout, ViewPool.Reusable {
-    private val recentsCoroutineScope: CoroutineScope = RecentsDependencies.get()
-    private val dispatcherProvider: DispatcherProvider = RecentsDependencies.get()
-
-    // This is initialised here and set in onAttachedToWindow because onLayout can be called before
-    // onAttachedToWindow so this property needs to be initialised as it is used below.
-    private lateinit var viewModel: TaskThumbnailViewModel
-
-    private lateinit var viewAttachedScope: CoroutineScope
-
     private val scrimView: View by lazy { findViewById(R.id.task_thumbnail_scrim) }
     private val liveTileView: LiveTileView by lazy { findViewById(R.id.task_thumbnail_live_tile) }
     private val thumbnailView: FixedSizeImageView by lazy { findViewById(R.id.task_thumbnail) }
@@ -70,10 +52,21 @@
     private val dimAlpha: MultiPropertyFactory<View> by lazy {
         MultiPropertyFactory(scrimView, VIEW_ALPHA, ScrimViewAlpha.entries.size, ::maxOf)
     }
+    private var onSizeChanged: ((width: Int, height: Int) -> Unit)? = null
 
     private var taskThumbnailViewHeader: TaskThumbnailViewHeader? = null
 
     private var uiState: TaskThumbnailUiState = Uninitialized
+
+    /**
+     * Sets the outline bounds of the view. Default to use view's bound as outline when set to null.
+     */
+    var outlineBounds: Rect? = null
+        set(value) {
+            field = value
+            invalidateOutline()
+        }
+
     private val bounds = Rect()
 
     var cornerRadius: Float = 0f
@@ -94,36 +87,24 @@
 
     override fun onFinishInflate() {
         super.onFinishInflate()
-
         maybeCreateHeader()
     }
 
     override fun onAttachedToWindow() {
         super.onAttachedToWindow()
-        viewAttachedScope =
-            CoroutineScope(
-                SupervisorJob() + Dispatchers.Main.immediate + CoroutineName("TaskThumbnailView")
-            )
-        viewModel = RecentsDependencies.get(this)
         clipToOutline = true
         outlineProvider =
             object : ViewOutlineProvider() {
                 override fun getOutline(view: View, outline: Outline) {
-                    outline.setRoundRect(bounds, cornerRadius)
+                    outline.setRoundRect(outlineBounds ?: bounds, cornerRadius)
                 }
             }
     }
 
-    // TODO(b/391842220): Cancel scope in onDetach instead of having a specific method for this.
-    fun destroyScopes() {
-        val scopeToCancel = viewAttachedScope
-        recentsCoroutineScope.launch(dispatcherProvider.background) {
-            scopeToCancel.cancel("TaskThumbnailView detaching from window")
-        }
-    }
-
     override fun onRecycle() {
         uiState = Uninitialized
+        onSizeChanged = null
+        outlineBounds = null
         resetViews()
     }
 
@@ -159,11 +140,13 @@
         splashIcon.alpha = value
     }
 
+    fun doOnSizeChange(action: (width: Int, height: Int) -> Unit) {
+        onSizeChanged = action
+    }
+
     override fun onSizeChanged(w: Int, h: Int, oldw: Int, oldh: Int) {
         super.onSizeChanged(w, h, oldw, oldh)
-        if (uiState is SnapshotSplash) {
-            setImageMatrix()
-        }
+        onSizeChanged?.invoke(width, height)
         bounds.set(0, 0, w, h)
         invalidateOutline()
     }
@@ -220,11 +203,12 @@
         drawBackground(snapshot.backgroundColor)
         thumbnailView.setImageBitmap(snapshot.bitmap)
         thumbnailView.isInvisible = false
-        setImageMatrix()
     }
 
-    private fun setImageMatrix() {
-        thumbnailView.imageMatrix = viewModel.getThumbnailPositionState(width, height, isLayoutRtl)
+    fun setImageMatrix(matrix: Matrix) {
+        if (uiState is SnapshotSplash) {
+            thumbnailView.imageMatrix = matrix
+        }
     }
 
     private fun logDebug(message: String) {
diff --git a/quickstep/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModel.kt b/quickstep/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModel.kt
index 81a904b..9bff3ac 100644
--- a/quickstep/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModel.kt
+++ b/quickstep/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModel.kt
@@ -19,9 +19,7 @@
 import android.graphics.Matrix
 import com.android.launcher3.util.coroutines.DispatcherProvider
 import com.android.quickstep.recents.data.RecentTasksRepository
-import com.android.quickstep.recents.usecase.GetThumbnailPositionUseCase
-import com.android.quickstep.recents.usecase.ThumbnailPositionState.MatrixScaling
-import com.android.quickstep.recents.usecase.ThumbnailPositionState.MissingThumbnail
+import com.android.quickstep.recents.domain.usecase.GetThumbnailPositionUseCase
 import com.android.quickstep.recents.viewmodel.RecentsViewData
 import com.android.quickstep.task.thumbnail.TaskOverlayUiState.Disabled
 import com.android.quickstep.task.thumbnail.TaskOverlayUiState.Enabled
@@ -36,7 +34,7 @@
     private val task: Task,
     recentsViewData: RecentsViewData,
     private val getThumbnailPositionUseCase: GetThumbnailPositionUseCase,
-    recentTasksRepository: RecentTasksRepository,
+    private val recentTasksRepository: RecentTasksRepository,
     dispatcherProvider: DispatcherProvider,
 ) {
     val overlayState =
@@ -60,22 +58,17 @@
             .flowOn(dispatcherProvider.background)
 
     fun getThumbnailPositionState(width: Int, height: Int, isRtl: Boolean): ThumbnailPositionState {
-        val matrix: Matrix
-        val isRotated: Boolean
-        when (
-            val thumbnailPositionState =
-                getThumbnailPositionUseCase.run(task.key.id, width, height, isRtl)
-        ) {
-            is MatrixScaling -> {
-                matrix = thumbnailPositionState.matrix
-                isRotated = thumbnailPositionState.isRotated
-            }
-            is MissingThumbnail -> {
-                matrix = Matrix.IDENTITY_MATRIX
-                isRotated = false
-            }
-        }
-        return ThumbnailPositionState(matrix, isRotated)
+        val thumbnailPositionState =
+            getThumbnailPositionUseCase(
+                thumbnailData = recentTasksRepository.getCurrentThumbnailById(task.key.id),
+                width = width,
+                height = height,
+                isRtl = isRtl,
+            )
+        return ThumbnailPositionState(
+            thumbnailPositionState.matrix,
+            thumbnailPositionState.isRotated,
+        )
     }
 
     data class ThumbnailPositionState(val matrix: Matrix, val isRotated: Boolean)
diff --git a/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModel.kt b/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModel.kt
deleted file mode 100644
index e641737..0000000
--- a/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModel.kt
+++ /dev/null
@@ -1,28 +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.quickstep.task.viewmodel
-
-import android.graphics.Matrix
-
-/** ViewModel for representing TaskThumbnails */
-interface TaskThumbnailViewModel {
-    /** Attaches this ViewModel to a specific task id for it to provide data from. */
-    fun bind(taskId: Int)
-
-    /** Returns a Matrix which can be applied to the snapshot */
-    fun getThumbnailPositionState(width: Int, height: Int, isRtl: Boolean): Matrix
-}
diff --git a/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModelImpl.kt b/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModelImpl.kt
deleted file mode 100644
index 94c40d1..0000000
--- a/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModelImpl.kt
+++ /dev/null
@@ -1,47 +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 goveryning permissions and
- * limitations under the License.
- */
-
-package com.android.quickstep.task.viewmodel
-
-import android.app.ActivityTaskManager.INVALID_TASK_ID
-import android.graphics.Matrix
-import android.util.Log
-import com.android.quickstep.recents.usecase.GetThumbnailPositionUseCase
-import com.android.quickstep.recents.usecase.ThumbnailPositionState
-
-class TaskThumbnailViewModelImpl(
-    private val getThumbnailPositionUseCase: GetThumbnailPositionUseCase
-) : TaskThumbnailViewModel {
-    private var taskId: Int = INVALID_TASK_ID
-
-    override fun bind(taskId: Int) {
-        Log.d(TAG, "bind taskId: $taskId")
-        this.taskId = taskId
-    }
-
-    override fun getThumbnailPositionState(width: Int, height: Int, isRtl: Boolean): Matrix =
-        when (
-            val thumbnailPositionState =
-                getThumbnailPositionUseCase.run(taskId, width, height, isRtl)
-        ) {
-            is ThumbnailPositionState.MatrixScaling -> thumbnailPositionState.matrix
-            is ThumbnailPositionState.MissingThumbnail -> Matrix.IDENTITY_MATRIX
-        }
-
-    private companion object {
-        const val TAG = "TaskThumbnailViewModel"
-    }
-}
diff --git a/quickstep/src/com/android/quickstep/util/AsyncClockEventDelegate.java b/quickstep/src/com/android/quickstep/util/AsyncClockEventDelegate.java
index 54f6443..207e482 100644
--- a/quickstep/src/com/android/quickstep/util/AsyncClockEventDelegate.java
+++ b/quickstep/src/com/android/quickstep/util/AsyncClockEventDelegate.java
@@ -59,8 +59,7 @@
 
     private final Context mContext;
     private final SettingsCache mSettingsCache;
-    private final SimpleBroadcastReceiver mReceiver =
-            new SimpleBroadcastReceiver(UI_HELPER_EXECUTOR, this::onClockEventReceived);
+    private final SimpleBroadcastReceiver mReceiver;
 
     private final ArrayMap<BroadcastReceiver, Handler> mTimeEventReceivers = new ArrayMap<>();
     private final List<ContentObserver> mFormatObservers = new ArrayList<>();
@@ -76,7 +75,9 @@
         super(context);
         mContext = context;
         mSettingsCache = settingsCache;
-        mReceiver.register(mContext, ACTION_TIME_CHANGED, ACTION_TIMEZONE_CHANGED);
+        mReceiver = new SimpleBroadcastReceiver(
+                context, UI_HELPER_EXECUTOR, this::onClockEventReceived);
+        mReceiver.register(ACTION_TIME_CHANGED, ACTION_TIMEZONE_CHANGED);
         tracker.addCloseable(this);
     }
 
@@ -138,6 +139,6 @@
     public void close() {
         mDestroyed = true;
         mSettingsCache.unregister(mFormatUri, this);
-        mReceiver.unregisterReceiverSafely(mContext);
+        mReceiver.unregisterReceiverSafely();
     }
 }
diff --git a/quickstep/src/com/android/quickstep/util/ContextualSearchStateManager.java b/quickstep/src/com/android/quickstep/util/ContextualSearchStateManager.java
index ed96399..a8d3c6d 100644
--- a/quickstep/src/com/android/quickstep/util/ContextualSearchStateManager.java
+++ b/quickstep/src/com/android/quickstep/util/ContextualSearchStateManager.java
@@ -74,8 +74,7 @@
             Settings.Secure.getUriFor(Settings.Secure.SEARCH_ALL_ENTRYPOINTS_ENABLED);
 
     private final Runnable mSysUiStateChangeListener = this::updateOverridesToSysUi;
-    private final SimpleBroadcastReceiver mContextualSearchPackageReceiver =
-            new SimpleBroadcastReceiver(UI_HELPER_EXECUTOR, (unused) -> requestUpdateProperties());
+    private final SimpleBroadcastReceiver mContextualSearchPackageReceiver;
     protected final EventLogArray mEventLogArray = new EventLogArray(TAG, MAX_DEBUG_EVENT_SIZE);
 
     // Cached value whether the ContextualSearch intent filter matched any enabled components.
@@ -95,6 +94,9 @@
             TopTaskTracker topTaskTracker,
             DaggerSingletonTracker lifeCycle) {
         mContext = context;
+        mContextualSearchPackageReceiver =
+                new SimpleBroadcastReceiver(context, UI_HELPER_EXECUTOR,
+                        (unused) -> requestUpdateProperties());
         mContextualSearchPackage = mContext.getResources().getString(
                 com.android.internal.R.string.config_defaultContextualSearchPackageName);
         mSystemUiProxy = systemUiProxy;
@@ -112,7 +114,7 @@
         requestUpdateProperties();
         registerSearchScreenSystemAction();
         mContextualSearchPackageReceiver.registerPkgActions(
-                context, mContextualSearchPackage, Intent.ACTION_PACKAGE_ADDED,
+                mContextualSearchPackage, Intent.ACTION_PACKAGE_ADDED,
                 Intent.ACTION_PACKAGE_CHANGED, Intent.ACTION_PACKAGE_REMOVED);
 
         SettingsCache.OnChangeListener settingChangedListener =
@@ -124,7 +126,7 @@
         systemUiProxy.addOnStateChangeListener(mSysUiStateChangeListener);
 
         lifeCycle.addCloseable(() -> {
-            mContextualSearchPackageReceiver.unregisterReceiverSafely(mContext);
+            mContextualSearchPackageReceiver.unregisterReceiverSafely();
             unregisterSearchScreenSystemAction();
             settingsCache.unregister(SEARCH_ALL_ENTRYPOINTS_ENABLED_URI, settingChangedListener);
             systemUiProxy.removeOnStateChangeListener(mSysUiStateChangeListener);
diff --git a/quickstep/src/com/android/quickstep/util/SplitAnimationController.kt b/quickstep/src/com/android/quickstep/util/SplitAnimationController.kt
index 88d14b7..a9dbbf2 100644
--- a/quickstep/src/com/android/quickstep/util/SplitAnimationController.kt
+++ b/quickstep/src/com/android/quickstep/util/SplitAnimationController.kt
@@ -127,7 +127,7 @@
                     val drawable = getDrawable(container.iconView, splitSelectSource)
                     return SplitAnimInitProps(
                         container.snapshotView,
-                        container.splitAnimationThumbnail,
+                        container.thumbnail,
                         drawable,
                         fadeWithThumbnail = true,
                         isStagedTask = true,
@@ -147,7 +147,7 @@
                 val drawable = getDrawable(it.iconView, splitSelectSource)
                 return SplitAnimInitProps(
                     it.snapshotView,
-                    it.splitAnimationThumbnail,
+                    it.thumbnail,
                     drawable,
                     fadeWithThumbnail = true,
                     isStagedTask = true,
@@ -215,13 +215,13 @@
         if (enableOverviewIconMenu()) {
             builder.add(
                 ObjectAnimator.ofFloat(
-                    (iconView as IconAppChipView).splitTranslationX,
+                    (iconView as IconAppChipView).getSplitTranslationX(),
                     MULTI_PROPERTY_VALUE,
                     0f,
                 )
             )
             builder.add(
-                ObjectAnimator.ofFloat(iconView.splitTranslationY, MULTI_PROPERTY_VALUE, 0f)
+                ObjectAnimator.ofFloat(iconView.getSplitTranslationY(), MULTI_PROPERTY_VALUE, 0f)
             )
         }
 
@@ -985,12 +985,11 @@
         val splitTree: Pair<Change, List<Change>>? = extractTopParentAndChildren(transitionInfo)
         check(splitTree != null) { "Could not find a split root candidate" }
         val rootCandidate = splitTree.first
-        val stageRootTaskIds: Set<Int> = splitTree.second
-            .map { it.taskInfo!!.taskId }
-            .toSet()
-        val leafTasks: List<Change> = transitionInfo.changes
-            .filter { it.taskInfo != null && it.taskInfo!!.parentTaskId in stageRootTaskIds}
-            .toList()
+        val stageRootTaskIds: Set<Int> = splitTree.second.map { it.taskInfo!!.taskId }.toSet()
+        val leafTasks: List<Change> =
+            transitionInfo.changes
+                .filter { it.taskInfo != null && it.taskInfo!!.parentTaskId in stageRootTaskIds }
+                .toList()
 
         // Starting position is a 34% size tile centered in the middle of the screen.
         // Ending position is the full device screen.
@@ -1031,8 +1030,13 @@
                         val endAbsBounds = leaf.endAbsBounds
 
                         t.setAlpha(leaf.leash, 1f)
-                        t.setCrop(leaf.leash, 0f, 0f,
-                            endAbsBounds.width().toFloat(), endAbsBounds.height().toFloat())
+                        t.setCrop(
+                            leaf.leash,
+                            0f,
+                            0f,
+                            endAbsBounds.width().toFloat(),
+                            endAbsBounds.height().toFloat(),
+                        )
                         t.setPosition(leaf.leash, 0f, 0f)
                     }
 
@@ -1040,10 +1044,18 @@
                         val endAbsBounds = stageRoot.endAbsBounds
 
                         t.setAlpha(stageRoot.leash, 1f)
-                        t.setCrop(stageRoot.leash, 0f, 0f,
-                            endAbsBounds.width().toFloat(), endAbsBounds.height().toFloat())
-                        t.setPosition(stageRoot.leash, endAbsBounds.left.toFloat(),
-                            endAbsBounds.top.toFloat())
+                        t.setCrop(
+                            stageRoot.leash,
+                            0f,
+                            0f,
+                            endAbsBounds.width().toFloat(),
+                            endAbsBounds.height().toFloat(),
+                        )
+                        t.setPosition(
+                            stageRoot.leash,
+                            endAbsBounds.left.toFloat(),
+                            endAbsBounds.top.toFloat(),
+                        )
                     }
                     t.apply()
                 }
diff --git a/quickstep/src/com/android/quickstep/util/SystemWindowManagerProxy.java b/quickstep/src/com/android/quickstep/util/SystemWindowManagerProxy.java
index 10ae7a3..d92cc86 100644
--- a/quickstep/src/com/android/quickstep/util/SystemWindowManagerProxy.java
+++ b/quickstep/src/com/android/quickstep/util/SystemWindowManagerProxy.java
@@ -73,8 +73,8 @@
     }
 
     @Override
-    public boolean isInDesktopMode() {
-        return mDesktopVisibilityController.areDesktopTasksVisible();
+    public boolean isInDesktopMode(int displayId) {
+        return mDesktopVisibilityController.isInDesktopMode(displayId);
     }
 
     @Override
diff --git a/quickstep/src/com/android/quickstep/util/TaskGridNavHelper.java b/quickstep/src/com/android/quickstep/util/TaskGridNavHelper.java
deleted file mode 100644
index da26622..0000000
--- a/quickstep/src/com/android/quickstep/util/TaskGridNavHelper.java
+++ /dev/null
@@ -1,140 +0,0 @@
-/*
- * Copyright (C) 2023 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.util;
-
-import static java.lang.annotation.RetentionPolicy.SOURCE;
-
-import androidx.annotation.IntDef;
-
-import com.android.launcher3.util.IntArray;
-
-import java.lang.annotation.Retention;
-import java.util.List;
-
-/**
- * Helper class for navigating RecentsView grid tasks via arrow keys and tab.
- */
-public class TaskGridNavHelper {
-    public static final int CLEAR_ALL_PLACEHOLDER_ID = -1;
-    public static final int ADD_DESK_PLACEHOLDER_ID = -2;
-
-    public static final int DIRECTION_UP = 0;
-    public static final int DIRECTION_DOWN = 1;
-    public static final int DIRECTION_LEFT = 2;
-    public static final int DIRECTION_RIGHT = 3;
-    public static final int DIRECTION_TAB = 4;
-
-    @Retention(SOURCE)
-    @IntDef({DIRECTION_UP, DIRECTION_DOWN, DIRECTION_LEFT, DIRECTION_RIGHT, DIRECTION_TAB})
-    public @interface TASK_NAV_DIRECTION {}
-
-    private final IntArray mOriginalTopRowIds;
-    private final IntArray mTopRowIds = new IntArray();
-    private final IntArray mBottomRowIds = new IntArray();
-
-    public TaskGridNavHelper(IntArray topIds, IntArray bottomIds,
-            List<Integer> largeTileIds, boolean hasAddDesktopButton) {
-        mOriginalTopRowIds = topIds.clone();
-        generateTaskViewIdGrid(topIds, bottomIds, largeTileIds, hasAddDesktopButton);
-    }
-
-    private void generateTaskViewIdGrid(IntArray topRowIdArray, IntArray bottomRowIdArray,
-            List<Integer> largeTileIds, boolean hasAddDesktopButton) {
-        // Add AddDesktopButton and lage tiles to both rows.
-        if (hasAddDesktopButton) {
-            mTopRowIds.add(ADD_DESK_PLACEHOLDER_ID);
-            mBottomRowIds.add(ADD_DESK_PLACEHOLDER_ID);
-        }
-        for (Integer tileId : largeTileIds) {
-            mTopRowIds.add(tileId);
-            mBottomRowIds.add(tileId);
-        }
-
-        // Add row ids to their respective rows.
-        mTopRowIds.addAll(topRowIdArray);
-        mBottomRowIds.addAll(bottomRowIdArray);
-
-        // Fill in the shorter array with the ids from the longer one.
-        while (mTopRowIds.size() > mBottomRowIds.size()) {
-            mBottomRowIds.add(mTopRowIds.get(mBottomRowIds.size()));
-        }
-        while (mBottomRowIds.size() > mTopRowIds.size()) {
-            mTopRowIds.add(mBottomRowIds.get(mTopRowIds.size()));
-        }
-
-        // Add the clear all button to the end of both arrays.
-        mTopRowIds.add(CLEAR_ALL_PLACEHOLDER_ID);
-        mBottomRowIds.add(CLEAR_ALL_PLACEHOLDER_ID);
-    }
-
-    /**
-     * Returns the id of the next page in the grid or -1 for the clear all button.
-     */
-    public int getNextGridPage(int currentPageTaskViewId, int delta,
-            @TASK_NAV_DIRECTION int direction, boolean cycle) {
-        boolean inTop = mTopRowIds.contains(currentPageTaskViewId);
-        int index = inTop ? mTopRowIds.indexOf(currentPageTaskViewId)
-                : mBottomRowIds.indexOf(currentPageTaskViewId);
-        int maxSize = Math.max(mTopRowIds.size(), mBottomRowIds.size());
-        int nextIndex = index + delta;
-
-        switch (direction) {
-            case DIRECTION_UP:
-            case DIRECTION_DOWN: {
-                return inTop ? mBottomRowIds.get(index) : mTopRowIds.get(index);
-            }
-            case DIRECTION_LEFT: {
-                int boundedIndex = cycle ? nextIndex % maxSize : Math.min(nextIndex, maxSize - 1);
-                return inTop ? mTopRowIds.get(boundedIndex)
-                        : mBottomRowIds.get(boundedIndex);
-            }
-            case DIRECTION_RIGHT: {
-                int boundedIndex =
-                        cycle ? (nextIndex < 0 ? maxSize - 1 : nextIndex)
-                                : Math.max(nextIndex, 0);
-                boolean inOriginalTop = mOriginalTopRowIds.contains(currentPageTaskViewId);
-                return inOriginalTop ? mTopRowIds.get(boundedIndex)
-                        : mBottomRowIds.get(boundedIndex);
-            }
-            case DIRECTION_TAB: {
-                int boundedIndex =
-                        cycle ? (nextIndex < 0 ? maxSize - 1 : nextIndex % maxSize)
-                                : Math.min(nextIndex, maxSize - 1);
-                if (delta >= 0) {
-                    return inTop && mTopRowIds.get(index) != mBottomRowIds.get(index)
-                            ? mBottomRowIds.get(index)
-                            : mTopRowIds.get(boundedIndex);
-                } else {
-                    if (mTopRowIds.contains(currentPageTaskViewId)) {
-                        if (boundedIndex < 0) {
-                            // If no cycling, always return the first task.
-                            return mTopRowIds.get(0);
-                        } else {
-                            return mBottomRowIds.get(boundedIndex);
-                        }
-                    } else {
-                        // Go up to top if there is task above
-                        return mTopRowIds.get(index) != mBottomRowIds.get(index)
-                                ? mTopRowIds.get(index)
-                                : mBottomRowIds.get(boundedIndex);
-                    }
-                }
-            }
-            default:
-                return currentPageTaskViewId;
-        }
-    }
-}
diff --git a/quickstep/src/com/android/quickstep/util/TaskGridNavHelper.kt b/quickstep/src/com/android/quickstep/util/TaskGridNavHelper.kt
new file mode 100644
index 0000000..0e78801
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/util/TaskGridNavHelper.kt
@@ -0,0 +1,159 @@
+/*
+ * Copyright (C) 2025 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.util
+
+import com.android.launcher3.util.IntArray
+import kotlin.math.abs
+import kotlin.math.max
+
+/** Helper class for navigating RecentsView grid tasks via arrow keys and tab. */
+class TaskGridNavHelper(
+    private val topIds: IntArray,
+    bottomIds: IntArray,
+    largeTileIds: List<Int>,
+    hasAddDesktopButton: Boolean,
+) {
+    private val topRowIds = mutableListOf<Int>()
+    private val bottomRowIds = mutableListOf<Int>()
+
+    init {
+        // Add AddDesktopButton and lage tiles to both rows.
+        if (hasAddDesktopButton) {
+            topRowIds += ADD_DESK_PLACEHOLDER_ID
+            bottomRowIds += ADD_DESK_PLACEHOLDER_ID
+        }
+        topRowIds += largeTileIds
+        bottomRowIds += largeTileIds
+
+        // Add row ids to their respective rows.
+        topRowIds += topIds
+        bottomRowIds += bottomIds
+
+        // Fill in the shorter array with the ids from the longer one.
+        topRowIds += bottomRowIds.takeLast(max(bottomRowIds.size - topRowIds.size, 0))
+        bottomRowIds += topRowIds.takeLast(max(topRowIds.size - bottomRowIds.size, 0))
+
+        // Add the clear all button to the end of both arrays.
+        topRowIds += CLEAR_ALL_PLACEHOLDER_ID
+        bottomRowIds += CLEAR_ALL_PLACEHOLDER_ID
+    }
+
+    /** Returns the id of the next page in the grid or -1 for the clear all button. */
+    fun getNextGridPage(
+        currentPageTaskViewId: Int,
+        delta: Int,
+        direction: TaskNavDirection,
+        cycle: Boolean,
+    ): Int {
+        val inTop = topRowIds.contains(currentPageTaskViewId)
+        val index =
+            if (inTop) topRowIds.indexOf(currentPageTaskViewId)
+            else bottomRowIds.indexOf(currentPageTaskViewId)
+        val maxSize = max(topRowIds.size, bottomRowIds.size)
+        val nextIndex = index + delta
+
+        return when (direction) {
+            TaskNavDirection.UP,
+            TaskNavDirection.DOWN -> {
+                if (inTop) bottomRowIds[index] else topRowIds[index]
+            }
+            TaskNavDirection.LEFT -> {
+                val boundedIndex =
+                    if (cycle) nextIndex % maxSize else nextIndex.coerceAtMost(maxSize - 1)
+                if (inTop) topRowIds[boundedIndex] else bottomRowIds[boundedIndex]
+            }
+            TaskNavDirection.RIGHT -> {
+                val boundedIndex =
+                    if (cycle) (if (nextIndex < 0) maxSize - 1 else nextIndex)
+                    else nextIndex.coerceAtLeast(0)
+                val inOriginalTop = topIds.contains(currentPageTaskViewId)
+                if (inOriginalTop) topRowIds[boundedIndex] else bottomRowIds[boundedIndex]
+            }
+            TaskNavDirection.TAB -> {
+                val boundedIndex =
+                    if (cycle) (if (nextIndex < 0) maxSize - 1 else nextIndex % maxSize)
+                    else nextIndex.coerceAtMost(maxSize - 1)
+                if (delta >= 0) {
+                    if (inTop && topRowIds[index] != bottomRowIds[index]) bottomRowIds[index]
+                    else topRowIds[boundedIndex]
+                } else {
+                    if (topRowIds.contains(currentPageTaskViewId)) {
+                        if (boundedIndex < 0) {
+                            // If no cycling, always return the first task.
+                            topRowIds[0]
+                        } else {
+                            bottomRowIds[boundedIndex]
+                        }
+                    } else {
+                        // Go up to top if there is task above
+                        if (topRowIds[index] != bottomRowIds[index]) topRowIds[index]
+                        else bottomRowIds[boundedIndex]
+                    }
+                }
+            }
+            else -> currentPageTaskViewId
+        }
+    }
+
+    /**
+     * Returns a sequence of pairs of (TaskView ID, offset) in the grid, ordered according to tab
+     * navigation, starting from the initial TaskView ID, towards the start or end of the grid.
+     *
+     * <p>A positive delta moves forward in the tab order towards the end of the grid, while a
+     * negative value moves backward towards the beginning. The offset is the distance between
+     * columns the tasks are in.
+     */
+    fun gridTaskViewIdOffsetPairInTabOrderSequence(
+        initialTaskViewId: Int,
+        towardsStart: Boolean,
+    ): Sequence<Pair<Int, Int>> = sequence {
+        val draggedTaskViewColumn = getColumn(initialTaskViewId)
+        var nextTaskViewId: Int = initialTaskViewId
+        var previousTaskViewId: Int = Int.MIN_VALUE
+        while (nextTaskViewId != previousTaskViewId && nextTaskViewId >= 0) {
+            previousTaskViewId = nextTaskViewId
+            nextTaskViewId =
+                getNextGridPage(
+                    nextTaskViewId,
+                    if (towardsStart) -1 else 1,
+                    TaskNavDirection.TAB,
+                    cycle = false,
+                )
+            if (nextTaskViewId >= 0 && nextTaskViewId != previousTaskViewId) {
+                val columnOffset = abs(getColumn(nextTaskViewId) - draggedTaskViewColumn)
+                yield(Pair(nextTaskViewId, columnOffset))
+            }
+        }
+    }
+
+    /** Returns the column of a task's id in the grid. */
+    private fun getColumn(taskViewId: Int): Int =
+        if (topRowIds.contains(taskViewId)) topRowIds.indexOf(taskViewId)
+        else bottomRowIds.indexOf(taskViewId)
+
+    enum class TaskNavDirection {
+        UP,
+        DOWN,
+        LEFT,
+        RIGHT,
+        TAB,
+    }
+
+    companion object {
+        const val CLEAR_ALL_PLACEHOLDER_ID: Int = -1
+        const val ADD_DESK_PLACEHOLDER_ID: Int = -2
+    }
+}
diff --git a/quickstep/src/com/android/quickstep/util/TransformParams.java b/quickstep/src/com/android/quickstep/util/TransformParams.java
index 1c1fbd8..716803a 100644
--- a/quickstep/src/com/android/quickstep/util/TransformParams.java
+++ b/quickstep/src/com/android/quickstep/util/TransformParams.java
@@ -222,7 +222,7 @@
         return null;
     }
 
-    // Pubic getters so outside packages can read the values.
+    // Public getters so outside packages can read the values.
 
     public float getProgress() {
         return mProgress;
diff --git a/quickstep/src/com/android/quickstep/views/DesktopTaskView.kt b/quickstep/src/com/android/quickstep/views/DesktopTaskView.kt
index da160f1..fbda3b3 100644
--- a/quickstep/src/com/android/quickstep/views/DesktopTaskView.kt
+++ b/quickstep/src/com/android/quickstep/views/DesktopTaskView.kt
@@ -21,6 +21,7 @@
 import android.graphics.Matrix
 import android.graphics.PointF
 import android.graphics.Rect
+import android.graphics.Rect.intersects
 import android.graphics.RectF
 import android.util.AttributeSet
 import android.util.Log
@@ -30,6 +31,7 @@
 import android.view.ViewStub
 import androidx.core.content.res.ResourcesCompat
 import androidx.core.view.updateLayoutParams
+import com.android.internal.hidden_from_bootclasspath.com.android.window.flags.Flags.enableDesktopRecentsTransitionsCornersBugfix
 import com.android.launcher3.Flags.enableDesktopExplodedView
 import com.android.launcher3.Flags.enableOverviewIconMenu
 import com.android.launcher3.Flags.enableRefactorTaskThumbnail
@@ -57,6 +59,7 @@
 import com.android.quickstep.task.thumbnail.TaskThumbnailView
 import com.android.quickstep.util.DesktopTask
 import com.android.quickstep.util.RecentsOrientedState
+import kotlin.math.roundToInt
 
 /** TaskView that contains all tasks that are part of the desktop. */
 class DesktopTaskView @JvmOverloads constructor(context: Context, attrs: AttributeSet? = null) :
@@ -102,7 +105,7 @@
      * Holds the default (user placed) positions of task windows. This can be moved into the
      * viewModel once RefactorTaskThumbnail has been launched.
      */
-    private var defaultTaskPositions: List<DesktopTaskBoundsData> = emptyList()
+    private var fullscreenTaskPositions: List<DesktopTaskBoundsData> = emptyList()
 
     /**
      * When enableDesktopExplodedView is enabled, this controls the gradual transition from the
@@ -172,42 +175,44 @@
 
         val thumbnailTopMarginPx = container.deviceProfile.overviewTaskThumbnailTopMarginPx
 
-        val containerWidth = layoutParams.width
-        val containerHeight = layoutParams.height - thumbnailTopMarginPx
+        val taskViewWidth = layoutParams.width
+        val taskViewHeight = layoutParams.height - thumbnailTopMarginPx
 
         BaseContainerInterface.getTaskDimension(mContext, container.deviceProfile, tempPointF)
 
-        val windowWidth = tempPointF.x.toInt()
-        val windowHeight = tempPointF.y.toInt()
-        val scaleWidth = containerWidth / windowWidth.toFloat()
-        val scaleHeight = containerHeight / windowHeight.toFloat()
+        val screenWidth = tempPointF.x.toInt()
+        val screenHeight = tempPointF.y.toInt()
+        val screenRect = Rect(0, 0, screenWidth, screenHeight)
+        val scaleWidth = taskViewWidth / screenWidth.toFloat()
+        val scaleHeight = taskViewHeight / screenHeight.toFloat()
 
         taskContainers.forEach {
             val taskId = it.task.key.id
-            val defaultPosition = defaultTaskPositions.firstOrNull { it.taskId == taskId } ?: return
-            val position =
+            val fullscreenTaskPosition =
+                fullscreenTaskPositions.firstOrNull { it.taskId == taskId } ?: return
+            val overviewTaskPosition =
                 if (enableDesktopExplodedView()) {
                     viewModel!!
                         .organizedDesktopTaskPositions
                         .firstOrNull { it.taskId == taskId }
                         ?.let { organizedPosition ->
-                            TEMP_RECT.apply {
+                            TEMP_OVERVIEW_TASK_POSITION.apply {
                                 lerpRect(
-                                    defaultPosition.bounds,
+                                    fullscreenTaskPosition.bounds,
                                     organizedPosition.bounds,
                                     explodeProgress,
                                 )
                             }
-                        } ?: defaultPosition.bounds
+                        } ?: fullscreenTaskPosition.bounds
                 } else {
-                    defaultPosition.bounds
+                    fullscreenTaskPosition.bounds
                 }
 
             if (enableDesktopExplodedView()) {
                 getRemoteTargetHandle(taskId)?.let { remoteTargetHandle ->
                     val fromRect =
-                        TEMP_RECTF1.apply {
-                            set(defaultPosition.bounds)
+                        TEMP_FROM_RECTF.apply {
+                            set(fullscreenTaskPosition.bounds)
                             scale(scaleWidth)
                             offset(
                                 lastComputedTaskSize.left.toFloat(),
@@ -215,8 +220,8 @@
                             )
                         }
                     val toRect =
-                        TEMP_RECTF2.apply {
-                            set(position)
+                        TEMP_TO_RECTF.apply {
+                            set(overviewTaskPosition)
                             scale(scaleWidth)
                             offset(
                                 lastComputedTaskSize.left.toFloat(),
@@ -230,10 +235,10 @@
                 }
             }
 
-            val taskLeft = position.left * scaleWidth
-            val taskTop = position.top * scaleHeight
-            val taskWidth = position.width() * scaleWidth
-            val taskHeight = position.height() * scaleHeight
+            val taskLeft = overviewTaskPosition.left * scaleWidth
+            val taskTop = overviewTaskPosition.top * scaleHeight
+            val taskWidth = overviewTaskPosition.width() * scaleWidth
+            val taskHeight = overviewTaskPosition.height() * scaleHeight
             // TODO(b/394660950): Revisit the choice to update the layout when explodeProgress == 1.
             // To run the explode animation in reverse, it may be simpler to use translation/scale
             // for all cases where the progress is non-zero.
@@ -254,6 +259,24 @@
                     leftMargin = taskLeft.toInt()
                     topMargin = taskTop.toInt()
                 }
+
+                if (
+                    enableDesktopRecentsTransitionsCornersBugfix() && enableRefactorTaskThumbnail()
+                ) {
+                    it.thumbnailView.outlineBounds =
+                        if (intersects(overviewTaskPosition, screenRect))
+                            Rect(overviewTaskPosition).apply {
+                                intersectUnchecked(screenRect)
+                                // Offset to 0,0 to transform into TaskThumbnailView's coordinate
+                                // system.
+                                offset(-overviewTaskPosition.left, -overviewTaskPosition.top)
+                                left = (left * scaleWidth).roundToInt()
+                                top = (top * scaleHeight).roundToInt()
+                                right = (right * scaleWidth).roundToInt()
+                                bottom = (bottom * scaleHeight).roundToInt()
+                            }
+                        else null
+                }
             } else {
                 // During the animation, apply translation and scale such that the view is
                 // transformed to where we want, without triggering layout.
@@ -346,13 +369,13 @@
         val desktopSize = Size(tempPointF.x.toInt(), tempPointF.y.toInt())
         DEFAULT_BOUNDS.set(0, 0, desktopSize.width / 4, desktopSize.height / 4)
 
-        defaultTaskPositions =
+        fullscreenTaskPositions =
             taskContainers.map {
                 DesktopTaskBoundsData(it.task.key.id, it.task.appBounds ?: DEFAULT_BOUNDS)
             }
 
         if (enableDesktopExplodedView()) {
-            viewModel?.organizeDesktopTasks(desktopSize, defaultTaskPositions)
+            viewModel?.organizeDesktopTasks(desktopSize, fullscreenTaskPositions)
         }
         positionTaskWindows()
     }
@@ -449,8 +472,8 @@
         private const val VIEW_POOL_INITIAL_SIZE = 0
         private val DEFAULT_BOUNDS = Rect()
         // Temporaries used for various purposes to avoid allocations.
-        private val TEMP_RECT = Rect()
-        private val TEMP_RECTF1 = RectF()
-        private val TEMP_RECTF2 = RectF()
+        private val TEMP_OVERVIEW_TASK_POSITION = Rect()
+        private val TEMP_FROM_RECTF = RectF()
+        private val TEMP_TO_RECTF = RectF()
     }
 }
diff --git a/quickstep/src/com/android/quickstep/views/GroupedTaskView.kt b/quickstep/src/com/android/quickstep/views/GroupedTaskView.kt
index 229c8f5..2abfb13 100644
--- a/quickstep/src/com/android/quickstep/views/GroupedTaskView.kt
+++ b/quickstep/src/com/android/quickstep/views/GroupedTaskView.kt
@@ -171,12 +171,10 @@
                 val iconMargins = (iconViewMarginStart + iconViewBackgroundMarginStart) * 2
                 // setMaxWidth() needs to be called before mIconView.setIconOrientation which is
                 // called in the super below.
-                (leftTopTaskContainer.iconView as IconAppChipView).setMaxWidth(
+                (leftTopTaskContainer.iconView as IconAppChipView).maxWidth =
                     groupedTaskViewSizes.first.x - iconMargins
-                )
-                (rightBottomTaskContainer.iconView as IconAppChipView).setMaxWidth(
+                (rightBottomTaskContainer.iconView as IconAppChipView).maxWidth =
                     groupedTaskViewSizes.second.x - iconMargins
-                )
             }
         }
         super.setOrientationState(orientationState)
diff --git a/quickstep/src/com/android/quickstep/views/IconAppChipView.java b/quickstep/src/com/android/quickstep/views/IconAppChipView.java
deleted file mode 100644
index 5270477..0000000
--- a/quickstep/src/com/android/quickstep/views/IconAppChipView.java
+++ /dev/null
@@ -1,471 +0,0 @@
-/*
- * Copyright (C) 2023 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 static com.android.app.animation.Interpolators.EMPHASIZED;
-import static com.android.app.animation.Interpolators.LINEAR;
-import static com.android.launcher3.LauncherAnimUtils.SCALE_PROPERTY;
-import static com.android.launcher3.LauncherAnimUtils.VIEW_TRANSLATE_X;
-import static com.android.launcher3.LauncherAnimUtils.VIEW_TRANSLATE_Y;
-
-import android.animation.Animator;
-import android.animation.AnimatorSet;
-import android.animation.ObjectAnimator;
-import android.animation.RectEvaluator;
-import android.animation.ValueAnimator;
-import android.content.Context;
-import android.content.res.Resources;
-import android.graphics.Outline;
-import android.graphics.Rect;
-import android.graphics.drawable.Drawable;
-import android.util.AttributeSet;
-import android.view.View;
-import android.view.ViewAnimationUtils;
-import android.view.ViewOutlineProvider;
-import android.widget.FrameLayout;
-import android.widget.ImageView;
-import android.widget.TextView;
-
-import androidx.annotation.Nullable;
-
-import com.android.launcher3.R;
-import com.android.launcher3.Utilities;
-import com.android.launcher3.util.MultiPropertyFactory;
-import com.android.launcher3.util.MultiValueAlpha;
-import com.android.quickstep.orientation.RecentsPagedOrientationHandler;
-import com.android.quickstep.util.RecentsOrientedState;
-
-/**
- * An icon app menu view which can be used in place of an IconView in overview TaskViews.
- */
-public class IconAppChipView extends FrameLayout implements TaskViewIcon {
-
-    private static final int MENU_BACKGROUND_REVEAL_DURATION = 417;
-    private static final int MENU_BACKGROUND_HIDE_DURATION = 333;
-
-    private static final int NUM_ALPHA_CHANNELS = 4;
-    private static final int INDEX_CONTENT_ALPHA = 0;
-    private static final int INDEX_COLOR_FILTER_ALPHA = 1;
-    private static final int INDEX_MODAL_ALPHA = 2;
-    /** Used to hide the app chip for 90:10 flex split. */
-    private static final int INDEX_MINIMUM_RATIO_ALPHA = 3;
-
-    private final MultiValueAlpha mMultiValueAlpha;
-
-    private View mMenuAnchorView;
-    private IconView mIconView;
-    // Two textview so we can ellipsize the collapsed view and crossfade on expand to the full name.
-    private TextView mIconTextCollapsedView;
-    private TextView mIconTextExpandedView;
-    private ImageView mIconArrowView;
-    private final Rect mBackgroundRelativeLtrLocation = new Rect();
-    final RectEvaluator mBackgroundAnimationRectEvaluator =
-            new RectEvaluator(mBackgroundRelativeLtrLocation);
-    private final int mCollapsedMenuDefaultWidth;
-    private final int mExpandedMenuDefaultWidth;
-    private final int mCollapsedMenuDefaultHeight;
-    private final int mExpandedMenuDefaultHeight;
-    private final int mIconMenuMarginTopStart;
-    private final int mMenuToChipGap;
-    private final int mBackgroundMarginTopStart;
-    private final int mAppNameHorizontalMargin;
-    private final int mIconViewMarginStart;
-    private final int mAppIconSize;
-    private final int mArrowSize;
-    private final int mIconViewDrawableExpandedSize;
-    private final int mArrowMarginEnd;
-    private AnimatorSet mAnimator;
-
-    private int mMaxWidth = Integer.MAX_VALUE;
-
-    private static final int INDEX_SPLIT_TRANSLATION = 0;
-    private static final int INDEX_MENU_TRANSLATION = 1;
-    private static final int INDEX_COUNT_TRANSLATION = 2;
-
-    private final MultiPropertyFactory<View> mViewTranslationX;
-    private final MultiPropertyFactory<View> mViewTranslationY;
-
-    /**
-     * Gets the view split x-axis translation
-     */
-    public MultiPropertyFactory<View>.MultiProperty getSplitTranslationX() {
-        return mViewTranslationX.get(INDEX_SPLIT_TRANSLATION);
-    }
-
-    /**
-     * Sets the view split x-axis translation
-     * @param translationX x-axis translation
-     */
-    public void setSplitTranslationX(float translationX) {
-        getSplitTranslationX().setValue(translationX);
-    }
-
-    /**
-     * Gets the view split y-axis translation
-     */
-    public MultiPropertyFactory<View>.MultiProperty getSplitTranslationY() {
-        return mViewTranslationY.get(INDEX_SPLIT_TRANSLATION);
-    }
-
-    /**
-     * Sets the view split y-axis translation
-     * @param translationY y-axis translation
-     */
-    public void setSplitTranslationY(float translationY) {
-        getSplitTranslationY().setValue(translationY);
-    }
-
-    /**
-     * Gets the menu x-axis translation for split task
-     */
-    public MultiPropertyFactory<View>.MultiProperty getMenuTranslationX() {
-        return mViewTranslationX.get(INDEX_MENU_TRANSLATION);
-    }
-
-    /**
-     * Gets the menu y-axis translation for split task
-     */
-    public MultiPropertyFactory<View>.MultiProperty getMenuTranslationY() {
-        return mViewTranslationY.get(INDEX_MENU_TRANSLATION);
-    }
-
-    public IconAppChipView(Context context) {
-        this(context, null);
-    }
-
-    public IconAppChipView(Context context, AttributeSet attrs) {
-        this(context, attrs, 0);
-    }
-
-    public IconAppChipView(Context context, AttributeSet attrs, int defStyleAttr) {
-        this(context, attrs, defStyleAttr, 0);
-    }
-
-    public IconAppChipView(Context context, @Nullable AttributeSet attrs, int defStyleAttr,
-            int defStyleRes) {
-        super(context, attrs, defStyleAttr, defStyleRes);
-        Resources res = getResources();
-        mMultiValueAlpha = new MultiValueAlpha(this, NUM_ALPHA_CHANNELS);
-        mMultiValueAlpha.setUpdateVisibility(/* updateVisibility= */ true);
-
-        // Menu dimensions
-        mCollapsedMenuDefaultWidth =
-                res.getDimensionPixelSize(R.dimen.task_thumbnail_icon_menu_collapsed_width);
-        mExpandedMenuDefaultWidth =
-                res.getDimensionPixelSize(R.dimen.task_thumbnail_icon_menu_expanded_width);
-        mCollapsedMenuDefaultHeight =
-                res.getDimensionPixelSize(R.dimen.task_thumbnail_icon_menu_collapsed_height);
-        mExpandedMenuDefaultHeight =
-                res.getDimensionPixelSize(R.dimen.task_thumbnail_icon_menu_expanded_height);
-        mIconMenuMarginTopStart = res.getDimensionPixelSize(
-                R.dimen.task_thumbnail_icon_menu_expanded_top_start_margin);
-        mMenuToChipGap = res.getDimensionPixelSize(
-                R.dimen.task_thumbnail_icon_menu_expanded_gap);
-
-        // Background dimensions
-        mBackgroundMarginTopStart = res.getDimensionPixelSize(
-                R.dimen.task_thumbnail_icon_menu_background_margin_top_start);
-
-        // Contents dimensions
-        mAppNameHorizontalMargin = res.getDimensionPixelSize(
-                R.dimen.task_thumbnail_icon_menu_app_name_margin_horizontal_collapsed);
-        mArrowMarginEnd = res.getDimensionPixelSize(R.dimen.task_thumbnail_icon_menu_arrow_margin);
-        mIconViewMarginStart = res.getDimensionPixelSize(
-                R.dimen.task_thumbnail_icon_view_start_margin);
-        mAppIconSize = res.getDimensionPixelSize(
-                R.dimen.task_thumbnail_icon_menu_app_icon_collapsed_size);
-        mArrowSize = res.getDimensionPixelSize(
-                R.dimen.task_thumbnail_icon_menu_arrow_size);
-        mIconViewDrawableExpandedSize = res.getDimensionPixelSize(
-                R.dimen.task_thumbnail_icon_menu_app_icon_expanded_size);
-
-        mViewTranslationX = new MultiPropertyFactory<>(this, VIEW_TRANSLATE_X,
-                INDEX_COUNT_TRANSLATION,
-                Float::sum);
-        mViewTranslationY = new MultiPropertyFactory<>(this, VIEW_TRANSLATE_Y,
-                INDEX_COUNT_TRANSLATION,
-                Float::sum);
-    }
-
-    @Override
-    protected void onFinishInflate() {
-        super.onFinishInflate();
-        mIconView = findViewById(R.id.icon_view);
-        mIconTextCollapsedView = findViewById(R.id.icon_text_collapsed);
-        mIconTextExpandedView = findViewById(R.id.icon_text_expanded);
-        mIconArrowView = findViewById(R.id.icon_arrow);
-        mMenuAnchorView = findViewById(R.id.icon_view_menu_anchor);
-    }
-
-    protected IconView getIconView() {
-        return mIconView;
-    }
-
-    @Override
-    public void setText(CharSequence text) {
-        if (mIconTextCollapsedView != null) {
-            mIconTextCollapsedView.setText(text);
-        }
-        if (mIconTextExpandedView != null) {
-            mIconTextExpandedView.setText(text);
-        }
-    }
-
-    @Override
-    public Drawable getDrawable() {
-        return mIconView == null ? null : mIconView.getDrawable();
-    }
-
-    @Override
-    public void setDrawable(Drawable icon) {
-        if (mIconView != null) {
-            mIconView.setDrawable(icon);
-        }
-    }
-
-    @Override
-    public void setDrawableSize(int iconWidth, int iconHeight) {
-        if (mIconView != null) {
-            mIconView.setDrawableSize(iconWidth, iconHeight);
-        }
-    }
-
-    /**
-     * Sets the maximum width of this Icon Menu. This is usually used when space is limited for
-     * split screen.
-     */
-    public void setMaxWidth(int maxWidth) {
-        // Width showing only the app icon and arrow. Max width should not be set to less than this.
-        int minimumMaxWidth = mIconViewMarginStart + mAppIconSize + mArrowSize + mArrowMarginEnd;
-        mMaxWidth = Math.max(maxWidth, minimumMaxWidth);
-    }
-
-    @Override
-    public void setIconOrientation(RecentsOrientedState orientationState, boolean isGridTask) {
-        RecentsPagedOrientationHandler orientationHandler =
-                orientationState.getOrientationHandler();
-        // Layout params for anchor view
-        LayoutParams anchorLayoutParams = (LayoutParams) mMenuAnchorView.getLayoutParams();
-        anchorLayoutParams.topMargin = mExpandedMenuDefaultHeight + mMenuToChipGap;
-        mMenuAnchorView.setLayoutParams(anchorLayoutParams);
-
-        // Layout Params for the Menu View (this)
-        LayoutParams iconMenuParams = (LayoutParams) getLayoutParams();
-        iconMenuParams.width = mExpandedMenuDefaultWidth;
-        iconMenuParams.height = mExpandedMenuDefaultHeight;
-        orientationHandler.setIconAppChipMenuParams(this, iconMenuParams, mIconMenuMarginTopStart,
-                mIconMenuMarginTopStart);
-        setLayoutParams(iconMenuParams);
-
-        // Layout params for the background
-        Rect collapsedBackgroundBounds = getCollapsedBackgroundLtrBounds();
-        mBackgroundRelativeLtrLocation.set(collapsedBackgroundBounds);
-        setOutlineProvider(new ViewOutlineProvider() {
-            final Rect mRtlAppliedOutlineBounds = new Rect();
-            @Override
-            public void getOutline(View view, Outline outline) {
-                mRtlAppliedOutlineBounds.set(mBackgroundRelativeLtrLocation);
-                if (isLayoutRtl()) {
-                    int width = getWidth();
-                    mRtlAppliedOutlineBounds.left = width - mBackgroundRelativeLtrLocation.right;
-                    mRtlAppliedOutlineBounds.right = width - mBackgroundRelativeLtrLocation.left;
-                }
-                outline.setRoundRect(
-                        mRtlAppliedOutlineBounds, mRtlAppliedOutlineBounds.height() / 2f);
-            }
-        });
-
-        // Layout Params for the Icon View
-        LayoutParams iconParams = (LayoutParams) mIconView.getLayoutParams();
-        int iconMarginStartRelativeToParent = mIconViewMarginStart + mBackgroundMarginTopStart;
-        orientationHandler.setIconAppChipChildrenParams(
-                iconParams, iconMarginStartRelativeToParent);
-
-        mIconView.setLayoutParams(iconParams);
-        mIconView.setDrawableSize(mAppIconSize, mAppIconSize);
-
-        // Layout Params for the collapsed Icon Text View
-        int textMarginStart =
-                iconMarginStartRelativeToParent + mAppIconSize + mAppNameHorizontalMargin;
-        LayoutParams iconTextCollapsedParams =
-                (LayoutParams) mIconTextCollapsedView.getLayoutParams();
-        orientationHandler.setIconAppChipChildrenParams(iconTextCollapsedParams, textMarginStart);
-        int collapsedTextWidth = collapsedBackgroundBounds.width() - mIconViewMarginStart
-                - mAppIconSize - mArrowSize - mAppNameHorizontalMargin - mArrowMarginEnd;
-        iconTextCollapsedParams.width = collapsedTextWidth;
-        mIconTextCollapsedView.setLayoutParams(iconTextCollapsedParams);
-        mIconTextCollapsedView.setAlpha(1f);
-
-        // Layout Params for the expanded Icon Text View
-        LayoutParams iconTextExpandedParams =
-                (LayoutParams) mIconTextExpandedView.getLayoutParams();
-        orientationHandler.setIconAppChipChildrenParams(iconTextExpandedParams, textMarginStart);
-        mIconTextExpandedView.setLayoutParams(iconTextExpandedParams);
-        mIconTextExpandedView.setAlpha(0f);
-        mIconTextExpandedView.setRevealClip(true, 0, mAppIconSize / 2f, collapsedTextWidth);
-
-        // Layout Params for the Icon Arrow View
-        LayoutParams iconArrowParams = (LayoutParams) mIconArrowView.getLayoutParams();
-        int arrowMarginStart = collapsedBackgroundBounds.right - mArrowMarginEnd - mArrowSize;
-        orientationHandler.setIconAppChipChildrenParams(iconArrowParams, arrowMarginStart);
-        mIconArrowView.setPivotY(iconArrowParams.height / 2f);
-        mIconArrowView.setLayoutParams(iconArrowParams);
-
-        // This method is called twice sometimes (like when rotating split tasks). It is called
-        // once before onMeasure and onLayout, and again after onMeasure but before onLayout with
-        // a new width. This happens because we update widths on rotation and on measure of
-        // grouped task views. Calling requestLayout() does not guarantee a call to onMeasure if
-        // it has just measured, so we explicitly call it here.
-        measure(MeasureSpec.makeMeasureSpec(getLayoutParams().width, MeasureSpec.EXACTLY),
-                MeasureSpec.makeMeasureSpec(getLayoutParams().height, MeasureSpec.EXACTLY));
-    }
-
-    @Override
-    public void setIconColorTint(int color, float amount) {
-        // RecentsView's COLOR_TINT animates between 0 and 0.5f, we want to hide the app chip menu.
-        float colorTintAlpha = Utilities.mapToRange(amount, 0f, 0.5f, 1f, 0f, LINEAR);
-        mMultiValueAlpha.get(INDEX_COLOR_FILTER_ALPHA).setValue(colorTintAlpha);
-    }
-
-    @Override
-    public void setContentAlpha(float alpha) {
-        mMultiValueAlpha.get(INDEX_CONTENT_ALPHA).setValue(alpha);
-    }
-
-    @Override
-    public void setModalAlpha(float alpha) {
-        mMultiValueAlpha.get(INDEX_MODAL_ALPHA).setValue(alpha);
-    }
-
-    @Override
-    public void setFlexSplitAlpha(float alpha) {
-        mMultiValueAlpha.get(INDEX_MINIMUM_RATIO_ALPHA).setValue(alpha);
-    }
-
-    @Override
-    public int getDrawableWidth() {
-        return mIconView == null ? 0 : mIconView.getDrawableWidth();
-    }
-
-    @Override
-    public int getDrawableHeight() {
-        return mIconView == null ? 0 : mIconView.getDrawableHeight();
-    }
-
-    protected void revealAnim(boolean isRevealing) {
-        cancelInProgressAnimations();
-        final Rect collapsedBackgroundBounds = getCollapsedBackgroundLtrBounds();
-        final Rect expandedBackgroundBounds = getExpandedBackgroundLtrBounds();
-        final Rect initialBackground = new Rect(mBackgroundRelativeLtrLocation);
-        mAnimator = new AnimatorSet();
-
-        if (isRevealing) {
-            boolean isRtl = isLayoutRtl();
-            bringToFront();
-            // Clip expanded text with reveal animation so it doesn't go beyond the edge of the menu
-            Animator expandedTextRevealAnim = ViewAnimationUtils.createCircularReveal(
-                    mIconTextExpandedView, 0, mIconTextExpandedView.getHeight() / 2,
-                    mIconTextCollapsedView.getWidth(), mIconTextExpandedView.getWidth());
-            // Animate background clipping
-            ValueAnimator backgroundAnimator = ValueAnimator.ofObject(
-                    mBackgroundAnimationRectEvaluator,
-                    initialBackground,
-                    expandedBackgroundBounds);
-            backgroundAnimator.addUpdateListener(valueAnimator -> invalidateOutline());
-
-            float iconViewScaling = mIconViewDrawableExpandedSize / (float) mAppIconSize;
-            float arrowTranslationX =
-                    expandedBackgroundBounds.right - collapsedBackgroundBounds.right;
-            float iconCenterToTextCollapsed = mAppIconSize / 2f + mAppNameHorizontalMargin;
-            float iconCenterToTextExpanded =
-                    mIconViewDrawableExpandedSize / 2f + mAppNameHorizontalMargin;
-            float textTranslationX = iconCenterToTextExpanded - iconCenterToTextCollapsed;
-
-            float textTranslationXWithRtl = isRtl ? -textTranslationX : textTranslationX;
-            float arrowTranslationWithRtl = isRtl ? -arrowTranslationX : arrowTranslationX;
-
-            mAnimator.playTogether(
-                    expandedTextRevealAnim,
-                    backgroundAnimator,
-                    ObjectAnimator.ofFloat(mIconView, SCALE_X, iconViewScaling),
-                    ObjectAnimator.ofFloat(mIconView, SCALE_Y, iconViewScaling),
-                    ObjectAnimator.ofFloat(mIconTextCollapsedView, TRANSLATION_X,
-                            textTranslationXWithRtl),
-                    ObjectAnimator.ofFloat(mIconTextExpandedView, TRANSLATION_X,
-                            textTranslationXWithRtl),
-                    ObjectAnimator.ofFloat(mIconTextCollapsedView, ALPHA, 0),
-                    ObjectAnimator.ofFloat(mIconTextExpandedView, ALPHA, 1),
-                    ObjectAnimator.ofFloat(mIconArrowView, TRANSLATION_X, arrowTranslationWithRtl),
-                    ObjectAnimator.ofFloat(mIconArrowView, SCALE_Y, -1));
-            mAnimator.setDuration(MENU_BACKGROUND_REVEAL_DURATION);
-        } else {
-            // Clip expanded text with reveal animation so it doesn't go beyond the edge of the menu
-            Animator expandedTextClipAnim = ViewAnimationUtils.createCircularReveal(
-                    mIconTextExpandedView, 0, mIconTextExpandedView.getHeight() / 2,
-                    mIconTextExpandedView.getWidth(), mIconTextCollapsedView.getWidth());
-
-            // Animate background clipping
-            ValueAnimator backgroundAnimator = ValueAnimator.ofObject(
-                    mBackgroundAnimationRectEvaluator,
-                    initialBackground,
-                    collapsedBackgroundBounds);
-            backgroundAnimator.addUpdateListener(valueAnimator -> invalidateOutline());
-
-            mAnimator.playTogether(
-                    expandedTextClipAnim,
-                    backgroundAnimator,
-                    ObjectAnimator.ofFloat(mIconView, SCALE_PROPERTY, 1),
-                    ObjectAnimator.ofFloat(mIconTextCollapsedView, TRANSLATION_X, 0),
-                    ObjectAnimator.ofFloat(mIconTextExpandedView, TRANSLATION_X, 0),
-                    ObjectAnimator.ofFloat(mIconTextCollapsedView, ALPHA, 1),
-                    ObjectAnimator.ofFloat(mIconTextExpandedView, ALPHA, 0),
-                    ObjectAnimator.ofFloat(mIconArrowView, TRANSLATION_X, 0),
-                    ObjectAnimator.ofFloat(mIconArrowView, SCALE_Y, 1));
-            mAnimator.setDuration(MENU_BACKGROUND_HIDE_DURATION);
-        }
-
-        mAnimator.setInterpolator(EMPHASIZED);
-        mAnimator.start();
-    }
-
-    private Rect getCollapsedBackgroundLtrBounds() {
-        Rect bounds = new Rect(
-                0,
-                0,
-                Math.min(mMaxWidth, mCollapsedMenuDefaultWidth),
-                mCollapsedMenuDefaultHeight);
-        bounds.offset(mBackgroundMarginTopStart, mBackgroundMarginTopStart);
-        return bounds;
-    }
-
-    private Rect getExpandedBackgroundLtrBounds() {
-        return new Rect(0, 0, mExpandedMenuDefaultWidth, mExpandedMenuDefaultHeight);
-    }
-
-    private void cancelInProgressAnimations() {
-        // We null the `AnimatorSet` because it holds references to the `Animators` which aren't
-        // expecting to be mutable and will cause a crash if they are re-used.
-        if (mAnimator != null && mAnimator.isStarted()) {
-            mAnimator.cancel();
-            mAnimator = null;
-        }
-    }
-
-    @Override
-    public View asView() {
-        return this;
-    }
-}
diff --git a/quickstep/src/com/android/quickstep/views/IconAppChipView.kt b/quickstep/src/com/android/quickstep/views/IconAppChipView.kt
new file mode 100644
index 0000000..85d14cc
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/views/IconAppChipView.kt
@@ -0,0 +1,438 @@
+/*
+ * Copyright (C) 2023 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.animation.AnimatorSet
+import android.animation.ObjectAnimator
+import android.animation.RectEvaluator
+import android.animation.ValueAnimator
+import android.content.Context
+import android.graphics.Outline
+import android.graphics.Rect
+import android.graphics.drawable.Drawable
+import android.util.AttributeSet
+import android.view.View
+import android.view.ViewAnimationUtils
+import android.view.ViewOutlineProvider
+import android.widget.FrameLayout
+import android.widget.ImageView
+import android.widget.TextView
+import com.android.app.animation.Interpolators
+import com.android.launcher3.LauncherAnimUtils.SCALE_PROPERTY
+import com.android.launcher3.LauncherAnimUtils.VIEW_TRANSLATE_X
+import com.android.launcher3.LauncherAnimUtils.VIEW_TRANSLATE_Y
+import com.android.launcher3.R
+import com.android.launcher3.Utilities
+import com.android.launcher3.util.MultiPropertyFactory
+import com.android.launcher3.util.MultiPropertyFactory.FloatBiFunction
+import com.android.launcher3.util.MultiValueAlpha
+import com.android.quickstep.util.RecentsOrientedState
+import kotlin.math.max
+import kotlin.math.min
+
+/** An icon app menu view which can be used in place of an IconView in overview TaskViews. */
+class IconAppChipView
+@JvmOverloads
+constructor(
+    context: Context,
+    attrs: AttributeSet? = null,
+    defStyleAttr: Int = 0,
+    defStyleRes: Int = 0,
+) : FrameLayout(context, attrs, defStyleAttr, defStyleRes), TaskViewIcon {
+
+    private var iconView: IconView? = null
+    private var iconArrowView: ImageView? = null
+    private var menuAnchorView: View? = null
+    // Two textview so we can ellipsize the collapsed view and crossfade on expand to the full name.
+    private var iconTextCollapsedView: TextView? = null
+    private var iconTextExpandedView: TextView? = null
+
+    private val backgroundRelativeLtrLocation = Rect()
+    private val backgroundAnimationRectEvaluator = RectEvaluator(backgroundRelativeLtrLocation)
+
+    // Menu dimensions
+    private val collapsedMenuDefaultWidth: Int =
+        resources.getDimensionPixelSize(R.dimen.task_thumbnail_icon_menu_collapsed_width)
+    private val expandedMenuDefaultWidth: Int =
+        resources.getDimensionPixelSize(R.dimen.task_thumbnail_icon_menu_expanded_width)
+    private val collapsedMenuDefaultHeight =
+        resources.getDimensionPixelSize(R.dimen.task_thumbnail_icon_menu_collapsed_height)
+    private val expandedMenuDefaultHeight =
+        resources.getDimensionPixelSize(R.dimen.task_thumbnail_icon_menu_expanded_height)
+    private val iconMenuMarginTopStart =
+        resources.getDimensionPixelSize(R.dimen.task_thumbnail_icon_menu_expanded_top_start_margin)
+    private val menuToChipGap: Int =
+        resources.getDimensionPixelSize(R.dimen.task_thumbnail_icon_menu_expanded_gap)
+
+    // Background dimensions
+    private val backgroundMarginTopStart: Int =
+        resources.getDimensionPixelSize(
+            R.dimen.task_thumbnail_icon_menu_background_margin_top_start
+        )
+
+    // Contents dimensions
+    private val appNameHorizontalMargin =
+        resources.getDimensionPixelSize(
+            R.dimen.task_thumbnail_icon_menu_app_name_margin_horizontal_collapsed
+        )
+    private val arrowMarginEnd =
+        resources.getDimensionPixelSize(R.dimen.task_thumbnail_icon_menu_arrow_margin)
+    private val iconViewMarginStart =
+        resources.getDimensionPixelSize(R.dimen.task_thumbnail_icon_view_start_margin)
+    private val appIconSize =
+        resources.getDimensionPixelSize(R.dimen.task_thumbnail_icon_menu_app_icon_collapsed_size)
+    private val arrowSize =
+        resources.getDimensionPixelSize(R.dimen.task_thumbnail_icon_menu_arrow_size)
+    private val iconViewDrawableExpandedSize =
+        resources.getDimensionPixelSize(R.dimen.task_thumbnail_icon_menu_app_icon_expanded_size)
+
+    private var animator: AnimatorSet? = null
+
+    private val multiValueAlpha: MultiValueAlpha =
+        MultiValueAlpha(this, NUM_ALPHA_CHANNELS).apply { setUpdateVisibility(true) }
+
+    private val viewTranslationX: MultiPropertyFactory<View> =
+        MultiPropertyFactory(this, VIEW_TRANSLATE_X, INDEX_COUNT_TRANSLATION, SUM_AGGREGATOR)
+
+    private val viewTranslationY: MultiPropertyFactory<View> =
+        MultiPropertyFactory(this, VIEW_TRANSLATE_Y, INDEX_COUNT_TRANSLATION, SUM_AGGREGATOR)
+
+    var maxWidth = Int.MAX_VALUE
+        /**
+         * Sets the maximum width of this Icon Menu. This is usually used when space is limited for
+         * split screen.
+         */
+        set(value) {
+            // Width showing only the app icon and arrow. Max width should not be set to less than
+            // this.
+            val minMaxWidth = iconViewMarginStart + appIconSize + arrowSize + arrowMarginEnd
+            field = max(value, minMaxWidth)
+        }
+
+    override fun onFinishInflate() {
+        super.onFinishInflate()
+        iconView = findViewById(R.id.icon_view)
+        iconTextCollapsedView = findViewById(R.id.icon_text_collapsed)
+        iconTextExpandedView = findViewById(R.id.icon_text_expanded)
+        iconArrowView = findViewById(R.id.icon_arrow)
+        menuAnchorView = findViewById(R.id.icon_view_menu_anchor)
+    }
+
+    override fun setText(text: CharSequence?) {
+        iconTextCollapsedView?.text = text
+        iconTextExpandedView?.text = text
+    }
+
+    override fun getDrawable(): Drawable? = iconView?.drawable
+
+    override fun setDrawable(icon: Drawable?) {
+        iconView?.drawable = icon
+    }
+
+    override fun setDrawableSize(iconWidth: Int, iconHeight: Int) {
+        iconView?.setDrawableSize(iconWidth, iconHeight)
+    }
+
+    override fun setIconOrientation(orientationState: RecentsOrientedState, isGridTask: Boolean) {
+        val orientationHandler = orientationState.orientationHandler
+        // Layout params for anchor view
+        val anchorLayoutParams = menuAnchorView!!.layoutParams as LayoutParams
+        anchorLayoutParams.topMargin = expandedMenuDefaultHeight + menuToChipGap
+        menuAnchorView!!.layoutParams = anchorLayoutParams
+
+        // Layout Params for the Menu View (this)
+        val iconMenuParams = layoutParams as LayoutParams
+        iconMenuParams.width = expandedMenuDefaultWidth
+        iconMenuParams.height = expandedMenuDefaultHeight
+        orientationHandler.setIconAppChipMenuParams(
+            this,
+            iconMenuParams,
+            iconMenuMarginTopStart,
+            iconMenuMarginTopStart,
+        )
+        layoutParams = iconMenuParams
+
+        // Layout params for the background
+        val collapsedBackgroundBounds = getCollapsedBackgroundLtrBounds()
+        backgroundRelativeLtrLocation.set(collapsedBackgroundBounds)
+        outlineProvider =
+            object : ViewOutlineProvider() {
+                val mRtlAppliedOutlineBounds: Rect = Rect()
+
+                override fun getOutline(view: View, outline: Outline) {
+                    mRtlAppliedOutlineBounds.set(backgroundRelativeLtrLocation)
+                    if (isLayoutRtl) {
+                        val width = width
+                        mRtlAppliedOutlineBounds.left = width - backgroundRelativeLtrLocation.right
+                        mRtlAppliedOutlineBounds.right = width - backgroundRelativeLtrLocation.left
+                    }
+                    outline.setRoundRect(
+                        mRtlAppliedOutlineBounds,
+                        mRtlAppliedOutlineBounds.height() / 2f,
+                    )
+                }
+            }
+
+        // Layout Params for the Icon View
+        val iconParams = iconView!!.layoutParams as LayoutParams
+        val iconMarginStartRelativeToParent = iconViewMarginStart + backgroundMarginTopStart
+        orientationHandler.setIconAppChipChildrenParams(iconParams, iconMarginStartRelativeToParent)
+
+        iconView!!.layoutParams = iconParams
+        iconView!!.setDrawableSize(appIconSize, appIconSize)
+
+        // Layout Params for the collapsed Icon Text View
+        val textMarginStart =
+            iconMarginStartRelativeToParent + appIconSize + appNameHorizontalMargin
+        val iconTextCollapsedParams = iconTextCollapsedView!!.layoutParams as LayoutParams
+        orientationHandler.setIconAppChipChildrenParams(iconTextCollapsedParams, textMarginStart)
+        val collapsedTextWidth =
+            (collapsedBackgroundBounds.width() -
+                iconViewMarginStart -
+                appIconSize -
+                arrowSize -
+                appNameHorizontalMargin -
+                arrowMarginEnd)
+        iconTextCollapsedParams.width = collapsedTextWidth
+        iconTextCollapsedView!!.layoutParams = iconTextCollapsedParams
+        iconTextCollapsedView!!.alpha = 1f
+
+        // Layout Params for the expanded Icon Text View
+        val iconTextExpandedParams = iconTextExpandedView!!.layoutParams as LayoutParams
+        orientationHandler.setIconAppChipChildrenParams(iconTextExpandedParams, textMarginStart)
+        iconTextExpandedView!!.layoutParams = iconTextExpandedParams
+        iconTextExpandedView!!.alpha = 0f
+        iconTextExpandedView!!.setRevealClip(
+            true,
+            0f,
+            appIconSize / 2f,
+            collapsedTextWidth.toFloat(),
+        )
+
+        // Layout Params for the Icon Arrow View
+        val iconArrowParams = iconArrowView!!.layoutParams as LayoutParams
+        val arrowMarginStart = collapsedBackgroundBounds.right - arrowMarginEnd - arrowSize
+        orientationHandler.setIconAppChipChildrenParams(iconArrowParams, arrowMarginStart)
+        iconArrowView!!.pivotY = iconArrowParams.height / 2f
+        iconArrowView!!.layoutParams = iconArrowParams
+
+        // This method is called twice sometimes (like when rotating split tasks). It is called
+        // once before onMeasure and onLayout, and again after onMeasure but before onLayout with
+        // a new width. This happens because we update widths on rotation and on measure of
+        // grouped task views. Calling requestLayout() does not guarantee a call to onMeasure if
+        // it has just measured, so we explicitly call it here.
+        measure(
+            MeasureSpec.makeMeasureSpec(layoutParams.width, MeasureSpec.EXACTLY),
+            MeasureSpec.makeMeasureSpec(layoutParams.height, MeasureSpec.EXACTLY),
+        )
+    }
+
+    override fun setIconColorTint(color: Int, amount: Float) {
+        // RecentsView's COLOR_TINT animates between 0 and 0.5f, we want to hide the app chip menu.
+        val colorTintAlpha = Utilities.mapToRange(amount, 0f, 0.5f, 1f, 0f, Interpolators.LINEAR)
+        multiValueAlpha[INDEX_COLOR_FILTER_ALPHA].value = colorTintAlpha
+    }
+
+    override fun setContentAlpha(alpha: Float) {
+        multiValueAlpha[INDEX_CONTENT_ALPHA].value = alpha
+    }
+
+    override fun setModalAlpha(alpha: Float) {
+        multiValueAlpha[INDEX_MODAL_ALPHA].value = alpha
+    }
+
+    override fun setFlexSplitAlpha(alpha: Float) {
+        multiValueAlpha[INDEX_MINIMUM_RATIO_ALPHA].value = alpha
+    }
+
+    override fun getDrawableWidth(): Int = iconView?.drawableWidth ?: 0
+
+    override fun getDrawableHeight(): Int = iconView?.drawableHeight ?: 0
+
+    /** Gets the view split x-axis translation */
+    fun getSplitTranslationX(): MultiPropertyFactory<View>.MultiProperty =
+        viewTranslationX.get(INDEX_SPLIT_TRANSLATION)
+
+    /**
+     * Sets the view split x-axis translation
+     *
+     * @param value x-axis translation
+     */
+    fun setSplitTranslationX(value: Float) {
+        getSplitTranslationX().value = value
+    }
+
+    /** Gets the view split y-axis translation */
+    fun getSplitTranslationY(): MultiPropertyFactory<View>.MultiProperty =
+        viewTranslationY[INDEX_SPLIT_TRANSLATION]
+
+    /**
+     * Sets the view split y-axis translation
+     *
+     * @param value y-axis translation
+     */
+    fun setSplitTranslationY(value: Float) {
+        getSplitTranslationY().value = value
+    }
+
+    /** Gets the menu x-axis translation for split task */
+    fun getMenuTranslationX(): MultiPropertyFactory<View>.MultiProperty =
+        viewTranslationX[INDEX_MENU_TRANSLATION]
+
+    /** Gets the menu y-axis translation for split task */
+    fun getMenuTranslationY(): MultiPropertyFactory<View>.MultiProperty =
+        viewTranslationY[INDEX_MENU_TRANSLATION]
+
+    internal fun revealAnim(isRevealing: Boolean) {
+        cancelInProgressAnimations()
+        val collapsedBackgroundBounds = getCollapsedBackgroundLtrBounds()
+        val expandedBackgroundBounds = getExpandedBackgroundLtrBounds()
+        val initialBackground = Rect(backgroundRelativeLtrLocation)
+        animator = AnimatorSet()
+
+        if (isRevealing) {
+            val isRtl = isLayoutRtl
+            bringToFront()
+            // Clip expanded text with reveal animation so it doesn't go beyond the edge of the menu
+            val expandedTextRevealAnim =
+                ViewAnimationUtils.createCircularReveal(
+                    iconTextExpandedView,
+                    0,
+                    iconTextExpandedView!!.height / 2,
+                    iconTextCollapsedView!!.width.toFloat(),
+                    iconTextExpandedView!!.width.toFloat(),
+                )
+            // Animate background clipping
+            val backgroundAnimator =
+                ValueAnimator.ofObject(
+                    backgroundAnimationRectEvaluator,
+                    initialBackground,
+                    expandedBackgroundBounds,
+                )
+            backgroundAnimator.addUpdateListener { invalidateOutline() }
+
+            val iconViewScaling = iconViewDrawableExpandedSize / appIconSize.toFloat()
+            val arrowTranslationX =
+                (expandedBackgroundBounds.right - collapsedBackgroundBounds.right).toFloat()
+            val iconCenterToTextCollapsed = appIconSize / 2f + appNameHorizontalMargin
+            val iconCenterToTextExpanded =
+                iconViewDrawableExpandedSize / 2f + appNameHorizontalMargin
+            val textTranslationX = iconCenterToTextExpanded - iconCenterToTextCollapsed
+
+            val textTranslationXWithRtl = if (isRtl) -textTranslationX else textTranslationX
+            val arrowTranslationWithRtl = if (isRtl) -arrowTranslationX else arrowTranslationX
+
+            animator!!.playTogether(
+                expandedTextRevealAnim,
+                backgroundAnimator,
+                ObjectAnimator.ofFloat(iconView, SCALE_X, iconViewScaling),
+                ObjectAnimator.ofFloat(iconView, SCALE_Y, iconViewScaling),
+                ObjectAnimator.ofFloat(
+                    iconTextCollapsedView,
+                    TRANSLATION_X,
+                    textTranslationXWithRtl,
+                ),
+                ObjectAnimator.ofFloat(
+                    iconTextExpandedView,
+                    TRANSLATION_X,
+                    textTranslationXWithRtl,
+                ),
+                ObjectAnimator.ofFloat(iconTextCollapsedView, ALPHA, 0f),
+                ObjectAnimator.ofFloat(iconTextExpandedView, ALPHA, 1f),
+                ObjectAnimator.ofFloat(iconArrowView, TRANSLATION_X, arrowTranslationWithRtl),
+                ObjectAnimator.ofFloat(iconArrowView, SCALE_Y, -1f),
+            )
+            animator!!.setDuration(MENU_BACKGROUND_REVEAL_DURATION.toLong())
+        } else {
+            // Clip expanded text with reveal animation so it doesn't go beyond the edge of the menu
+            val expandedTextClipAnim =
+                ViewAnimationUtils.createCircularReveal(
+                    iconTextExpandedView,
+                    0,
+                    iconTextExpandedView!!.height / 2,
+                    iconTextExpandedView!!.width.toFloat(),
+                    iconTextCollapsedView!!.width.toFloat(),
+                )
+
+            // Animate background clipping
+            val backgroundAnimator =
+                ValueAnimator.ofObject(
+                    backgroundAnimationRectEvaluator,
+                    initialBackground,
+                    collapsedBackgroundBounds,
+                )
+            backgroundAnimator.addUpdateListener { valueAnimator: ValueAnimator? ->
+                invalidateOutline()
+            }
+
+            animator!!.playTogether(
+                expandedTextClipAnim,
+                backgroundAnimator,
+                ObjectAnimator.ofFloat(iconView, SCALE_PROPERTY, 1f),
+                ObjectAnimator.ofFloat(iconTextCollapsedView, TRANSLATION_X, 0f),
+                ObjectAnimator.ofFloat(iconTextExpandedView, TRANSLATION_X, 0f),
+                ObjectAnimator.ofFloat(iconTextCollapsedView, ALPHA, 1f),
+                ObjectAnimator.ofFloat(iconTextExpandedView, ALPHA, 0f),
+                ObjectAnimator.ofFloat(iconArrowView, TRANSLATION_X, 0f),
+                ObjectAnimator.ofFloat(iconArrowView, SCALE_Y, 1f),
+            )
+            animator!!.setDuration(MENU_BACKGROUND_HIDE_DURATION.toLong())
+        }
+
+        animator!!.interpolator = Interpolators.EMPHASIZED
+        animator!!.start()
+    }
+
+    private fun getCollapsedBackgroundLtrBounds(): Rect {
+        val bounds =
+            Rect(0, 0, min(maxWidth, collapsedMenuDefaultWidth), collapsedMenuDefaultHeight)
+        bounds.offset(backgroundMarginTopStart, backgroundMarginTopStart)
+        return bounds
+    }
+
+    private fun getExpandedBackgroundLtrBounds() =
+        Rect(0, 0, expandedMenuDefaultWidth, expandedMenuDefaultHeight)
+
+    private fun cancelInProgressAnimations() {
+        // We null the `AnimatorSet` because it holds references to the `Animators` which aren't
+        // expecting to be mutable and will cause a crash if they are re-used.
+        if (animator != null && animator!!.isStarted) {
+            animator!!.cancel()
+            animator = null
+        }
+    }
+
+    override fun asView(): View = this
+
+    private companion object {
+        private val SUM_AGGREGATOR = FloatBiFunction { a: Float, b: Float -> a + b }
+
+        private const val MENU_BACKGROUND_REVEAL_DURATION = 417
+        private const val MENU_BACKGROUND_HIDE_DURATION = 333
+
+        private const val NUM_ALPHA_CHANNELS = 4
+        private const val INDEX_CONTENT_ALPHA = 0
+        private const val INDEX_COLOR_FILTER_ALPHA = 1
+        private const val INDEX_MODAL_ALPHA = 2
+
+        /** Used to hide the app chip for 90:10 flex split. */
+        private const val INDEX_MINIMUM_RATIO_ALPHA = 3
+
+        private const val INDEX_SPLIT_TRANSLATION = 0
+        private const val INDEX_MENU_TRANSLATION = 1
+        private const val INDEX_COUNT_TRANSLATION = 2
+    }
+}
diff --git a/quickstep/src/com/android/quickstep/views/LauncherRecentsView.java b/quickstep/src/com/android/quickstep/views/LauncherRecentsView.java
index c6bd677..0f1c294 100644
--- a/quickstep/src/com/android/quickstep/views/LauncherRecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/LauncherRecentsView.java
@@ -275,7 +275,8 @@
         boolean showDesktopApps = false;
         GestureState.GestureEndTarget endTarget = mCurrentGestureEndTarget;
         if (endTarget == GestureState.GestureEndTarget.LAST_TASK
-                && desktopVisibilityController.areDesktopTasksVisibleAndNotInOverview()) {
+                && desktopVisibilityController.isInDesktopModeAndNotInOverview(
+                        mContainer.getDisplayId())) {
             // Recents gesture was cancelled and we are returning to the previous task.
             // After super class has handled clean up, show desktop apps on top again
             showDesktopApps = true;
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 88850ab..5bfcf43 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -65,11 +65,6 @@
 import static com.android.quickstep.BaseContainerInterface.getTaskDimension;
 import static com.android.quickstep.TaskUtils.checkCurrentOrManagedUserId;
 import static com.android.quickstep.util.LogUtils.splitFailureMessage;
-import static com.android.quickstep.util.TaskGridNavHelper.DIRECTION_DOWN;
-import static com.android.quickstep.util.TaskGridNavHelper.DIRECTION_LEFT;
-import static com.android.quickstep.util.TaskGridNavHelper.DIRECTION_RIGHT;
-import static com.android.quickstep.util.TaskGridNavHelper.DIRECTION_TAB;
-import static com.android.quickstep.util.TaskGridNavHelper.DIRECTION_UP;
 import static com.android.quickstep.views.ClearAllButton.DISMISS_ALPHA;
 import static com.android.quickstep.views.OverviewActionsView.HIDDEN_ACTIONS_IN_MENU;
 import static com.android.quickstep.views.OverviewActionsView.HIDDEN_DESKTOP;
@@ -132,6 +127,7 @@
 import android.widget.Toast;
 import android.window.DesktopModeFlags;
 import android.window.PictureInPictureSurfaceTransaction;
+import android.window.TransitionInfo;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
@@ -1395,12 +1391,13 @@
         RemoteAnimationTargets targets = params.getTargetSet();
         if (targets != null && targets.findTask(taskId) != null) {
             launchSideTaskInLiveTileMode(taskId, targets.apps, targets.wallpapers,
-                    targets.nonApps);
+                    targets.nonApps, /* transitionInfo= */ null);
         }
     }
 
     public void launchSideTaskInLiveTileMode(int taskId, RemoteAnimationTarget[] apps,
-            RemoteAnimationTarget[] wallpaper, RemoteAnimationTarget[] nonApps) {
+            RemoteAnimationTarget[] wallpaper, RemoteAnimationTarget[] nonApps,
+            @Nullable TransitionInfo transitionInfo) {
         AnimatorSet anim = new AnimatorSet();
         TaskView taskView = getTaskViewByTaskId(taskId);
         if (taskView == null || !isTaskViewVisible(taskView)) {
@@ -1452,7 +1449,7 @@
         } else {
             TaskViewUtils.composeRecentsLaunchAnimator(anim, taskView, apps, wallpaper, nonApps,
                     true /* launcherClosing */, getStateManager(), this,
-                    getDepthController());
+                    getDepthController(), transitionInfo);
         }
         anim.start();
     }
@@ -2694,8 +2691,6 @@
         }
         if (enableRefactorTaskThumbnail()) {
             mRecentsViewModel.onReset();
-            // TODO(b/391842220) Remove TaskViews rather than calling specific logic to cancel scope
-            getTaskViews().forEach(TaskView::destroyScopes);
         }
     }
 
@@ -4538,7 +4533,7 @@
     }
 
     private boolean snapToPageRelative(int delta, boolean cycle,
-            @TaskGridNavHelper.TASK_NAV_DIRECTION int direction) {
+            TaskGridNavHelper.TaskNavDirection direction) {
         // Set next page if scroll animation is still running, otherwise cannot snap to the
         // next page on successive key presses. Setting the current page aborts the scroll.
         if (!mScroller.isFinished()) {
@@ -4557,7 +4552,7 @@
         return true;
     }
 
-    private int getNextPageInternal(int delta, @TaskGridNavHelper.TASK_NAV_DIRECTION int direction,
+    private int getNextPageInternal(int delta, TaskGridNavHelper.TaskNavDirection direction,
             boolean cycle) {
         if (!showAsGrid()) {
             return getNextPage() + delta;
@@ -4645,15 +4640,19 @@
         switch (event.getKeyCode()) {
             case KeyEvent.KEYCODE_TAB:
                 return snapToPageRelative(event.isShiftPressed() ? -1 : 1, true /* cycle */,
-                        DIRECTION_TAB);
+                        TaskGridNavHelper.TaskNavDirection.TAB);
             case KeyEvent.KEYCODE_DPAD_RIGHT:
-                return snapToPageRelative(mIsRtl ? -1 : 1, true /* cycle */, DIRECTION_RIGHT);
+                return snapToPageRelative(mIsRtl ? -1 : 1, true /* cycle */,
+                        TaskGridNavHelper.TaskNavDirection.RIGHT);
             case KeyEvent.KEYCODE_DPAD_LEFT:
-                return snapToPageRelative(mIsRtl ? 1 : -1, true /* cycle */, DIRECTION_LEFT);
+                return snapToPageRelative(mIsRtl ? 1 : -1, true /* cycle */,
+                        TaskGridNavHelper.TaskNavDirection.LEFT);
             case KeyEvent.KEYCODE_DPAD_UP:
-                return snapToPageRelative(1, false /* cycle */, DIRECTION_UP);
+                return snapToPageRelative(1, false /* cycle */,
+                        TaskGridNavHelper.TaskNavDirection.UP);
             case KeyEvent.KEYCODE_DPAD_DOWN:
-                return snapToPageRelative(1, false /* cycle */, DIRECTION_DOWN);
+                return snapToPageRelative(1, false /* cycle */,
+                        TaskGridNavHelper.TaskNavDirection.DOWN);
             case KeyEvent.KEYCODE_DEL:
             case KeyEvent.KEYCODE_FORWARD_DEL:
                 dismissCurrentTask();
diff --git a/quickstep/src/com/android/quickstep/views/RecentsViewUtils.kt b/quickstep/src/com/android/quickstep/views/RecentsViewUtils.kt
index d37a3f9..d3549fb 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsViewUtils.kt
+++ b/quickstep/src/com/android/quickstep/views/RecentsViewUtils.kt
@@ -17,6 +17,7 @@
 package com.android.quickstep.views
 
 import android.graphics.Rect
+import android.os.VibrationAttributes
 import android.view.View
 import androidx.core.view.children
 import androidx.dynamicanimation.animation.FloatPropertyCompat
@@ -26,14 +27,18 @@
 import com.android.launcher3.Flags.enableLargeDesktopWindowingTile
 import com.android.launcher3.Flags.enableSeparateExternalDisplayTasks
 import com.android.launcher3.R
+import com.android.launcher3.Utilities.boundToRange
 import com.android.launcher3.touch.SingleAxisSwipeDetector
 import com.android.launcher3.util.DynamicResource
 import com.android.launcher3.util.IntArray
+import com.android.launcher3.util.MSDLPlayerWrapper
 import com.android.quickstep.util.GroupTask
 import com.android.quickstep.util.TaskGridNavHelper
 import com.android.quickstep.util.isExternalDisplay
 import com.android.quickstep.views.RecentsView.RUNNING_TASK_ATTACH_ALPHA
 import com.android.systemui.shared.recents.model.ThumbnailData
+import com.google.android.msdl.data.model.MSDLToken
+import com.google.android.msdl.domain.InteractionProperties
 import java.util.function.BiConsumer
 import kotlin.math.abs
 
@@ -378,7 +383,7 @@
             // negative displacement to positive displacement). We do not check for an exact value
             // to compare to, as the update listener does not necessarily hit every value (e.g. a
             // value of zero). Do not check again once it has started settling, as a spring can
-            // bounce past the origin multiple times depending on the stifness and damping ratio.
+            // bounce past the origin multiple times depending on the stiffness and damping ratio.
             if (startSettling) return@addUpdateListener
             if (lastPosition < 0 && value >= 0) {
                 startSettling = true
@@ -386,75 +391,79 @@
             lastPosition = value
             if (startSettling) {
                 neighborsToSettle.setStartVelocity(velocity).animateToFinalPosition(0f)
+                playDismissSettlingHaptic(velocity)
             }
         }
 
         // Add tasks before dragged index, fanning out from the dragged task.
         // The order they are added matters, as each spring drives the next.
         var previousNeighbor = neighborsToSettle
-        getTasksAdjacentToDraggedTask(draggedTaskView, towardsStart = true).forEach {
-            previousNeighbor = createNeighboringTaskViewSpringAnimation(it, previousNeighbor)
+        getTasksOffsetPairAdjacentToDraggedTask(draggedTaskView, towardsStart = true).forEach {
+            (taskView, offset) ->
+            previousNeighbor =
+                createNeighboringTaskViewSpringAnimation(
+                    taskView,
+                    offset * ADDITIONAL_DISMISS_DAMPING_RATIO,
+                    previousNeighbor,
+                )
         }
         // Add tasks after dragged index, fanning out from the dragged task.
         // The order they are added matters, as each spring drives the next.
         previousNeighbor = neighborsToSettle
-        getTasksAdjacentToDraggedTask(draggedTaskView, towardsStart = false).forEach {
-            previousNeighbor = createNeighboringTaskViewSpringAnimation(it, previousNeighbor)
+        getTasksOffsetPairAdjacentToDraggedTask(draggedTaskView, towardsStart = false).forEach {
+            (taskView, offset) ->
+            previousNeighbor =
+                createNeighboringTaskViewSpringAnimation(
+                    taskView,
+                    offset * ADDITIONAL_DISMISS_DAMPING_RATIO,
+                    previousNeighbor,
+                )
         }
     }
 
-    /** Gets adjacent tasks either before or after the dragged task in visual order. */
-    private fun getTasksAdjacentToDraggedTask(
+    /**
+     * Gets pairs of (TaskView, offset) adjacent the dragged task in visual order.
+     *
+     * <p>Gets tasks either before or after the dragged task along with their offset from it. The
+     * offset is the distance between indices for carousels, or distance between columns for grids.
+     */
+    private fun getTasksOffsetPairAdjacentToDraggedTask(
         draggedTaskView: TaskView,
         towardsStart: Boolean,
-    ): Sequence<TaskView> {
+    ): Sequence<Pair<TaskView, Int>> {
         if (recentsView.showAsGrid()) {
-            return gridTaskViewInTabOrderSequence(draggedTaskView, towardsStart)
+            val taskGridNavHelper =
+                TaskGridNavHelper(
+                    recentsView.topRowIdArray,
+                    recentsView.bottomRowIdArray,
+                    getLargeTaskViewIds(),
+                    hasAddDesktopButton = false,
+                )
+            return taskGridNavHelper
+                .gridTaskViewIdOffsetPairInTabOrderSequence(
+                    draggedTaskView.taskViewId,
+                    towardsStart,
+                )
+                .mapNotNull { (taskViewId, columnOffset) ->
+                    recentsView.getTaskViewFromTaskViewId(taskViewId)?.let { taskView ->
+                        Pair(taskView, columnOffset)
+                    }
+                }
         } else {
             val taskViewList = taskViews.toList()
             val draggedTaskViewIndex = taskViewList.indexOf(draggedTaskView)
 
             return if (towardsStart) {
-                taskViewList.take(draggedTaskViewIndex).reversed().asSequence()
+                taskViewList
+                    .take(draggedTaskViewIndex)
+                    .reversed()
+                    .mapIndexed { index, taskView -> Pair(taskView, index + 1) }
+                    .asSequence()
             } else {
-                taskViewList.takeLast(taskViewList.size - draggedTaskViewIndex - 1).asSequence()
-            }
-        }
-    }
-
-    /**
-     * Returns a sequence of TaskViews in the grid, ordered according to tab navigation, starting
-     * from the dragged TaskView, in the direction of the provided delta.
-     *
-     * <p>A positive delta moves forward in the tab order towards the end of the grid, while a
-     * negative value moves backward towards the beginning.
-     */
-    private fun gridTaskViewInTabOrderSequence(
-        draggedTaskView: TaskView,
-        towardsStart: Boolean,
-    ): Sequence<TaskView> = sequence {
-        val taskGridNavHelper =
-            TaskGridNavHelper(
-                recentsView.topRowIdArray,
-                recentsView.bottomRowIdArray,
-                getLargeTaskViewIds(),
-                /* hasAddDesktopButton= */ false,
-            )
-        var nextTaskView: TaskView? = draggedTaskView
-        var previousTaskView: TaskView? = null
-        while (nextTaskView != previousTaskView && nextTaskView != null) {
-            previousTaskView = nextTaskView
-            nextTaskView =
-                recentsView.getTaskViewFromTaskViewId(
-                    taskGridNavHelper.getNextGridPage(
-                        nextTaskView.taskViewId,
-                        if (towardsStart) -1 else 1,
-                        TaskGridNavHelper.DIRECTION_TAB,
-                        /* cycle = */ false,
-                    )
-                )
-            if (nextTaskView != null && nextTaskView != previousTaskView) {
-                yield(nextTaskView)
+                taskViewList
+                    .takeLast(taskViewList.size - draggedTaskViewIndex - 1)
+                    .mapIndexed { index, taskView -> Pair(taskView, index + 1) }
+                    .asSequence()
             }
         }
     }
@@ -462,6 +471,7 @@
     /** Creates a neighboring task view spring, driven by the spring of its neighbor. */
     private fun createNeighboringTaskViewSpringAnimation(
         taskView: TaskView,
+        dampingOffsetRatio: Float,
         previousNeighborSpringAnimation: SpringAnimation,
     ): SpringAnimation {
         val neighboringTaskViewSpringAnimation =
@@ -471,7 +481,7 @@
                         taskView.secondaryDismissTranslationProperty
                     ),
                 )
-                .setSpring(createExpressiveDismissSpringForce())
+                .setSpring(createExpressiveDismissSpringForce(dampingOffsetRatio))
         // Update live tile on spring animation.
         if (taskView.isRunningTask && recentsView.enableDrawingLiveTile) {
             neighboringTaskViewSpringAnimation.addUpdateListener { _, _, _ ->
@@ -489,18 +499,43 @@
         return neighboringTaskViewSpringAnimation
     }
 
-    private fun createExpressiveDismissSpringForce(): SpringForce {
+    private fun createExpressiveDismissSpringForce(dampingRatioOffset: Float = 0f): SpringForce {
         val resourceProvider = DynamicResource.provider(recentsView.mContainer)
         return SpringForce()
             .setDampingRatio(
-                resourceProvider.getFloat(R.dimen.expressive_dismiss_task_trans_y_damping_ratio)
+                resourceProvider.getFloat(R.dimen.expressive_dismiss_task_trans_y_damping_ratio) +
+                    dampingRatioOffset
             )
             .setStiffness(
                 resourceProvider.getFloat(R.dimen.expressive_dismiss_task_trans_y_stiffness)
             )
     }
 
+    /**
+     * Plays a haptic as the dragged task view settles back into its rest state.
+     *
+     * <p>Haptic intensity is proportional to velocity.
+     */
+    private fun playDismissSettlingHaptic(velocity: Float) {
+        val maxDismissSettlingVelocity =
+            recentsView.pagedOrientationHandler.getSecondaryDimension(recentsView)
+        MSDLPlayerWrapper.INSTANCE.get(recentsView.context)
+            .playToken(
+                MSDLToken.CANCEL,
+                InteractionProperties.DynamicVibrationScale(
+                    boundToRange(velocity / maxDismissSettlingVelocity, 0f, 1f),
+                    VibrationAttributes.Builder()
+                        .setUsage(VibrationAttributes.USAGE_TOUCH)
+                        .setFlags(VibrationAttributes.FLAG_PIPELINED_EFFECT)
+                        .build(),
+                ),
+            )
+    }
+
     companion object {
         val TEMP_RECT = Rect()
+
+        // The additional damping to apply to tasks further from the dismissed task.
+        const val ADDITIONAL_DISMISS_DAMPING_RATIO = 0.15f
     }
 }
diff --git a/quickstep/src/com/android/quickstep/views/TaskContainer.kt b/quickstep/src/com/android/quickstep/views/TaskContainer.kt
index bbe1af4..ce20bf9 100644
--- a/quickstep/src/com/android/quickstep/views/TaskContainer.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskContainer.kt
@@ -17,6 +17,7 @@
 package com.android.quickstep.views
 
 import android.graphics.Bitmap
+import android.graphics.Matrix
 import android.view.View
 import com.android.launcher3.Flags.enableRefactorTaskThumbnail
 import com.android.launcher3.model.data.TaskViewItemInfo
@@ -26,11 +27,9 @@
 import com.android.quickstep.ViewUtils.addAccessibleChildToList
 import com.android.quickstep.recents.di.RecentsDependencies
 import com.android.quickstep.recents.di.getScope
-import com.android.quickstep.recents.di.inject
 import com.android.quickstep.recents.ui.mapper.TaskUiStateMapper
 import com.android.quickstep.recents.ui.viewmodel.TaskData
 import com.android.quickstep.task.thumbnail.TaskThumbnailView
-import com.android.quickstep.task.viewmodel.TaskThumbnailViewModel
 import com.android.systemui.shared.recents.model.Task
 import com.android.systemui.shared.recents.model.ThumbnailData
 
@@ -55,10 +54,6 @@
 ) {
     val overlay: TaskOverlayFactory.TaskOverlay<*> = taskOverlayFactory.createOverlay(this)
 
-    // TODO(b/390581380): Remove this after this bug is fixed
-    private val taskThumbnailViewModel: TaskThumbnailViewModel by
-        RecentsDependencies.inject(snapshotView)
-
     init {
         if (enableRefactorTaskThumbnail()) {
             require(snapshotView is TaskThumbnailView)
@@ -75,7 +70,10 @@
     }
 
     internal var thumbnailData: ThumbnailData? = null
-    val splitAnimationThumbnail: Bitmap?
+        private set
+
+    val thumbnail: Bitmap?
+        /** If possible don't use this. It should be replaced as part of b/331753115. */
         get() =
             if (enableRefactorTaskThumbnail()) thumbnailData?.thumbnail
             else thumbnailViewDeprecated.thumbnail
@@ -106,9 +104,7 @@
 
     fun bind() {
         digitalWellBeingToast?.bind(task, taskView, snapshotView, stagePosition)
-        if (enableRefactorTaskThumbnail()) {
-            taskThumbnailViewModel.bind(task.key.id)
-        } else {
+        if (!enableRefactorTaskThumbnail()) {
             thumbnailViewDeprecated.bind(task, overlay, taskView)
         }
         overlay.init()
@@ -128,11 +124,6 @@
         }
     }
 
-    // TODO(b/391842220): Cancel scope in onDetach instead of having a specific method for this.
-    fun destroyScopes() {
-        thumbnailView.destroyScopes()
-    }
-
     fun setOverlayEnabled(enabled: Boolean) {
         if (!enableRefactorTaskThumbnail()) {
             thumbnailViewDeprecated.setOverlayEnabled(enabled)
@@ -189,4 +180,8 @@
             thumbnailViewDeprecated.setSplashAlpha(progress)
         }
     }
+
+    fun updateThumbnailMatrix(matrix: Matrix) {
+        thumbnailView.setImageMatrix(matrix)
+    }
 }
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.kt b/quickstep/src/com/android/quickstep/views/TaskView.kt
index 609262f..49ec31d 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskView.kt
@@ -41,7 +41,6 @@
 import android.widget.Toast
 import androidx.annotation.IntDef
 import androidx.annotation.VisibleForTesting
-import androidx.core.view.doOnLayout
 import androidx.core.view.updateLayoutParams
 import com.android.app.animation.Interpolators
 import com.android.launcher3.Flags.enableCursorHoverStates
@@ -671,11 +670,6 @@
         taskContainers.forEach { it.destroy() }
     }
 
-    fun destroyScopes() {
-        // TODO(b/391842220): Cancel scope in onDetach instead of having a specific method for this.
-        taskContainers.forEach { it.destroyScopes() }
-    }
-
     // TODO: Clip-out the icon region from the thumbnail, since they are overlapping.
     override fun hasOverlappingRendering() = false
 
@@ -778,9 +772,14 @@
             container.setState(
                 state = containerState,
                 liveTile = state.isLiveTile,
-                hasHeader = type == TaskViewType.DESKTOP,
+                hasHeader = state.hasHeader,
             )
             updateThumbnailValidity(container)
+            updateThumbnailMatrix(
+                container = container,
+                width = container.thumbnailView.width,
+                height = container.thumbnailView.height,
+            )
 
             if (enableOverviewIconMenu()) {
                 setIconState(container, containerState)
@@ -798,6 +797,23 @@
         applyThumbnailSplashAlpha()
     }
 
+    /**
+     * Updates the thumbnail's transformation matrix and rotation state within a TaskContainer.
+     *
+     * This function is called to reposition the thumbnail in the following scenarios:
+     * - When the TTV's size changes (onSizeChanged), and it's displaying a SnapshotSplash.
+     * - When drawing a snapshot (drawSnapshot).
+     *
+     * @param container The TaskContainer holding the thumbnail to be updated.
+     * @param width The desired width of the thumbnail's container.
+     * @param height The desired height of the thumbnail's container.
+     */
+    private fun updateThumbnailMatrix(container: TaskContainer, width: Int, height: Int) {
+        val thumbnailPosition =
+            viewModel!!.getThumbnailPosition(container.thumbnailData, width, height, isLayoutRtl)
+        container.updateThumbnailMatrix(thumbnailPosition.matrix)
+    }
+
     override fun onDetachedFromWindow() {
         super.onDetachedFromWindow()
         if (enableRefactorTaskThumbnail()) {
@@ -843,6 +859,7 @@
                         getTaskUseCase = RecentsDependencies.get(),
                         getSysUiStatusNavFlagsUseCase = RecentsDependencies.get(),
                         isThumbnailValidUseCase = RecentsDependencies.get(),
+                        getThumbnailPositionUseCase = RecentsDependencies.get(),
                         dispatcherProvider = RecentsDependencies.get(),
                     )
                     .apply { bind(*taskIds) }
@@ -852,7 +869,10 @@
             container.bind()
             if (enableRefactorTaskThumbnail()) {
                 container.thumbnailView.cornerRadius = thumbnailFullscreenParams.currentCornerRadius
-                container.thumbnailView.doOnLayout { updateThumbnailValidity(container) }
+                container.thumbnailView.doOnSizeChange { width, height ->
+                    updateThumbnailValidity(container)
+                    updateThumbnailMatrix(container, width, height)
+                }
             }
         }
         setOrientationState(orientedState)
@@ -1220,6 +1240,7 @@
                 recentsView.stateManager,
                 recentsView,
                 recentsView.depthController,
+                /* transitionInfo= */ null,
             )
             addListener(
                 object : AnimatorListenerAdapter() {
@@ -1396,7 +1417,7 @@
             container.task,
             container.iconView.drawable,
             container.snapshotView,
-            container.splitAnimationThumbnail,
+            container.thumbnail,
             /* intent */ null,
             /* user */ null,
             container.itemInfo,
diff --git a/quickstep/tests/multivalentScreenshotTests/src/com/android/quickstep/task/thumbnail/FakeTaskThumbnailViewModel.kt b/quickstep/tests/multivalentScreenshotTests/src/com/android/quickstep/task/thumbnail/FakeTaskThumbnailViewModel.kt
deleted file mode 100644
index 1a2b1c3..0000000
--- a/quickstep/tests/multivalentScreenshotTests/src/com/android/quickstep/task/thumbnail/FakeTaskThumbnailViewModel.kt
+++ /dev/null
@@ -1,29 +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.quickstep.task.thumbnail
-
-import android.graphics.Matrix
-import com.android.quickstep.task.viewmodel.TaskThumbnailViewModel
-
-class FakeTaskThumbnailViewModel : TaskThumbnailViewModel {
-    override fun bind(taskId: Int) {
-        // no-op
-    }
-
-    override fun getThumbnailPositionState(width: Int, height: Int, isRtl: Boolean) =
-        Matrix.IDENTITY_MATRIX
-}
diff --git a/quickstep/tests/multivalentScreenshotTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewScreenshotTest.kt b/quickstep/tests/multivalentScreenshotTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewScreenshotTest.kt
index 232a08a..b9d2fed 100644
--- a/quickstep/tests/multivalentScreenshotTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewScreenshotTest.kt
+++ b/quickstep/tests/multivalentScreenshotTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewScreenshotTest.kt
@@ -19,10 +19,8 @@
 import android.graphics.Color
 import android.view.LayoutInflater
 import com.android.launcher3.R
-import com.android.quickstep.recents.di.RecentsDependencies
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.BackgroundOnly
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.Uninitialized
-import com.android.quickstep.task.viewmodel.TaskThumbnailViewModel
 import com.google.android.apps.nexuslauncher.imagecomparison.goldenpathmanager.ViewScreenshotGoldenPathManager
 import org.junit.Rule
 import org.junit.Test
@@ -45,8 +43,6 @@
             ViewScreenshotGoldenPathManager(getEmulatedDevicePathConfig(emulationSpec)),
         )
 
-    private val taskThumbnailViewModel = FakeTaskThumbnailViewModel()
-
     @Test
     fun taskThumbnailView_uninitializedByDefault() {
         screenshotRule.screenshotTest("taskThumbnailView_uninitialized") { activity ->
@@ -119,14 +115,10 @@
     }
 
     private fun createTaskThumbnailView(context: Context): TaskThumbnailView {
-        val di = RecentsDependencies.initialize(context)
         val taskThumbnailView =
             LayoutInflater.from(context).inflate(R.layout.task_thumbnail, null, false)
                 as TaskThumbnailView
         taskThumbnailView.cornerRadius = CORNER_RADIUS
-        val ttvDiScopeId = di.getScope(taskThumbnailView).scopeId
-        di.provide(TaskThumbnailViewModel::class.java, ttvDiScopeId) { taskThumbnailViewModel }
-
         return taskThumbnailView
     }
 
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/model/data/TaskViewItemInfoTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/model/data/TaskViewItemInfoTest.kt
index de0da64..adfbca5 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/model/data/TaskViewItemInfoTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/model/data/TaskViewItemInfoTest.kt
@@ -47,6 +47,7 @@
 import com.google.common.truth.Truth.assertThat
 import dagger.BindsInstance
 import dagger.Component
+import org.junit.After
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -78,6 +79,11 @@
         RecentsDependencies.initialize(context)
     }
 
+    @After
+    fun tearDown() {
+        RecentsDependencies.destroy()
+    }
+
     @Test
     fun singleTask() {
         val taskContainers = listOf(createTaskContainer(createTask(1)))
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/util/SettingsCacheSandbox.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/util/SettingsCacheSandbox.kt
index dcd5352..52238c8 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/util/SettingsCacheSandbox.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/util/SettingsCacheSandbox.kt
@@ -17,19 +17,22 @@
 package com.android.launcher3.util
 
 import android.net.Uri
+import com.android.launcher3.util.SettingsCache.OnChangeListener
 import org.mockito.kotlin.any
 import org.mockito.kotlin.doAnswer
 import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
 
-/**
- * Provides a sandboxed [SettingsCache] for testing.
- *
- * Note that listeners registered to [cache] will never be invoked.
- */
+/** Provides [SettingsCache] sandboxed from system settings for testing. */
 class SettingsCacheSandbox {
     private val values = mutableMapOf<Uri, Int>()
+    private val listeners = mutableMapOf<Uri, MutableSet<OnChangeListener>>()
 
-    /** Fake cache that delegates [SettingsCache.getValue] to [values]. */
+    /**
+     * Fake cache that delegates:
+     * - [SettingsCache.getValue] to [values]
+     * - [SettingsCache.mListenerMap] to [listeners].
+     */
     val cache =
         mock<SettingsCache> {
             on { getValue(any<Uri>()) } doAnswer { mock.getValue(it.getArgument(0), 1) }
@@ -37,11 +40,22 @@
                 {
                     values.getOrDefault(it.getArgument(0), it.getArgument(1)) == 1
                 }
+
+            doAnswer {
+                    listeners.getOrPut(it.getArgument(0)) { mutableSetOf() }.add(it.getArgument(1))
+                }
+                .whenever(mock)
+                .register(any(), any())
+            doAnswer { listeners[it.getArgument(0)]?.remove(it.getArgument(1)) }
+                .whenever(mock)
+                .unregister(any(), any())
         }
 
     operator fun get(key: Uri): Int? = values[key]
 
     operator fun set(key: Uri, value: Int) {
+        if (value == values[key]) return
         values[key] = value
+        listeners[key]?.forEach { it.onSettingsChanged(value == 1) }
     }
 }
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/AllAppsActionManagerTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/AllAppsActionManagerTest.kt
index 73b35e8..a1bd107 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/AllAppsActionManagerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/AllAppsActionManagerTest.kt
@@ -18,32 +18,59 @@
 
 import android.app.PendingIntent
 import android.content.IIntentSender
+import android.provider.Settings
+import android.provider.Settings.Secure.USER_SETUP_COMPLETE
 import androidx.test.ext.junit.runners.AndroidJUnit4
-import androidx.test.platform.app.InstrumentationRegistry
+import com.android.launcher3.dagger.LauncherAppComponent
+import com.android.launcher3.dagger.LauncherAppSingleton
+import com.android.launcher3.util.AllModulesForTest
 import com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR
+import com.android.launcher3.util.SandboxApplication
+import com.android.launcher3.util.SettingsCache
+import com.android.launcher3.util.SettingsCacheSandbox
 import com.android.launcher3.util.TestUtil
 import com.google.common.truth.Truth.assertThat
+import dagger.BindsInstance
+import dagger.Component
 import java.util.concurrent.Semaphore
 import java.util.concurrent.TimeUnit.SECONDS
+import org.junit.After
+import org.junit.Before
+import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 
 private const val TIMEOUT = 5L
+private val USER_SETUP_COMPLETE_URI = Settings.Secure.getUriFor(USER_SETUP_COMPLETE)
 
 @RunWith(AndroidJUnit4::class)
 class AllAppsActionManagerTest {
     private val callbackSemaphore = Semaphore(0)
     private val bgExecutor = UI_HELPER_EXECUTOR
 
-    private val allAppsActionManager =
-        AllAppsActionManager(
-            InstrumentationRegistry.getInstrumentation().targetContext,
-            bgExecutor,
-        ) {
-            callbackSemaphore.release()
-            PendingIntent(IIntentSender.Default())
+    @get:Rule val context = SandboxApplication()
+
+    private val settingsCacheSandbox =
+        SettingsCacheSandbox().also { it[USER_SETUP_COMPLETE_URI] = 1 }
+
+    private val allAppsActionManager by
+        lazy(LazyThreadSafetyMode.NONE) {
+            AllAppsActionManager(context, bgExecutor) {
+                callbackSemaphore.release()
+                PendingIntent(IIntentSender.Default())
+            }
         }
 
+    @Before
+    fun initDaggerComponent() {
+        context.initDaggerComponent(
+            DaggerAllAppsActionManagerTestComponent.builder()
+                .bindSettingsCache(settingsCacheSandbox.cache)
+        )
+    }
+
+    @After fun destroyManager() = allAppsActionManager.onDestroy()
+
     @Test
     fun taskbarPresent_actionRegistered() {
         allAppsActionManager.isTaskbarPresent = true
@@ -88,4 +115,50 @@
         assertThat(callbackSemaphore.tryAcquire(TIMEOUT, SECONDS)).isTrue()
         assertThat(allAppsActionManager.isActionRegistered).isTrue()
     }
+
+    @Test
+    fun taskbarPresent_userSetupIncomplete_actionUnregistered() {
+        settingsCacheSandbox[USER_SETUP_COMPLETE_URI] = 0
+        allAppsActionManager.isTaskbarPresent = true
+        assertThat(allAppsActionManager.isActionRegistered).isFalse()
+    }
+
+    @Test
+    fun taskbarPresent_setupUiVisible_actionUnregistered() {
+        allAppsActionManager.isSetupUiVisible = true
+        allAppsActionManager.isTaskbarPresent = true
+        assertThat(allAppsActionManager.isActionRegistered).isFalse()
+    }
+
+    @Test
+    fun taskbarPresent_userSetupCompleted_actionRegistered() {
+        settingsCacheSandbox[USER_SETUP_COMPLETE_URI] = 0
+        allAppsActionManager.isTaskbarPresent = true
+
+        settingsCacheSandbox[USER_SETUP_COMPLETE_URI] = 1
+        assertThat(callbackSemaphore.tryAcquire(TIMEOUT, SECONDS)).isTrue()
+        assertThat(allAppsActionManager.isActionRegistered).isTrue()
+    }
+
+    @Test
+    fun taskbarPresent_setupUiDismissed_actionRegistered() {
+        allAppsActionManager.isSetupUiVisible = true
+        allAppsActionManager.isTaskbarPresent = true
+
+        allAppsActionManager.isSetupUiVisible = false
+        assertThat(callbackSemaphore.tryAcquire(TIMEOUT, SECONDS)).isTrue()
+        assertThat(allAppsActionManager.isActionRegistered).isTrue()
+    }
+}
+
+@LauncherAppSingleton
+@Component(modules = [AllModulesForTest::class])
+interface AllAppsActionManagerTestComponent : LauncherAppComponent {
+
+    @Component.Builder
+    interface Builder : LauncherAppComponent.Builder {
+        @BindsInstance fun bindSettingsCache(settingsCache: SettingsCache): Builder
+
+        override fun build(): AllAppsActionManagerTestComponent
+    }
 }
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/usecase/GetThumbnailPositionUseCaseTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/domain/usecase/GetThumbnailPositionUseCaseTest.kt
similarity index 61%
rename from quickstep/tests/multivalentTests/src/com/android/quickstep/recents/usecase/GetThumbnailPositionUseCaseTest.kt
rename to quickstep/tests/multivalentTests/src/com/android/quickstep/recents/domain/usecase/GetThumbnailPositionUseCaseTest.kt
index bd7d970..a253280 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/usecase/GetThumbnailPositionUseCaseTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/domain/usecase/GetThumbnailPositionUseCaseTest.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2024 The Android Open Source Project
+ * Copyright (C) 2025 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.
@@ -14,22 +14,15 @@
  * limitations under the License.
  */
 
-package com.android.quickstep.recents.usecase
+package com.android.quickstep.recents.domain.usecase
 
-import android.content.ComponentName
-import android.content.Intent
 import android.graphics.Bitmap
-import android.graphics.Color
 import android.graphics.Matrix
 import android.graphics.Rect
 import android.view.Surface.ROTATION_90
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.quickstep.recents.data.FakeRecentsDeviceProfileRepository
 import com.android.quickstep.recents.data.FakeRecentsRotationStateRepository
-import com.android.quickstep.recents.data.FakeTasksRepository
-import com.android.quickstep.recents.usecase.ThumbnailPositionState.MatrixScaling
-import com.android.quickstep.recents.usecase.ThumbnailPositionState.MissingThumbnail
-import com.android.systemui.shared.recents.model.Task
 import com.android.systemui.shared.recents.model.ThumbnailData
 import com.android.systemui.shared.recents.utilities.PreviewPositionHelper
 import com.google.common.truth.Truth.assertThat
@@ -43,55 +36,34 @@
 /** Test for [GetThumbnailPositionUseCase] */
 @RunWith(AndroidJUnit4::class)
 class GetThumbnailPositionUseCaseTest {
-    private val task =
-        Task(Task.TaskKey(TASK_ID, 0, Intent(), ComponentName("", ""), 0, 2000)).apply {
-            colorBackground = Color.BLACK
-        }
-    private val thumbnailData =
-        ThumbnailData(
-            thumbnail =
-                mock<Bitmap>().apply {
-                    whenever(width).thenReturn(THUMBNAIL_WIDTH)
-                    whenever(height).thenReturn(THUMBNAIL_HEIGHT)
-                }
-        )
-
     private val deviceProfileRepository = FakeRecentsDeviceProfileRepository()
     private val rotationStateRepository = FakeRecentsRotationStateRepository()
-    private val tasksRepository = FakeTasksRepository()
     private val previewPositionHelper = mock<PreviewPositionHelper>()
 
     private val systemUnderTest =
         GetThumbnailPositionUseCase(
             deviceProfileRepository,
             rotationStateRepository,
-            tasksRepository,
             previewPositionHelper,
         )
 
     @Test
-    fun invisibleTask_returnsIdentityMatrix() = runTest {
-        tasksRepository.seedTasks(listOf(task))
-
-        assertThat(systemUnderTest.run(TASK_ID, CANVAS_WIDTH, CANVAS_HEIGHT, isRtl = true))
-            .isInstanceOf(MissingThumbnail::class.java)
+    fun nullThumbnailData_returnsIdentityMatrix() = runTest {
+        val expectedResult = ThumbnailPosition(Matrix.IDENTITY_MATRIX, false)
+        val result = systemUnderTest.invoke(null, CANVAS_WIDTH, CANVAS_HEIGHT, isRtl = true)
+        assertThat(result).isEqualTo(expectedResult)
     }
 
     @Test
-    fun visibleTaskWithoutThumbnailData_returnsIdentityMatrix() = runTest {
-        tasksRepository.seedTasks(listOf(task))
-        tasksRepository.setVisibleTasks(setOf(TASK_ID))
-
-        assertThat(systemUnderTest.run(TASK_ID, CANVAS_WIDTH, CANVAS_HEIGHT, isRtl = true))
-            .isInstanceOf(MissingThumbnail::class.java)
+    fun withoutThumbnail_returnsIdentityMatrix() = runTest {
+        val expectedResult = ThumbnailPosition(Matrix.IDENTITY_MATRIX, false)
+        val result =
+            systemUnderTest.invoke(ThumbnailData(), CANVAS_WIDTH, CANVAS_HEIGHT, isRtl = true)
+        assertThat(result).isEqualTo(expectedResult)
     }
 
     @Test
     fun visibleTaskWithThumbnailData_returnsTransformedMatrix() = runTest {
-        tasksRepository.seedThumbnailData(mapOf(TASK_ID to thumbnailData))
-        tasksRepository.seedTasks(listOf(task))
-        tasksRepository.setVisibleTasks(setOf(TASK_ID))
-
         val isLargeScreen = true
         deviceProfileRepository.setRecentsDeviceProfile(
             deviceProfileRepository.getRecentsDeviceProfile().copy(isLargeScreen = isLargeScreen)
@@ -108,13 +80,14 @@
         whenever(previewPositionHelper.matrix).thenReturn(MATRIX)
         whenever(previewPositionHelper.isOrientationChanged).thenReturn(isRotated)
 
-        assertThat(systemUnderTest.run(TASK_ID, CANVAS_WIDTH, CANVAS_HEIGHT, isRtl))
-            .isEqualTo(MatrixScaling(MATRIX, isRotated))
+        val result = systemUnderTest.invoke(THUMBNAIL_DATA, CANVAS_WIDTH, CANVAS_HEIGHT, isRtl)
+        val expectedResult = ThumbnailPosition(MATRIX, isRotated)
+        assertThat(result).isEqualTo(expectedResult)
 
         verify(previewPositionHelper)
             .updateThumbnailMatrix(
                 Rect(0, 0, THUMBNAIL_WIDTH, THUMBNAIL_HEIGHT),
-                thumbnailData,
+                THUMBNAIL_DATA,
                 CANVAS_WIDTH,
                 CANVAS_HEIGHT,
                 isLargeScreen,
@@ -123,8 +96,7 @@
             )
     }
 
-    companion object {
-        const val TASK_ID = 2
+    private companion object {
         const val THUMBNAIL_WIDTH = 100
         const val THUMBNAIL_HEIGHT = 200
         const val CANVAS_WIDTH = 300
@@ -133,5 +105,14 @@
             Matrix().apply {
                 setValues(floatArrayOf(2.3f, 4.5f, 2.6f, 7.4f, 3.4f, 2.3f, 2.5f, 6.0f, 3.4f))
             }
+
+        val THUMBNAIL_DATA =
+            ThumbnailData(
+                thumbnail =
+                    mock<Bitmap>().apply {
+                        whenever(width).thenReturn(THUMBNAIL_WIDTH)
+                        whenever(height).thenReturn(THUMBNAIL_HEIGHT)
+                    }
+            )
     }
 }
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModelTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModelTest.kt
index 08e459b..a97ef0c 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModelTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModelTest.kt
@@ -29,6 +29,7 @@
 import com.android.quickstep.recents.domain.model.TaskModel
 import com.android.quickstep.recents.domain.usecase.GetSysUiStatusNavFlagsUseCase
 import com.android.quickstep.recents.domain.usecase.GetTaskUseCase
+import com.android.quickstep.recents.domain.usecase.GetThumbnailPositionUseCase
 import com.android.quickstep.recents.domain.usecase.IsThumbnailValidUseCase
 import com.android.quickstep.recents.viewmodel.RecentsViewData
 import com.android.quickstep.views.TaskViewType
@@ -58,6 +59,7 @@
 
     private val recentsViewData = RecentsViewData()
     private val getTaskUseCase = mock<GetTaskUseCase>()
+    private val getThumbnailPositionUseCase = mock<GetThumbnailPositionUseCase>()
     private val isThumbnailValidUseCase =
         spy(IsThumbnailValidUseCase(FakeRecentsRotationStateRepository()))
     private lateinit var sut: TaskViewModel
@@ -71,6 +73,7 @@
                 getTaskUseCase = getTaskUseCase,
                 getSysUiStatusNavFlagsUseCase = GetSysUiStatusNavFlagsUseCase(),
                 isThumbnailValidUseCase = isThumbnailValidUseCase,
+                getThumbnailPositionUseCase = getThumbnailPositionUseCase,
                 dispatcherProvider = TestDispatcherProvider(unconfinedTestDispatcher),
             )
         whenever(getTaskUseCase.invoke(TASK_MODEL_1.id)).thenReturn(flow { emit(TASK_MODEL_1) })
@@ -112,6 +115,7 @@
                         getTaskUseCase = getTaskUseCase,
                         getSysUiStatusNavFlagsUseCase = GetSysUiStatusNavFlagsUseCase(),
                         isThumbnailValidUseCase = isThumbnailValidUseCase,
+                        getThumbnailPositionUseCase = getThumbnailPositionUseCase,
                         dispatcherProvider = TestDispatcherProvider(unconfinedTestDispatcher),
                     )
                 sut.bind(TASK_MODEL_1.id)
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/usecase/GetThumbnailUseCaseTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/usecase/GetThumbnailUseCaseTest.kt
index 0044631..90ef61d 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/usecase/GetThumbnailUseCaseTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/usecase/GetThumbnailUseCaseTest.kt
@@ -21,7 +21,6 @@
 import android.graphics.Bitmap
 import android.graphics.Color
 import com.android.quickstep.recents.data.FakeTasksRepository
-import com.android.quickstep.task.viewmodel.TaskOverlayViewModelTest
 import com.android.systemui.shared.recents.model.Task
 import com.android.systemui.shared.recents.model.ThumbnailData
 import com.google.common.truth.Truth.assertThat
@@ -62,7 +61,7 @@
     @Test
     fun taskNotVisible_returnsNull() {
         tasksRepository.seedTasks(listOf(task))
-        tasksRepository.seedThumbnailData(mapOf(TaskOverlayViewModelTest.TASK_ID to thumbnailData))
+        tasksRepository.seedThumbnailData(mapOf(TASK_ID to thumbnailData))
 
         assertThat(systemUnderTest.run(TASK_ID)).isNull()
     }
@@ -70,8 +69,8 @@
     @Test
     fun taskVisible_returnsThumbnail() {
         tasksRepository.seedTasks(listOf(task))
-        tasksRepository.seedThumbnailData(mapOf(TaskOverlayViewModelTest.TASK_ID to thumbnailData))
-        tasksRepository.setVisibleTasks(setOf(TaskOverlayViewModelTest.TASK_ID))
+        tasksRepository.seedThumbnailData(mapOf(TASK_ID to thumbnailData))
+        tasksRepository.setVisibleTasks(setOf(TASK_ID))
 
         assertThat(systemUnderTest.run(TASK_ID)).isEqualTo(thumbnailData.thumbnail)
     }
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelImplTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelImplTest.kt
deleted file mode 100644
index 4b4e2eb..0000000
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelImplTest.kt
+++ /dev/null
@@ -1,88 +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.quickstep.task.thumbnail
-
-import android.graphics.Matrix
-import android.platform.test.flag.junit.SetFlagsRule
-import androidx.test.ext.junit.runners.AndroidJUnit4
-import com.android.quickstep.recents.usecase.GetThumbnailPositionUseCase
-import com.android.quickstep.recents.usecase.ThumbnailPositionState.MatrixScaling
-import com.android.quickstep.recents.usecase.ThumbnailPositionState.MissingThumbnail
-import com.android.quickstep.task.viewmodel.TaskThumbnailViewModelImpl
-import com.google.common.truth.Truth.assertThat
-import kotlinx.coroutines.test.StandardTestDispatcher
-import kotlinx.coroutines.test.TestScope
-import kotlinx.coroutines.test.runTest
-import org.junit.Rule
-import org.junit.Test
-import org.junit.runner.RunWith
-import org.mockito.kotlin.mock
-import org.mockito.kotlin.whenever
-
-/** Test for [TaskThumbnailView] */
-@RunWith(AndroidJUnit4::class)
-class TaskThumbnailViewModelImplTest {
-    @get:Rule val setFlagsRule = SetFlagsRule()
-
-    private val dispatcher = StandardTestDispatcher()
-    private val testScope = TestScope(dispatcher)
-
-    private val mGetThumbnailPositionUseCase = mock<GetThumbnailPositionUseCase>()
-
-    private val systemUnderTest by lazy { TaskThumbnailViewModelImpl(mGetThumbnailPositionUseCase) }
-
-    @Test
-    fun getSnapshotMatrix_MissingThumbnail() =
-        testScope.runTest {
-            val taskId = 2
-            val isRtl = true
-
-            whenever(mGetThumbnailPositionUseCase.run(taskId, CANVAS_WIDTH, CANVAS_HEIGHT, isRtl))
-                .thenReturn(MissingThumbnail)
-
-            systemUnderTest.bind(taskId)
-            assertThat(
-                    systemUnderTest.getThumbnailPositionState(CANVAS_WIDTH, CANVAS_HEIGHT, isRtl)
-                )
-                .isEqualTo(Matrix.IDENTITY_MATRIX)
-        }
-
-    @Test
-    fun getSnapshotMatrix_MatrixScaling() =
-        testScope.runTest {
-            val taskId = 2
-            val isRtl = true
-
-            whenever(mGetThumbnailPositionUseCase.run(taskId, CANVAS_WIDTH, CANVAS_HEIGHT, isRtl))
-                .thenReturn(MatrixScaling(MATRIX, isRotated = false))
-
-            systemUnderTest.bind(taskId)
-            assertThat(
-                    systemUnderTest.getThumbnailPositionState(CANVAS_WIDTH, CANVAS_HEIGHT, isRtl)
-                )
-                .isEqualTo(MATRIX)
-        }
-
-    private companion object {
-        const val CANVAS_WIDTH = 300
-        const val CANVAS_HEIGHT = 600
-        val MATRIX =
-            Matrix().apply {
-                setValues(floatArrayOf(2.3f, 4.5f, 2.6f, 7.4f, 3.4f, 2.3f, 2.5f, 6.0f, 3.4f))
-            }
-    }
-}
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModelTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModelTest.kt
deleted file mode 100644
index 95504af..0000000
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModelTest.kt
+++ /dev/null
@@ -1,195 +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.quickstep.task.viewmodel
-
-import android.content.ComponentName
-import android.content.Intent
-import android.graphics.Bitmap
-import android.graphics.Color
-import android.graphics.Matrix
-import androidx.test.ext.junit.runners.AndroidJUnit4
-import com.android.launcher3.util.TestDispatcherProvider
-import com.android.quickstep.recents.data.FakeTasksRepository
-import com.android.quickstep.recents.usecase.GetThumbnailPositionUseCase
-import com.android.quickstep.recents.usecase.ThumbnailPositionState.MatrixScaling
-import com.android.quickstep.recents.usecase.ThumbnailPositionState.MissingThumbnail
-import com.android.quickstep.recents.viewmodel.RecentsViewData
-import com.android.quickstep.task.thumbnail.TaskOverlayUiState.Disabled
-import com.android.quickstep.task.thumbnail.TaskOverlayUiState.Enabled
-import com.android.quickstep.task.viewmodel.TaskOverlayViewModel.ThumbnailPositionState
-import com.android.systemui.shared.recents.model.Task
-import com.android.systemui.shared.recents.model.ThumbnailData
-import com.google.common.truth.Truth.assertThat
-import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.flow.first
-import kotlinx.coroutines.test.TestScope
-import kotlinx.coroutines.test.UnconfinedTestDispatcher
-import kotlinx.coroutines.test.runTest
-import org.junit.Test
-import org.junit.runner.RunWith
-import org.mockito.kotlin.mock
-import org.mockito.kotlin.whenever
-
-/** Test for [TaskOverlayViewModel] */
-@OptIn(ExperimentalCoroutinesApi::class)
-@RunWith(AndroidJUnit4::class)
-class TaskOverlayViewModelTest {
-    private val task =
-        Task(Task.TaskKey(TASK_ID, 0, Intent(), ComponentName("", ""), 0, 2000)).apply {
-            colorBackground = Color.BLACK
-        }
-    private val thumbnailData =
-        ThumbnailData(
-            thumbnail =
-                mock<Bitmap>().apply {
-                    whenever(width).thenReturn(THUMBNAIL_WIDTH)
-                    whenever(height).thenReturn(THUMBNAIL_HEIGHT)
-                }
-        )
-    private val recentsViewData = RecentsViewData()
-    private val tasksRepository = FakeTasksRepository()
-    private val mGetThumbnailPositionUseCase = mock<GetThumbnailPositionUseCase>()
-    private val dispatcher = UnconfinedTestDispatcher()
-    private val testScope = TestScope(dispatcher)
-    private val systemUnderTest =
-        TaskOverlayViewModel(
-            task,
-            recentsViewData,
-            mGetThumbnailPositionUseCase,
-            tasksRepository,
-            TestDispatcherProvider(dispatcher),
-        )
-
-    @Test
-    fun initialStateIsDisabled() =
-        testScope.runTest { assertThat(systemUnderTest.overlayState.first()).isEqualTo(Disabled) }
-
-    @Test
-    fun recentsViewOverlayDisabled_Disabled() =
-        testScope.runTest {
-            recentsViewData.overlayEnabled.value = false
-            recentsViewData.settledFullyVisibleTaskIds.value = setOf(TASK_ID)
-
-            assertThat(systemUnderTest.overlayState.first()).isEqualTo(Disabled)
-        }
-
-    @Test
-    fun taskNotFullyVisible_Disabled() =
-        testScope.runTest {
-            recentsViewData.overlayEnabled.value = true
-            recentsViewData.settledFullyVisibleTaskIds.value = setOf()
-
-            assertThat(systemUnderTest.overlayState.first()).isEqualTo(Disabled)
-        }
-
-    @Test
-    fun noThumbnail_Enabled() =
-        testScope.runTest {
-            recentsViewData.overlayEnabled.value = true
-            recentsViewData.settledFullyVisibleTaskIds.value = setOf(TASK_ID)
-            task.isLocked = false
-
-            assertThat(systemUnderTest.overlayState.first())
-                .isEqualTo(Enabled(isRealSnapshot = false, thumbnail = null))
-        }
-
-    @Test
-    fun withThumbnail_RealSnapshot_NotLocked_Enabled() =
-        testScope.runTest {
-            recentsViewData.overlayEnabled.value = true
-            recentsViewData.settledFullyVisibleTaskIds.value = setOf(TASK_ID)
-            tasksRepository.seedTasks(listOf(task))
-            tasksRepository.seedThumbnailData(mapOf(TASK_ID to thumbnailData))
-            tasksRepository.setVisibleTasks(setOf(TASK_ID))
-            thumbnailData.isRealSnapshot = true
-            task.isLocked = false
-
-            assertThat(systemUnderTest.overlayState.first())
-                .isEqualTo(Enabled(isRealSnapshot = true, thumbnail = thumbnailData.thumbnail))
-        }
-
-    @Test
-    fun withThumbnail_RealSnapshot_Locked_Enabled() =
-        testScope.runTest {
-            recentsViewData.overlayEnabled.value = true
-            recentsViewData.settledFullyVisibleTaskIds.value = setOf(TASK_ID)
-            tasksRepository.seedTasks(listOf(task))
-            tasksRepository.seedThumbnailData(mapOf(TASK_ID to thumbnailData))
-            tasksRepository.setVisibleTasks(setOf(TASK_ID))
-            thumbnailData.isRealSnapshot = true
-            task.isLocked = true
-
-            assertThat(systemUnderTest.overlayState.first())
-                .isEqualTo(Enabled(isRealSnapshot = false, thumbnail = thumbnailData.thumbnail))
-        }
-
-    @Test
-    fun withThumbnail_FakeSnapshot_Enabled() =
-        testScope.runTest {
-            recentsViewData.overlayEnabled.value = true
-            recentsViewData.settledFullyVisibleTaskIds.value = setOf(TASK_ID)
-            tasksRepository.seedTasks(listOf(task))
-            tasksRepository.seedThumbnailData(mapOf(TASK_ID to thumbnailData))
-            tasksRepository.setVisibleTasks(setOf(TASK_ID))
-            thumbnailData.isRealSnapshot = false
-            task.isLocked = false
-
-            assertThat(systemUnderTest.overlayState.first())
-                .isEqualTo(Enabled(isRealSnapshot = false, thumbnail = thumbnailData.thumbnail))
-        }
-
-    @Test
-    fun getThumbnailMatrix_MissingThumbnail() =
-        testScope.runTest {
-            val isRtl = true
-
-            whenever(mGetThumbnailPositionUseCase.run(TASK_ID, CANVAS_WIDTH, CANVAS_HEIGHT, isRtl))
-                .thenReturn(MissingThumbnail)
-
-            assertThat(
-                    systemUnderTest.getThumbnailPositionState(CANVAS_WIDTH, CANVAS_HEIGHT, isRtl)
-                )
-                .isEqualTo(ThumbnailPositionState(Matrix.IDENTITY_MATRIX, isRotated = false))
-        }
-
-    @Test
-    fun getThumbnailMatrix_MatrixScaling() =
-        testScope.runTest {
-            val isRtl = true
-            val isRotated = true
-
-            whenever(mGetThumbnailPositionUseCase.run(TASK_ID, CANVAS_WIDTH, CANVAS_HEIGHT, isRtl))
-                .thenReturn(MatrixScaling(MATRIX, isRotated))
-
-            assertThat(
-                    systemUnderTest.getThumbnailPositionState(CANVAS_WIDTH, CANVAS_HEIGHT, isRtl)
-                )
-                .isEqualTo(ThumbnailPositionState(MATRIX, isRotated))
-        }
-
-    companion object {
-        const val TASK_ID = 0
-        const val THUMBNAIL_WIDTH = 100
-        const val THUMBNAIL_HEIGHT = 200
-        const val CANVAS_WIDTH = 300
-        const val CANVAS_HEIGHT = 600
-        val MATRIX =
-            Matrix().apply {
-                setValues(floatArrayOf(2.3f, 4.5f, 2.6f, 7.4f, 3.4f, 2.3f, 2.5f, 6.0f, 3.4f))
-            }
-    }
-}
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/taskbar/controllers/TaskbarPinningControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/taskbar/controllers/TaskbarPinningControllerTest.kt
index 5b42d6c..26418d8 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/taskbar/controllers/TaskbarPinningControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/taskbar/controllers/TaskbarPinningControllerTest.kt
@@ -71,8 +71,9 @@
         whenever(taskbarActivityContext.dragLayer).thenReturn(taskbarDragLayer)
         whenever(taskbarActivityContext.statsLogManager).thenReturn(statsLogManager)
         whenever(
-                taskbarControllers.taskbarDesktopModeController
-                    .areDesktopTasksVisibleAndNotInOverview
+                taskbarControllers.taskbarDesktopModeController.isInDesktopModeAndNotInOverview(
+                    taskbarActivityContext.displayId
+                )
             )
             .thenAnswer { _ -> isInDesktopMode }
         pinningController = spy(TaskbarPinningController(taskbarActivityContext))
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/SplitAnimationControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/SplitAnimationControllerTest.kt
index 5051251..9d000a4 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/SplitAnimationControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/SplitAnimationControllerTest.kt
@@ -89,7 +89,7 @@
     @Before
     fun setup() {
         whenever(mockTaskContainer.snapshotView).thenReturn(mockSnapshotView)
-        whenever(mockTaskContainer.splitAnimationThumbnail).thenReturn(mockBitmap)
+        whenever(mockTaskContainer.thumbnail).thenReturn(mockBitmap)
         whenever(mockTaskContainer.iconView).thenReturn(mockIconView)
         whenever(mockTaskContainer.task).thenReturn(mockTask)
         whenever(mockIconView.drawable).thenReturn(mockTaskViewDrawable)
@@ -117,7 +117,7 @@
         assertEquals(
             "Did not fallback to use splitSource icon drawable",
             mockSplitSourceDrawable,
-            splitAnimInitProps.iconDrawable
+            splitAnimInitProps.iconDrawable,
         )
     }
 
@@ -133,7 +133,7 @@
         assertEquals(
             "Did not use taskView icon drawable",
             mockTaskViewDrawable,
-            splitAnimInitProps.iconDrawable
+            splitAnimInitProps.iconDrawable,
         )
     }
 
@@ -152,7 +152,7 @@
         assertEquals(
             "Did not use taskView icon drawable",
             mockTaskViewDrawable,
-            splitAnimInitProps.iconDrawable
+            splitAnimInitProps.iconDrawable,
         )
     }
 
@@ -168,7 +168,7 @@
         assertEquals(
             "Did not use splitSource icon drawable",
             mockSplitSourceDrawable,
-            splitAnimInitProps.iconDrawable
+            splitAnimInitProps.iconDrawable,
         )
     }
 
@@ -190,13 +190,13 @@
         val splitAnimInitProps: SplitAnimationController.Companion.SplitAnimInitProps =
             splitAnimationController.getFirstAnimInitViews(
                 { mockGroupedTaskView },
-                { splitSelectSource }
+                { splitSelectSource },
             )
 
         assertEquals(
             "Did not use splitSource icon drawable",
             mockSplitSourceDrawable,
-            splitAnimInitProps.iconDrawable
+            splitAnimInitProps.iconDrawable,
         )
     }
 
@@ -214,7 +214,7 @@
                 any(),
                 any(),
                 any(),
-                any()
+                any(),
             )
 
         spySplitAnimationController.playSplitLaunchAnimation(
@@ -230,7 +230,7 @@
             null /* info */,
             null /* t */,
             {} /* finishCallback */,
-            1f /* cornerRadius */
+            1f, /* cornerRadius */
         )
 
         verify(spySplitAnimationController)
@@ -243,7 +243,7 @@
                 any(),
                 any(),
                 any(),
-                any()
+                any(),
             )
     }
 
@@ -267,7 +267,7 @@
             transitionInfo,
             transaction,
             {} /* finishCallback */,
-            1f /* cornerRadius */
+            1f, /* cornerRadius */
         )
 
         verify(spySplitAnimationController)
@@ -296,7 +296,7 @@
             transitionInfo,
             transaction,
             {} /* finishCallback */,
-            1f /* cornerRadius */
+            1f, /* cornerRadius */
         )
 
         verify(spySplitAnimationController)
@@ -325,7 +325,7 @@
             transitionInfo,
             transaction,
             {} /* finishCallback */,
-            1f /* cornerRadius */
+            1f, /* cornerRadius */
         )
 
         verify(spySplitAnimationController)
@@ -353,7 +353,7 @@
             transitionInfo,
             transaction,
             {} /* finishCallback */,
-            1f /* cornerRadius */
+            1f, /* cornerRadius */
         )
 
         verify(spySplitAnimationController)
@@ -381,7 +381,7 @@
             transitionInfo,
             transaction,
             {} /* finishCallback */,
-            1f /* cornerRadius */
+            1f, /* cornerRadius */
         )
 
         verify(spySplitAnimationController)
@@ -408,7 +408,7 @@
             transitionInfo,
             transaction,
             {} /* finishCallback */,
-            1f /* cornerRadius */
+            1f, /* cornerRadius */
         )
 
         verify(spySplitAnimationController)
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/TaskGridNavHelperTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/TaskGridNavHelperTest.kt
index f2fa0c5..cb088fd 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/TaskGridNavHelperTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/TaskGridNavHelperTest.kt
@@ -16,13 +16,13 @@
 package com.android.quickstep.util
 
 import com.android.launcher3.util.IntArray
-import com.android.quickstep.util.TaskGridNavHelper.ADD_DESK_PLACEHOLDER_ID
-import com.android.quickstep.util.TaskGridNavHelper.CLEAR_ALL_PLACEHOLDER_ID
-import com.android.quickstep.util.TaskGridNavHelper.DIRECTION_DOWN
-import com.android.quickstep.util.TaskGridNavHelper.DIRECTION_LEFT
-import com.android.quickstep.util.TaskGridNavHelper.DIRECTION_RIGHT
-import com.android.quickstep.util.TaskGridNavHelper.DIRECTION_TAB
-import com.android.quickstep.util.TaskGridNavHelper.DIRECTION_UP
+import com.android.quickstep.util.TaskGridNavHelper.Companion.ADD_DESK_PLACEHOLDER_ID
+import com.android.quickstep.util.TaskGridNavHelper.Companion.CLEAR_ALL_PLACEHOLDER_ID
+import com.android.quickstep.util.TaskGridNavHelper.TaskNavDirection.DOWN
+import com.android.quickstep.util.TaskGridNavHelper.TaskNavDirection.LEFT
+import com.android.quickstep.util.TaskGridNavHelper.TaskNavDirection.RIGHT
+import com.android.quickstep.util.TaskGridNavHelper.TaskNavDirection.TAB
+import com.android.quickstep.util.TaskGridNavHelper.TaskNavDirection.UP
 import com.google.common.truth.Truth.assertThat
 import org.junit.Test
 
@@ -35,8 +35,7 @@
     */
     @Test
     fun equalLengthRows_noFocused_onTop_pressDown_goesToBottom() {
-        assertThat(getNextGridPage(currentPageTaskViewId = 1, DIRECTION_DOWN, delta = 1))
-            .isEqualTo(2)
+        assertThat(getNextGridPage(currentPageTaskViewId = 1, DOWN, delta = 1)).isEqualTo(2)
     }
 
     /*                      ↑----→
@@ -46,7 +45,7 @@
     */
     @Test
     fun equalLengthRows_noFocused_onTop_pressUp_goesToBottom() {
-        assertThat(getNextGridPage(currentPageTaskViewId = 1, DIRECTION_UP, delta = 1)).isEqualTo(2)
+        assertThat(getNextGridPage(currentPageTaskViewId = 1, UP, delta = 1)).isEqualTo(2)
     }
 
     /*                      ↓----↑
@@ -57,8 +56,7 @@
     */
     @Test
     fun equalLengthRows_noFocused_onBottom_pressDown_goesToTop() {
-        assertThat(getNextGridPage(currentPageTaskViewId = 2, DIRECTION_DOWN, delta = 1))
-            .isEqualTo(1)
+        assertThat(getNextGridPage(currentPageTaskViewId = 2, DOWN, delta = 1)).isEqualTo(1)
     }
 
     /*
@@ -68,7 +66,7 @@
     */
     @Test
     fun equalLengthRows_noFocused_onBottom_pressUp_goesToTop() {
-        assertThat(getNextGridPage(currentPageTaskViewId = 2, DIRECTION_UP, delta = 1)).isEqualTo(1)
+        assertThat(getNextGridPage(currentPageTaskViewId = 2, UP, delta = 1)).isEqualTo(1)
     }
 
     /*
@@ -78,8 +76,7 @@
     */
     @Test
     fun equalLengthRows_noFocused_onTop_pressLeft_goesLeft() {
-        assertThat(getNextGridPage(currentPageTaskViewId = 1, DIRECTION_LEFT, delta = 1))
-            .isEqualTo(3)
+        assertThat(getNextGridPage(currentPageTaskViewId = 1, LEFT, delta = 1)).isEqualTo(3)
     }
 
     /*
@@ -89,8 +86,7 @@
     */
     @Test
     fun equalLengthRows_noFocused_onBottom_pressLeft_goesLeft() {
-        assertThat(getNextGridPage(currentPageTaskViewId = 2, DIRECTION_LEFT, delta = 1))
-            .isEqualTo(4)
+        assertThat(getNextGridPage(currentPageTaskViewId = 2, LEFT, delta = 1)).isEqualTo(4)
     }
 
     /*
@@ -100,8 +96,7 @@
     */
     @Test
     fun equalLengthRows_noFocused_onTop_secondItem_pressRight_goesRight() {
-        assertThat(getNextGridPage(currentPageTaskViewId = 3, DIRECTION_RIGHT, delta = -1))
-            .isEqualTo(1)
+        assertThat(getNextGridPage(currentPageTaskViewId = 3, RIGHT, delta = -1)).isEqualTo(1)
     }
 
     /*
@@ -111,8 +106,7 @@
     */
     @Test
     fun equalLengthRows_noFocused_onBottom_secondItem_pressRight_goesRight() {
-        assertThat(getNextGridPage(currentPageTaskViewId = 4, DIRECTION_RIGHT, delta = -1))
-            .isEqualTo(2)
+        assertThat(getNextGridPage(currentPageTaskViewId = 4, RIGHT, delta = -1)).isEqualTo(2)
     }
 
     /*
@@ -124,7 +118,7 @@
     */
     @Test
     fun equalLengthRows_noFocused_onTop_pressRight_cycleToClearAll() {
-        assertThat(getNextGridPage(currentPageTaskViewId = 1, DIRECTION_RIGHT, delta = -1))
+        assertThat(getNextGridPage(currentPageTaskViewId = 1, RIGHT, delta = -1))
             .isEqualTo(CLEAR_ALL_PLACEHOLDER_ID)
     }
 
@@ -137,7 +131,7 @@
     */
     @Test
     fun equalLengthRows_noFocused_onBottom_pressRight_cycleToClearAll() {
-        assertThat(getNextGridPage(currentPageTaskViewId = 2, DIRECTION_RIGHT, delta = -1))
+        assertThat(getNextGridPage(currentPageTaskViewId = 2, RIGHT, delta = -1))
             .isEqualTo(CLEAR_ALL_PLACEHOLDER_ID)
     }
 
@@ -149,7 +143,7 @@
     */
     @Test
     fun equalLengthRows_noFocused_onTop_lastItem_pressLeft_toClearAll() {
-        assertThat(getNextGridPage(currentPageTaskViewId = 5, DIRECTION_LEFT, delta = 1))
+        assertThat(getNextGridPage(currentPageTaskViewId = 5, LEFT, delta = 1))
             .isEqualTo(CLEAR_ALL_PLACEHOLDER_ID)
     }
 
@@ -161,7 +155,7 @@
     */
     @Test
     fun equalLengthRows_noFocused_onBottom_lastItem_pressLeft_toClearAll() {
-        assertThat(getNextGridPage(currentPageTaskViewId = 6, DIRECTION_LEFT, delta = 1))
+        assertThat(getNextGridPage(currentPageTaskViewId = 6, LEFT, delta = 1))
             .isEqualTo(CLEAR_ALL_PLACEHOLDER_ID)
     }
 
@@ -176,11 +170,7 @@
     @Test
     fun equalLengthRows_noFocused_onClearAll_pressLeft_cycleToFirst() {
         assertThat(
-                getNextGridPage(
-                    currentPageTaskViewId = CLEAR_ALL_PLACEHOLDER_ID,
-                    DIRECTION_LEFT,
-                    delta = 1,
-                )
+                getNextGridPage(currentPageTaskViewId = CLEAR_ALL_PLACEHOLDER_ID, LEFT, delta = 1)
             )
             .isEqualTo(1)
     }
@@ -194,11 +184,7 @@
     @Test
     fun equalLengthRows_noFocused_onClearAll_pressRight_toLastInBottom() {
         assertThat(
-                getNextGridPage(
-                    currentPageTaskViewId = CLEAR_ALL_PLACEHOLDER_ID,
-                    DIRECTION_RIGHT,
-                    delta = -1,
-                )
+                getNextGridPage(currentPageTaskViewId = CLEAR_ALL_PLACEHOLDER_ID, RIGHT, delta = -1)
             )
             .isEqualTo(6)
     }
@@ -214,7 +200,7 @@
         assertThat(
                 getNextGridPage(
                     currentPageTaskViewId = FOCUSED_TASK_ID,
-                    DIRECTION_LEFT,
+                    LEFT,
                     delta = 1,
                     largeTileIds = listOf(FOCUSED_TASK_ID),
                 )
@@ -233,7 +219,7 @@
         assertThat(
                 getNextGridPage(
                     currentPageTaskViewId = FOCUSED_TASK_ID,
-                    DIRECTION_UP,
+                    UP,
                     delta = 1,
                     largeTileIds = listOf(FOCUSED_TASK_ID),
                 )
@@ -254,7 +240,7 @@
         assertThat(
                 getNextGridPage(
                     currentPageTaskViewId = FOCUSED_TASK_ID,
-                    DIRECTION_DOWN,
+                    DOWN,
                     delta = 1,
                     largeTileIds = listOf(FOCUSED_TASK_ID),
                 )
@@ -275,7 +261,7 @@
         assertThat(
                 getNextGridPage(
                     currentPageTaskViewId = FOCUSED_TASK_ID,
-                    DIRECTION_RIGHT,
+                    RIGHT,
                     delta = -1,
                     largeTileIds = listOf(FOCUSED_TASK_ID),
                 )
@@ -297,7 +283,7 @@
         assertThat(
                 getNextGridPage(
                     currentPageTaskViewId = CLEAR_ALL_PLACEHOLDER_ID,
-                    DIRECTION_LEFT,
+                    LEFT,
                     delta = 1,
                     largeTileIds = listOf(FOCUSED_TASK_ID),
                 )
@@ -316,7 +302,7 @@
         assertThat(
                 getNextGridPage(
                     currentPageTaskViewId = 7,
-                    DIRECTION_DOWN,
+                    DOWN,
                     delta = 1,
                     topIds = IntArray.wrap(1, 3, 5, 7),
                 )
@@ -335,7 +321,7 @@
         assertThat(
                 getNextGridPage(
                     /* topIds = */ currentPageTaskViewId = 7,
-                    DIRECTION_UP,
+                    UP,
                     delta = 1,
                     topIds = IntArray.wrap(1, 3, 5, 7),
                 )
@@ -353,7 +339,7 @@
         assertThat(
                 getNextGridPage(
                     /* topIds = */ currentPageTaskViewId = 6,
-                    DIRECTION_LEFT,
+                    LEFT,
                     delta = 1,
                     topIds = IntArray.wrap(1, 3, 5, 7),
                 )
@@ -372,7 +358,7 @@
         assertThat(
                 getNextGridPage(
                     /* topIds = */ currentPageTaskViewId = CLEAR_ALL_PLACEHOLDER_ID,
-                    DIRECTION_RIGHT,
+                    RIGHT,
                     delta = -1,
                     topIds = IntArray.wrap(1, 3, 5, 7),
                 )
@@ -391,7 +377,7 @@
         assertThat(
                 getNextGridPage(
                     currentPageTaskViewId = CLEAR_ALL_PLACEHOLDER_ID,
-                    DIRECTION_RIGHT,
+                    RIGHT,
                     delta = -1,
                     bottomIds = IntArray.wrap(2, 4, 6, 7),
                 )
@@ -406,8 +392,7 @@
     */
     @Test
     fun equalLengthRows_noFocused_onTop_pressTab_goesToBottom() {
-        assertThat(getNextGridPage(currentPageTaskViewId = 1, DIRECTION_TAB, delta = 1))
-            .isEqualTo(2)
+        assertThat(getNextGridPage(currentPageTaskViewId = 1, TAB, delta = 1)).isEqualTo(2)
     }
 
     /*
@@ -418,8 +403,7 @@
     */
     @Test
     fun equalLengthRows_noFocused_onBottom_pressTab_goesToNextTop() {
-        assertThat(getNextGridPage(currentPageTaskViewId = 2, DIRECTION_TAB, delta = 1))
-            .isEqualTo(3)
+        assertThat(getNextGridPage(currentPageTaskViewId = 2, TAB, delta = 1)).isEqualTo(3)
     }
 
     /*
@@ -431,8 +415,7 @@
     */
     @Test
     fun equalLengthRows_noFocused_onTop_pressTabWithShift_goesToPreviousBottom() {
-        assertThat(getNextGridPage(currentPageTaskViewId = 3, DIRECTION_TAB, delta = -1))
-            .isEqualTo(2)
+        assertThat(getNextGridPage(currentPageTaskViewId = 3, TAB, delta = -1)).isEqualTo(2)
     }
 
     /*
@@ -442,8 +425,7 @@
     */
     @Test
     fun equalLengthRows_noFocused_onBottom_pressTabWithShift_goesToTop() {
-        assertThat(getNextGridPage(currentPageTaskViewId = 2, DIRECTION_TAB, delta = -1))
-            .isEqualTo(1)
+        assertThat(getNextGridPage(currentPageTaskViewId = 2, TAB, delta = -1)).isEqualTo(1)
     }
 
     /*
@@ -453,9 +435,7 @@
     */
     @Test
     fun equalLengthRows_noFocused_onTop_pressTabWithShift_noCycle_staysOnTop() {
-        assertThat(
-                getNextGridPage(currentPageTaskViewId = 1, DIRECTION_TAB, delta = -1, cycle = false)
-            )
+        assertThat(getNextGridPage(currentPageTaskViewId = 1, TAB, delta = -1, cycle = false))
             .isEqualTo(1)
     }
 
@@ -469,7 +449,7 @@
         assertThat(
                 getNextGridPage(
                     currentPageTaskViewId = CLEAR_ALL_PLACEHOLDER_ID,
-                    DIRECTION_TAB,
+                    TAB,
                     delta = 1,
                     cycle = false,
                 )
@@ -487,7 +467,7 @@
         assertThat(
                 getNextGridPage(
                     currentPageTaskViewId = DESKTOP_TASK_ID,
-                    DIRECTION_LEFT,
+                    LEFT,
                     delta = 1,
                     largeTileIds = listOf(DESKTOP_TASK_ID, FOCUSED_TASK_ID),
                 )
@@ -505,7 +485,7 @@
         assertThat(
                 getNextGridPage(
                     currentPageTaskViewId = FOCUSED_TASK_ID,
-                    DIRECTION_RIGHT,
+                    RIGHT,
                     delta = -1,
                     largeTileIds = listOf(DESKTOP_TASK_ID, FOCUSED_TASK_ID),
                 )
@@ -525,7 +505,7 @@
         assertThat(
                 getNextGridPage(
                     currentPageTaskViewId = DESKTOP_TASK_ID,
-                    DIRECTION_RIGHT,
+                    RIGHT,
                     delta = -1,
                     largeTileIds = listOf(DESKTOP_TASK_ID, FOCUSED_TASK_ID),
                 )
@@ -546,7 +526,7 @@
         assertThat(
                 getNextGridPage(
                     currentPageTaskViewId = CLEAR_ALL_PLACEHOLDER_ID,
-                    DIRECTION_LEFT,
+                    LEFT,
                     delta = 1,
                     largeTileIds = listOf(DESKTOP_TASK_ID, FOCUSED_TASK_ID),
                 )
@@ -565,7 +545,7 @@
         assertThat(
                 getNextGridPage(
                     currentPageTaskViewId = 2,
-                    DIRECTION_RIGHT,
+                    RIGHT,
                     delta = -1,
                     largeTileIds = listOf(DESKTOP_TASK_ID, FOCUSED_TASK_ID),
                 )
@@ -584,7 +564,7 @@
         assertThat(
                 getNextGridPage(
                     currentPageTaskViewId = 1,
-                    DIRECTION_RIGHT,
+                    RIGHT,
                     delta = -1,
                     largeTileIds = listOf(DESKTOP_TASK_ID, FOCUSED_TASK_ID),
                 )
@@ -603,7 +583,7 @@
         assertThat(
                 getNextGridPage(
                     currentPageTaskViewId = DESKTOP_TASK_ID,
-                    DIRECTION_TAB,
+                    TAB,
                     delta = 1,
                     largeTileIds = listOf(DESKTOP_TASK_ID, FOCUSED_TASK_ID),
                 )
@@ -621,7 +601,7 @@
         assertThat(
                 getNextGridPage(
                     currentPageTaskViewId = FOCUSED_TASK_ID,
-                    DIRECTION_LEFT,
+                    LEFT,
                     delta = 1,
                     topIds = IntArray(),
                     bottomIds = IntArray.wrap(2),
@@ -643,7 +623,7 @@
         assertThat(
                 getNextGridPage(
                     currentPageTaskViewId = DESKTOP_TASK_ID,
-                    DIRECTION_TAB,
+                    TAB,
                     delta = -1,
                     largeTileIds = listOf(DESKTOP_TASK_ID, FOCUSED_TASK_ID),
                 )
@@ -662,7 +642,7 @@
         assertThat(
                 getNextGridPage(
                     currentPageTaskViewId = 1,
-                    DIRECTION_RIGHT,
+                    RIGHT,
                     delta = -1,
                     hasAddDesktopButton = true,
                 )
@@ -681,7 +661,7 @@
         assertThat(
                 getNextGridPage(
                     currentPageTaskViewId = 2,
-                    DIRECTION_RIGHT,
+                    RIGHT,
                     delta = -1,
                     hasAddDesktopButton = true,
                 )
@@ -701,7 +681,7 @@
         assertThat(
                 getNextGridPage(
                     currentPageTaskViewId = ADD_DESK_PLACEHOLDER_ID,
-                    DIRECTION_RIGHT,
+                    RIGHT,
                     delta = -1,
                     hasAddDesktopButton = true,
                 )
@@ -722,7 +702,7 @@
         assertThat(
                 getNextGridPage(
                     currentPageTaskViewId = CLEAR_ALL_PLACEHOLDER_ID,
-                    DIRECTION_LEFT,
+                    LEFT,
                     delta = 1,
                     hasAddDesktopButton = true,
                 )
@@ -741,7 +721,7 @@
         assertThat(
                 getNextGridPage(
                     currentPageTaskViewId = ADD_DESK_PLACEHOLDER_ID,
-                    DIRECTION_UP,
+                    UP,
                     delta = 1,
                     hasAddDesktopButton = true,
                 )
@@ -760,7 +740,7 @@
         assertThat(
                 getNextGridPage(
                     currentPageTaskViewId = ADD_DESK_PLACEHOLDER_ID,
-                    DIRECTION_DOWN,
+                    DOWN,
                     delta = 1,
                     hasAddDesktopButton = true,
                 )
@@ -778,7 +758,7 @@
         assertThat(
                 getNextGridPage(
                     currentPageTaskViewId = ADD_DESK_PLACEHOLDER_ID,
-                    DIRECTION_LEFT,
+                    LEFT,
                     delta = 1,
                     largeTileIds = listOf(DESKTOP_TASK_ID),
                     hasAddDesktopButton = true,
@@ -797,7 +777,7 @@
         assertThat(
                 getNextGridPage(
                     currentPageTaskViewId = DESKTOP_TASK_ID,
-                    DIRECTION_RIGHT,
+                    RIGHT,
                     delta = -1,
                     largeTileIds = listOf(DESKTOP_TASK_ID),
                     hasAddDesktopButton = true,
@@ -806,9 +786,43 @@
             .isEqualTo(ADD_DESK_PLACEHOLDER_ID)
     }
 
+    // Col offset:  0   1   2
+    //             -----------
+    // ID grid:     4   2   0  start
+    //         end [5]  3   1
+    @Test
+    fun gridTaskViewIdOffsetPairInTabOrderSequence_towardsStart() {
+        val expected = listOf(Pair(4, 0), Pair(3, 1), Pair(2, 1), Pair(1, 2), Pair(0, 2))
+        assertThat(
+                gridTaskViewIdOffsetPairInTabOrderSequence(
+                        initialTaskViewId = 5,
+                        towardsStart = true,
+                    )
+                    .toList()
+            )
+            .isEqualTo(expected)
+    }
+
+    // Col offset:  2   1   0
+    //             -----------
+    // ID grid:     4   2  [0] start
+    //          end 5   3   1
+    @Test
+    fun gridTaskViewIdOffsetPairInTabOrderSequence_towardsEnd() {
+        val expected = listOf(Pair(1, 0), Pair(2, 1), Pair(3, 1), Pair(4, 2), Pair(5, 2))
+        assertThat(
+                gridTaskViewIdOffsetPairInTabOrderSequence(
+                        initialTaskViewId = 0,
+                        towardsStart = false,
+                    )
+                    .toList()
+            )
+            .isEqualTo(expected)
+    }
+
     private fun getNextGridPage(
         currentPageTaskViewId: Int,
-        direction: Int,
+        direction: TaskGridNavHelper.TaskNavDirection,
         delta: Int,
         topIds: IntArray = IntArray.wrap(1, 3, 5),
         bottomIds: IntArray = IntArray.wrap(2, 4, 6),
@@ -821,6 +835,22 @@
         return taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle)
     }
 
+    private fun gridTaskViewIdOffsetPairInTabOrderSequence(
+        initialTaskViewId: Int,
+        towardsStart: Boolean,
+        topIds: IntArray = IntArray.wrap(0, 2, 4),
+        bottomIds: IntArray = IntArray.wrap(1, 3, 5),
+        largeTileIds: List<Int> = emptyList(),
+        hasAddDesktopButton: Boolean = false,
+    ): Sequence<Pair<Int, Int>> {
+        val taskGridNavHelper =
+            TaskGridNavHelper(topIds, bottomIds, largeTileIds, hasAddDesktopButton)
+        return taskGridNavHelper.gridTaskViewIdOffsetPairInTabOrderSequence(
+            initialTaskViewId,
+            towardsStart,
+        )
+    }
+
     private companion object {
         const val FOCUSED_TASK_ID = 99
         const val DESKTOP_TASK_ID = 100
diff --git a/src/com/android/launcher3/BubbleTextView.java b/src/com/android/launcher3/BubbleTextView.java
index fb847f9..03dd943 100644
--- a/src/com/android/launcher3/BubbleTextView.java
+++ b/src/com/android/launcher3/BubbleTextView.java
@@ -75,8 +75,8 @@
 import com.android.launcher3.dragndrop.DragOptions.PreDragCondition;
 import com.android.launcher3.dragndrop.DraggableView;
 import com.android.launcher3.folder.FolderIcon;
-import com.android.launcher3.graphics.IconShape;
 import com.android.launcher3.graphics.PreloadIconDrawable;
+import com.android.launcher3.graphics.ShapeDelegate;
 import com.android.launcher3.icons.DotRenderer;
 import com.android.launcher3.icons.FastBitmapDrawable;
 import com.android.launcher3.icons.IconCache.ItemInfoUpdateReceiver;
@@ -724,7 +724,7 @@
         if (!mForceHideDot && (hasDot() || mDotParams.scale > 0)) {
             getIconBounds(mDotParams.iconBounds);
             Utilities.scaleRectAboutCenter(mDotParams.iconBounds,
-                    IconShape.INSTANCE.get(getContext()).getNormalizationScale());
+                    ShapeDelegate.getNormalizationScale());
             final int scrollX = getScrollX();
             final int scrollY = getScrollY();
             canvas.translate(scrollX, scrollY);
@@ -775,7 +775,7 @@
         getIconBounds(mRunningAppIconBounds);
         Utilities.scaleRectAboutCenter(
                 mRunningAppIconBounds,
-                IconShape.INSTANCE.get(getContext()).getNormalizationScale());
+                ShapeDelegate.getNormalizationScale());
 
         final boolean isMinimized = mRunningAppState == RunningAppState.MINIMIZED;
         final int indicatorTop = mRunningAppIconBounds.bottom + mRunningAppIndicatorTopMargin;
@@ -964,7 +964,8 @@
 
     @Override
     public void setTextColor(ColorStateList colors) {
-        mTextColor = shouldDrawAppContrastTile() ? PillColorProvider.getInstance(
+        mTextColor = (shouldDrawAppContrastTile() && !TextUtils.isEmpty(getText()))
+                ? PillColorProvider.getInstance(
                 getContext()).getAppTitleTextPaint().getColor()
                 : colors.getDefaultColor();
         mTextColorStateList = colors;
@@ -989,7 +990,7 @@
     public boolean shouldDrawAppContrastTile() {
         return mDisplay == DISPLAY_WORKSPACE && shouldTextBeVisible()
                 && PillColorProvider.getInstance(getContext()).isMatchaEnabled()
-                && enableContrastTiles() && !TextUtils.isEmpty(getText());
+                && enableContrastTiles();
     }
 
     public void setTextVisibility(boolean visible) {
diff --git a/src/com/android/launcher3/DeviceProfile.java b/src/com/android/launcher3/DeviceProfile.java
index 813d8f1..f34add8 100644
--- a/src/com/android/launcher3/DeviceProfile.java
+++ b/src/com/android/launcher3/DeviceProfile.java
@@ -52,7 +52,7 @@
 import com.android.launcher3.CellLayout.ContainerType;
 import com.android.launcher3.DevicePaddings.DevicePadding;
 import com.android.launcher3.folder.ClippedFolderIconLayoutRule;
-import com.android.launcher3.graphics.IconShape;
+import com.android.launcher3.graphics.ThemeManager;
 import com.android.launcher3.icons.DotRenderer;
 import com.android.launcher3.icons.IconNormalizer;
 import com.android.launcher3.model.data.ItemInfo;
@@ -387,7 +387,8 @@
     }
 
     /** TODO: Once we fully migrate to staged split, remove "isMultiWindowMode" */
-    DeviceProfile(Context context, InvariantDeviceProfile inv, Info info, WindowBounds windowBounds,
+    DeviceProfile(Context context, InvariantDeviceProfile inv, Info info,
+            WindowManagerProxy wmProxy, ThemeManager themeManager, WindowBounds windowBounds,
             SparseArray<DotRenderer> dotRendererCache, boolean isMultiWindowMode,
             boolean transposeLayoutWithOrientation, boolean isMultiDisplay, boolean isGestureMode,
             @NonNull final ViewScaleProvider viewScaleProvider,
@@ -420,7 +421,7 @@
         isPhone = !isTablet;
         isTwoPanels = isTablet && isMultiDisplay;
         isTaskbarPresent = (isTablet || (enableTinyTaskbar() && isGestureMode))
-                && WindowManagerProxy.INSTANCE.get(context).isTaskbarDrawnInProcess();
+                && wmProxy.isTaskbarDrawnInProcess();
 
         // Some more constants.
         context = getContext(context, info, inv.isFixedLandscape
@@ -845,8 +846,8 @@
         dimensionOverrideProvider.accept(this);
 
         // This is done last, after iconSizePx is calculated above.
-        mDotRendererWorkSpace = createDotRenderer(context, iconSizePx, dotRendererCache);
-        mDotRendererAllApps = createDotRenderer(context, allAppsIconSizePx, dotRendererCache);
+        mDotRendererWorkSpace = createDotRenderer(themeManager, iconSizePx, dotRendererCache);
+        mDotRendererAllApps = createDotRenderer(themeManager, allAppsIconSizePx, dotRendererCache);
     }
 
     /**
@@ -868,12 +869,12 @@
     }
 
     private static DotRenderer createDotRenderer(
-            @NonNull Context context, int size, @NonNull SparseArray<DotRenderer> cache) {
+            @NonNull ThemeManager themeManager, int size, @NonNull SparseArray<DotRenderer> cache) {
         DotRenderer renderer = cache.get(size);
         if (renderer == null) {
             renderer = new DotRenderer(
                     size,
-                    IconShape.INSTANCE.get(context).getShape().getPath(DEFAULT_DOT_SIZE),
+                    themeManager.getIconShape().getPath(DEFAULT_DOT_SIZE),
                     DEFAULT_DOT_SIZE);
             cache.put(size, renderer);
         }
@@ -1090,7 +1091,7 @@
         dotRendererCache.put(iconSizePx, mDotRendererWorkSpace);
         dotRendererCache.put(allAppsIconSizePx, mDotRendererAllApps);
 
-        return new Builder(context, inv, mInfo)
+        return inv.newDPBuilder(context, mInfo)
                 .setWindowBounds(bounds)
                 .setIsMultiDisplay(isMultiDisplay)
                 .setMultiWindowMode(isMultiWindowMode)
@@ -2473,9 +2474,11 @@
     }
 
     public static class Builder {
-        private Context mContext;
-        private InvariantDeviceProfile mInv;
-        private Info mInfo;
+        private final Context mContext;
+        private final InvariantDeviceProfile mInv;
+        private final Info mInfo;
+        private final WindowManagerProxy mWMProxy;
+        private final ThemeManager mThemeManager;
 
         private WindowBounds mWindowBounds;
         private boolean mIsMultiDisplay;
@@ -2491,10 +2494,13 @@
 
         private boolean mIsTransientTaskbar;
 
-        public Builder(Context context, InvariantDeviceProfile inv, Info info) {
+        public Builder(Context context, InvariantDeviceProfile inv, Info info,
+                WindowManagerProxy wmProxy, ThemeManager themeManager) {
             mContext = context;
             mInv = inv;
             mInfo = info;
+            mWMProxy = wmProxy;
+            mThemeManager = themeManager;
             mIsTransientTaskbar = info.isTransientTaskbar();
         }
 
@@ -2575,7 +2581,8 @@
             if (mOverrideProvider == null) {
                 mOverrideProvider = DEFAULT_DIMENSION_PROVIDER;
             }
-            return new DeviceProfile(mContext, mInv, mInfo, mWindowBounds, mDotRendererCache,
+            return new DeviceProfile(mContext, mInv, mInfo, mWMProxy, mThemeManager,
+                    mWindowBounds, mDotRendererCache,
                     mIsMultiWindowMode, mTransposeLayoutWithOrientation, mIsMultiDisplay,
                     mIsGestureMode, mViewScaleProvider, mOverrideProvider, mIsTransientTaskbar);
         }
diff --git a/src/com/android/launcher3/InvariantDeviceProfile.java b/src/com/android/launcher3/InvariantDeviceProfile.java
index 56c2b8e..c2d6df5 100644
--- a/src/com/android/launcher3/InvariantDeviceProfile.java
+++ b/src/com/android/launcher3/InvariantDeviceProfile.java
@@ -30,7 +30,6 @@
 import static com.android.launcher3.util.DisplayController.CHANGE_TASKBAR_PINNING;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 
-import android.annotation.TargetApi;
 import android.content.Context;
 import android.content.Intent;
 import android.content.res.Resources;
@@ -46,7 +45,6 @@
 import android.util.Log;
 import android.util.SparseArray;
 import android.util.Xml;
-import android.view.Display;
 
 import androidx.annotation.DimenRes;
 import androidx.annotation.IntDef;
@@ -56,18 +54,21 @@
 import androidx.core.content.res.ResourcesCompat;
 
 import com.android.launcher3.config.FeatureFlags;
+import com.android.launcher3.dagger.ApplicationContext;
+import com.android.launcher3.dagger.LauncherAppComponent;
+import com.android.launcher3.dagger.LauncherAppSingleton;
+import com.android.launcher3.graphics.ThemeManager;
 import com.android.launcher3.icons.DotRenderer;
 import com.android.launcher3.logging.FileLog;
 import com.android.launcher3.model.DeviceGridState;
 import com.android.launcher3.provider.RestoreDbTask;
 import com.android.launcher3.testing.shared.ResourceUtils;
+import com.android.launcher3.util.DaggerSingletonObject;
+import com.android.launcher3.util.DaggerSingletonTracker;
 import com.android.launcher3.util.DisplayController;
 import com.android.launcher3.util.DisplayController.Info;
-import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.util.Partner;
 import com.android.launcher3.util.ResourceHelper;
-import com.android.launcher3.util.RunnableList;
-import com.android.launcher3.util.SafeCloseable;
 import com.android.launcher3.util.SimpleBroadcastReceiver;
 import com.android.launcher3.util.WindowBounds;
 import com.android.launcher3.util.window.CachedDisplayInfo;
@@ -86,12 +87,15 @@
 import java.util.Objects;
 import java.util.stream.Collectors;
 
-public class InvariantDeviceProfile implements SafeCloseable {
+import javax.inject.Inject;
+
+@LauncherAppSingleton
+public class InvariantDeviceProfile {
 
     public static final String TAG = "IDP";
     // We do not need any synchronization for this variable as its only written on UI thread.
-    public static final MainThreadInitializedObject<InvariantDeviceProfile> INSTANCE =
-            new MainThreadInitializedObject<>(InvariantDeviceProfile::new);
+    public static final DaggerSingletonObject<InvariantDeviceProfile> INSTANCE =
+            new DaggerSingletonObject<>(LauncherAppComponent::getIDP);
 
     public static final String GRID_NAME_PREFS_KEY = "idp_grid_name";
     public static final String NON_FIXED_LANDSCAPE_GRID_NAME_PREFS_KEY =
@@ -129,7 +133,10 @@
     private static final String RES_GRID_NUM_COLUMNS = "grid_num_columns";
     private static final String RES_GRID_ICON_SIZE_DP = "grid_icon_size_dp";
 
-    private final RunnableList mCloseActions = new RunnableList();
+    private final DisplayController mDisplayController;
+    private final WindowManagerProxy mWMProxy;
+    private final LauncherPrefs mPrefs;
+    private final ThemeManager mThemeManager;
 
     /**
      * Number of icons per row and column in the workspace.
@@ -247,16 +254,22 @@
 
     private final ArrayList<OnIDPChangeListener> mChangeListeners = new ArrayList<>();
 
-    @VisibleForTesting
-    public InvariantDeviceProfile() {
-    }
+    @Inject
+    InvariantDeviceProfile(
+            @ApplicationContext Context context,
+            LauncherPrefs prefs,
+            DisplayController dc,
+            WindowManagerProxy wmProxy,
+            ThemeManager themeManager,
+            DaggerSingletonTracker lifeCycle) {
+        mDisplayController = dc;
+        mWMProxy = wmProxy;
+        mPrefs = prefs;
+        mThemeManager = themeManager;
 
-    @TargetApi(23)
-    private InvariantDeviceProfile(Context context) {
-        String gridName = getCurrentGridName(context);
+        String gridName = prefs.get(GRID_NAME);
         initGrid(context, gridName);
 
-        DisplayController dc = DisplayController.INSTANCE.get(context);
         dc.setPriorityListener(
                 (displayContext, info, flags) -> {
                     if ((flags & (CHANGE_DENSITY | CHANGE_SUPPORTED_BOUNDS
@@ -265,126 +278,46 @@
                         onConfigChanged(displayContext);
                     }
                 });
-        mCloseActions.add(() -> dc.setPriorityListener(null));
+        lifeCycle.addCloseable(() -> dc.setPriorityListener(null));
 
         LauncherPrefChangeListener prefListener = key -> {
             if (FIXED_LANDSCAPE_MODE.getSharedPrefKey().equals(key)
-                    && isFixedLandscape != FIXED_LANDSCAPE_MODE.get(context)) {
+                    && isFixedLandscape != prefs.get(FIXED_LANDSCAPE_MODE)) {
                 Trace.beginSection("InvariantDeviceProfile#setFixedLandscape");
                 if (isFixedLandscape) {
-                    setCurrentGrid(
-                            context, LauncherPrefs.get(context).get(NON_FIXED_LANDSCAPE_GRID_NAME));
+                    setCurrentGrid(context, prefs.get(NON_FIXED_LANDSCAPE_GRID_NAME));
                 } else {
-                    LauncherPrefs.get(context)
-                            .put(NON_FIXED_LANDSCAPE_GRID_NAME, getCurrentGridName(context));
+                    prefs.put(NON_FIXED_LANDSCAPE_GRID_NAME, mPrefs.get(GRID_NAME));
                     onConfigChanged(context);
                 }
                 Trace.endSection();
             } else if (ENABLE_TWOLINE_ALLAPPS_TOGGLE.getSharedPrefKey().equals(key)
-                    && enableTwoLinesInAllApps != ENABLE_TWOLINE_ALLAPPS_TOGGLE.get(context)) {
+                    && enableTwoLinesInAllApps != prefs.get(ENABLE_TWOLINE_ALLAPPS_TOGGLE)) {
                 onConfigChanged(context);
             }
         };
-        LauncherPrefs prefs = LauncherPrefs.INSTANCE.get(context);
         prefs.addListener(prefListener, FIXED_LANDSCAPE_MODE, ENABLE_TWOLINE_ALLAPPS_TOGGLE);
-        mCloseActions.add(() -> prefs.removeListener(prefListener,
+        lifeCycle.addCloseable(() -> prefs.removeListener(prefListener,
                 FIXED_LANDSCAPE_MODE, ENABLE_TWOLINE_ALLAPPS_TOGGLE));
 
-        SimpleBroadcastReceiver localeReceiver = new SimpleBroadcastReceiver(
+        SimpleBroadcastReceiver localeReceiver = new SimpleBroadcastReceiver(context,
                 MAIN_EXECUTOR, i -> onConfigChanged(context));
-        localeReceiver.register(context, Intent.ACTION_LOCALE_CHANGED);
-        mCloseActions.add(() -> localeReceiver.unregisterReceiverSafely(context));
-    }
-
-    /**
-     * This constructor should NOT have any monitors by design.
-     */
-    public InvariantDeviceProfile(Context context, String gridName) {
-        String newName = initGrid(context, gridName);
-        if (newName == null || !newName.equals(gridName)) {
-            throw new IllegalArgumentException("Unknown grid name: " + gridName);
-        }
-    }
-
-    /**
-     * This constructor should NOT have any monitors by design.
-     */
-    public InvariantDeviceProfile(Context context, Display display) {
-        // Ensure that the main device profile is initialized
-        INSTANCE.get(context);
-        String gridName = getCurrentGridName(context);
-
-        // Get the display info based on default display and interpolate it to existing display
-        Info defaultInfo = DisplayController.INSTANCE.get(context).getInfo();
-        @DeviceType int defaultDeviceType = defaultInfo.getDeviceType();
-        DisplayOption defaultDisplayOption = invDistWeightedInterpolate(
-                defaultInfo,
-                getPredefinedDeviceProfiles(
-                        context,
-                        gridName,
-                        defaultInfo,
-                        /*allowDisabledGrid=*/false,
-                        FIXED_LANDSCAPE_MODE.get(context)
-                ),
-                defaultDeviceType);
-
-        Context displayContext = context.createDisplayContext(display);
-        Info myInfo = new Info(displayContext);
-        @DeviceType int deviceType = myInfo.getDeviceType();
-        DisplayOption myDisplayOption = invDistWeightedInterpolate(
-                myInfo,
-                getPredefinedDeviceProfiles(
-                        context,
-                        gridName,
-                        myInfo,
-                        /*allowDisabledGrid=*/false,
-                        FIXED_LANDSCAPE_MODE.get(context)
-                ),
-                deviceType);
-
-        DisplayOption result = new DisplayOption(defaultDisplayOption.grid)
-                .add(myDisplayOption);
-        result.iconSizes[INDEX_DEFAULT] =
-                defaultDisplayOption.iconSizes[INDEX_DEFAULT];
-        for (int i = 1; i < COUNT_SIZES; i++) {
-            result.iconSizes[i] = Math.min(
-                    defaultDisplayOption.iconSizes[i], myDisplayOption.iconSizes[i]);
-        }
-
-        System.arraycopy(defaultDisplayOption.minCellSize, 0, result.minCellSize, 0,
-                COUNT_SIZES);
-        System.arraycopy(defaultDisplayOption.borderSpaces, 0, result.borderSpaces, 0,
-                COUNT_SIZES);
-
-        initGrid(context, myInfo, result);
-    }
-
-    @Override
-    public void close() {
-        mCloseActions.executeAllAndDestroy();
-    }
-
-    public static String getCurrentGridName(Context context) {
-        return LauncherPrefs.get(context).get(GRID_NAME);
+        localeReceiver.register(Intent.ACTION_LOCALE_CHANGED);
+        lifeCycle.addCloseable(() -> localeReceiver.unregisterReceiverSafely());
     }
 
     private String initGrid(Context context, String gridName) {
-        FileLog.d(TAG, "Before initGrid:"
-                + "gridName:" + gridName
-                + ", dbFile:" + dbFile
-                + ", LauncherPrefs GRID_NAME:" + LauncherPrefs.get(context).get(GRID_NAME)
-                + ", LauncherPrefs DB_FILE:" + LauncherPrefs.get(context).get(DB_FILE));
-        Info displayInfo = DisplayController.INSTANCE.get(context).getInfo();
+        Info displayInfo = mDisplayController.getInfo();
         List<DisplayOption> allOptions = getPredefinedDeviceProfiles(
                 context,
                 gridName,
                 displayInfo,
-                RestoreDbTask.isPending(context),
-                FIXED_LANDSCAPE_MODE.get(context)
+                RestoreDbTask.isPending(mPrefs),
+                mPrefs.get(FIXED_LANDSCAPE_MODE)
         );
 
         // Filter out options that don't have the same number of columns as the grid
-        DeviceGridState deviceGridState = new DeviceGridState(context);
+        DeviceGridState deviceGridState = new DeviceGridState(mPrefs);
         List<DisplayOption> allOptionsFilteredByColCount =
                 filterByColumnCount(allOptions, deviceGridState.getColumns());
 
@@ -395,15 +328,15 @@
                         displayInfo.getDeviceType());
 
         if (!displayOption.grid.name.equals(gridName)) {
-            LauncherPrefs.get(context).put(GRID_NAME, displayOption.grid.name);
+            mPrefs.put(GRID_NAME, displayOption.grid.name);
         }
 
         initGrid(context, displayInfo, displayOption);
         FileLog.d(TAG, "After initGrid:"
                 + "gridName:" + gridName
                 + ", dbFile:" + dbFile
-                + ", LauncherPrefs GRID_NAME:" + LauncherPrefs.get(context).get(GRID_NAME)
-                + ", LauncherPrefs DB_FILE:" + LauncherPrefs.get(context).get(DB_FILE));
+                + ", LauncherPrefs GRID_NAME:" + mPrefs.get(GRID_NAME)
+                + ", LauncherPrefs DB_FILE:" + mPrefs.get(DB_FILE));
         return displayOption.grid.name;
     }
 
@@ -420,18 +353,13 @@
      */
     @Deprecated
     public void reset(Context context) {
-        initGrid(context, getCurrentGridName(context));
-    }
-
-    @VisibleForTesting
-    public static String getDefaultGridName(Context context) {
-        return new InvariantDeviceProfile().initGrid(context, null);
+        initGrid(context, mPrefs.get(GRID_NAME));
     }
 
     private void initGrid(Context context, Info displayInfo, DisplayOption displayOption) {
         enableTwoLinesInAllApps = Flags.enableTwolineToggle()
                 && Utilities.isEnglishLanguage(context)
-                && ENABLE_TWOLINE_ALLAPPS_TOGGLE.get(context);
+                && mPrefs.get(ENABLE_TWOLINE_ALLAPPS_TOGGLE);
         mLocale = context.getResources().getConfiguration().locale.toString();
 
         DisplayMetrics metrics = context.getResources().getDisplayMetrics();
@@ -522,7 +450,7 @@
         defaultWallpaperSize = new Point(displayInfo.currentSize);
         SparseArray<DotRenderer> dotRendererCache = new SparseArray<>();
         for (WindowBounds bounds : displayInfo.supportedBounds) {
-            localSupportedProfiles.add(new DeviceProfile.Builder(context, this, displayInfo)
+            localSupportedProfiles.add(newDPBuilder(context, displayInfo)
                     .setIsMultiDisplay(deviceType == TYPE_MULTI_DISPLAY)
                     .setWindowBounds(bounds)
                     .setDotRendererCache(dotRendererCache)
@@ -561,6 +489,10 @@
                 });
     }
 
+    DeviceProfile.Builder newDPBuilder(Context context, Info info) {
+        return new DeviceProfile.Builder(context, this, info, mWMProxy, mThemeManager);
+    }
+
     public void addOnChangeListener(OnIDPChangeListener listener) {
         mChangeListeners.add(listener);
     }
@@ -574,7 +506,7 @@
      * migration.
      */
     public void setCurrentGrid(Context context, String newGridName) {
-        LauncherPrefs.get(context).put(GRID_NAME, newGridName);
+        mPrefs.put(GRID_NAME, newGridName);
         MAIN_EXECUTOR.execute(() -> {
             Trace.beginSection("InvariantDeviceProfile#setCurrentGrid");
             onConfigChanged(context.getApplicationContext());
@@ -594,8 +526,7 @@
         Object[] oldState = toModelState();
 
         // Re-init grid
-        String gridName = getCurrentGridName(context);
-        initGrid(context, gridName);
+        initGrid(context, mPrefs.get(GRID_NAME));
 
         boolean modelPropsChanged = !Arrays.equals(oldState, toModelState());
         for (OnIDPChangeListener listener : mChangeListeners) {
@@ -912,11 +843,20 @@
         return out;
     }
 
-    public DeviceProfile getDeviceProfile(Context context) {
-        WindowManagerProxy windowManagerProxy = WindowManagerProxy.INSTANCE.get(context);
-        Rect bounds = windowManagerProxy.getCurrentBounds(context);
-        int rotation = windowManagerProxy.getRotation(context);
+    public DeviceProfile createDeviceProfileForSecondaryDisplay(Context displayContext) {
+        // Disable transpose layout and use multi-window mode so that the icons are scaled properly
+        return newDPBuilder(displayContext, new Info(displayContext))
+                .setIsMultiDisplay(false)
+                .setMultiWindowMode(true)
+                .setWindowBounds(mWMProxy.getRealBounds(
+                        displayContext, mWMProxy.getDisplayInfo(displayContext)))
+                .setTransposeLayoutWithOrientation(false)
+                .build();
+    }
 
+    public DeviceProfile getDeviceProfile(Context context) {
+        Rect bounds = mWMProxy.getCurrentBounds(context);
+        int rotation = mWMProxy.getRotation(context);
         return getBestMatch(bounds.width(), bounds.height(), rotation);
     }
 
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index db79662..3edba99 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -71,6 +71,7 @@
 import static com.android.launcher3.LauncherState.NO_SCALE;
 import static com.android.launcher3.LauncherState.SPRING_LOADED;
 import static com.android.launcher3.Utilities.postAsyncCallback;
+import static com.android.launcher3.anim.AnimatorListeners.forEndCallback;
 import static com.android.launcher3.config.FeatureFlags.FOLDABLE_SINGLE_PAGE;
 import static com.android.launcher3.config.FeatureFlags.MULTI_SELECT_EDIT_MODE;
 import static com.android.launcher3.logging.KeyboardStateManager.KeyboardState.HIDE;
@@ -1466,7 +1467,10 @@
             }
 
             getModelWriter().addItemToDatabase(info, container, screenId, cellXY[0], cellXY[1]);
-            mWorkspace.addInScreen(view, info);
+            AnimatorSet anim = new AnimatorSet();
+            anim.addListener(forEndCallback(() ->
+                    view.sendAccessibilityEvent(AccessibilityEvent.TYPE_VIEW_FOCUSED)));
+            bindInflatedItems(Collections.singletonList(Pair.create(info, view)), anim);
         } else {
             // Adding a shortcut to a Folder.
             FolderIcon folderIcon = findFolderIcon(container);
diff --git a/src/com/android/launcher3/LauncherAppState.java b/src/com/android/launcher3/LauncherAppState.java
index e560a14..71013c3 100644
--- a/src/com/android/launcher3/LauncherAppState.java
+++ b/src/com/android/launcher3/LauncherAppState.java
@@ -119,14 +119,13 @@
         }
 
         SimpleBroadcastReceiver modelChangeReceiver =
-                new SimpleBroadcastReceiver(UI_HELPER_EXECUTOR, mModel::onBroadcastIntent);
+                new SimpleBroadcastReceiver(context, UI_HELPER_EXECUTOR, mModel::onBroadcastIntent);
         modelChangeReceiver.register(
-                mContext,
                 ACTION_DEVICE_POLICY_RESOURCE_UPDATED);
         if (BuildConfig.IS_STUDIO_BUILD) {
-            modelChangeReceiver.register(mContext, RECEIVER_EXPORTED, ACTION_FORCE_ROLOAD);
+            modelChangeReceiver.register(RECEIVER_EXPORTED, ACTION_FORCE_ROLOAD);
         }
-        mOnTerminateCallback.add(() -> modelChangeReceiver.unregisterReceiverSafely(mContext));
+        mOnTerminateCallback.add(() -> modelChangeReceiver.unregisterReceiverSafely());
 
         SafeCloseable userChangeListener = UserCache.INSTANCE.get(mContext)
                 .addUserEventListener(mModel::onUserEvent);
diff --git a/src/com/android/launcher3/dagger/LauncherBaseAppComponent.java b/src/com/android/launcher3/dagger/LauncherBaseAppComponent.java
index 7bd7c3e..31d0da0 100644
--- a/src/com/android/launcher3/dagger/LauncherBaseAppComponent.java
+++ b/src/com/android/launcher3/dagger/LauncherBaseAppComponent.java
@@ -18,9 +18,10 @@
 
 import android.content.Context;
 
+import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherPrefs;
-import com.android.launcher3.graphics.IconShape;
 import com.android.launcher3.graphics.ThemeManager;
+import com.android.launcher3.icons.LauncherIcons.IconPool;
 import com.android.launcher3.model.ItemInstallQueue;
 import com.android.launcher3.pm.InstallSessionHelper;
 import com.android.launcher3.pm.UserCache;
@@ -55,7 +56,6 @@
     ApiWrapper getApiWrapper();
     CustomWidgetManager getCustomWidgetManager();
     DynamicResource getDynamicResource();
-    IconShape getIconShape();
     InstallSessionHelper getInstallSessionHelper();
     ItemInstallQueue getItemInstallQueue();
     RefreshRateTracker getRefreshRateTracker();
@@ -72,6 +72,8 @@
     DisplayController getDisplayController();
     WallpaperColorHints getWallpaperColorHints();
     LockedUserState getLockedUserState();
+    InvariantDeviceProfile getIDP();
+    IconPool getIconPool();
 
     /** Builder for LauncherBaseAppComponent. */
     interface Builder {
diff --git a/src/com/android/launcher3/dagger/LauncherComponentProvider.kt b/src/com/android/launcher3/dagger/LauncherComponentProvider.kt
index 5015e54..71e3354 100644
--- a/src/com/android/launcher3/dagger/LauncherComponentProvider.kt
+++ b/src/com/android/launcher3/dagger/LauncherComponentProvider.kt
@@ -47,6 +47,10 @@
             .component
     }
 
+    /** Extension method easily access LauncherAppComponent */
+    val Context.appComponent: LauncherAppComponent
+        get() = get(this)
+
     private data class Holder(
         val component: LauncherAppComponent,
         private val filter: LayoutInflater.Filter?,
diff --git a/src/com/android/launcher3/dragndrop/DragView.java b/src/com/android/launcher3/dragndrop/DragView.java
index 072673d..bff323c 100644
--- a/src/com/android/launcher3/dragndrop/DragView.java
+++ b/src/com/android/launcher3/dragndrop/DragView.java
@@ -60,7 +60,7 @@
 import com.android.app.animation.Interpolators;
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
-import com.android.launcher3.graphics.IconShape;
+import com.android.launcher3.graphics.ThemeManager;
 import com.android.launcher3.icons.FastBitmapDrawable;
 import com.android.launcher3.icons.LauncherIcons;
 import com.android.launcher3.model.data.ItemInfo;
@@ -274,9 +274,9 @@
                 Utilities.scaleRectAboutCenter(shrunkBounds, 0.98f);
                 adaptiveIcon.setBounds(shrunkBounds);
 
-                IconShape iconShape = IconShape.INSTANCE.get(getContext());
+                ThemeManager themeManager = ThemeManager.INSTANCE.get(getContext());
                 final Path mask = (adaptiveIcon instanceof FolderAdaptiveIcon
-                        ? iconShape.getFolderShape() : iconShape.getShape())
+                        ? themeManager.getFolderShape() : themeManager.getIconShape())
                         .getPath(shrunkBounds);
 
                 mTranslateX = new SpringFloatValue(DragView.this,
diff --git a/src/com/android/launcher3/folder/FolderAnimationManager.java b/src/com/android/launcher3/folder/FolderAnimationManager.java
index 2157610..d2354c1 100644
--- a/src/com/android/launcher3/folder/FolderAnimationManager.java
+++ b/src/com/android/launcher3/folder/FolderAnimationManager.java
@@ -46,8 +46,8 @@
 import com.android.launcher3.anim.PropertyResetListener;
 import com.android.launcher3.apppairs.AppPairIcon;
 import com.android.launcher3.celllayout.CellLayoutLayoutParams;
-import com.android.launcher3.graphics.IconShape;
-import com.android.launcher3.graphics.IconShape.ShapeDelegate;
+import com.android.launcher3.graphics.ShapeDelegate;
+import com.android.launcher3.graphics.ThemeManager;
 import com.android.launcher3.util.Themes;
 import com.android.launcher3.views.BaseDragLayer;
 
@@ -237,7 +237,7 @@
         }
         play(a, getAnimator(mFolder.mFooter, ALPHA, 0, 1f), footerStartDelay, footerAlphaDuration);
 
-        ShapeDelegate shapeDelegate = IconShape.INSTANCE.get(mContext).getFolderShape();
+        ShapeDelegate shapeDelegate = ThemeManager.INSTANCE.get(mContext).getFolderShape();
         // Create reveal animator for the folder background
         play(a, shapeDelegate.createRevealAnimator(
                 mFolder, startRect, endRect, finalRadius, !mIsOpening));
diff --git a/src/com/android/launcher3/folder/PreviewBackground.java b/src/com/android/launcher3/folder/PreviewBackground.java
index 77fa355..ba8a290 100644
--- a/src/com/android/launcher3/folder/PreviewBackground.java
+++ b/src/com/android/launcher3/folder/PreviewBackground.java
@@ -47,8 +47,8 @@
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.R;
 import com.android.launcher3.celllayout.DelegatedCellDrawing;
-import com.android.launcher3.graphics.IconShape;
-import com.android.launcher3.graphics.IconShape.ShapeDelegate;
+import com.android.launcher3.graphics.ShapeDelegate;
+import com.android.launcher3.graphics.ThemeManager;
 import com.android.launcher3.util.Themes;
 import com.android.launcher3.views.ActivityContext;
 
@@ -260,7 +260,7 @@
     }
 
     private ShapeDelegate getShape() {
-        return IconShape.INSTANCE.get(mContext).getFolderShape();
+        return ThemeManager.INSTANCE.get(mContext).getFolderShape();
     }
 
     public void drawShadow(Canvas canvas) {
diff --git a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
index 911064c..03f0582 100644
--- a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
+++ b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
@@ -24,6 +24,7 @@
 import static com.android.launcher3.BubbleTextView.DISPLAY_WORKSPACE;
 import static com.android.launcher3.DeviceProfile.DEFAULT_SCALE;
 import static com.android.launcher3.Hotseat.ALPHA_CHANNEL_PREVIEW_RENDERER;
+import static com.android.launcher3.LauncherPrefs.GRID_NAME;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT_PREDICTION;
 import static com.android.launcher3.Utilities.SHOULD_SHOW_FIRST_PAGE_WIDGET;
@@ -137,14 +138,14 @@
 
         private final String mPrefName;
 
-        public PreviewContext(Context base, InvariantDeviceProfile idp) {
+        public PreviewContext(Context base, String gridName) {
             super(base);
             mPrefName = "preview-" + UUID.randomUUID().toString();
-            initDaggerComponent(DaggerLauncherPreviewRenderer_PreviewAppComponent.builder()
-                    .bindPrefs(new ProxyPrefs(
-                            this, getSharedPreferences(mPrefName, MODE_PRIVATE))));
-
-            putObject(InvariantDeviceProfile.INSTANCE, idp);
+            LauncherPrefs prefs =
+                    new ProxyPrefs(this, getSharedPreferences(mPrefName, MODE_PRIVATE));
+            prefs.put(GRID_NAME, gridName);
+            initDaggerComponent(
+                    DaggerLauncherPreviewRenderer_PreviewAppComponent.builder().bindPrefs(prefs));
             putObject(LauncherAppState.INSTANCE,
                     new LauncherAppState(this, null /* iconCacheFileName */));
         }
@@ -192,8 +193,8 @@
                 this::getAppWidgetScale).build();
         if (context instanceof PreviewContext) {
             Context tempContext = ((PreviewContext) context).getBaseContext();
-            mDpOrig = new InvariantDeviceProfile(tempContext, InvariantDeviceProfile
-                    .getCurrentGridName(tempContext)).getDeviceProfile(tempContext)
+            mDpOrig = InvariantDeviceProfile.INSTANCE.get(tempContext)
+                    .getDeviceProfile(tempContext)
                     .copy(tempContext);
         } else {
             mDpOrig = mDp;
diff --git a/src/com/android/launcher3/graphics/PreloadIconDrawable.java b/src/com/android/launcher3/graphics/PreloadIconDrawable.java
index 50d6d1c..3bd9fb5 100644
--- a/src/com/android/launcher3/graphics/PreloadIconDrawable.java
+++ b/src/com/android/launcher3/graphics/PreloadIconDrawable.java
@@ -119,7 +119,7 @@
                 IconPalette.getPreloadProgressColor(context, info.bitmap.color),
                 getPreloadColors(context),
                 Utilities.isDarkTheme(context),
-                IconShape.INSTANCE.get(context).getShape().getPath(DEFAULT_PATH_SIZE)
+                ThemeManager.INSTANCE.get(context).getIconShape().getPath(DEFAULT_PATH_SIZE)
         );
     }
 
diff --git a/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java b/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
index 7a60814..f0d670e 100644
--- a/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
+++ b/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
@@ -20,6 +20,7 @@
 import static android.content.res.Configuration.UI_MODE_NIGHT_YES;
 import static android.view.Display.DEFAULT_DISPLAY;
 
+import static com.android.launcher3.LauncherPrefs.GRID_NAME;
 import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
 import static com.android.launcher3.graphics.ThemeManager.PREF_ICON_SHAPE;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
@@ -39,6 +40,7 @@
 import android.util.SparseIntArray;
 import android.view.ContextThemeWrapper;
 import android.view.Display;
+import android.view.Surface;
 import android.view.SurfaceControlViewHost;
 import android.view.SurfaceControlViewHost.SurfacePackage;
 import android.view.View;
@@ -60,7 +62,6 @@
 import com.android.launcher3.model.BaseLauncherBinder;
 import com.android.launcher3.model.BgDataModel;
 import com.android.launcher3.model.BgDataModel.Callbacks;
-import com.android.launcher3.model.GridSizeMigrationDBController;
 import com.android.launcher3.model.LoaderTask;
 import com.android.launcher3.model.ModelDbController;
 import com.android.launcher3.provider.LauncherDbUtils;
@@ -117,7 +118,7 @@
         mGridName = bundle.getString("name");
         bundle.remove("name");
         if (mGridName == null) {
-            mGridName = InvariantDeviceProfile.getCurrentGridName(context);
+            mGridName = LauncherPrefs.get(context).get(GRID_NAME);
         }
         mWallpaperColors = bundle.getParcelable(KEY_COLORS);
         if (Flags.newCustomizationPickerUi()) {
@@ -286,6 +287,20 @@
             }
             context = context.createConfigurationContext(configuration);
         }
+        if (InvariantDeviceProfile.INSTANCE.get(context).isFixedLandscape) {
+            Configuration configuration = new Configuration(
+                    context.getResources().getConfiguration()
+            );
+            int width = configuration.screenWidthDp;
+            int height = configuration.screenHeightDp;
+            if (configuration.screenHeightDp > configuration.screenWidthDp) {
+                configuration.screenWidthDp = height;
+                configuration.screenHeightDp = width;
+                configuration.orientation = Surface.ROTATION_90;
+            }
+            context = context.createConfigurationContext(configuration);
+        }
+
         if (Flags.newCustomizationPickerUi()) {
             if (mPreviewColorOverride != null) {
                 LocalColorExtractor.newInstance(context)
@@ -316,11 +331,10 @@
     @WorkerThread
     private void loadModelData() {
         final Context inflationContext = getPreviewContext();
-        final InvariantDeviceProfile idp = new InvariantDeviceProfile(inflationContext, mGridName);
-        if (GridSizeMigrationDBController.needsToMigrate(inflationContext, idp)
+        if (!mGridName.equals(LauncherPrefs.INSTANCE.get(mContext).get(GRID_NAME))
                 || mShapeKey != null) {
             // Start the migration
-            PreviewContext previewContext = new PreviewContext(inflationContext, idp);
+            PreviewContext previewContext = new PreviewContext(inflationContext, mGridName);
             if (mShapeKey != null) {
                 LauncherPrefs.INSTANCE.get(previewContext).put(PREF_ICON_SHAPE, mShapeKey);
             }
@@ -348,6 +362,7 @@
 
                 @Override
                 public void run() {
+                    InvariantDeviceProfile idp = LauncherAppState.getIDP(previewContext);
                     DeviceProfile deviceProfile = idp.getDeviceProfile(previewContext);
                     String query =
                             LauncherSettings.Favorites.SCREEN + " = " + Workspace.FIRST_SCREEN_ID
@@ -371,7 +386,7 @@
             LauncherAppState.getInstance(inflationContext).getModel().loadAsync(dataModel -> {
                 if (dataModel != null) {
                     MAIN_EXECUTOR.execute(() -> renderView(inflationContext, dataModel, null,
-                            null, idp));
+                            null, LauncherAppState.getIDP(inflationContext)));
                 } else {
                     Log.e(TAG, "Model loading failed");
                 }
@@ -396,15 +411,28 @@
         }
         renderer.hideBottomRow(mHideQsb);
         View view = renderer.getRenderedView(dataModel, widgetProviderInfoMap);
-        // This aspect scales the view to fit in the surface and centers it
-        final float scale = Math.min(mWidth / (float) view.getMeasuredWidth(),
-                mHeight / (float) view.getMeasuredHeight());
-        view.setScaleX(scale);
-        view.setScaleY(scale);
+
         view.setPivotX(0);
         view.setPivotY(0);
-        view.setTranslationX((mWidth - scale * view.getWidth()) / 2);
-        view.setTranslationY((mHeight - scale * view.getHeight()) / 2);
+        if (idp.isFixedLandscape) {
+            final float scale = Math.min(mHeight / (float) view.getMeasuredWidth(),
+                    mWidth / (float) view.getMeasuredHeight());
+            view.setScaleX(scale);
+            view.setScaleY(scale);
+            view.setRotation(90);
+            view.setTranslationX((mHeight - scale * view.getWidth()) / 2 + mWidth);
+            view.setTranslationY((mWidth - scale * view.getHeight()) / 2);
+        } else {
+            // This aspect scales the view to fit in the surface and centers it
+            final float scale = Math.min(mWidth / (float) view.getMeasuredWidth(),
+                    mHeight / (float) view.getMeasuredHeight());
+            view.setScaleX(scale);
+            view.setScaleY(scale);
+            view.setTranslationX((mWidth - scale * view.getWidth()) / 2);
+            view.setTranslationY((mHeight - scale * view.getHeight()) / 2);
+        }
+
+
         if (!Flags.newCustomizationPickerUi()) {
             view.setAlpha(0);
             view.animate().alpha(1)
diff --git a/src/com/android/launcher3/graphics/IconShape.kt b/src/com/android/launcher3/graphics/ShapeDelegate.kt
similarity index 81%
rename from src/com/android/launcher3/graphics/IconShape.kt
rename to src/com/android/launcher3/graphics/ShapeDelegate.kt
index eac3440..df0c8f9 100644
--- a/src/com/android/launcher3/graphics/IconShape.kt
+++ b/src/com/android/launcher3/graphics/ShapeDelegate.kt
@@ -42,78 +42,38 @@
 import androidx.graphics.shapes.toPath
 import androidx.graphics.shapes.transformed
 import com.android.launcher3.anim.RoundedRectRevealOutlineProvider
-import com.android.launcher3.dagger.LauncherAppComponent
-import com.android.launcher3.dagger.LauncherAppSingleton
-import com.android.launcher3.graphics.ThemeManager.ThemeChangeListener
 import com.android.launcher3.icons.GraphicsUtils
 import com.android.launcher3.icons.IconNormalizer.normalizeAdaptiveIcon
-import com.android.launcher3.util.DaggerSingletonObject
-import com.android.launcher3.util.DaggerSingletonTracker
 import com.android.launcher3.views.ClipPathView
-import javax.inject.Inject
 
 /** Abstract representation of the shape of an icon shape */
-@LauncherAppSingleton
-class IconShape
-@Inject
-constructor(private val themeManager: ThemeManager, lifeCycle: DaggerSingletonTracker) {
+interface ShapeDelegate {
 
-    val normalizationScale =
-        normalizeAdaptiveIcon(
-            AdaptiveIconDrawable(null, ColorDrawable(Color.BLACK)),
-            AREA_CALC_SIZE,
-        )
+    fun getPath(pathSize: Float = DEFAULT_PATH_SIZE) =
+        Path().apply { addToPath(this, 0f, 0f, pathSize / 2) }
 
-    var shape: ShapeDelegate = pickBestShape(themeManager.iconState.iconMask)
-        private set
-
-    var folderShape: ShapeDelegate =
-        themeManager.iconState.run {
-            if (folderShapeMask == iconMask || folderShapeMask.isEmpty()) shape
-            else pickBestShape(folderShapeMask)
+    fun getPath(bounds: Rect) =
+        Path().apply {
+            addToPath(
+                this,
+                bounds.left.toFloat(),
+                bounds.top.toFloat(),
+                // Radius is half of the average size of the icon
+                (bounds.width() + bounds.height()) / 4f,
+            )
         }
-        private set
 
-    init {
-        val changeListener = ThemeChangeListener {
-            shape = pickBestShape(themeManager.iconState.iconMask)
-            folderShape =
-                themeManager.iconState.run {
-                    if (folderShapeMask == iconMask || folderShapeMask.isEmpty()) shape
-                    else pickBestShape(folderShapeMask)
-                }
-        }
-        themeManager.addChangeListener(changeListener)
-        lifeCycle.addCloseable { themeManager.removeChangeListener(changeListener) }
-    }
+    fun drawShape(canvas: Canvas, offsetX: Float, offsetY: Float, radius: Float, paint: Paint)
 
-    interface ShapeDelegate {
-        fun getPath(pathSize: Float = DEFAULT_PATH_SIZE) =
-            Path().apply { addToPath(this, 0f, 0f, pathSize / 2) }
+    fun addToPath(path: Path, offsetX: Float, offsetY: Float, radius: Float)
 
-        fun getPath(bounds: Rect) =
-            Path().apply {
-                addToPath(
-                    this,
-                    bounds.left.toFloat(),
-                    bounds.top.toFloat(),
-                    // Radius is half of the average size of the icon
-                    (bounds.width() + bounds.height()) / 4f,
-                )
-            }
-
-        fun drawShape(canvas: Canvas, offsetX: Float, offsetY: Float, radius: Float, paint: Paint)
-
-        fun addToPath(path: Path, offsetX: Float, offsetY: Float, radius: Float)
-
-        fun <T> createRevealAnimator(
-            target: T,
-            startRect: Rect,
-            endRect: Rect,
-            endRadius: Float,
-            isReversed: Boolean,
-        ): ValueAnimator where T : View, T : ClipPathView
-    }
+    fun <T> createRevealAnimator(
+        target: T,
+        startRect: Rect,
+        endRect: Rect,
+        endRadius: Float,
+        isReversed: Boolean,
+    ): ValueAnimator where T : View, T : ClipPathView
 
     class Circle : RoundedSquare(1f) {
 
@@ -179,10 +139,15 @@
                 }
                 .createRevealAnimator(target, isReversed)
         }
+
+        override fun equals(other: Any?) =
+            other is RoundedSquare && other.radiusRatio == radiusRatio
+
+        override fun hashCode() = radiusRatio.hashCode()
     }
 
     /** Generic shape delegate with pathString in bounds [0, 0, 100, 100] */
-    class GenericPathShape(pathString: String) : ShapeDelegate {
+    data class GenericPathShape(private val pathString: String) : ShapeDelegate {
         private val poly =
             RoundedPolygon(
                 features = SvgPathParser.parseFeatures(pathString),
@@ -287,7 +252,6 @@
     }
 
     companion object {
-        @JvmField var INSTANCE = DaggerSingletonObject(LauncherAppComponent::getIconShape)
 
         const val TAG = "IconShape"
         const val DEFAULT_PATH_SIZE = 100f
@@ -312,7 +276,6 @@
             }
         }
 
-        @VisibleForTesting
         fun pickBestShape(shapeStr: String): ShapeDelegate {
             val baseShape =
                 if (shapeStr.isNotEmpty()) {
@@ -332,7 +295,6 @@
             return pickBestShape(baseShape, shapeStr)
         }
 
-        @VisibleForTesting
         fun pickBestShape(baseShape: Path, shapeStr: String): ShapeDelegate {
             val calcAreaDiff = areaDiffCalculator(baseShape)
 
@@ -380,5 +342,12 @@
                 centerY = (bottom - top) / 2,
                 rounding = CornerRounding(cornerR),
             )
+
+        @JvmStatic
+        val normalizationScale =
+            normalizeAdaptiveIcon(
+                AdaptiveIconDrawable(null, ColorDrawable(Color.BLACK)),
+                AREA_CALC_SIZE,
+            )
     }
 }
diff --git a/src/com/android/launcher3/graphics/ThemeManager.kt b/src/com/android/launcher3/graphics/ThemeManager.kt
index 242220a..4a0ff8c 100644
--- a/src/com/android/launcher3/graphics/ThemeManager.kt
+++ b/src/com/android/launcher3/graphics/ThemeManager.kt
@@ -25,6 +25,7 @@
 import com.android.launcher3.dagger.ApplicationContext
 import com.android.launcher3.dagger.LauncherAppComponent
 import com.android.launcher3.dagger.LauncherAppSingleton
+import com.android.launcher3.graphics.ShapeDelegate.Companion.pickBestShape
 import com.android.launcher3.icons.IconThemeController
 import com.android.launcher3.icons.mono.MonoIconThemeController
 import com.android.launcher3.shapes.ShapesProvider
@@ -46,24 +47,30 @@
 ) {
 
     /** Representation of the current icon state */
-    var iconState = parseIconState()
+    var iconState = parseIconState(null)
         private set
 
     var isMonoThemeEnabled
         set(value) = prefs.put(THEMED_ICONS, value)
         get() = prefs.get(THEMED_ICONS)
 
-    val themeController: IconThemeController?
+    val themeController
         get() = iconState.themeController
 
-    val isIconThemeEnabled: Boolean
+    val isIconThemeEnabled
         get() = themeController != null
 
+    val iconShape
+        get() = iconState.iconShape
+
+    val folderShape
+        get() = iconState.folderShape
+
     private val listeners = CopyOnWriteArrayList<ThemeChangeListener>()
 
     init {
-        val receiver = SimpleBroadcastReceiver(MAIN_EXECUTOR) { verifyIconState() }
-        receiver.registerPkgActions(context, "android", ACTION_OVERLAY_CHANGED)
+        val receiver = SimpleBroadcastReceiver(context, MAIN_EXECUTOR) { verifyIconState() }
+        receiver.registerPkgActions("android", ACTION_OVERLAY_CHANGED)
 
         val prefListener = LauncherPrefChangeListener { key ->
             when (key) {
@@ -74,13 +81,13 @@
         prefs.addListener(prefListener, THEMED_ICONS, PREF_ICON_SHAPE)
 
         lifecycle.addCloseable {
-            receiver.unregisterReceiverSafely(context)
+            receiver.unregisterReceiverSafely()
             prefs.removeListener(prefListener)
         }
     }
 
     protected fun verifyIconState() {
-        val newState = parseIconState()
+        val newState = parseIconState(iconState)
         if (newState == iconState) return
         iconState = newState
 
@@ -91,7 +98,7 @@
 
     fun removeChangeListener(listener: ThemeChangeListener) = listeners.remove(listener)
 
-    private fun parseIconState(): IconState {
+    private fun parseIconState(oldState: IconState?): IconState {
         val shapeModel =
             prefs.get(PREF_ICON_SHAPE).let { shapeOverride ->
                 ShapesProvider.iconShapes.values.firstOrNull { it.key == shapeOverride }
@@ -102,10 +109,27 @@
                 CONFIG_ICON_MASK_RES_ID == Resources.ID_NULL -> ""
                 else -> context.resources.getString(CONFIG_ICON_MASK_RES_ID)
             }
+
+        val iconShape =
+            if (oldState != null && oldState.iconMask == iconMask) oldState.iconShape
+            else pickBestShape(iconMask)
+
+        val folderShapeMask = shapeModel?.folderPathString ?: iconMask
+        val folderShape =
+            when {
+                oldState != null && oldState.folderShapeMask == folderShapeMask ->
+                    oldState.folderShape
+                folderShapeMask == iconMask || folderShapeMask.isEmpty() -> iconShape
+                else -> pickBestShape(folderShapeMask)
+            }
+
         return IconState(
             iconMask = iconMask,
-            folderShapeMask = shapeModel?.folderPathString ?: iconMask,
+            folderShapeMask = folderShapeMask,
             themeController = createThemeController(),
+            iconScale = shapeModel?.iconScale ?: 1f,
+            iconShape = iconShape,
+            folderShape = folderShape,
         )
     }
 
@@ -118,6 +142,9 @@
         val folderShapeMask: String,
         val themeController: IconThemeController?,
         val themeCode: String = themeController?.themeID ?: "no-theme",
+        val iconScale: Float = 1f,
+        val iconShape: ShapeDelegate,
+        val folderShape: ShapeDelegate,
     ) {
         fun toUniqueId() = "${iconMask.hashCode()},$themeCode"
     }
diff --git a/src/com/android/launcher3/icons/LauncherIconProvider.java b/src/com/android/launcher3/icons/LauncherIconProvider.java
index 482360c..836b7d1 100644
--- a/src/com/android/launcher3/icons/LauncherIconProvider.java
+++ b/src/com/android/launcher3/icons/LauncherIconProvider.java
@@ -30,7 +30,7 @@
 
 import com.android.launcher3.R;
 import com.android.launcher3.config.FeatureFlags;
-import com.android.launcher3.graphics.IconShape;
+import com.android.launcher3.graphics.ShapeDelegate;
 import com.android.launcher3.graphics.ThemeManager;
 import com.android.launcher3.util.ApiWrapper;
 
@@ -54,13 +54,13 @@
     private Map<String, ThemeData> mThemedIconMap;
 
     private final ApiWrapper mApiWrapper;
-    private final IconShape mIconShape;
+    private final ThemeManager mThemeManager;
 
     public LauncherIconProvider(Context context) {
         super(context);
-        setIconThemeSupported(ThemeManager.INSTANCE.get(context).isMonoThemeEnabled());
+        mThemeManager = ThemeManager.INSTANCE.get(context);
         mApiWrapper = ApiWrapper.INSTANCE.get(context);
-        mIconShape = IconShape.INSTANCE.get(context);
+        setIconThemeSupported(mThemeManager.isMonoThemeEnabled());
     }
 
     /**
@@ -91,7 +91,7 @@
     @Override
     protected Drawable loadAppInfoIcon(ApplicationInfo info, Resources resources, int density) {
         // Tries to load the round icon res, if the app defines it as an adaptive icon
-        if (mIconShape.getShape() instanceof IconShape.Circle) {
+        if (mThemeManager.getIconShape() instanceof ShapeDelegate.Circle) {
             int roundIconRes = mApiWrapper.getRoundIconRes(info);
             if (roundIconRes != 0 && roundIconRes != info.icon) {
                 try {
diff --git a/src/com/android/launcher3/icons/LauncherIcons.java b/src/com/android/launcher3/icons/LauncherIcons.java
deleted file mode 100644
index 5c6debe..0000000
--- a/src/com/android/launcher3/icons/LauncherIcons.java
+++ /dev/null
@@ -1,164 +0,0 @@
-/*
- * Copyright (C) 2016 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.icons;
-
-import static android.graphics.Color.BLACK;
-
-import static com.android.launcher3.graphics.ThemeManager.PREF_ICON_SHAPE;
-
-import android.content.Context;
-import android.graphics.Canvas;
-import android.graphics.Path;
-import android.graphics.Rect;
-import android.graphics.drawable.AdaptiveIconDrawable;
-import android.os.UserHandle;
-
-import androidx.annotation.NonNull;
-
-import com.android.launcher3.Flags;
-import com.android.launcher3.InvariantDeviceProfile;
-import com.android.launcher3.LauncherPrefs;
-import com.android.launcher3.graphics.IconShape;
-import com.android.launcher3.graphics.ThemeManager;
-import com.android.launcher3.pm.UserCache;
-import com.android.launcher3.util.MainThreadInitializedObject;
-import com.android.launcher3.util.SafeCloseable;
-import com.android.launcher3.util.UserIconInfo;
-
-import java.util.concurrent.ConcurrentLinkedQueue;
-
-/**
- * Wrapper class to provide access to {@link BaseIconFactory} and also to provide pool of this class
- * that are threadsafe.
- */
-public class LauncherIcons extends BaseIconFactory implements AutoCloseable {
-
-    private static final float SEVEN_SIDED_COOKIE_SCALE = 72f / 80f;
-    private static final float FOUR_SIDED_COOKIE_SCALE = 72f / 83.4f;
-    private static final float VERY_SUNNY_SCALE = 72f / 92f;
-    private static final float DEFAULT_ICON_SCALE = 1f;
-
-
-    private static final MainThreadInitializedObject<Pool> POOL =
-            new MainThreadInitializedObject<>(Pool::new);
-
-    /**
-     * Return a new Message instance from the global pool. Allows us to
-     * avoid allocating new objects in many cases.
-     */
-    public static LauncherIcons obtain(Context context) {
-        return POOL.get(context).obtain();
-    }
-
-    public static void clearPool(Context context) {
-        POOL.get(context).close();
-    }
-
-    private final ConcurrentLinkedQueue<LauncherIcons> mPool;
-
-    protected LauncherIcons(Context context, int fillResIconDpi, int iconBitmapSize,
-            ConcurrentLinkedQueue<LauncherIcons> pool) {
-        super(context, fillResIconDpi, iconBitmapSize);
-        mThemeController = ThemeManager.INSTANCE.get(context).getThemeController();
-        mPool = pool;
-    }
-
-    /**
-     * Recycles a LauncherIcons that may be in-use.
-     */
-    public void recycle() {
-        clear();
-        mPool.add(this);
-    }
-
-    @NonNull
-    @Override
-    protected UserIconInfo getUserInfo(@NonNull UserHandle user) {
-        return UserCache.INSTANCE.get(mContext).getUserInfo(user);
-    }
-
-    @NonNull
-    @Override
-    public Path getShapePath(AdaptiveIconDrawable drawable, Rect iconBounds) {
-        if (!Flags.enableLauncherIconShapes()) return drawable.getIconMask();
-        return IconShape.INSTANCE.get(mContext).getShape().getPath(iconBounds);
-    }
-
-    @Override
-    protected void drawAdaptiveIcon(
-            @NonNull Canvas canvas,
-            @NonNull AdaptiveIconDrawable drawable,
-            @NonNull Path overridePath
-    ) {
-        if (!Flags.enableLauncherIconShapes()) {
-            super.drawAdaptiveIcon(canvas, drawable, overridePath);
-            return;
-        }
-        String shapeKey = LauncherPrefs.get(mContext).get(PREF_ICON_SHAPE);
-        float iconScale = switch (shapeKey) {
-            case "seven_sided_cookie" -> SEVEN_SIDED_COOKIE_SCALE;
-            case "four_sided_cookie" -> FOUR_SIDED_COOKIE_SCALE;
-            case "sunny" -> VERY_SUNNY_SCALE;
-            default -> DEFAULT_ICON_SCALE;
-        };
-        canvas.clipPath(overridePath);
-        canvas.drawColor(BLACK);
-        canvas.save();
-        canvas.scale(iconScale, iconScale, canvas.getWidth() / 2f, canvas.getHeight() / 2f);
-        if (drawable.getBackground() != null) {
-            drawable.getBackground().draw(canvas);
-        }
-        if (drawable.getForeground() != null) {
-            drawable.getForeground().draw(canvas);
-        }
-        canvas.restore();
-    }
-
-    @Override
-    public void close() {
-        recycle();
-    }
-
-    private static class Pool implements SafeCloseable {
-
-        private final Context mContext;
-
-        @NonNull
-        private ConcurrentLinkedQueue<LauncherIcons> mPool = new ConcurrentLinkedQueue<>();
-
-        private Pool(Context context) {
-            mContext = context;
-        }
-
-        public LauncherIcons obtain() {
-            ConcurrentLinkedQueue<LauncherIcons> pool = mPool;
-            LauncherIcons m = pool.poll();
-
-            if (m == null) {
-                InvariantDeviceProfile idp = InvariantDeviceProfile.INSTANCE.get(mContext);
-                return new LauncherIcons(mContext, idp.fillResIconDpi, idp.iconBitmapSize, pool);
-            } else {
-                return m;
-            }
-        }
-
-        @Override
-        public void close() {
-            mPool = new ConcurrentLinkedQueue<>();
-        }
-    }
-}
diff --git a/src/com/android/launcher3/icons/LauncherIcons.kt b/src/com/android/launcher3/icons/LauncherIcons.kt
new file mode 100644
index 0000000..6c018e8
--- /dev/null
+++ b/src/com/android/launcher3/icons/LauncherIcons.kt
@@ -0,0 +1,127 @@
+/*
+ * Copyright (C) 2016 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.icons
+
+import android.content.Context
+import android.graphics.Canvas
+import android.graphics.Color
+import android.graphics.Path
+import android.graphics.Rect
+import android.graphics.drawable.AdaptiveIconDrawable
+import android.os.UserHandle
+import com.android.launcher3.Flags
+import com.android.launcher3.InvariantDeviceProfile
+import com.android.launcher3.dagger.ApplicationContext
+import com.android.launcher3.dagger.LauncherAppSingleton
+import com.android.launcher3.dagger.LauncherComponentProvider.appComponent
+import com.android.launcher3.graphics.ThemeManager
+import com.android.launcher3.pm.UserCache
+import com.android.launcher3.util.UserIconInfo
+import dagger.assisted.Assisted
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
+import java.util.concurrent.ConcurrentLinkedQueue
+import javax.inject.Inject
+
+/**
+ * Wrapper class to provide access to [BaseIconFactory] and also to provide pool of this class that
+ * are threadsafe.
+ */
+class LauncherIcons
+@AssistedInject
+internal constructor(
+    @ApplicationContext context: Context,
+    idp: InvariantDeviceProfile,
+    private var themeManager: ThemeManager,
+    private var userCache: UserCache,
+    @Assisted private val pool: ConcurrentLinkedQueue<LauncherIcons>,
+) : BaseIconFactory(context, idp.fillResIconDpi, idp.iconBitmapSize), AutoCloseable {
+
+    private val iconScale = themeManager.iconState.iconScale
+
+    init {
+        mThemeController = themeManager.themeController
+    }
+
+    /** Recycles a LauncherIcons that may be in-use. */
+    fun recycle() {
+        clear()
+        pool.add(this)
+    }
+
+    override fun getUserInfo(user: UserHandle): UserIconInfo {
+        return userCache.getUserInfo(user)
+    }
+
+    public override fun getShapePath(drawable: AdaptiveIconDrawable, iconBounds: Rect): Path {
+        if (!Flags.enableLauncherIconShapes()) return drawable.iconMask
+        return themeManager.iconShape.getPath(iconBounds)
+    }
+
+    override fun drawAdaptiveIcon(
+        canvas: Canvas,
+        drawable: AdaptiveIconDrawable,
+        overridePath: Path,
+    ) {
+        if (!Flags.enableLauncherIconShapes()) {
+            super.drawAdaptiveIcon(canvas, drawable, overridePath)
+            return
+        }
+        canvas.clipPath(overridePath)
+        canvas.drawColor(Color.BLACK)
+        canvas.save()
+        canvas.scale(iconScale, iconScale, canvas.width / 2f, canvas.height / 2f)
+        if (drawable.background != null) {
+            drawable.background.draw(canvas)
+        }
+        if (drawable.foreground != null) {
+            drawable.foreground.draw(canvas)
+        }
+        canvas.restore()
+    }
+
+    override fun close() {
+        recycle()
+    }
+
+    @AssistedFactory
+    internal interface LauncherIconsFactory {
+        fun create(pool: ConcurrentLinkedQueue<LauncherIcons>): LauncherIcons
+    }
+
+    @LauncherAppSingleton
+    class IconPool @Inject internal constructor(private val factory: LauncherIconsFactory) {
+        private var pool = ConcurrentLinkedQueue<LauncherIcons>()
+
+        fun obtain(): LauncherIcons = pool.let { it.poll() ?: factory.create(it) }
+
+        fun clear() {
+            pool = ConcurrentLinkedQueue()
+        }
+    }
+
+    companion object {
+
+        /**
+         * Return a new LauncherIcons instance from the global pool. Allows us to avoid allocating
+         * new objects in many cases.
+         */
+        @JvmStatic
+        fun obtain(context: Context): LauncherIcons = context.appComponent.iconPool.obtain()
+
+        @JvmStatic fun clearPool(context: Context) = context.appComponent.iconPool.clear()
+    }
+}
diff --git a/src/com/android/launcher3/model/DeviceGridState.java b/src/com/android/launcher3/model/DeviceGridState.java
index 90af215..6f12c97 100644
--- a/src/com/android/launcher3/model/DeviceGridState.java
+++ b/src/com/android/launcher3/model/DeviceGridState.java
@@ -68,7 +68,10 @@
     }
 
     public DeviceGridState(Context context) {
-        LauncherPrefs lp = LauncherPrefs.get(context);
+        this(LauncherPrefs.get(context));
+    }
+
+    public DeviceGridState(LauncherPrefs lp) {
         mGridSizeString = lp.get(WORKSPACE_SIZE);
         mNumHotseat = lp.get(HOTSEAT_COUNT);
         mDeviceType = lp.get(DEVICE_TYPE);
diff --git a/src/com/android/launcher3/pm/UserCache.java b/src/com/android/launcher3/pm/UserCache.java
index 0b18a87..20c0ecc 100644
--- a/src/com/android/launcher3/pm/UserCache.java
+++ b/src/com/android/launcher3/pm/UserCache.java
@@ -81,10 +81,7 @@
     }
 
     private final List<BiConsumer<UserHandle, String>> mUserEventListeners = new ArrayList<>();
-    private final SimpleBroadcastReceiver mUserChangeReceiver =
-            new SimpleBroadcastReceiver(MODEL_EXECUTOR, this::onUsersChanged);
-
-    private final Context mContext;
+    private final SimpleBroadcastReceiver mUserChangeReceiver;
     private final ApiWrapper mApiWrapper;
 
     @NonNull
@@ -99,16 +96,17 @@
             DaggerSingletonTracker tracker,
             ApiWrapper apiWrapper
     ) {
-        mContext = context;
         mApiWrapper = apiWrapper;
+        mUserChangeReceiver = new SimpleBroadcastReceiver(context,
+                MODEL_EXECUTOR, this::onUsersChanged);
         mUserToSerialMap = Collections.emptyMap();
         MODEL_EXECUTOR.execute(this::initAsync);
-        tracker.addCloseable(() -> mUserChangeReceiver.unregisterReceiverSafely(mContext));
+        tracker.addCloseable(() -> mUserChangeReceiver.unregisterReceiverSafely());
     }
 
     @WorkerThread
     private void initAsync() {
-        mUserChangeReceiver.register(mContext,
+        mUserChangeReceiver.register(
                 Intent.ACTION_MANAGED_PROFILE_AVAILABLE,
                 Intent.ACTION_MANAGED_PROFILE_UNAVAILABLE,
                 Intent.ACTION_MANAGED_PROFILE_REMOVED,
diff --git a/src/com/android/launcher3/provider/RestoreDbTask.java b/src/com/android/launcher3/provider/RestoreDbTask.java
index dc42920..34c9117 100644
--- a/src/com/android/launcher3/provider/RestoreDbTask.java
+++ b/src/com/android/launcher3/provider/RestoreDbTask.java
@@ -415,7 +415,11 @@
     }
 
     public static boolean isPending(Context context) {
-        return LauncherPrefs.get(context).has(RESTORE_DEVICE);
+        return isPending(LauncherPrefs.get(context));
+    }
+
+    public static boolean isPending(LauncherPrefs prefs) {
+        return prefs.has(RESTORE_DEVICE);
     }
 
     /**
diff --git a/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncher.java b/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncher.java
index df27b54..c20d655 100644
--- a/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncher.java
+++ b/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncher.java
@@ -115,15 +115,9 @@
         if (mDragLayer != null) {
             return;
         }
-        InvariantDeviceProfile currentDisplayIdp = new InvariantDeviceProfile(
-                this, getWindow().getDecorView().getDisplay());
 
-        // Disable transpose layout and use multi-window mode so that the icons are scaled properly
-        mDeviceProfile = currentDisplayIdp.getDeviceProfile(this)
-                .toBuilder(this)
-                .setMultiWindowMode(true)
-                .setTransposeLayoutWithOrientation(false)
-                .build();
+        mDeviceProfile = InvariantDeviceProfile.INSTANCE.get(this)
+                .createDeviceProfileForSecondaryDisplay(this);
         mDeviceProfile.autoResizeAllAppsCells();
 
         setContentView(R.layout.secondary_launcher);
diff --git a/src/com/android/launcher3/shapes/IconShapeModel.kt b/src/com/android/launcher3/shapes/IconShapeModel.kt
index dd6c432..fc49adc 100644
--- a/src/com/android/launcher3/shapes/IconShapeModel.kt
+++ b/src/com/android/launcher3/shapes/IconShapeModel.kt
@@ -21,4 +21,5 @@
     val title: String,
     val pathString: String,
     val folderPathString: String = pathString,
+    val iconScale: Float = 1f,
 )
diff --git a/src/com/android/launcher3/shapes/ShapesProvider.kt b/src/com/android/launcher3/shapes/ShapesProvider.kt
index 7e1f640..463c816 100644
--- a/src/com/android/launcher3/shapes/ShapesProvider.kt
+++ b/src/com/android/launcher3/shapes/ShapesProvider.kt
@@ -174,6 +174,7 @@
                         pathString =
                             "M39.888,4.517C46.338 7.319 53.662 7.319 60.112 4.517L63.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 3Z",
                         folderPathString = folderShapes["complexClover"]!!,
+                        iconScale = 72f / 83.4f,
                     ),
                 "seven_sided_cookie" to
                     IconShapeModel(
@@ -182,6 +183,7 @@
                         pathString =
                             "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",
                         folderPathString = folderShapes["clover"]!!,
+                        iconScale = 72f / 80f,
                     ),
                 "arch" to
                     IconShapeModel(
@@ -198,6 +200,7 @@
                         pathString =
                             "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",
                         folderPathString = folderShapes["clover"]!!,
+                        iconScale = 72f / 92f,
                     ),
             )
         } else {
diff --git a/src/com/android/launcher3/util/DisplayController.java b/src/com/android/launcher3/util/DisplayController.java
index ee1af81..376a61e 100644
--- a/src/com/android/launcher3/util/DisplayController.java
+++ b/src/com/android/launcher3/util/DisplayController.java
@@ -107,7 +107,6 @@
     private static final String ACTION_OVERLAY_CHANGED = "android.intent.action.OVERLAY_CHANGED";
     private static final String TARGET_OVERLAY_PACKAGE = "android";
 
-    private final Context mContext;
     private final WindowManagerProxy mWMProxy;
 
     // Null for SDK < S
@@ -120,8 +119,7 @@
 
     // We will register broadcast receiver on main thread to ensure not missing changes on
     // TARGET_OVERLAY_PACKAGE and ACTION_OVERLAY_CHANGED.
-    private final SimpleBroadcastReceiver mReceiver =
-            new SimpleBroadcastReceiver(MAIN_EXECUTOR, this::onIntent);
+    private final SimpleBroadcastReceiver mReceiver;
 
     private Info mInfo;
     private boolean mDestroyed = false;
@@ -131,7 +129,6 @@
             WindowManagerProxy wmProxy,
             LauncherPrefs prefs,
             DaggerSingletonTracker lifecycle) {
-        mContext = context;
         mWMProxy = wmProxy;
 
         if (enableTaskbarPinning()) {
@@ -155,11 +152,12 @@
 
         Display display = context.getSystemService(DisplayManager.class)
                 .getDisplay(DEFAULT_DISPLAY);
-        mWindowContext = mContext.createWindowContext(display, TYPE_APPLICATION, null);
+        mWindowContext = context.createWindowContext(display, TYPE_APPLICATION, null);
         mWindowContext.registerComponentCallbacks(this);
 
         // Initialize navigation mode change listener
-        mReceiver.registerPkgActions(mContext, TARGET_OVERLAY_PACKAGE, ACTION_OVERLAY_CHANGED);
+        mReceiver = new SimpleBroadcastReceiver(context, MAIN_EXECUTOR, this::onIntent);
+        mReceiver.registerPkgActions(TARGET_OVERLAY_PACKAGE, ACTION_OVERLAY_CHANGED);
 
         mInfo = new Info(mWindowContext, wmProxy,
                 wmProxy.estimateInternalDisplayBounds(mWindowContext));
@@ -169,7 +167,7 @@
         lifecycle.addCloseable(() -> {
             mDestroyed = true;
             mWindowContext.unregisterComponentCallbacks(this);
-            mReceiver.unregisterReceiverSafely(mContext);
+            mReceiver.unregisterReceiverSafely();
             wmProxy.unregisterDesktopVisibilityListener(this);
         });
     }
@@ -229,8 +227,10 @@
     }
 
     @Override
-    public void onDesktopVisibilityChanged(boolean visible) {
-        notifyConfigChange();
+    public void onIsInDesktopModeChanged(int displayId, boolean isInDesktopModeAndNotInOverview) {
+        if (DEFAULT_DISPLAY == displayId) {
+            notifyConfigChange();
+        }
     }
 
     /**
@@ -448,7 +448,7 @@
             mIsTaskbarPinned = LauncherPrefs.get(displayInfoContext).get(TASKBAR_PINNING);
             mIsTaskbarPinnedInDesktopMode = LauncherPrefs.get(displayInfoContext).get(
                     TASKBAR_PINNING_IN_DESKTOP_MODE);
-            mIsInDesktopMode = wmProxy.isInDesktopMode();
+            mIsInDesktopMode = wmProxy.isInDesktopMode(DEFAULT_DISPLAY);
             mShowLockedTaskbarOnHome = wmProxy.showLockedTaskbarOnHome(displayInfoContext);
             mShowDesktopTaskbarForFreeformDisplay = wmProxy.showDesktopTaskbarForFreeformDisplay(
                     displayInfoContext);
diff --git a/src/com/android/launcher3/util/LockedUserState.kt b/src/com/android/launcher3/util/LockedUserState.kt
index a6a6ceb..742a327 100644
--- a/src/com/android/launcher3/util/LockedUserState.kt
+++ b/src/com/android/launcher3/util/LockedUserState.kt
@@ -44,7 +44,7 @@
 
     @VisibleForTesting
     val userUnlockedReceiver =
-        SimpleBroadcastReceiver(UI_HELPER_EXECUTOR) {
+        SimpleBroadcastReceiver(context, UI_HELPER_EXECUTOR) {
             if (Intent.ACTION_USER_UNLOCKED == it.action) {
                 isUserUnlocked = true
             }
@@ -61,7 +61,6 @@
         isUserUnlockedAtLauncherStartup = isUserUnlocked
         if (!isUserUnlocked) {
             userUnlockedReceiver.register(
-                context,
                 {
                     // If user is unlocked while registering broadcast receiver, we should update
                     // [isUserUnlocked], which will call [notifyUserUnlocked] in setter
@@ -72,7 +71,7 @@
                 Intent.ACTION_USER_UNLOCKED,
             )
         }
-        lifeCycle.addCloseable { userUnlockedReceiver.unregisterReceiverSafely(context) }
+        lifeCycle.addCloseable { userUnlockedReceiver.unregisterReceiverSafely() }
     }
 
     private fun checkIsUserUnlocked() =
@@ -80,7 +79,7 @@
 
     private fun notifyUserUnlocked() {
         mUserUnlockedActions.executeAllAndDestroy()
-        userUnlockedReceiver.unregisterReceiverSafely(context)
+        userUnlockedReceiver.unregisterReceiverSafely()
     }
 
     /**
diff --git a/src/com/android/launcher3/util/ScreenOnTracker.java b/src/com/android/launcher3/util/ScreenOnTracker.java
index 50be98b..8ffe9ea 100644
--- a/src/com/android/launcher3/util/ScreenOnTracker.java
+++ b/src/com/android/launcher3/util/ScreenOnTracker.java
@@ -46,34 +46,31 @@
     private final SimpleBroadcastReceiver mReceiver;
     private final CopyOnWriteArrayList<ScreenOnListener> mListeners = new CopyOnWriteArrayList<>();
 
-    private final Context mContext;
     private boolean mIsScreenOn;
 
     @Inject
     ScreenOnTracker(@ApplicationContext Context context, DaggerSingletonTracker tracker) {
         // Assume that the screen is on to begin with
-        mContext = context;
-        mReceiver = new SimpleBroadcastReceiver(UI_HELPER_EXECUTOR, this::onReceive);
+        mReceiver = new SimpleBroadcastReceiver(context, UI_HELPER_EXECUTOR, this::onReceive);
         init(tracker);
     }
 
     @VisibleForTesting
     ScreenOnTracker(@ApplicationContext Context context, SimpleBroadcastReceiver receiver,
             DaggerSingletonTracker tracker) {
-        mContext = context;
         mReceiver = receiver;
         init(tracker);
     }
 
     private void init(DaggerSingletonTracker tracker) {
         mIsScreenOn = true;
-        mReceiver.register(mContext, ACTION_SCREEN_ON, ACTION_SCREEN_OFF, ACTION_USER_PRESENT);
+        mReceiver.register(ACTION_SCREEN_ON, ACTION_SCREEN_OFF, ACTION_USER_PRESENT);
         tracker.addCloseable(this);
     }
 
     @Override
     public void close() {
-        mReceiver.unregisterReceiverSafely(mContext);
+        mReceiver.unregisterReceiverSafely();
     }
 
     @VisibleForTesting
diff --git a/src/com/android/launcher3/util/SimpleBroadcastReceiver.java b/src/com/android/launcher3/util/SimpleBroadcastReceiver.java
index 539a7cb..7a40abe 100644
--- a/src/com/android/launcher3/util/SimpleBroadcastReceiver.java
+++ b/src/com/android/launcher3/util/SimpleBroadcastReceiver.java
@@ -25,22 +25,29 @@
 import android.text.TextUtils;
 
 import androidx.annotation.AnyThread;
+import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
 import java.util.function.Consumer;
 
 public class SimpleBroadcastReceiver extends BroadcastReceiver {
+    public static final String TAG = "SimpleBroadcastReceiver";
+    // Keeps a strong reference to the context.
+    private final Context mContext;
 
     private final Consumer<Intent> mIntentConsumer;
 
     // Handler to register/unregister broadcast receiver
     private final Handler mHandler;
 
-    public SimpleBroadcastReceiver(LooperExecutor looperExecutor, Consumer<Intent> intentConsumer) {
-        this(looperExecutor.getHandler(), intentConsumer);
+    public SimpleBroadcastReceiver(@NonNull Context context, LooperExecutor looperExecutor,
+            Consumer<Intent> intentConsumer) {
+        this(context, looperExecutor.getHandler(), intentConsumer);
     }
 
-    public SimpleBroadcastReceiver(Handler handler, Consumer<Intent> intentConsumer) {
+    public SimpleBroadcastReceiver(@NonNull Context context, Handler handler,
+            Consumer<Intent> intentConsumer) {
+        mContext = context;
         mIntentConsumer = intentConsumer;
         mHandler = handler;
     }
@@ -50,18 +57,18 @@
         mIntentConsumer.accept(intent);
     }
 
-    /** Calls {@link #register(Context, Runnable, String...)} with null completionCallback. */
+    /** Calls {@link #register(Runnable, String...)} with null completionCallback. */
     @AnyThread
-    public void register(Context context, String... actions) {
-        register(context, null, actions);
+    public void register(String... actions) {
+        register(null, actions);
     }
 
     /**
-     * Calls {@link #register(Context, Runnable, int, String...)} with null completionCallback.
+     * Calls {@link #register(Runnable, int, String...)} with null completionCallback.
      */
     @AnyThread
-    public void register(Context context, int flags, String... actions) {
-        register(context, null, flags, actions);
+    public void register(int flags, String... actions) {
+        register(null, flags, actions);
     }
 
     /**
@@ -74,19 +81,18 @@
      *                           while registerReceiver() is executed on a binder call.
      */
     @AnyThread
-    public void register(
-            Context context, @Nullable Runnable completionCallback, String... actions) {
+    public void register(@Nullable Runnable completionCallback, String... actions) {
         if (Looper.myLooper() == mHandler.getLooper()) {
-            registerInternal(context, completionCallback, actions);
+            registerInternal(mContext, completionCallback, actions);
         } else {
-            mHandler.post(() -> registerInternal(context, completionCallback, actions));
+            mHandler.post(() -> registerInternal(mContext, completionCallback, actions));
         }
     }
 
     /** Register broadcast receiver and run completion callback if passed. */
     @AnyThread
     private void registerInternal(
-            Context context, @Nullable Runnable completionCallback, String... actions) {
+            @NonNull Context context, @Nullable Runnable completionCallback, String... actions) {
         context.registerReceiver(this, getFilter(actions));
         if (completionCallback != null) {
             completionCallback.run();
@@ -94,37 +100,37 @@
     }
 
     /**
-     * Same as {@link #register(Context, Runnable, String...)} above but with additional flags
-     * params.
+     * Same as {@link #register(Runnable, String...)} above but with additional flags
+     * params utilizine the original {@link Context}.
      */
     @AnyThread
-    public void register(
-            Context context, @Nullable Runnable completionCallback, int flags, String... actions) {
+    public void register(@Nullable Runnable completionCallback, int flags, String... actions) {
         if (Looper.myLooper() == mHandler.getLooper()) {
-            registerInternal(context, completionCallback, flags, actions);
+            registerInternal(mContext, completionCallback, flags, actions);
         } else {
-            mHandler.post(() -> registerInternal(context, completionCallback, flags, actions));
+            mHandler.post(() -> registerInternal(mContext, completionCallback, flags, actions));
         }
     }
 
     /** Register broadcast receiver and run completion callback if passed. */
     @AnyThread
     private void registerInternal(
-            Context context, @Nullable Runnable completionCallback, int flags, String... actions) {
+            @NonNull Context context, @Nullable Runnable completionCallback, int flags,
+            String... actions) {
         context.registerReceiver(this, getFilter(actions), flags);
         if (completionCallback != null) {
             completionCallback.run();
         }
     }
 
-    /** Same as {@link #register(Context, Runnable, String...)} above but with pkg name. */
+    /** Same as {@link #register(Runnable, String...)} above but with pkg name. */
     @AnyThread
-    public void registerPkgActions(Context context, @Nullable String pkg, String... actions) {
+    public void registerPkgActions(@Nullable String pkg, String... actions) {
         if (Looper.myLooper() == mHandler.getLooper()) {
-            context.registerReceiver(this, getPackageFilter(pkg, actions));
+            mContext.registerReceiver(this, getPackageFilter(pkg, actions));
         } else {
             mHandler.post(() -> {
-                context.registerReceiver(this, getPackageFilter(pkg, actions));
+                mContext.registerReceiver(this, getPackageFilter(pkg, actions));
             });
         }
     }
@@ -135,19 +141,19 @@
      * unregister happens on {@link #mHandler}'s looper.
      */
     @AnyThread
-    public void unregisterReceiverSafely(Context context) {
+    public void unregisterReceiverSafely() {
         if (Looper.myLooper() == mHandler.getLooper()) {
-            unregisterReceiverSafelyInternal(context);
+            unregisterReceiverSafelyInternal(mContext);
         } else {
             mHandler.post(() -> {
-                unregisterReceiverSafelyInternal(context);
+                unregisterReceiverSafelyInternal(mContext);
             });
         }
     }
 
     /** Unregister broadcast receiver ignoring any errors. */
     @AnyThread
-    private void unregisterReceiverSafelyInternal(Context context) {
+    private void unregisterReceiverSafelyInternal(@NonNull Context context) {
         try {
             context.unregisterReceiver(this);
         } catch (IllegalArgumentException e) {
diff --git a/src/com/android/launcher3/util/WallpaperOffsetInterpolator.java b/src/com/android/launcher3/util/WallpaperOffsetInterpolator.java
index f8cbe0d..26a04a5 100644
--- a/src/com/android/launcher3/util/WallpaperOffsetInterpolator.java
+++ b/src/com/android/launcher3/util/WallpaperOffsetInterpolator.java
@@ -31,8 +31,7 @@
     // Don't use all the wallpaper for parallax until you have at least this many pages
     private static final int MIN_PARALLAX_PAGE_SPAN = 4;
 
-    private final SimpleBroadcastReceiver mWallpaperChangeReceiver =
-            new SimpleBroadcastReceiver(UI_HELPER_EXECUTOR, i -> onWallpaperChanged());
+    private final SimpleBroadcastReceiver mWallpaperChangeReceiver;
     private final Workspace<?> mWorkspace;
     private final boolean mIsRtl;
     private final Handler mHandler;
@@ -46,6 +45,8 @@
 
     public WallpaperOffsetInterpolator(Workspace<?> workspace) {
         mWorkspace = workspace;
+        mWallpaperChangeReceiver = new SimpleBroadcastReceiver(
+                workspace.getContext(), UI_HELPER_EXECUTOR, i -> onWallpaperChanged());
         mIsRtl = Utilities.isRtl(workspace.getResources());
         mHandler = new OffsetHandler(workspace.getContext());
     }
@@ -198,11 +199,10 @@
     public void setWindowToken(IBinder token) {
         mWindowToken = token;
         if (mWindowToken == null && mRegistered) {
-            mWallpaperChangeReceiver.unregisterReceiverSafely(mWorkspace.getContext());
+            mWallpaperChangeReceiver.unregisterReceiverSafely();
             mRegistered = false;
         } else if (mWindowToken != null && !mRegistered) {
-            mWallpaperChangeReceiver.register(
-                    mWorkspace.getContext(), ACTION_WALLPAPER_CHANGED);
+            mWallpaperChangeReceiver.register(ACTION_WALLPAPER_CHANGED);
             onWallpaperChanged();
             mRegistered = true;
         }
diff --git a/src/com/android/launcher3/util/window/WindowManagerProxy.java b/src/com/android/launcher3/util/window/WindowManagerProxy.java
index f511ef2..647d170 100644
--- a/src/com/android/launcher3/util/window/WindowManagerProxy.java
+++ b/src/com/android/launcher3/util/window/WindowManagerProxy.java
@@ -109,7 +109,7 @@
     /**
      * Returns if we are in desktop mode or not.
      */
-    public boolean isInDesktopMode() {
+    public boolean isInDesktopMode(int displayId) {
         return false;
     }
 
@@ -503,11 +503,13 @@
     /** A listener for when the user enters/exits Desktop Mode.  */
     public interface DesktopVisibilityListener {
         /**
-         * Callback for when the user enters or exits Desktop Mode
+         * Called when the desktop mode state on the display whose ID is `displayId` changes.
          *
-         * @param visible whether Desktop Mode is now visible
+         * @param displayId The ID of the display for which this notification is triggering.
+         * @param isInDesktopModeAndNotInOverview True if a desktop is currently active on the given
+         *                                        display, and Overview is currently inactive.
          */
-        void onDesktopVisibilityChanged(boolean visible);
+        void onIsInDesktopModeChanged(int displayId, boolean isInDesktopModeAndNotInOverview);
     }
 
 }
diff --git a/src/com/android/launcher3/views/ClipIconView.java b/src/com/android/launcher3/views/ClipIconView.java
index f90a3e4..ddf18df 100644
--- a/src/com/android/launcher3/views/ClipIconView.java
+++ b/src/com/android/launcher3/views/ClipIconView.java
@@ -47,7 +47,8 @@
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.dragndrop.FolderAdaptiveIcon;
-import com.android.launcher3.graphics.IconShape;
+import com.android.launcher3.graphics.ShapeDelegate;
+import com.android.launcher3.graphics.ThemeManager;
 
 /**
  * A view used to draw both layers of an {@link AdaptiveIconDrawable}.
@@ -174,7 +175,7 @@
         if (mIsAdaptiveIcon) {
             if (!isOpening && progress >= shapeProgressStart) {
                 if (mRevealAnimator == null) {
-                    mRevealAnimator = IconShape.INSTANCE.get(getContext()).getShape()
+                    mRevealAnimator = ThemeManager.INSTANCE.get(getContext()).getIconShape()
                             .createRevealAnimator(this, mStartRevealRect,
                                     mOutline, mTaskCornerRadius, !isOpening);
                     mRevealAnimator.addListener(forEndCallback(() -> mRevealAnimator = null));
@@ -259,7 +260,7 @@
 
             if (!isFolderIcon) {
                 Utilities.scaleRectAboutCenter(mStartRevealRect,
-                        IconShape.INSTANCE.get(getContext()).getNormalizationScale());
+                        ShapeDelegate.getNormalizationScale());
             }
 
             if (dp.isLandscape) {
diff --git a/src/com/android/launcher3/views/DoubleShadowBubbleTextView.java b/src/com/android/launcher3/views/DoubleShadowBubbleTextView.java
index 392d9a7..05bc4d8 100644
--- a/src/com/android/launcher3/views/DoubleShadowBubbleTextView.java
+++ b/src/com/android/launcher3/views/DoubleShadowBubbleTextView.java
@@ -25,6 +25,7 @@
 import android.os.Build;
 import android.text.Spannable;
 import android.text.SpannableString;
+import android.text.TextUtils;
 import android.text.style.ImageSpan;
 import android.util.AttributeSet;
 import android.util.Log;
@@ -102,7 +103,7 @@
 
     @Override
     public void onDraw(Canvas canvas) {
-        if (shouldDrawAppContrastTile()) {
+        if (shouldDrawAppContrastTile() && !TextUtils.isEmpty(getText())) {
             drawAppContrastTile(canvas);
         }
         // If text is transparent or shadow alpha is 0, don't draw any shadow
diff --git a/tests/multivalentTests/src/com/android/launcher3/AbstractDeviceProfileTest.kt b/tests/multivalentTests/src/com/android/launcher3/AbstractDeviceProfileTest.kt
index 9c64ec9..9d42e1b 100644
--- a/tests/multivalentTests/src/com/android/launcher3/AbstractDeviceProfileTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/AbstractDeviceProfileTest.kt
@@ -28,11 +28,13 @@
 import android.view.Surface
 import androidx.test.core.app.ApplicationProvider.getApplicationContext
 import androidx.test.platform.app.InstrumentationRegistry
+import com.android.launcher3.LauncherPrefs.Companion.GRID_NAME
 import com.android.launcher3.dagger.LauncherAppComponent
 import com.android.launcher3.dagger.LauncherAppSingleton
 import com.android.launcher3.testing.shared.ResourceUtils
 import com.android.launcher3.util.AllModulesMinusWMProxy
 import com.android.launcher3.util.DisplayController
+import com.android.launcher3.util.FakePrefsModule
 import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext
 import com.android.launcher3.util.NavigationMode
 import com.android.launcher3.util.WindowBounds
@@ -70,7 +72,7 @@
     protected open val runningContext: Context = getApplicationContext()
     private val displayController: DisplayController = mock()
     private val windowManagerProxy: WindowManagerProxy = mock()
-    private val launcherPrefs: LauncherPrefs = mock()
+    private lateinit var launcherPrefs: LauncherPrefs
 
     @get:Rule val setFlagsRule = SetFlagsRule(SetFlagsRule.DefaultInitValueType.DEVICE_DEFAULT)
 
@@ -132,6 +134,7 @@
         isGestureMode: Boolean = true,
         isVerticalBar: Boolean = false,
         isFixedLandscape: Boolean = false,
+        gridName: String? = GRID_NAME.defaultValue,
     ) {
         val (naturalX, naturalY) = deviceSpec.naturalSize
         val windowsBounds = phoneWindowsBounds(deviceSpec, isGestureMode, naturalX, naturalY)
@@ -145,6 +148,7 @@
             isGestureMode,
             densityDpi = deviceSpec.densityDpi,
             isFixedLandscape = isFixedLandscape,
+            gridName = gridName,
         )
     }
 
@@ -152,6 +156,7 @@
         deviceSpec: DeviceSpec,
         isLandscape: Boolean = false,
         isGestureMode: Boolean = true,
+        gridName: String? = GRID_NAME.defaultValue,
     ) {
         val (naturalX, naturalY) = deviceSpec.naturalSize
         val windowsBounds = tabletWindowsBounds(deviceSpec, naturalX, naturalY)
@@ -164,6 +169,7 @@
             rotation = if (isLandscape) Surface.ROTATION_0 else Surface.ROTATION_90,
             isGestureMode,
             densityDpi = deviceSpec.densityDpi,
+            gridName = gridName,
         )
     }
 
@@ -173,6 +179,7 @@
         isLandscape: Boolean = false,
         isGestureMode: Boolean = true,
         isFolded: Boolean = false,
+        gridName: String? = GRID_NAME.defaultValue,
     ) {
         val (unfoldedNaturalX, unfoldedNaturalY) = deviceSpecUnfolded.naturalSize
         val unfoldedWindowsBounds =
@@ -199,6 +206,7 @@
                 rotation = if (isLandscape) Surface.ROTATION_90 else Surface.ROTATION_0,
                 isGestureMode = isGestureMode,
                 densityDpi = deviceSpecFolded.densityDpi,
+                gridName = gridName,
             )
         } else {
             initializeCommonVars(
@@ -207,6 +215,7 @@
                 rotation = if (isLandscape) Surface.ROTATION_0 else Surface.ROTATION_90,
                 isGestureMode = isGestureMode,
                 densityDpi = deviceSpecUnfolded.densityDpi,
+                gridName = gridName,
             )
         }
     }
@@ -282,6 +291,7 @@
         isGestureMode: Boolean = true,
         densityDpi: Int,
         isFixedLandscape: Boolean = false,
+        gridName: String? = GRID_NAME.defaultValue,
     ) {
         setFlagsRule.setFlags(true, Flags.FLAG_ENABLE_TWOLINE_TOGGLE)
         val windowsBounds = perDisplayBoundsCache[displayInfo]!!
@@ -311,18 +321,23 @@
         context.initDaggerComponent(
             DaggerAbsDPTestSandboxComponent.builder()
                 .bindWMProxy(windowManagerProxy)
-                .bindLauncherPrefs(launcherPrefs)
                 .bindDisplayController(displayController)
         )
+        launcherPrefs = context.appComponent.launcherPrefs
+        launcherPrefs.put(
+            LauncherPrefs.TASKBAR_PINNING.to(false),
+            LauncherPrefs.TASKBAR_PINNING_IN_DESKTOP_MODE.to(true),
+            LauncherPrefs.FIXED_LANDSCAPE_MODE.to(isFixedLandscape),
+            LauncherPrefs.HOTSEAT_COUNT.to(-1),
+            LauncherPrefs.DEVICE_TYPE.to(-1),
+            LauncherPrefs.WORKSPACE_SIZE.to(""),
+            LauncherPrefs.DB_FILE.to(""),
+            LauncherPrefs.ENABLE_TWOLINE_ALLAPPS_TOGGLE.to(true),
+        )
+        if (gridName != null) {
+            launcherPrefs.put(GRID_NAME, gridName)
+        }
 
-        whenever(launcherPrefs.get(LauncherPrefs.TASKBAR_PINNING)).thenReturn(false)
-        whenever(launcherPrefs.get(LauncherPrefs.TASKBAR_PINNING_IN_DESKTOP_MODE)).thenReturn(true)
-        whenever(launcherPrefs.get(LauncherPrefs.FIXED_LANDSCAPE_MODE)).thenReturn(isFixedLandscape)
-        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("")
-        whenever(launcherPrefs.get(LauncherPrefs.ENABLE_TWOLINE_ALLAPPS_TOGGLE)).thenReturn(true)
         val info = spy(DisplayController.Info(context, windowManagerProxy, perDisplayBoundsCache))
         whenever(displayController.info).thenReturn(info)
         whenever(info.isTransientTaskbar).thenReturn(isGestureMode)
@@ -365,15 +380,13 @@
 }
 
 @LauncherAppSingleton
-@Component(modules = [AllModulesMinusWMProxy::class])
+@Component(modules = [AllModulesMinusWMProxy::class, FakePrefsModule::class])
 interface AbsDPTestSandboxComponent : LauncherAppComponent {
 
     @Component.Builder
     interface Builder : LauncherAppComponent.Builder {
         @BindsInstance fun bindWMProxy(proxy: WindowManagerProxy): Builder
 
-        @BindsInstance fun bindLauncherPrefs(prefs: LauncherPrefs): Builder
-
         @BindsInstance fun bindDisplayController(displayController: DisplayController): Builder
 
         override fun build(): AbsDPTestSandboxComponent
diff --git a/tests/multivalentTests/src/com/android/launcher3/FakeInvariantDeviceProfileTest.kt b/tests/multivalentTests/src/com/android/launcher3/FakeInvariantDeviceProfileTest.kt
index bfbdb18..f855c51 100644
--- a/tests/multivalentTests/src/com/android/launcher3/FakeInvariantDeviceProfileTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/FakeInvariantDeviceProfileTest.kt
@@ -15,7 +15,6 @@
  */
 package com.android.launcher3
 
-import android.content.Context
 import android.graphics.PointF
 import android.graphics.Rect
 import android.platform.test.rule.AllowedDevices
@@ -23,10 +22,11 @@
 import android.platform.test.rule.IgnoreLimit
 import android.platform.test.rule.LimitDevicesRule
 import android.util.SparseArray
-import androidx.test.core.app.ApplicationProvider
 import com.android.launcher3.DeviceProfile.DEFAULT_DIMENSION_PROVIDER
 import com.android.launcher3.DeviceProfile.DEFAULT_PROVIDER
+import com.android.launcher3.testing.shared.ResourceUtils.INVALID_RESOURCE_HANDLE
 import com.android.launcher3.util.DisplayController.Info
+import com.android.launcher3.util.SandboxApplication
 import com.android.launcher3.util.WindowBounds
 import java.io.PrintWriter
 import java.io.StringWriter
@@ -46,7 +46,8 @@
 @IgnoreLimit(ignoreLimit = BuildConfig.IS_STUDIO_BUILD)
 abstract class FakeInvariantDeviceProfileTest {
 
-    protected lateinit var context: Context
+    @get:Rule val context = SandboxApplication()
+
     protected lateinit var inv: InvariantDeviceProfile
     protected val info = mock<Info>()
     protected lateinit var windowBounds: WindowBounds
@@ -59,7 +60,6 @@
 
     @Before
     open fun setUp() {
-        context = ApplicationProvider.getApplicationContext()
         // make sure to reset values
         useTwoPanels = false
         isGestureMode = true
@@ -70,6 +70,8 @@
             context,
             inv,
             info,
+            context.appComponent.wmProxy,
+            context.appComponent.themeManager,
             windowBounds,
             SparseArray(),
             /*isMultiWindowMode=*/ false,
@@ -107,7 +109,7 @@
         transposeLayoutWithOrientation = true
 
         inv =
-            InvariantDeviceProfile().apply {
+            context.appComponent.idp.apply {
                 numRows = 5
                 numColumns = 4
                 numSearchContainerColumns = 4
@@ -169,6 +171,14 @@
                 inlineQsb = BooleanArray(4) { false }
 
                 devicePaddingId = R.xml.paddings_handhelds
+
+                isFixedLandscape = false
+                workspaceSpecsId = INVALID_RESOURCE_HANDLE
+                allAppsSpecsId = INVALID_RESOURCE_HANDLE
+                folderSpecsId = INVALID_RESOURCE_HANDLE
+                hotseatSpecsId = INVALID_RESOURCE_HANDLE
+                workspaceCellSpecsId = INVALID_RESOURCE_HANDLE
+                allAppsCellSpecsId = INVALID_RESOURCE_HANDLE
             }
     }
 
@@ -189,7 +199,7 @@
         useTwoPanels = false
 
         inv =
-            InvariantDeviceProfile().apply {
+            context.appComponent.idp.apply {
                 numRows = 5
                 numColumns = 6
                 numSearchContainerColumns = 3
@@ -252,6 +262,14 @@
                 inlineQsb = booleanArrayOf(false, true, false, false)
 
                 devicePaddingId = R.xml.paddings_handhelds
+
+                isFixedLandscape = false
+                workspaceSpecsId = INVALID_RESOURCE_HANDLE
+                allAppsSpecsId = INVALID_RESOURCE_HANDLE
+                folderSpecsId = INVALID_RESOURCE_HANDLE
+                hotseatSpecsId = INVALID_RESOURCE_HANDLE
+                workspaceCellSpecsId = INVALID_RESOURCE_HANDLE
+                allAppsCellSpecsId = INVALID_RESOURCE_HANDLE
             }
     }
 
@@ -274,7 +292,7 @@
         useTwoPanels = true
 
         inv =
-            InvariantDeviceProfile().apply {
+            context.appComponent.idp.apply {
                 numRows = rows
                 numColumns = cols
                 numSearchContainerColumns = cols
@@ -332,6 +350,14 @@
                 inlineQsb = booleanArrayOf(false, false, false, false)
 
                 devicePaddingId = R.xml.paddings_handhelds
+
+                isFixedLandscape = false
+                workspaceSpecsId = INVALID_RESOURCE_HANDLE
+                allAppsSpecsId = INVALID_RESOURCE_HANDLE
+                folderSpecsId = INVALID_RESOURCE_HANDLE
+                hotseatSpecsId = INVALID_RESOURCE_HANDLE
+                workspaceCellSpecsId = INVALID_RESOURCE_HANDLE
+                allAppsCellSpecsId = INVALID_RESOURCE_HANDLE
             }
     }
 
diff --git a/tests/multivalentTests/src/com/android/launcher3/FakeLauncherPrefs.kt b/tests/multivalentTests/src/com/android/launcher3/FakeLauncherPrefs.kt
index 5e1e548..4d01d4d 100644
--- a/tests/multivalentTests/src/com/android/launcher3/FakeLauncherPrefs.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/FakeLauncherPrefs.kt
@@ -21,19 +21,24 @@
 import android.content.SharedPreferences
 import com.android.launcher3.dagger.ApplicationContext
 import com.android.launcher3.dagger.LauncherAppSingleton
-import java.io.File
+import com.android.launcher3.util.DaggerSingletonTracker
+import java.util.UUID
 import javax.inject.Inject
 
 /** Emulates Launcher preferences for a test environment. */
 @LauncherAppSingleton
-class FakeLauncherPrefs @Inject constructor(@ApplicationContext context: Context) :
+class FakeLauncherPrefs
+@Inject
+constructor(@ApplicationContext context: Context, lifeCycle: DaggerSingletonTracker) :
     LauncherPrefs(context) {
 
-    private val backingPrefs =
-        context.getSharedPreferences(
-            File.createTempFile("fake-pref", ".xml", context.filesDir),
-            MODE_PRIVATE,
-        )
+    private val prefName = "fake-pref-" + UUID.randomUUID().toString()
+
+    private val backingPrefs = context.getSharedPreferences(prefName, MODE_PRIVATE)
+
+    init {
+        lifeCycle.addCloseable { context.deleteSharedPreferences(prefName) }
+    }
 
     override fun getSharedPrefs(item: Item): SharedPreferences = backingPrefs
 }
diff --git a/tests/multivalentTests/src/com/android/launcher3/FakeLauncherPrefsTest.kt b/tests/multivalentTests/src/com/android/launcher3/FakeLauncherPrefsTest.kt
index c57c86f..0941c79 100644
--- a/tests/multivalentTests/src/com/android/launcher3/FakeLauncherPrefsTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/FakeLauncherPrefsTest.kt
@@ -18,10 +18,14 @@
 
 import androidx.test.core.app.ApplicationProvider.getApplicationContext
 import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
+import com.android.launcher3.util.DaggerSingletonTracker
 import com.android.launcher3.util.LauncherMultivalentJUnit
 import com.google.common.truth.Truth.assertThat
+import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.MockitoAnnotations
 
 private val TEST_CONSTANT_ITEM = LauncherPrefs.nonRestorableItem("TEST_BOOLEAN_ITEM", false)
 
@@ -36,7 +40,15 @@
 
 @RunWith(LauncherMultivalentJUnit::class)
 class FakeLauncherPrefsTest {
-    private val launcherPrefs = FakeLauncherPrefs(getApplicationContext())
+
+    @Mock lateinit var lifeCycle: DaggerSingletonTracker
+    private lateinit var launcherPrefs: FakeLauncherPrefs
+
+    @Before
+    fun setup() {
+        MockitoAnnotations.initMocks(this)
+        launcherPrefs = FakeLauncherPrefs(getApplicationContext(), lifeCycle)
+    }
 
     @Test
     fun testGet_constantItemNotInPrefs_returnsDefaultValue() {
diff --git a/tests/multivalentTests/src/com/android/launcher3/graphics/IconShapeTest.kt b/tests/multivalentTests/src/com/android/launcher3/graphics/ShapeDelegateTest.kt
similarity index 93%
rename from tests/multivalentTests/src/com/android/launcher3/graphics/IconShapeTest.kt
rename to tests/multivalentTests/src/com/android/launcher3/graphics/ShapeDelegateTest.kt
index 311676a..7e38f0e 100644
--- a/tests/multivalentTests/src/com/android/launcher3/graphics/IconShapeTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/graphics/ShapeDelegateTest.kt
@@ -28,13 +28,13 @@
 import androidx.core.graphics.PathParser
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
-import com.android.launcher3.graphics.IconShape.Circle
-import com.android.launcher3.graphics.IconShape.Companion.AREA_CALC_SIZE
-import com.android.launcher3.graphics.IconShape.Companion.AREA_DIFF_THRESHOLD
-import com.android.launcher3.graphics.IconShape.Companion.areaDiffCalculator
-import com.android.launcher3.graphics.IconShape.Companion.pickBestShape
-import com.android.launcher3.graphics.IconShape.GenericPathShape
-import com.android.launcher3.graphics.IconShape.RoundedSquare
+import com.android.launcher3.graphics.ShapeDelegate.Circle
+import com.android.launcher3.graphics.ShapeDelegate.Companion.AREA_CALC_SIZE
+import com.android.launcher3.graphics.ShapeDelegate.Companion.AREA_DIFF_THRESHOLD
+import com.android.launcher3.graphics.ShapeDelegate.Companion.areaDiffCalculator
+import com.android.launcher3.graphics.ShapeDelegate.Companion.pickBestShape
+import com.android.launcher3.graphics.ShapeDelegate.GenericPathShape
+import com.android.launcher3.graphics.ShapeDelegate.RoundedSquare
 import com.android.launcher3.icons.GraphicsUtils
 import com.android.launcher3.views.ClipPathView
 import com.google.common.truth.Truth.assertThat
@@ -43,7 +43,7 @@
 
 @SmallTest
 @RunWith(AndroidJUnit4::class)
-class IconShapeTest {
+class ShapeDelegateTest {
 
     @Test
     fun `areaDiffCalculator increases with outwards shape`() {
diff --git a/tests/multivalentTests/src/com/android/launcher3/shapes/ShapesProviderTest.kt b/tests/multivalentTests/src/com/android/launcher3/shapes/ShapesProviderTest.kt
index 2b8896e..508c9a4 100644
--- a/tests/multivalentTests/src/com/android/launcher3/shapes/ShapesProviderTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/shapes/ShapesProviderTest.kt
@@ -22,7 +22,7 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.launcher3.Flags.FLAG_ENABLE_LAUNCHER_ICON_SHAPES
-import com.android.launcher3.graphics.IconShape.GenericPathShape
+import com.android.launcher3.graphics.ShapeDelegate.GenericPathShape
 import com.android.systemui.shared.Flags.FLAG_NEW_CUSTOMIZATION_PICKER_UI
 import org.junit.Rule
 import org.junit.Test
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/DisplayControllerTest.kt b/tests/multivalentTests/src/com/android/launcher3/util/DisplayControllerTest.kt
index 588a668..aa1451b 100644
--- a/tests/multivalentTests/src/com/android/launcher3/util/DisplayControllerTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/util/DisplayControllerTest.kt
@@ -203,7 +203,7 @@
     fun testTaskbarPinningChangeInLockedTaskbarChange() {
         whenever(windowManagerProxy.showLockedTaskbarOnHome(any())).thenReturn(true)
         whenever(windowManagerProxy.isHomeVisible(any())).thenReturn(true)
-        whenever(windowManagerProxy.isInDesktopMode()).thenReturn(false)
+        whenever(windowManagerProxy.isInDesktopMode(any())).thenReturn(false)
         whenever(launcherPrefs.get(TASKBAR_PINNING)).thenReturn(false)
         DisplayController.enableTaskbarModePreferenceForTests(true)
 
@@ -219,7 +219,7 @@
     fun testLockedTaskbarChangeOnConfigurationChanged() {
         whenever(windowManagerProxy.showLockedTaskbarOnHome(any())).thenReturn(true)
         whenever(windowManagerProxy.isHomeVisible(any())).thenReturn(true)
-        whenever(windowManagerProxy.isInDesktopMode()).thenReturn(false)
+        whenever(windowManagerProxy.isInDesktopMode(any())).thenReturn(false)
         whenever(launcherPrefs.get(TASKBAR_PINNING)).thenReturn(false)
         DisplayController.enableTaskbarModePreferenceForTests(true)
         assertTrue(displayController.getInfo().isTransientTaskbar())
@@ -237,7 +237,7 @@
         whenever(windowManagerProxy.showDesktopTaskbarForFreeformDisplay(any())).thenReturn(true)
         whenever(launcherPrefs.get(TASKBAR_PINNING)).thenReturn(false)
         whenever(launcherPrefs.get(TASKBAR_PINNING_IN_DESKTOP_MODE)).thenReturn(false)
-        whenever(windowManagerProxy.isInDesktopMode()).thenReturn(true)
+        whenever(windowManagerProxy.isInDesktopMode(any())).thenReturn(true)
         whenever(windowManagerProxy.isHomeVisible(any())).thenReturn(false)
         DisplayController.enableTaskbarModePreferenceForTests(true)
 
@@ -256,7 +256,7 @@
         whenever(windowManagerProxy.showDesktopTaskbarForFreeformDisplay(any())).thenReturn(true)
         whenever(launcherPrefs.get(TASKBAR_PINNING)).thenReturn(false)
         whenever(launcherPrefs.get(TASKBAR_PINNING_IN_DESKTOP_MODE)).thenReturn(false)
-        whenever(windowManagerProxy.isInDesktopMode()).thenReturn(false)
+        whenever(windowManagerProxy.isInDesktopMode(any())).thenReturn(false)
         whenever(windowManagerProxy.isHomeVisible(any())).thenReturn(false)
         DisplayController.enableTaskbarModePreferenceForTests(true)
 
@@ -275,7 +275,7 @@
         whenever(windowManagerProxy.showDesktopTaskbarForFreeformDisplay(any())).thenReturn(true)
         whenever(launcherPrefs.get(TASKBAR_PINNING)).thenReturn(false)
         whenever(launcherPrefs.get(TASKBAR_PINNING_IN_DESKTOP_MODE)).thenReturn(false)
-        whenever(windowManagerProxy.isInDesktopMode()).thenReturn(false)
+        whenever(windowManagerProxy.isInDesktopMode(any())).thenReturn(false)
         whenever(windowManagerProxy.isHomeVisible(any())).thenReturn(true)
         DisplayController.enableTaskbarModePreferenceForTests(true)
 
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/ScreenOnTrackerTest.kt b/tests/multivalentTests/src/com/android/launcher3/util/ScreenOnTrackerTest.kt
index 45cc19c..9c3f223 100644
--- a/tests/multivalentTests/src/com/android/launcher3/util/ScreenOnTrackerTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/util/ScreenOnTrackerTest.kt
@@ -51,7 +51,7 @@
 
     @Test
     fun test_default_state() {
-        verify(receiver).register(context, ACTION_SCREEN_ON, ACTION_SCREEN_OFF, ACTION_USER_PRESENT)
+        verify(receiver).register(ACTION_SCREEN_ON, ACTION_SCREEN_OFF, ACTION_USER_PRESENT)
         assertThat(underTest.isScreenOn).isTrue()
     }
 
@@ -59,7 +59,7 @@
     fun close_unregister_receiver() {
         underTest.close()
 
-        verify(receiver).unregisterReceiverSafely(context)
+        verify(receiver).unregisterReceiverSafely()
     }
 
     @Test
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/SimpleBroadcastReceiverTest.kt b/tests/multivalentTests/src/com/android/launcher3/util/SimpleBroadcastReceiverTest.kt
index d3e27b6..17933f2 100644
--- a/tests/multivalentTests/src/com/android/launcher3/util/SimpleBroadcastReceiverTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/util/SimpleBroadcastReceiverTest.kt
@@ -52,7 +52,7 @@
     @Before
     fun setUp() {
         MockitoAnnotations.initMocks(this)
-        underTest = SimpleBroadcastReceiver(UI_HELPER_EXECUTOR, intentConsumer)
+        underTest = SimpleBroadcastReceiver(context, UI_HELPER_EXECUTOR, intentConsumer)
         if (Looper.getMainLooper() == null) {
             Looper.prepareMainLooper()
         }
@@ -60,7 +60,7 @@
 
     @Test
     fun async_register() {
-        underTest.register(context, "test_action_1", "test_action_2")
+        underTest.register("test_action_1", "test_action_2")
         awaitTasksCompleted()
 
         verify(context).registerReceiver(same(underTest), intentFilterCaptor.capture())
@@ -72,7 +72,7 @@
 
     @Test
     fun async_register_withCompletionRunnable() {
-        underTest.register(context, completionRunnable, "test_action_1", "test_action_2")
+        underTest.register(completionRunnable, "test_action_1", "test_action_2")
         awaitTasksCompleted()
 
         verify(context).registerReceiver(same(underTest), intentFilterCaptor.capture())
@@ -85,7 +85,7 @@
 
     @Test
     fun async_register_withCompletionRunnable_and_flag() {
-        underTest.register(context, completionRunnable, 1, "test_action_1", "test_action_2")
+        underTest.register(completionRunnable, 1, "test_action_1", "test_action_2")
         awaitTasksCompleted()
 
         verify(context).registerReceiver(same(underTest), intentFilterCaptor.capture(), eq(1))
@@ -98,7 +98,7 @@
 
     @Test
     fun async_register_with_package() {
-        underTest.registerPkgActions(context, "pkg", "test_action_1", "test_action_2")
+        underTest.registerPkgActions("pkg", "test_action_1", "test_action_2")
 
         awaitTasksCompleted()
         verify(context).registerReceiver(same(underTest), intentFilterCaptor.capture())
@@ -112,9 +112,10 @@
 
     @Test
     fun sync_register_withCompletionRunnable_and_flag() {
-        underTest = SimpleBroadcastReceiver(Handler(Looper.getMainLooper()), intentConsumer)
+        underTest =
+            SimpleBroadcastReceiver(context, Handler(Looper.getMainLooper()), intentConsumer)
 
-        underTest.register(context, completionRunnable, 1, "test_action_1", "test_action_2")
+        underTest.register(completionRunnable, 1, "test_action_1", "test_action_2")
         getInstrumentation().waitForIdleSync()
 
         verify(context).registerReceiver(same(underTest), intentFilterCaptor.capture(), eq(1))
@@ -127,7 +128,7 @@
 
     @Test
     fun async_unregister() {
-        underTest.unregisterReceiverSafely(context)
+        underTest.unregisterReceiverSafely()
 
         awaitTasksCompleted()
         verify(context).unregisterReceiver(same(underTest))
@@ -135,9 +136,10 @@
 
     @Test
     fun sync_unregister() {
-        underTest = SimpleBroadcastReceiver(Handler(Looper.getMainLooper()), intentConsumer)
+        underTest =
+            SimpleBroadcastReceiver(context, Handler(Looper.getMainLooper()), intentConsumer)
 
-        underTest.unregisterReceiverSafely(context)
+        underTest.unregisterReceiverSafely()
         getInstrumentation().waitForIdleSync()
 
         verify(context).unregisterReceiver(same(underTest))
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/GeneratedPreviewTest.kt b/tests/multivalentTests/src/com/android/launcher3/widget/GeneratedPreviewTest.kt
index b92582c..6af0950 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/GeneratedPreviewTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/GeneratedPreviewTest.kt
@@ -17,7 +17,6 @@
 import androidx.test.filters.SmallTest
 import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
 import com.android.launcher3.Flags.FLAG_ENABLE_GENERATED_PREVIEWS
-import com.android.launcher3.InvariantDeviceProfile
 import com.android.launcher3.R
 import com.android.launcher3.icons.IconCache
 import com.android.launcher3.model.WidgetItem
@@ -100,7 +99,7 @@
 
     private fun createWidgetItem() {
         Executors.MODEL_EXECUTOR.submit {
-                val idp = InvariantDeviceProfile()
+                val idp = context.appComponent.idp
                 widgetItem = WidgetItem(appWidgetProviderInfo, idp, iconCache, context)
             }
             .get()
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/LauncherAppWidgetProviderInfoTest.java b/tests/multivalentTests/src/com/android/launcher3/widget/LauncherAppWidgetProviderInfoTest.java
index 48cf3df..b3fd0f7 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/LauncherAppWidgetProviderInfoTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/LauncherAppWidgetProviderInfoTest.java
@@ -15,15 +15,12 @@
  */
 package com.android.launcher3.widget;
 
-import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
-
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doAnswer;
 
 import android.appwidget.AppWidgetHostView;
-import android.content.Context;
 import android.graphics.Point;
 import android.graphics.Rect;
 
@@ -32,16 +29,14 @@
 
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.InvariantDeviceProfile;
-import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.util.SandboxApplication;
 
-import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mockito;
 
-import java.util.ArrayList;
 import java.util.Collections;
-import java.util.List;
 
 @SmallTest
 @RunWith(AndroidJUnit4.class)
@@ -52,12 +47,7 @@
     private static final int NUM_OF_COLS = 4;
     private static final int NUM_OF_ROWS = 5;
 
-    private Context mContext;
-
-    @Before
-    public void setUp() {
-        mContext = getApplicationContext();
-    }
+    @Rule public SandboxApplication mContext = new SandboxApplication();
 
     @Test
     public void initSpans_minWidthSmallerThanCellWidth_shouldInitializeSpansToOne() {
@@ -256,8 +246,9 @@
     }
 
     private InvariantDeviceProfile createIDP() {
-        DeviceProfile dp = LauncherAppState.getIDP(mContext)
-                .getDeviceProfile(mContext).copy(mContext);
+        InvariantDeviceProfile idp = InvariantDeviceProfile.INSTANCE.get(mContext);
+
+        DeviceProfile dp = idp.getDeviceProfile(mContext).copy(mContext);
         DeviceProfile profile = Mockito.spy(dp);
         doAnswer(i -> {
             ((Point) i.getArgument(0)).set(CELL_SIZE, CELL_SIZE);
@@ -267,10 +258,7 @@
         profile.cellLayoutBorderSpacePx = new Point(SPACE_SIZE, SPACE_SIZE);
         profile.widgetPadding.setEmpty();
 
-        InvariantDeviceProfile idp = new InvariantDeviceProfile();
-        List<DeviceProfile> supportedProfiles = new ArrayList<>(idp.supportedProfiles);
-        supportedProfiles.add(profile);
-        idp.supportedProfiles = Collections.unmodifiableList(supportedProfiles);
+        idp.supportedProfiles = Collections.singletonList(profile);
         idp.numColumns = NUM_OF_COLS;
         idp.numRows = NUM_OF_ROWS;
         return idp;
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetRecommendationCategoryProviderTest.java b/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetRecommendationCategoryProviderTest.java
index ac67d2b..2fbeaf1 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetRecommendationCategoryProviderTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetRecommendationCategoryProviderTest.java
@@ -25,8 +25,6 @@
 import static android.content.pm.ApplicationInfo.CATEGORY_VIDEO;
 import static android.content.pm.ApplicationInfo.FLAG_INSTALLED;
 
-import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
-
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.ArgumentMatchers.any;
@@ -35,11 +33,9 @@
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.spy;
-import static org.mockito.Mockito.when;
 
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.ComponentName;
-import android.content.Context;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.LauncherApps;
 import android.os.Process;
@@ -53,12 +49,14 @@
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.util.Executors;
+import com.android.launcher3.util.SandboxApplication;
 import com.android.launcher3.util.WidgetUtils;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 
 import com.google.common.collect.ImmutableMap;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
@@ -92,22 +90,22 @@
 
     private final ApplicationInfo mTestAppInfo = ApplicationInfoBuilder.newBuilder().setPackageName(
             TEST_PACKAGE).setName(TEST_APP_NAME).build();
-    private Context mContext;
+
+    @Rule public SandboxApplication mContext = new SandboxApplication();
     @Mock
     private IconCache mIconCache;
 
     private WidgetItem mTestWidgetItem;
-    @Mock
+
     private LauncherApps mLauncherApps;
     private InvariantDeviceProfile mTestProfile;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
-        mContext = spy(getInstrumentation().getTargetContext());
-        doReturn(mLauncherApps).when(mContext).getSystemService(LauncherApps.class);
+        mLauncherApps = mContext.spyService(LauncherApps.class);
         mTestAppInfo.flags = FLAG_INSTALLED;
-        mTestProfile = new InvariantDeviceProfile();
+        mTestProfile = InvariantDeviceProfile.INSTANCE.get(mContext);
         mTestProfile.numRows = 5;
         mTestProfile.numColumns = 5;
         createTestWidgetItem();
@@ -128,10 +126,10 @@
                 testCategories.entrySet()) {
 
             mTestAppInfo.category = testCategory.getKey();
-            when(mLauncherApps.getApplicationInfo(/*packageName=*/ eq(TEST_PACKAGE),
+            doReturn(mTestAppInfo).when(mLauncherApps).getApplicationInfo(
+                    /*packageName=*/ eq(TEST_PACKAGE),
                     /*flags=*/ anyInt(),
-                    /*user=*/ eq(Process.myUserHandle())))
-                    .thenReturn(mTestAppInfo);
+                    /*user=*/ eq(Process.myUserHandle()));
 
             WidgetRecommendationCategory category = Executors.MODEL_EXECUTOR.submit(() ->
                     new WidgetRecommendationCategoryProvider().getWidgetRecommendationCategory(
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java b/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
index c9b6d4f..767ab63 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetsListHeaderViewHolderBinderTest.java
@@ -15,8 +15,6 @@
  */
 package com.android.launcher3.widget.picker;
 
-import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
-
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.ArgumentMatchers.any;
@@ -48,11 +46,13 @@
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.util.ActivityContextWrapper;
 import com.android.launcher3.util.PackageUserKey;
+import com.android.launcher3.util.SandboxApplication;
 import com.android.launcher3.util.WidgetUtils;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
@@ -67,6 +67,7 @@
     private static final String TEST_PACKAGE = "com.google.test";
     private static final String APP_NAME = "Test app";
 
+    @Rule public SandboxApplication app = new SandboxApplication();
     private Context mContext;
     private WidgetsListHeaderViewHolderBinder mViewHolderBinder;
     private InvariantDeviceProfile mTestProfile;
@@ -80,9 +81,9 @@
     public void setUp() {
         MockitoAnnotations.initMocks(this);
 
-        mContext = new ActivityContextWrapper(new ContextThemeWrapper(getApplicationContext(),
-                R.style.WidgetContainerTheme));
-        mTestProfile = new InvariantDeviceProfile();
+        mContext = new ActivityContextWrapper(new ContextThemeWrapper(
+                app, R.style.WidgetContainerTheme));
+        mTestProfile = InvariantDeviceProfile.INSTANCE.get(app);
         mTestProfile.numRows = 5;
         mTestProfile.numColumns = 5;
 
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java b/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
index 86bbcc1..e6f13a6 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinderTest.java
@@ -15,7 +15,6 @@
  */
 package com.android.launcher3.widget.picker;
 
-import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
 import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
 
 import static com.google.common.truth.Truth.assertThat;
@@ -50,6 +49,7 @@
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.util.ActivityContextWrapper;
+import com.android.launcher3.util.SandboxApplication;
 import com.android.launcher3.util.WidgetUtils;
 import com.android.launcher3.widget.DatabaseWidgetPreviewLoader;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
@@ -57,6 +57,7 @@
 import com.android.launcher3.widget.model.WidgetsListContentEntry;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
@@ -71,6 +72,7 @@
     private static final String TEST_PACKAGE = "com.google.test";
     private static final String APP_NAME = "Test app";
 
+    @Rule public SandboxApplication app = new SandboxApplication();
     private Context mContext;
     private WidgetsListTableViewHolderBinder mViewHolderBinder;
     private InvariantDeviceProfile mTestProfile;
@@ -85,8 +87,8 @@
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
-        mContext = new ActivityContextWrapper(getApplicationContext());
-        mTestProfile = new InvariantDeviceProfile();
+        mContext = new ActivityContextWrapper(app);
+        mTestProfile = InvariantDeviceProfile.INSTANCE.get(app);
         mTestProfile.numRows = 5;
         mTestProfile.numColumns = 5;
 
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java b/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java
index 6088c8e..bd34de6 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/picker/model/WidgetsListContentEntryTest.java
@@ -37,10 +37,12 @@
 import com.android.launcher3.icons.cache.CachedObject;
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.model.data.PackageItemInfo;
+import com.android.launcher3.util.SandboxApplication;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.model.WidgetsListContentEntry;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
@@ -64,6 +66,7 @@
     private final ComponentName mWidget3 = ComponentName.createRelative(PACKAGE_NAME, ".mWidget3");
     private final Map<ComponentName, String> mWidgetsToLabels = new HashMap();
 
+    @Rule public SandboxApplication app = new SandboxApplication();
     @Mock private IconCache mIconCache;
 
     private InvariantDeviceProfile mTestProfile;
@@ -76,7 +79,7 @@
         mWidgetsToLabels.put(mWidget2, "Dog");
         mWidgetsToLabels.put(mWidget3, "Bird");
 
-        mTestProfile = new InvariantDeviceProfile();
+        mTestProfile = InvariantDeviceProfile.INSTANCE.get(app);
         mTestProfile.numRows = 5;
         mTestProfile.numColumns = 5;
 
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java b/tests/multivalentTests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java
index 59f352b..0cdda3a 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java
@@ -46,6 +46,7 @@
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.search.SearchCallback;
+import com.android.launcher3.util.SandboxApplication;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.model.WidgetsListBaseEntry;
 import com.android.launcher3.widget.model.WidgetsListContentEntry;
@@ -53,6 +54,7 @@
 import com.android.launcher3.widget.picker.search.WidgetsSearchBar.WidgetsSearchDataProvider;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
@@ -66,6 +68,7 @@
 @RunWith(AndroidJUnit4.class)
 public class SimpleWidgetsSearchAlgorithmTest {
 
+    @Rule public SandboxApplication app = new SandboxApplication();
     @Mock private IconCache mIconCache;
 
     private InvariantDeviceProfile mTestProfile;
@@ -90,7 +93,7 @@
             CachedObject componentWithLabel = invocation.getArgument(0);
             return componentWithLabel.getComponent().getShortClassName();
         }).when(mIconCache).getTitleNoCache(any());
-        mTestProfile = new InvariantDeviceProfile();
+        mTestProfile = InvariantDeviceProfile.INSTANCE.get(app);
         mTestProfile.numRows = 5;
         mTestProfile.numColumns = 5;
         mContext = getApplicationContext();
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/picker/util/WidgetPreviewContainerSizesTest.kt b/tests/multivalentTests/src/com/android/launcher3/widget/picker/util/WidgetPreviewContainerSizesTest.kt
index 7a858e4..2452a88 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/picker/util/WidgetPreviewContainerSizesTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/picker/util/WidgetPreviewContainerSizesTest.kt
@@ -25,6 +25,7 @@
 import com.android.launcher3.DeviceProfile
 import com.android.launcher3.InvariantDeviceProfile
 import com.android.launcher3.LauncherAppState
+import com.android.launcher3.dagger.LauncherComponentProvider.appComponent
 import com.android.launcher3.icons.IconCache
 import com.android.launcher3.model.WidgetItem
 import com.android.launcher3.util.ActivityContextWrapper
@@ -53,7 +54,7 @@
         context = ActivityContextWrapper(ApplicationProvider.getApplicationContext())
         testInvariantProfile = LauncherAppState.getIDP(context)
         widgetItemInvariantProfile =
-            InvariantDeviceProfile().apply {
+            context.appComponent.idp.apply {
                 numRows = TEST_GRID_SIZE
                 numColumns = TEST_GRID_SIZE
             }
@@ -143,13 +144,13 @@
             widgetSize: Point,
             context: Context,
             invariantDeviceProfile: InvariantDeviceProfile,
-            iconCache: IconCache
+            iconCache: IconCache,
         ): WidgetItem {
             val providerInfo =
                 createAppWidgetProviderInfo(
                     ComponentName.createRelative(
                         TEST_PACKAGE,
-                        /*cls=*/ ".WidgetProvider_" + widgetSize.x + "x" + widgetSize.y
+                        /*cls=*/ ".WidgetProvider_" + widgetSize.x + "x" + widgetSize.y,
                     )
                 )
             val widgetInfo =
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/picker/util/WidgetsTableUtilsTest.java b/tests/multivalentTests/src/com/android/launcher3/widget/picker/util/WidgetsTableUtilsTest.java
index 2f5fcfe..a17e472 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/picker/util/WidgetsTableUtilsTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/picker/util/WidgetsTableUtilsTest.java
@@ -15,8 +15,6 @@
  */
 package com.android.launcher3.widget.picker.util;
 
-import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
-
 import static com.android.launcher3.util.WidgetUtils.createAppWidgetProviderInfo;
 
 import static com.google.common.truth.Truth.assertThat;
@@ -44,10 +42,12 @@
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.pm.ShortcutConfigActivityInfo;
 import com.android.launcher3.util.ActivityContextWrapper;
+import com.android.launcher3.util.SandboxApplication;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.util.WidgetsTableUtils;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
@@ -68,6 +68,8 @@
     private static final int NUM_OF_COLS = 5;
     private static final int NUM_OF_ROWS = 5;
 
+    @Rule public SandboxApplication app = new SandboxApplication();
+
     @Mock
     private IconCache mIconCache;
 
@@ -89,9 +91,8 @@
     public void setUp() {
         MockitoAnnotations.initMocks(this);
 
-        mContext = new ActivityContextWrapper(getApplicationContext());
-
-        mTestInvariantProfile = new InvariantDeviceProfile();
+        mContext = new ActivityContextWrapper(app);
+        mTestInvariantProfile = InvariantDeviceProfile.INSTANCE.get(app);
         mTestInvariantProfile.numColumns = NUM_OF_COLS;
         mTestInvariantProfile.numRows = NUM_OF_ROWS;
 
diff --git a/tests/src/com/android/launcher3/model/LoaderTaskTest.kt b/tests/src/com/android/launcher3/model/LoaderTaskTest.kt
index cdb45fc..8f64e84 100644
--- a/tests/src/com/android/launcher3/model/LoaderTaskTest.kt
+++ b/tests/src/com/android/launcher3/model/LoaderTaskTest.kt
@@ -15,7 +15,6 @@
 import androidx.test.filters.SmallTest
 import com.android.dx.mockito.inline.extended.ExtendedMockito
 import com.android.launcher3.Flags
-import com.android.launcher3.InvariantDeviceProfile
 import com.android.launcher3.LauncherAppState
 import com.android.launcher3.LauncherModel
 import com.android.launcher3.LauncherModel.LoaderTransaction
@@ -120,12 +119,11 @@
                 .mockStatic(FirstScreenBroadcastHelper::class.java)
                 .startMocking()
         val idp =
-            InvariantDeviceProfile().apply {
+            context.appComponent.idp.apply {
                 numRows = 5
                 numColumns = 6
                 numDatabaseHotseatIcons = 5
             }
-        context.putObject(InvariantDeviceProfile.INSTANCE, idp)
         context.putObject(LauncherAppState.INSTANCE, app)
 
         doReturn(TestViewHelpers.findWidgetProvider(false))
diff --git a/tests/src/com/android/launcher3/nonquickstep/DeviceProfileDumpTest.kt b/tests/src/com/android/launcher3/nonquickstep/DeviceProfileDumpTest.kt
index 2e2b6cd..f6e45a3 100644
--- a/tests/src/com/android/launcher3/nonquickstep/DeviceProfileDumpTest.kt
+++ b/tests/src/com/android/launcher3/nonquickstep/DeviceProfileDumpTest.kt
@@ -19,7 +19,6 @@
 import com.android.launcher3.AbstractDeviceProfileTest
 import com.android.launcher3.DeviceProfile
 import com.android.launcher3.Flags
-import com.android.launcher3.InvariantDeviceProfile
 import com.android.launcher3.util.rule.setFlags
 import org.junit.Before
 import org.junit.Test
@@ -41,12 +40,13 @@
             Flags.FLAG_ENABLE_SCALING_REVEAL_HOME_ANIMATION,
         )
         setFlagsRule.setFlags(false, Flags.FLAG_ONE_GRID_SPECS)
+        setFlagsRule.setFlags(false, Flags.FLAG_ENABLE_LAUNCHER_ICON_SHAPES)
     }
 
     @Test
     fun dumpPortraitGesture() {
         initializeDevice(instance.deviceName, isGestureMode = true, isLandscape = false)
-        val dp = getDeviceProfileForGrid(instance.gridName)
+        val dp = context.appComponent.idp.getDeviceProfile(context)
         dp.isTaskbarPresentInApps = instance.isTaskbarPresentInApps
 
         assertDump(dp, instance.filename("Portrait"))
@@ -55,7 +55,7 @@
     @Test
     fun dumpPortrait3Button() {
         initializeDevice(instance.deviceName, isGestureMode = false, isLandscape = false)
-        val dp = getDeviceProfileForGrid(instance.gridName)
+        val dp = context.appComponent.idp.getDeviceProfile(context)
         dp.isTaskbarPresentInApps = instance.isTaskbarPresentInApps
 
         assertDump(dp, instance.filename("Portrait3Button"))
@@ -64,7 +64,7 @@
     @Test
     fun dumpLandscapeGesture() {
         initializeDevice(instance.deviceName, isGestureMode = true, isLandscape = true)
-        val dp = getDeviceProfileForGrid(instance.gridName)
+        val dp = context.appComponent.idp.getDeviceProfile(context)
         dp.isTaskbarPresentInApps = instance.isTaskbarPresentInApps
 
         val testName =
@@ -79,7 +79,7 @@
     @Test
     fun dumpLandscape3Button() {
         initializeDevice(instance.deviceName, isGestureMode = false, isLandscape = true)
-        val dp = getDeviceProfileForGrid(instance.gridName)
+        val dp = context.appComponent.idp.getDeviceProfile(context)
         dp.isTaskbarPresentInApps = instance.isTaskbarPresentInApps
 
         val testName =
@@ -101,26 +101,25 @@
                     deviceSpecFolded = deviceSpecs["twopanel-phone"]!!,
                     isLandscape = isLandscape,
                     isGestureMode = isGestureMode,
+                    gridName = instance.gridName,
                 )
             "tablet" ->
                 initializeVarsForTablet(
                     deviceSpec = deviceSpec,
                     isLandscape = isLandscape,
                     isGestureMode = isGestureMode,
+                    gridName = instance.gridName,
                 )
             else ->
                 initializeVarsForPhone(
                     deviceSpec = deviceSpec,
                     isVerticalBar = isLandscape,
                     isGestureMode = isGestureMode,
+                    gridName = instance.gridName,
                 )
         }
     }
 
-    private fun getDeviceProfileForGrid(gridName: String): DeviceProfile {
-        return InvariantDeviceProfile(context, gridName).getDeviceProfile(context)
-    }
-
     private fun assertDump(dp: DeviceProfile, filename: String) {
         assertDump(dp, folderName, filename)
     }