Merge "Add feature flag for enable_overview_on_connected_displays" into main
diff --git a/quickstep/res/values/strings.xml b/quickstep/res/values/strings.xml
index 0474000..d2a7029 100644
--- a/quickstep/res/values/strings.xml
+++ b/quickstep/res/values/strings.xml
@@ -304,10 +304,6 @@
     <string name="taskbar_a11y_shown_with_bubbles_left_title">Taskbar &#38; bubbles left shown</string>
     <!-- Accessibility title for the Taskbar window appearing together with bubble bar on right. [CHAR_LIMIT=30] -->
     <string name="taskbar_a11y_shown_with_bubbles_right_title">Taskbar &#38; bubbles right shown</string>
-    <!-- Accessibility title for the Taskbar window being closed. [CHAR_LIMIT=30] -->
-    <string name="taskbar_a11y_hidden_title">Taskbar hidden</string>
-    <!-- Accessibility title for the Taskbar window being closed together with bubble bar. [CHAR_LIMIT=30] -->
-    <string name="taskbar_a11y_hidden_with_bubbles_title">Taskbar &#38; bubbles hidden</string>
     <!-- Accessibility title for the Taskbar window on phones. [CHAR_LIMIT=NONE] -->
     <string name="taskbar_phone_a11y_title">Navigation bar</string>
     <!-- Text in popup dialog for user to switch between always showing Taskbar or not. [CHAR LIMIT=30] -->
diff --git a/quickstep/src/com/android/launcher3/hybridhotseat/HotseatPredictionModel.java b/quickstep/src/com/android/launcher3/hybridhotseat/HotseatPredictionModel.java
index 56945ba..70868c5 100644
--- a/quickstep/src/com/android/launcher3/hybridhotseat/HotseatPredictionModel.java
+++ b/quickstep/src/com/android/launcher3/hybridhotseat/HotseatPredictionModel.java
@@ -17,7 +17,6 @@
 
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT_PREDICTION;
 import static com.android.launcher3.model.PredictionHelper.getAppTargetFromItemInfo;
-import static com.android.launcher3.model.PredictionHelper.isTrackedForHotseatPrediction;
 import static com.android.launcher3.model.PredictionHelper.wrapAppTargetWithItemLocation;
 
 import android.app.prediction.AppTarget;
@@ -27,9 +26,12 @@
 
 import com.android.launcher3.model.BgDataModel;
 import com.android.launcher3.model.BgDataModel.FixedContainerItems;
+import com.android.launcher3.model.PredictionHelper;
 import com.android.launcher3.model.data.ItemInfo;
 
 import java.util.ArrayList;
+import java.util.Objects;
+import java.util.stream.Collectors;
 
 /**
  * Model helper for app predictions in workspace
@@ -43,13 +45,18 @@
      */
     public static Bundle convertDataModelToAppTargetBundle(Context context, BgDataModel dataModel) {
         Bundle bundle = new Bundle();
-        ArrayList<AppTargetEvent> events = new ArrayList<>();
-        ArrayList<ItemInfo> workspaceItems = dataModel.getAllWorkspaceItems();
-        for (ItemInfo item : workspaceItems) {
-            AppTarget target = getAppTargetFromItemInfo(context, item);
-            if (target != null && !isTrackedForHotseatPrediction(item)) continue;
-            events.add(wrapAppTargetWithItemLocation(target, AppTargetEvent.ACTION_PIN, item));
-        }
+        ArrayList<AppTargetEvent> events = dataModel.itemsIdMap
+                .stream()
+                .filter(PredictionHelper::isTrackedForHotseatPrediction)
+                .map(item -> {
+                    AppTarget target = getAppTargetFromItemInfo(context, item);
+                    return target != null
+                            ? wrapAppTargetWithItemLocation(target, AppTargetEvent.ACTION_PIN, item)
+                            : null;
+                })
+                .filter(Objects::nonNull)
+                .collect(Collectors.toCollection(ArrayList::new));
+
         ArrayList<AppTarget> currentTargets = new ArrayList<>();
         FixedContainerItems hotseatItems = dataModel.extraItems.get(CONTAINER_HOTSEAT_PREDICTION);
         if (hotseatItems != null) {
diff --git a/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java b/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
index 25e1813..40e8fc2 100644
--- a/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
+++ b/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
@@ -462,7 +462,7 @@
     private Bundle getBundleForWidgetsOnWorkspace(Context context, BgDataModel dataModel) {
         Bundle bundle = new Bundle();
         ArrayList<AppTargetEvent> widgetEvents =
-                dataModel.getAllWorkspaceItems().stream()
+                dataModel.itemsIdMap.stream()
                         .filter(PredictionHelper::isTrackedForWidgetPrediction)
                         .map(item -> {
                             AppTarget target = getAppTargetFromItemInfo(context, item);
diff --git a/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java b/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
index 9d9054e..40e1c10 100644
--- a/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
+++ b/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
@@ -18,6 +18,7 @@
 import static com.android.launcher3.Flags.enableCategorizedWidgetSuggestions;
 import static com.android.launcher3.Flags.enableTieredWidgetsByDefaultInPicker;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION;
+import static com.android.launcher3.model.ModelUtils.WIDGET_FILTER;
 
 import static java.util.stream.Collectors.groupingBy;
 import static java.util.stream.Collectors.toMap;
@@ -69,9 +70,11 @@
             @NonNull AllAppsList apps) {
         Predicate<WidgetItem> predictedWidgetsFilter = enableTieredWidgetsByDefaultInPicker()
                 ? dataModel.widgetsModel.getPredictedWidgetsFilter() : null;
-        Set<ComponentKey> widgetsInWorkspace = dataModel.appWidgets.stream().map(
-                widget -> new ComponentKey(widget.providerName, widget.user)).collect(
-                Collectors.toSet());
+        Set<ComponentKey> widgetsInWorkspace = dataModel.itemsIdMap
+                .stream()
+                .filter(WIDGET_FILTER)
+                .map(item -> new ComponentKey(item.getTargetComponent(), item.user))
+                .collect(Collectors.toSet());
 
         // Widgets (excluding shortcuts & already added widgets) that belong to apps eligible for
         // being in predictions.
diff --git a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchController.java b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchController.java
index b928de0..3736e6d 100644
--- a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchController.java
@@ -33,6 +33,7 @@
 import com.android.quickstep.util.DesktopTask;
 import com.android.quickstep.util.GroupTask;
 import com.android.quickstep.util.LayoutUtils;
+import com.android.quickstep.util.SingleTask;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.recents.model.ThumbnailData;
 import com.android.systemui.shared.system.ActivityManagerWrapper;
@@ -275,7 +276,7 @@
 
         if (desktopTask != null) {
             mTasks = desktopTask.getTasks().stream()
-                    .map(GroupTask::new)
+                    .map(SingleTask::new)
                     .filter(task -> !shouldExcludeTask(task, taskIdsToExclude))
                     .collect(Collectors.toList());
             // All other tasks, apart from the grouped desktop task, are hidden
diff --git a/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java b/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
index 289e720..927254d 100644
--- a/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
@@ -41,8 +41,9 @@
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_A11Y_BUTTON_LONG_CLICKABLE;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_BACK_DISABLED;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_HOME_DISABLED;
-import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_SHOWING;
-import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_SWITCHER_BUTTON_SHOWING;
+import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_ALT_BACK;
+import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_VISIBLE;
+import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_SWITCHER_BUTTON_VISIBLE;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_NAV_BAR_HIDDEN;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_NOTIFICATION_PANEL_EXPANDED;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_OVERVIEW_DISABLED;
@@ -133,7 +134,8 @@
     private static final int FLAG_IME_SWITCHER_BUTTON_VISIBLE = 1 << 0;
     /** Whether the IME is visible. */
     private static final int FLAG_IME_VISIBLE = 1 << 1;
-    private static final int FLAG_ROTATION_BUTTON_VISIBLE = 1 << 2;
+    /** Whether the back button is adjusted for the IME. */
+    private static final int FLAG_IME_ALT_BACK = 1 << 2;
     private static final int FLAG_A11Y_VISIBLE = 1 << 3;
     private static final int FLAG_ONLY_BACK_FOR_BOUNCER_VISIBLE = 1 << 4;
     private static final int FLAG_KEYGUARD_VISIBLE = 1 << 5;
@@ -303,8 +305,7 @@
                     isThreeButtonNav ? mStartContextualContainer : mEndContextualContainer,
                     mControllers.navButtonController, R.id.ime_switcher);
             mPropertyHolders.add(new StatePropertyHolder(mImeSwitcherButton,
-                    flags -> ((flags & FLAG_IME_SWITCHER_BUTTON_VISIBLE) != 0)
-                            && ((flags & FLAG_ROTATION_BUTTON_VISIBLE) == 0)));
+                    flags -> ((flags & FLAG_IME_SWITCHER_BUTTON_VISIBLE) != 0)));
         }
 
         mPropertyHolders.add(new StatePropertyHolder(
@@ -333,7 +334,7 @@
 
         // Make sure to remove nav bar buttons translation when any of the following occur:
         // - Notification shade is expanded
-        // - IME is showing (add separate translation for IME)
+        // - IME is visible (add separate translation for IME)
         // - VoiceInteractionWindow (assistant) is showing
         // - Keyboard shortcuts helper is showing
         if (!mContext.isPhoneMode()) {
@@ -449,7 +450,7 @@
                     flags -> (flags & FLAG_IME_VISIBLE) == 0));
         }
         mPropertyHolders.add(new StatePropertyHolder(mBackButton,
-                flags -> (flags & FLAG_IME_VISIBLE) != 0,
+                flags -> (flags & FLAG_IME_ALT_BACK) != 0,
                 ROTATION_DRAWABLE_PERCENT, 1f, 0f));
         // Translate back button to be at end/start of other buttons for keyguard (only after SUW
         // since it is laid to align with SUW actions while in that state)
@@ -498,8 +499,7 @@
                 endContainer, navButtonController, R.id.accessibility_button,
                 R.layout.taskbar_contextual_button);
         mPropertyHolders.add(new StatePropertyHolder(mA11yButton,
-                flags -> (flags & FLAG_A11Y_VISIBLE) != 0
-                        && (flags & FLAG_ROTATION_BUTTON_VISIBLE) == 0));
+                flags -> (flags & FLAG_A11Y_VISIBLE) != 0));
 
         mSpace = new Space(mNavButtonsView.getContext());
         mSpace.setOnClickListener(view -> navButtonController.onButtonClick(BUTTON_SPACE, view));
@@ -510,8 +510,9 @@
     private void parseSystemUiFlags(@SystemUiStateFlags long sysUiStateFlags) {
         mSysuiStateFlags = sysUiStateFlags;
         boolean isImeSwitcherButtonVisible =
-                (sysUiStateFlags & SYSUI_STATE_IME_SWITCHER_BUTTON_SHOWING) != 0;
-        boolean isImeVisible = (sysUiStateFlags & SYSUI_STATE_IME_SHOWING) != 0;
+                (sysUiStateFlags & SYSUI_STATE_IME_SWITCHER_BUTTON_VISIBLE) != 0;
+        boolean isImeVisible = (sysUiStateFlags & SYSUI_STATE_IME_VISIBLE) != 0;
+        boolean useImeAltBack = (sysUiStateFlags & SYSUI_STATE_IME_ALT_BACK) != 0;
         boolean a11yVisible = (sysUiStateFlags & SYSUI_STATE_A11Y_BUTTON_CLICKABLE) != 0;
         boolean isHomeDisabled = (sysUiStateFlags & SYSUI_STATE_HOME_DISABLED) != 0;
         boolean isRecentsDisabled = (sysUiStateFlags & SYSUI_STATE_OVERVIEW_DISABLED) != 0;
@@ -527,6 +528,7 @@
 
         updateStateForFlag(FLAG_IME_SWITCHER_BUTTON_VISIBLE, isImeSwitcherButtonVisible);
         updateStateForFlag(FLAG_IME_VISIBLE, isImeVisible);
+        updateStateForFlag(FLAG_IME_ALT_BACK, useImeAltBack);
         updateStateForFlag(FLAG_A11Y_VISIBLE, a11yVisible);
         updateStateForFlag(FLAG_DISABLE_HOME, isHomeDisabled);
         updateStateForFlag(FLAG_DISABLE_RECENTS, isRecentsDisabled);
@@ -1231,7 +1233,7 @@
         appendFlag(str, flags, FLAG_IME_SWITCHER_BUTTON_VISIBLE,
                 "FLAG_IME_SWITCHER_BUTTON_VISIBLE");
         appendFlag(str, flags, FLAG_IME_VISIBLE, "FLAG_IME_VISIBLE");
-        appendFlag(str, flags, FLAG_ROTATION_BUTTON_VISIBLE, "FLAG_ROTATION_BUTTON_VISIBLE");
+        appendFlag(str, flags, FLAG_IME_ALT_BACK, "FLAG_IME_ALT_BACK");
         appendFlag(str, flags, FLAG_A11Y_VISIBLE, "FLAG_A11Y_VISIBLE");
         appendFlag(str, flags, FLAG_ONLY_BACK_FOR_BOUNCER_VISIBLE,
                 "FLAG_ONLY_BACK_FOR_BOUNCER_VISIBLE");
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
index b46b0dc..ebd4ee5 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
@@ -723,6 +723,11 @@
     public void onDisplayRemoved(int displayId) {
     }
 
+    /**
+     * Signal from SysUI indicating that system decorations should be removed from the display.
+     */
+    public void onDisplayRemoveSystemDecorations(int displayId) {}
+
     private void removeActivityCallbacksAndListeners() {
         if (mActivity != null) {
             mActivity.removeOnDeviceProfileChangeListener(mDebugActivityDeviceProfileChanged);
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarRecentAppsController.kt b/quickstep/src/com/android/launcher3/taskbar/TaskbarRecentAppsController.kt
index 5eb92d8..6047999 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarRecentAppsController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarRecentAppsController.kt
@@ -30,6 +30,7 @@
 import com.android.quickstep.RecentsModel
 import com.android.quickstep.util.DesktopTask
 import com.android.quickstep.util.GroupTask
+import com.android.quickstep.util.SingleTask
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
 import java.io.PrintWriter
 
@@ -218,7 +219,7 @@
         // Kind of hacky, we wrap each single task in the Desktop as a GroupTask.
         val orderFromId = orderedRunningTaskIds.withIndex().associate { (index, id) -> id to index }
         val sortedTasks = tasks.sortedWith(compareBy(nullsLast()) { orderFromId[it.key.id] })
-        return sortedTasks.map { GroupTask(it) }
+        return sortedTasks.map { SingleTask(it) }
     }
 
     private fun reloadRecentTasksIfNeeded() {
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
index 68114f1..1ca3dfb 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
@@ -34,8 +34,7 @@
 import static com.android.quickstep.util.SystemUiFlagUtils.isTaskbarHidden;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_BUBBLES_EXPANDED;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_DIALOG_SHOWING;
-import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_SHOWING;
-import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_SWITCHER_BUTTON_SHOWING;
+import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_VISIBLE;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_NOTIFICATION_PANEL_VISIBLE;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_SCREEN_PINNING;
 
@@ -260,9 +259,7 @@
     private @Nullable AnimatorSet mAnimator;
     private boolean mIsSystemGestureInProgress;
     /** Whether the IME is visible. */
-    private boolean mIsImeShowing;
-    /** Whether the IME Switcher button is visible. */
-    private boolean mIsImeSwitcherButtonShowing;
+    private boolean mIsImeVisible;
 
     private final Alarm mTimeoutAlarm = new Alarm();
     private boolean mEnableBlockingTimeoutDuringTests = false;
@@ -1122,7 +1119,7 @@
      */
     @VisibleForTesting
     long getTaskbarStashStartDelayForIme() {
-        if (mIsImeShowing) {
+        if (mIsImeVisible) {
             // Only delay when IME is exiting, not entering.
             return 0;
         }
@@ -1148,9 +1145,7 @@
         updateStateForFlag(FLAG_STASHED_DEVICE_LOCKED,
                 SystemUiFlagUtils.isLocked(systemUiStateFlags));
 
-        mIsImeShowing = hasAnyFlag(systemUiStateFlags, SYSUI_STATE_IME_SHOWING);
-        mIsImeSwitcherButtonShowing =
-                hasAnyFlag(systemUiStateFlags, SYSUI_STATE_IME_SWITCHER_BUTTON_SHOWING);
+        mIsImeVisible = hasAnyFlag(systemUiStateFlags, SYSUI_STATE_IME_VISIBLE);
         if (updateStateForFlag(FLAG_STASHED_IME, shouldStashForIme())) {
             animDuration = TASKBAR_STASH_DURATION_FOR_IME;
             startDelay = getTaskbarStashStartDelayForIme();
@@ -1166,7 +1161,7 @@
     }
 
     /**
-     * We stash when IME or IME switcher is showing.
+     * We stash when the IME is visible.
      *
      * <p>Do not stash if in small screen, with 3 button nav, and in landscape (or seascape).
      * <p>Do not stash if taskbar is transient.
@@ -1202,7 +1197,7 @@
             return false;
         }
 
-        return mIsImeShowing || mIsImeSwitcherButtonShowing;
+        return mIsImeVisible;
     }
 
     /**
@@ -1377,8 +1372,7 @@
         pw.println(prefix + "\tappliedState=" + getStateString(mStatePropertyHolder.mPrevFlags));
         pw.println(prefix + "\tmState=" + getStateString(mState));
         pw.println(prefix + "\tmIsSystemGestureInProgress=" + mIsSystemGestureInProgress);
-        pw.println(prefix + "\tmIsImeShowing=" + mIsImeShowing);
-        pw.println(prefix + "\tmIsImeSwitcherButtonShowing=" + mIsImeSwitcherButtonShowing);
+        pw.println(prefix + "\tmIsImeVisible=" + mIsImeVisible);
     }
 
     private static String getStateString(long flags) {
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
index e4e97e5..457ba3d 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
@@ -33,7 +33,6 @@
 import android.graphics.Canvas;
 import android.graphics.Rect;
 import android.graphics.drawable.Drawable;
-import android.os.Bundle;
 import android.util.ArraySet;
 import android.util.AttributeSet;
 import android.view.DisplayCutout;
@@ -41,7 +40,6 @@
 import android.view.LayoutInflater;
 import android.view.MotionEvent;
 import android.view.View;
-import android.view.accessibility.AccessibilityNodeInfo;
 import android.widget.FrameLayout;
 
 import androidx.annotation.LayoutRes;
@@ -311,16 +309,6 @@
         mShouldTryStartAlign = mActivityContext.shouldStartAlignTaskbar();
     }
 
-    @Override
-    public boolean performAccessibilityActionInternal(int action, Bundle arguments) {
-        if (action == AccessibilityNodeInfo.ACTION_ACCESSIBILITY_FOCUS) {
-            announceTaskbarShown();
-        } else if (action == AccessibilityNodeInfo.ACTION_CLEAR_ACCESSIBILITY_FOCUS) {
-            announceTaskbarHidden();
-        }
-        return super.performAccessibilityActionInternal(action, arguments);
-    }
-
     private void announceTaskbarShown() {
         BubbleBarLocation bubbleBarLocation = mControllerCallbacks.getBubbleBarLocationIfVisible();
         if (bubbleBarLocation == null) {
@@ -334,21 +322,12 @@
         }
     }
 
-    private void announceTaskbarHidden() {
-        BubbleBarLocation bubbleBarLocation = mControllerCallbacks.getBubbleBarLocationIfVisible();
-        if (bubbleBarLocation == null) {
-            announceForAccessibility(mContext.getString(R.string.taskbar_a11y_hidden_title));
-        } else {
-            announceForAccessibility(
-                    mContext.getString(R.string.taskbar_a11y_hidden_with_bubbles_title));
-        }
-    }
-
     protected void announceAccessibilityChanges() {
-        this.performAccessibilityAction(
-                isVisibleToUser() ? AccessibilityNodeInfo.ACTION_ACCESSIBILITY_FOCUS
-                        : AccessibilityNodeInfo.ACTION_CLEAR_ACCESSIBILITY_FOCUS, null);
-
+        // Only announce taskbar window shown. Window disappearing is generally not announce.
+        // This also aligns with talkback guidelines and unnecessary announcement to users.
+        if (isVisibleToUser()) {
+            announceTaskbarShown();
+        }
         ActivityContext.lookupContext(getContext()).getDragLayer()
                 .sendAccessibilityEvent(TYPE_WINDOW_CONTENT_CHANGED);
     }
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
index aff879e..4e029e3 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
@@ -19,8 +19,7 @@
 
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_BOUNCER_SHOWING;
-import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_SHOWING;
-import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_SWITCHER_BUTTON_SHOWING;
+import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_VISIBLE;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_NOTIFICATION_PANEL_EXPANDED;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_QUICK_SETTINGS_EXPANDED;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_STATUS_BAR_KEYGUARD_SHOWING;
@@ -92,10 +91,9 @@
     private static final long MASK_HIDE_BUBBLE_BAR = SYSUI_STATE_BOUNCER_SHOWING
             | SYSUI_STATE_STATUS_BAR_KEYGUARD_SHOWING
             | SYSUI_STATE_STATUS_BAR_KEYGUARD_SHOWING_OCCLUDED
-            | SYSUI_STATE_IME_SHOWING
+            | SYSUI_STATE_IME_VISIBLE
             | SYSUI_STATE_NOTIFICATION_PANEL_EXPANDED
-            | SYSUI_STATE_QUICK_SETTINGS_EXPANDED
-            | SYSUI_STATE_IME_SWITCHER_BUTTON_SHOWING;
+            | SYSUI_STATE_QUICK_SETTINGS_EXPANDED;
 
     private static final long MASK_HIDE_HANDLE_VIEW = SYSUI_STATE_BOUNCER_SHOWING
             | SYSUI_STATE_STATUS_BAR_KEYGUARD_SHOWING
@@ -239,7 +237,7 @@
 
         boolean sysuiLocked = (flags & MASK_SYSUI_LOCKED) != 0;
         mBubbleStashController.setSysuiLocked(sysuiLocked);
-        mIsImeVisible = (flags & SYSUI_STATE_IME_SHOWING) != 0;
+        mIsImeVisible = (flags & SYSUI_STATE_IME_VISIBLE) != 0;
         if (mIsImeVisible) {
             mBubbleBarViewController.onImeVisible();
         }
diff --git a/quickstep/src/com/android/quickstep/DesktopSystemShortcut.kt b/quickstep/src/com/android/quickstep/DesktopSystemShortcut.kt
index d60dab6..914855b 100644
--- a/quickstep/src/com/android/quickstep/DesktopSystemShortcut.kt
+++ b/quickstep/src/com/android/quickstep/DesktopSystemShortcut.kt
@@ -26,6 +26,7 @@
 import com.android.quickstep.views.RecentsViewContainer
 import com.android.quickstep.views.TaskContainer
 import com.android.systemui.shared.system.InteractionJankMonitorWrapper
+import com.android.wm.shell.shared.desktopmode.DesktopModeCompatPolicy
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
 import com.android.wm.shell.shared.desktopmode.DesktopModeTransitionSource
 
@@ -70,16 +71,31 @@
                     container: RecentsViewContainer,
                     taskContainer: TaskContainer,
                 ): List<DesktopSystemShortcut>? {
-                    return if (!DesktopModeStatus.canEnterDesktopMode(container.asContext())) null
-                    else if (!taskContainer.task.isDockable) null
-                    else
-                        listOf(
-                            DesktopSystemShortcut(
-                                container,
-                                taskContainer,
-                                abstractFloatingViewHelper,
+                    val context = container.asContext()
+                    val taskKey = taskContainer.task.key
+                    val desktopModeCompatPolicy = DesktopModeCompatPolicy(context)
+                    return when {
+                        !DesktopModeStatus.canEnterDesktopMode(context) -> null
+
+                        desktopModeCompatPolicy.isTopActivityExemptFromDesktopWindowing(
+                            taskKey.baseActivity?.packageName,
+                            taskKey.numActivities,
+                            taskKey.isTopActivityNoDisplay,
+                            taskKey.isActivityStackTransparent,
+                        ) -> null
+
+                        !taskContainer.task.isDockable -> null
+
+                        else -> {
+                            listOf(
+                                DesktopSystemShortcut(
+                                    container,
+                                    taskContainer,
+                                    abstractFloatingViewHelper,
+                                )
                             )
-                        )
+                        }
+                    }
                 }
 
                 override fun showForGroupedTask() = true
diff --git a/quickstep/src/com/android/quickstep/ExternalDisplaySystemShortcut.kt b/quickstep/src/com/android/quickstep/ExternalDisplaySystemShortcut.kt
index 46c4f36..f97cf9c 100644
--- a/quickstep/src/com/android/quickstep/ExternalDisplaySystemShortcut.kt
+++ b/quickstep/src/com/android/quickstep/ExternalDisplaySystemShortcut.kt
@@ -25,6 +25,7 @@
 import com.android.quickstep.views.RecentsViewContainer
 import com.android.quickstep.views.TaskContainer
 import com.android.window.flags.Flags
+import com.android.wm.shell.shared.desktopmode.DesktopModeCompatPolicy
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
 
 /** A menu item that allows the user to move the current app into external display. */
@@ -66,18 +67,31 @@
                     container: RecentsViewContainer,
                     taskContainer: TaskContainer,
                 ): List<ExternalDisplaySystemShortcut>? {
-                    return if (
-                        DesktopModeStatus.canEnterDesktopMode(container.asContext()) &&
-                            Flags.moveToExternalDisplayShortcut()
-                    )
-                        listOf(
-                            ExternalDisplaySystemShortcut(
-                                container,
-                                abstractFloatingViewHelper,
-                                taskContainer,
+                    val context = container.asContext()
+                    val taskKey = taskContainer.task.key
+                    val desktopModeCompatPolicy = DesktopModeCompatPolicy(context)
+                    return when {
+                        !DesktopModeStatus.canEnterDesktopMode(context) -> null
+
+                        !Flags.moveToExternalDisplayShortcut() -> null
+
+                        desktopModeCompatPolicy.isTopActivityExemptFromDesktopWindowing(
+                            taskKey.baseActivity?.packageName,
+                            taskKey.numActivities,
+                            taskKey.isTopActivityNoDisplay,
+                            taskKey.isActivityStackTransparent,
+                        ) -> null
+
+                        else -> {
+                            listOf(
+                                ExternalDisplaySystemShortcut(
+                                    container,
+                                    abstractFloatingViewHelper,
+                                    taskContainer,
+                                )
                             )
-                        )
-                    else null
+                        }
+                    }
                 }
 
                 override fun showForGroupedTask() = true
diff --git a/quickstep/src/com/android/quickstep/OverviewCommandHelper.kt b/quickstep/src/com/android/quickstep/OverviewCommandHelper.kt
index 66f307c..6ad9a2c 100644
--- a/quickstep/src/com/android/quickstep/OverviewCommandHelper.kt
+++ b/quickstep/src/com/android/quickstep/OverviewCommandHelper.kt
@@ -28,7 +28,9 @@
 import androidx.annotation.UiThread
 import androidx.annotation.VisibleForTesting
 import com.android.internal.jank.Cuj
+import com.android.launcher3.Flags.enableFallbackOverviewInWindow
 import com.android.launcher3.Flags.enableLargeDesktopWindowingTile
+import com.android.launcher3.Flags.enableLauncherOverviewInWindow
 import com.android.launcher3.Flags.enableOverviewCommandHelperTimeout
 import com.android.launcher3.PagedView
 import com.android.launcher3.logger.LauncherAtom
@@ -344,9 +346,12 @@
             return false
         }
 
-        val activity = containerInterface.getCreatedContainer()
-        if (activity != null) {
-            InteractionJankMonitorWrapper.begin(activity.rootView, Cuj.CUJ_LAUNCHER_QUICK_SWITCH)
+        val recentsInWindowFlagSet =
+            enableFallbackOverviewInWindow() || enableLauncherOverviewInWindow()
+        if (!recentsInWindowFlagSet) {
+            containerInterface.getCreatedContainer()?.rootView?.let { view ->
+                InteractionJankMonitorWrapper.begin(view, Cuj.CUJ_LAUNCHER_QUICK_SWITCH)
+            }
         }
 
         val gestureState =
@@ -373,6 +378,12 @@
                     transitionInfo: TransitionInfo,
                 ) {
                     Log.d(TAG, "recents animation started: $command")
+                    if (recentsInWindowFlagSet) {
+                        containerInterface.getCreatedContainer()?.rootView?.let { view ->
+                            InteractionJankMonitorWrapper.begin(view, Cuj.CUJ_LAUNCHER_QUICK_SWITCH)
+                        }
+                    }
+
                     updateRecentsViewFocus(command)
                     logShowOverviewFrom(command.type)
                     containerInterface.runOnInitBackgroundStateUI {
diff --git a/quickstep/src/com/android/quickstep/RecentTasksList.java b/quickstep/src/com/android/quickstep/RecentTasksList.java
index 379d71b4..01ced75 100644
--- a/quickstep/src/com/android/quickstep/RecentTasksList.java
+++ b/quickstep/src/com/android/quickstep/RecentTasksList.java
@@ -39,6 +39,8 @@
 import com.android.launcher3.util.SplitConfigurationOptions;
 import com.android.quickstep.util.DesktopTask;
 import com.android.quickstep.util.GroupTask;
+import com.android.quickstep.util.SingleTask;
+import com.android.quickstep.util.SplitTask;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.wm.shell.Flags;
 import com.android.wm.shell.recents.IRecentTasksListener;
@@ -50,6 +52,7 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.Objects;
 import java.util.function.Consumer;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
@@ -360,22 +363,19 @@
                 final Task.TaskKey task1Key = new Task.TaskKey(taskInfo1);
                 final Task task1 = Task.from(task1Key, taskInfo1,
                         tmpLockedUsers.get(task1Key.userId) /* isLocked */);
-                final Task task2;
-                final SplitConfigurationOptions.SplitBounds launcherSplitBounds;
 
                 if (rawTask.isBaseType(TYPE_SPLIT)) {
                     final TaskInfo taskInfo2 = rawTask.getBaseGroupedTask().getTaskInfo2();
                     final Task.TaskKey task2Key = new Task.TaskKey(taskInfo2);
-                    task2 = Task.from(task2Key, taskInfo2,
+                    final Task task2 = Task.from(task2Key, taskInfo2,
                             tmpLockedUsers.get(task2Key.userId) /* isLocked */);
-                    launcherSplitBounds =
+                    final SplitConfigurationOptions.SplitBounds launcherSplitBounds =
                             convertShellSplitBoundsToLauncher(
                                     rawTask.getBaseGroupedTask().getSplitBounds());
+                    allTasks.add(new SplitTask(task1, task2, launcherSplitBounds));
                 } else {
-                    task2 = null;
-                    launcherSplitBounds = null;
+                    allTasks.add(new SingleTask(task1));
                 }
-                allTasks.add(new GroupTask(task1, task2, launcherSplitBounds));
             } else {
                 TaskInfo taskInfo1 = rawTask.getTaskInfo1();
                 TaskInfo taskInfo2 = rawTask.getTaskInfo2();
@@ -407,9 +407,14 @@
                 if (taskInfo1.isVisible) {
                     numVisibleTasks++;
                 }
-                final SplitConfigurationOptions.SplitBounds launcherSplitBounds =
-                        convertShellSplitBoundsToLauncher(rawTask.getSplitBounds());
-                allTasks.add(new GroupTask(task1, task2, launcherSplitBounds));
+                if (task2 != null) {
+                    Objects.requireNonNull(rawTask.getSplitBounds());
+                    final SplitConfigurationOptions.SplitBounds launcherSplitBounds =
+                            convertShellSplitBoundsToLauncher(rawTask.getSplitBounds());
+                    allTasks.add(new SplitTask(task1, task2, launcherSplitBounds));
+                } else {
+                    allTasks.add(new SingleTask(task1));
+                }
             }
         }
 
diff --git a/quickstep/src/com/android/quickstep/RecentsAnimationDeviceState.java b/quickstep/src/com/android/quickstep/RecentsAnimationDeviceState.java
index 0c89a80..e1e962a 100644
--- a/quickstep/src/com/android/quickstep/RecentsAnimationDeviceState.java
+++ b/quickstep/src/com/android/quickstep/RecentsAnimationDeviceState.java
@@ -37,7 +37,7 @@
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_DISABLE_GESTURE_PIP_ANIMATING;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_DISABLE_GESTURE_SPLIT_INVOCATION;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_HOME_DISABLED;
-import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_SHOWING;
+import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_VISIBLE;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_MAGNIFICATION_OVERLAP;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_NAV_BAR_HIDDEN;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_NOTIFICATION_PANEL_EXPANDED;
@@ -586,7 +586,7 @@
     /** Returns whether IME is rendering nav buttons, and IME is currently showing. */
     public boolean isImeRenderingNavButtons() {
         return mCanImeRenderGesturalNavButtons && mMode == NO_BUTTON
-                && ((mSystemUiStateFlags & SYSUI_STATE_IME_SHOWING) != 0);
+                && ((mSystemUiStateFlags & SYSUI_STATE_IME_VISIBLE) != 0);
     }
 
     /**
diff --git a/quickstep/src/com/android/quickstep/RemoteTargetGluer.java b/quickstep/src/com/android/quickstep/RemoteTargetGluer.java
index 8edbacb..ef63b9b 100644
--- a/quickstep/src/com/android/quickstep/RemoteTargetGluer.java
+++ b/quickstep/src/com/android/quickstep/RemoteTargetGluer.java
@@ -142,7 +142,9 @@
         if (mSplitBounds == null) {
             SplitBounds shellSplitBounds = targets.extras.getParcelable(KEY_EXTRA_SPLIT_BOUNDS,
                     SplitBounds.class);
-            mSplitBounds = convertShellSplitBoundsToLauncher(shellSplitBounds);
+            if (shellSplitBounds != null) {
+                mSplitBounds = convertShellSplitBoundsToLauncher(shellSplitBounds);
+            }
         }
 
         boolean containsSplitTargets = mSplitBounds != null;
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index 15f320d..33ba780 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -317,6 +317,17 @@
 
         @BinderThread
         @Override
+        public void onDisplayRemoveSystemDecorations(int displayId) {
+            // TODO(b/391786915): Replace all
+            // `executeForTouchInteractionService(executeForTaskbarManager())` with just
+            // `executeForTaskbarManager` directly (since `tis` is unused).
+            MAIN_EXECUTOR.execute(() -> executeForTouchInteractionService(
+                    tis -> executeForTaskbarManager(taskbarManager -> taskbarManager
+                            .onDisplayRemoveSystemDecorations(displayId))));
+        }
+
+        @BinderThread
+        @Override
         public void updateWallpaperVisibility(int displayId, boolean visible) {
             MAIN_EXECUTOR.execute(() -> executeForTouchInteractionService(
                     tis -> executeForTaskbarManager(
diff --git a/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java b/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java
index d9209bf..fff7e9b 100644
--- a/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java
+++ b/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java
@@ -33,7 +33,6 @@
 import com.android.launcher3.Flags;
 import com.android.launcher3.anim.AnimatorPlaybackController;
 import com.android.launcher3.anim.PendingAnimation;
-import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.desktop.DesktopRecentsTransitionController;
 import com.android.launcher3.logging.StatsLogManager;
 import com.android.launcher3.statehandlers.DesktopVisibilityController;
@@ -47,6 +46,7 @@
 import com.android.quickstep.GestureState;
 import com.android.quickstep.fallback.window.RecentsDisplayModel;
 import com.android.quickstep.util.GroupTask;
+import com.android.quickstep.util.SingleTask;
 import com.android.quickstep.util.SplitSelectStateController;
 import com.android.quickstep.util.TaskViewSimulator;
 import com.android.quickstep.views.OverviewActionsView;
@@ -210,7 +210,7 @@
             if (!found) {
                 ArrayList<GroupTask> newList = new ArrayList<>(taskGroups.size() + 1);
                 newList.addAll(taskGroups);
-                newList.add(new GroupTask(mHomeTask, null, null));
+                newList.add(new SingleTask(mHomeTask));
                 taskGroups = newList;
             }
         }
diff --git a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
index 9bfe71f..b2a30ca 100644
--- a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
+++ b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
@@ -31,15 +31,15 @@
 import static com.android.launcher3.logger.LauncherAtomExtensions.ExtendedContainers.ContainerCase.DEVICE_SEARCH_RESULT_CONTAINER;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_WORKSPACE_SNAPSHOT;
 import static com.android.systemui.shared.system.SysUiStatsLog.LAUNCHER_UICHANGED__DISPLAY_ROTATION__ROTATION_0;
-import static com.android.systemui.shared.system.SysUiStatsLog.LAUNCHER_UICHANGED__DISPLAY_ROTATION__ROTATION_90;
 import static com.android.systemui.shared.system.SysUiStatsLog.LAUNCHER_UICHANGED__DISPLAY_ROTATION__ROTATION_180;
 import static com.android.systemui.shared.system.SysUiStatsLog.LAUNCHER_UICHANGED__DISPLAY_ROTATION__ROTATION_270;
+import static com.android.systemui.shared.system.SysUiStatsLog.LAUNCHER_UICHANGED__DISPLAY_ROTATION__ROTATION_90;
 import static com.android.systemui.shared.system.SysUiStatsLog.LAUNCHER_UICHANGED__DST_STATE__ALLAPPS;
 import static com.android.systemui.shared.system.SysUiStatsLog.LAUNCHER_UICHANGED__DST_STATE__BACKGROUND;
 import static com.android.systemui.shared.system.SysUiStatsLog.LAUNCHER_UICHANGED__DST_STATE__HOME;
 import static com.android.systemui.shared.system.SysUiStatsLog.LAUNCHER_UICHANGED__DST_STATE__OVERVIEW;
-import static com.android.systemui.shared.system.SysUiStatsLog.LAUNCHER_UICHANGED__RECENTS_ORIENTATION_HANDLER__PORTRAIT;
 import static com.android.systemui.shared.system.SysUiStatsLog.LAUNCHER_UICHANGED__RECENTS_ORIENTATION_HANDLER__LANDSCAPE;
+import static com.android.systemui.shared.system.SysUiStatsLog.LAUNCHER_UICHANGED__RECENTS_ORIENTATION_HANDLER__PORTRAIT;
 import static com.android.systemui.shared.system.SysUiStatsLog.LAUNCHER_UICHANGED__RECENTS_ORIENTATION_HANDLER__SEASCAPE;
 
 import android.content.Context;
@@ -69,6 +69,7 @@
 import com.android.launcher3.logger.LauncherAtomExtensions.ExtendedContainers;
 import com.android.launcher3.logging.InstanceId;
 import com.android.launcher3.logging.StatsLogManager;
+import com.android.launcher3.model.data.CollectionInfo;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.util.DisplayController;
 import com.android.launcher3.util.Executors;
@@ -386,7 +387,8 @@
                 // and then write to StatsLog.
                 app.getModel().enqueueModelUpdateTask((taskController, dataModel, apps) ->
                         write(event, applyOverwrites(mItemInfo.buildProto(
-                                dataModel.collections.get(mItemInfo.container), mContext))));
+                                (CollectionInfo) dataModel.itemsIdMap.get(mItemInfo.container),
+                                mContext))));
             })) {
                 // Write log on the model thread so that logs do not go out of order
                 // (for eg: drop comes after drag)
diff --git a/quickstep/src/com/android/quickstep/util/DesktopTask.kt b/quickstep/src/com/android/quickstep/util/DesktopTask.kt
index 1cee2d2..5463cf7 100644
--- a/quickstep/src/com/android/quickstep/util/DesktopTask.kt
+++ b/quickstep/src/com/android/quickstep/util/DesktopTask.kt
@@ -17,7 +17,6 @@
 
 import com.android.quickstep.views.TaskViewType
 import com.android.systemui.shared.recents.model.Task
-import java.util.Objects
 
 /**
  * A [Task] container that can contain N number of tasks that are part of the desktop in recent
@@ -39,9 +38,6 @@
     override fun equals(o: Any?): Boolean {
         if (this === o) return true
         if (o !is DesktopTask) return false
-        if (!super.equals(o)) return false
-        return tasks == o.tasks
+        return super.equals(o)
     }
-
-    override fun hashCode() = Objects.hash(super.hashCode(), tasks)
 }
diff --git a/quickstep/src/com/android/quickstep/util/GroupTask.kt b/quickstep/src/com/android/quickstep/util/GroupTask.kt
index 0bee5f6..d5bbcd3 100644
--- a/quickstep/src/com/android/quickstep/util/GroupTask.kt
+++ b/quickstep/src/com/android/quickstep/util/GroupTask.kt
@@ -22,10 +22,10 @@
 import java.util.Objects
 
 /**
- * A [Task] container that can contain one or two tasks, depending on if the two tasks are
- * represented as an app-pair in the recents task list.
+ * An abstract class for creating [Task] containers that can be [SingleTask]s, [SplitTask]s, or
+ * [DesktopTask]s in the recent tasks list.
  */
-open class GroupTask
+abstract class GroupTask
 @VisibleForTesting(otherwise = VisibleForTesting.PROTECTED)
 constructor(
     @Deprecated("Prefer using `getTasks()` instead") @JvmField val task1: Task,
@@ -33,13 +33,16 @@
     @JvmField val mSplitBounds: SplitConfigurationOptions.SplitBounds?,
     @JvmField val taskViewType: TaskViewType,
 ) {
-    constructor(task: Task) : this(task, null, null)
-
-    constructor(
-        t1: Task,
-        t2: Task?,
+    protected constructor(
+        task1: Task,
+        task2: Task?,
         splitBounds: SplitConfigurationOptions.SplitBounds?,
-    ) : this(t1, t2, splitBounds, if (t2 != null) TaskViewType.GROUPED else TaskViewType.SINGLE)
+    ) : this(
+        task1,
+        task2,
+        splitBounds,
+        if (task2 != null) TaskViewType.GROUPED else TaskViewType.SINGLE,
+    )
 
     open fun containsTask(taskId: Int) =
         task1.key.id == taskId || (task2 != null && task2.key.id == taskId)
@@ -59,18 +62,50 @@
         get() = listOfNotNull(task1, task2)
 
     /** Creates a copy of this instance */
-    open fun copy() = GroupTask(Task(task1), if (task2 != null) Task(task2) else null, mSplitBounds)
+    abstract fun copy(): GroupTask
 
     override fun toString() = "type=$taskViewType task1=$task1 task2=$task2"
 
     override fun equals(o: Any?): Boolean {
         if (this === o) return true
         if (o !is GroupTask) return false
-        return taskViewType == o.taskViewType &&
-            task1 == o.task1 &&
-            task2 == o.task2 &&
-            mSplitBounds == o.mSplitBounds
+        return taskViewType == o.taskViewType && tasks == o.tasks
     }
 
-    override fun hashCode() = Objects.hash(task1, task2, mSplitBounds, taskViewType)
+    override fun hashCode() = Objects.hash(tasks, taskViewType)
+}
+
+/** A [Task] container that must contain exactly one task in the recent tasks list. */
+class SingleTask(task: Task) :
+    GroupTask(task, task2 = null, mSplitBounds = null, TaskViewType.SINGLE) {
+    override fun copy() = SingleTask(task1)
+
+    override fun toString() = "type=$taskViewType task=$task1"
+
+    override fun equals(o: Any?): Boolean {
+        if (this === o) return true
+        if (o !is SingleTask) return false
+        return super.equals(o)
+    }
+}
+
+/**
+ * A [Task] container that must contain exactly two tasks and split bounds to represent an app-pair
+ * in the recent tasks list.
+ */
+class SplitTask(task1: Task, task2: Task, splitBounds: SplitConfigurationOptions.SplitBounds) :
+    GroupTask(task1, task2, splitBounds, TaskViewType.GROUPED) {
+
+    override fun copy() = SplitTask(task1, task2!!, mSplitBounds!!)
+
+    override fun toString() = "type=$taskViewType task1=$task1 task2=$task2"
+
+    override fun equals(o: Any?): Boolean {
+        if (this === o) return true
+        if (o !is SplitTask) return false
+        if (mSplitBounds!! != o.mSplitBounds!!) return false
+        return super.equals(o)
+    }
+
+    override fun hashCode() = Objects.hash(super.hashCode(), mSplitBounds)
 }
diff --git a/quickstep/src/com/android/quickstep/util/SplitScreenUtils.kt b/quickstep/src/com/android/quickstep/util/SplitScreenUtils.kt
index d982e81..4005c5a 100644
--- a/quickstep/src/com/android/quickstep/util/SplitScreenUtils.kt
+++ b/quickstep/src/com/android/quickstep/util/SplitScreenUtils.kt
@@ -33,38 +33,34 @@
         // TODO(b/254378592): Remove these methods when the two classes are reunited
         /** Converts the shell version of SplitBounds to the launcher version */
         @JvmStatic
-        fun convertShellSplitBoundsToLauncher(
-            shellSplitBounds: SplitBounds?
-        ): SplitConfigurationOptions.SplitBounds? {
-            return if (shellSplitBounds == null) {
-                null
-            } else {
-                SplitConfigurationOptions.SplitBounds(
-                    shellSplitBounds.leftTopBounds,
-                    shellSplitBounds.rightBottomBounds,
-                    shellSplitBounds.leftTopTaskId,
-                    shellSplitBounds.rightBottomTaskId,
-                    shellSplitBounds.snapPosition
-                )
-            }
-        }
+        fun convertShellSplitBoundsToLauncher(shellSplitBounds: SplitBounds) =
+            SplitConfigurationOptions.SplitBounds(
+                shellSplitBounds.leftTopBounds,
+                shellSplitBounds.rightBottomBounds,
+                shellSplitBounds.leftTopTaskId,
+                shellSplitBounds.rightBottomTaskId,
+                shellSplitBounds.snapPosition,
+            )
 
         /**
          * Given a TransitionInfo, generates the tree structure for those changes and extracts out
-         * the top most root and it's two immediate children.
-         * Changes can be provided in any order.
+         * the top most root and it's two immediate children. Changes can be provided in any order.
          *
-         * @return a [Pair] where first -> top most split root,
-         *         second -> [List] of 2, leftTop/bottomRight stage roots
+         * @return a [Pair] where first -> top most split root, second -> [List] of 2,
+         *   leftTop/bottomRight stage roots
          */
-        fun extractTopParentAndChildren(transitionInfo: TransitionInfo):
-                Pair<Change, List<Change>>? {
+        fun extractTopParentAndChildren(
+            transitionInfo: TransitionInfo
+        ): Pair<Change, List<Change>>? {
             val parentToChildren = mutableMapOf<Change, MutableList<Change>>()
             val hasParent = mutableSetOf<Change>()
             // filter out anything that isn't opening and the divider
-            val taskChanges: List<Change> = transitionInfo.changes
-                    .filter { change -> (change.mode == TRANSIT_OPEN ||
-                            change.mode == TRANSIT_TO_FRONT) && change.flags < FLAG_FIRST_CUSTOM}
+            val taskChanges: List<Change> =
+                transitionInfo.changes
+                    .filter { change ->
+                        (change.mode == TRANSIT_OPEN || change.mode == TRANSIT_TO_FRONT) &&
+                            change.flags < FLAG_FIRST_CUSTOM
+                    }
                     .toList()
 
             // 1. Build Parent-Child Relationships
@@ -73,8 +69,8 @@
                 //  startAnimation() and we can know the precise taskIds of launching tasks.
                 change.parent?.let { parent ->
                     parentToChildren
-                            .getOrPut(transitionInfo.getChange(parent)!!) { mutableListOf() }
-                            .add(change)
+                        .getOrPut(transitionInfo.getChange(parent)!!) { mutableListOf() }
+                        .add(change)
                     hasParent.add(change)
                 }
             }
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarRecentAppsControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarRecentAppsControllerTest.kt
index 6eccc36..c792783 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarRecentAppsControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarRecentAppsControllerTest.kt
@@ -21,6 +21,7 @@
 import android.content.Context
 import android.content.Intent
 import android.content.res.Resources
+import android.graphics.Rect
 import android.os.Process
 import android.os.UserHandle
 import android.platform.test.annotations.EnableFlags
@@ -36,12 +37,16 @@
 import com.android.launcher3.model.data.WorkspaceItemInfo
 import com.android.launcher3.taskbar.TaskbarRecentAppsController.TaskState
 import com.android.launcher3.util.LauncherMultivalentJUnit
+import com.android.launcher3.util.SplitConfigurationOptions
 import com.android.quickstep.RecentsModel
 import com.android.quickstep.RecentsModel.RecentTasksChangedListener
 import com.android.quickstep.TaskIconCache
 import com.android.quickstep.util.DesktopTask
 import com.android.quickstep.util.GroupTask
+import com.android.quickstep.util.SingleTask
+import com.android.quickstep.util.SplitTask
 import com.android.systemui.shared.recents.model.Task
+import com.android.wm.shell.shared.split.SplitScreenConstants
 import com.google.common.truth.Truth.assertThat
 import java.util.function.Consumer
 import org.junit.Before
@@ -913,15 +918,21 @@
         return packageNames.map { packageName ->
             if (packageName.startsWith("split")) {
                 val splitPackages = packageName.split("_")
-                GroupTask(
+                SplitTask(
                     createTask(100, splitPackages[0]),
                     createTask(101, splitPackages[1]),
-                    /* splitBounds = */ null,
+                    SplitConfigurationOptions.SplitBounds(
+                        /* leftTopBounds = */ Rect(),
+                        /* rightBottomBounds = */ Rect(),
+                        /* leftTopTaskId = */ -1,
+                        /* rightBottomTaskId = */ -1,
+                        /* snapPosition = */ SplitScreenConstants.SNAP_TO_2_50_50,
+                    ),
                 )
             } else {
                 // Use the number at the end of the test packageName as the id.
                 val id = 1000 + packageName[packageName.length - 1].code
-                GroupTask(createTask(id, packageName))
+                SingleTask(createTask(id, packageName))
             }
         }
     }
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarStashControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarStashControllerTest.kt
index 588c22c..021e1e4 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarStashControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarStashControllerTest.kt
@@ -53,7 +53,7 @@
 import com.android.launcher3.util.LauncherMultivalentJUnit
 import com.android.launcher3.util.LauncherMultivalentJUnit.EmulatedDevices
 import com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_BUBBLES_EXPANDED
-import com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_SHOWING
+import com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_VISIBLE
 import com.android.wm.shell.Flags.FLAG_ENABLE_BUBBLE_BAR
 import com.google.common.truth.Truth.assertThat
 import com.google.common.truth.TruthJUnit.assume
@@ -542,7 +542,7 @@
         assume().that(activityContext.isHardwareKeyboard).isFalse()
 
         getInstrumentation().runOnMainSync {
-            stashController.updateStateForSysuiFlags(SYSUI_STATE_IME_SHOWING, false)
+            stashController.updateStateForSysuiFlags(SYSUI_STATE_IME_VISIBLE, false)
             animatorTestRule.advanceTimeBy(TASKBAR_STASH_DURATION_FOR_IME)
         }
         assertThat(viewController.areIconsVisible()).isFalse()
@@ -555,7 +555,7 @@
         assume().that(activityContext.isHardwareKeyboard).isFalse()
 
         getInstrumentation().runOnMainSync {
-            stashController.updateStateForSysuiFlags(SYSUI_STATE_IME_SHOWING, true)
+            stashController.updateStateForSysuiFlags(SYSUI_STATE_IME_VISIBLE, true)
             animatorTestRule.advanceTimeBy(0)
         }
 
@@ -574,7 +574,7 @@
 
         // Start with IME shown.
         getInstrumentation().runOnMainSync {
-            stashController.updateStateForSysuiFlags(SYSUI_STATE_IME_SHOWING, true)
+            stashController.updateStateForSysuiFlags(SYSUI_STATE_IME_VISIBLE, true)
             animatorTestRule.advanceTimeBy(0)
         }
 
@@ -600,7 +600,7 @@
         assume().that(activityContext.isHardwareKeyboard).isFalse()
 
         getInstrumentation().runOnMainSync {
-            stashController.updateStateForSysuiFlags(SYSUI_STATE_IME_SHOWING, false)
+            stashController.updateStateForSysuiFlags(SYSUI_STATE_IME_VISIBLE, false)
             animatorTestRule.advanceTimeBy(TASKBAR_STASH_DURATION_FOR_IME)
         }
         assertThat(viewController.areIconsVisible()).isFalse()
@@ -613,7 +613,7 @@
         assume().that(activityContext.isHardwareKeyboard).isFalse()
 
         getInstrumentation().runOnMainSync {
-            stashController.updateStateForSysuiFlags(SYSUI_STATE_IME_SHOWING, false)
+            stashController.updateStateForSysuiFlags(SYSUI_STATE_IME_VISIBLE, false)
             animatorTestRule.advanceTimeBy(TASKBAR_STASH_DURATION_FOR_IME)
         }
 
@@ -633,7 +633,7 @@
         assume().that(activityContext.isHardwareKeyboard).isFalse()
 
         getInstrumentation().runOnMainSync {
-            stashController.updateStateForSysuiFlags(SYSUI_STATE_IME_SHOWING, true)
+            stashController.updateStateForSysuiFlags(SYSUI_STATE_IME_VISIBLE, true)
             animatorTestRule.advanceTimeBy(0)
         }
 
@@ -653,7 +653,7 @@
 
         getInstrumentation().runOnMainSync {
             stashController.updateStateForFlag(FLAG_IN_APP, true)
-            stashController.updateStateForSysuiFlags(SYSUI_STATE_IME_SHOWING, true)
+            stashController.updateStateForSysuiFlags(SYSUI_STATE_IME_VISIBLE, true)
         }
 
         assertThat(stashController.isStashed).isFalse()
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarViewTestUtil.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarViewTestUtil.kt
index e7f3523..df70b10 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarViewTestUtil.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarViewTestUtil.kt
@@ -28,6 +28,7 @@
 import com.android.launcher3.taskbar.TaskbarIconType.OVERFLOW
 import com.android.launcher3.taskbar.TaskbarIconType.RECENT
 import com.android.quickstep.util.GroupTask
+import com.android.quickstep.util.SingleTask
 import com.android.systemui.shared.recents.model.Task
 import com.android.systemui.shared.recents.model.Task.TaskKey
 import com.google.common.truth.FailureMetadata
@@ -56,7 +57,7 @@
     /** Creates a list of fake recent tasks. */
     fun createRecents(size: Int): List<GroupTask> {
         return List(size) {
-            GroupTask(
+            SingleTask(
                 Task().apply {
                     key =
                         TaskKey(
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/AbsSwipeUpHandlerTestCase.java b/quickstep/tests/multivalentTests/src/com/android/quickstep/AbsSwipeUpHandlerTestCase.java
index 542eb64..0c74610 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/AbsSwipeUpHandlerTestCase.java
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/AbsSwipeUpHandlerTestCase.java
@@ -18,6 +18,8 @@
 
 import static com.android.quickstep.AbsSwipeUpHandler.STATE_HANDLER_INVALIDATED;
 import static com.android.wm.shell.shared.ShellSharedConstants.KEY_EXTRA_SHELL_CAN_HAND_OFF_ANIMATION;
+import static com.android.wm.shell.shared.split.SplitBounds.KEY_EXTRA_SPLIT_BOUNDS;
+import static com.android.wm.shell.shared.split.SplitScreenConstants.SNAP_TO_2_50_50;
 
 import static junit.framework.TestCase.assertFalse;
 import static junit.framework.TestCase.assertTrue;
@@ -68,6 +70,7 @@
 import com.android.quickstep.views.RecentsViewContainer;
 import com.android.systemui.shared.Flags;
 import com.android.systemui.shared.system.InputConsumerController;
+import com.android.wm.shell.shared.split.SplitBounds;
 
 import com.google.android.msdl.data.model.MSDLToken;
 
@@ -140,6 +143,12 @@
     public void setUpAnimationTargets() {
         Bundle extras = new Bundle();
         extras.putBoolean(KEY_EXTRA_SHELL_CAN_HAND_OFF_ANIMATION, true);
+        extras.putParcelable(KEY_EXTRA_SPLIT_BOUNDS, new SplitBounds(
+                /* leftTopBounds = */ new Rect(),
+                /* rightBottomBounds = */ new Rect(),
+                /* leftTopTaskId = */ -1,
+                /* rightBottomTaskId = */ -1,
+                /* snapPosition = */ SNAP_TO_2_50_50));
         mRecentsAnimationTargets = new RecentsAnimationTargets(
                 new RemoteAnimationTarget[] {mRemoteAnimationTarget},
                 new RemoteAnimationTarget[] {mRemoteAnimationTarget},
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/RecentTasksListTest.java b/quickstep/tests/multivalentTests/src/com/android/quickstep/RecentTasksListTest.java
index 4a7c537..cad3b99 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/RecentTasksListTest.java
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/RecentTasksListTest.java
@@ -35,6 +35,7 @@
 import android.app.TaskInfo;
 import android.content.Context;
 import android.content.res.Resources;
+import android.graphics.Rect;
 
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
@@ -45,6 +46,8 @@
 import com.android.quickstep.views.TaskViewType;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.wm.shell.shared.GroupedTaskInfo;
+import com.android.wm.shell.shared.split.SplitBounds;
+import com.android.wm.shell.shared.split.SplitScreenConstants;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -99,7 +102,12 @@
     @Test
     public void loadTasksInBackground_onlyKeys_noValidTaskDescription() throws Exception  {
         GroupedTaskInfo recentTaskInfos = GroupedTaskInfo.forSplitTasks(
-                new RecentTaskInfo(), new RecentTaskInfo(), null);
+                new RecentTaskInfo(), new RecentTaskInfo(), new SplitBounds(
+                        /* leftTopBounds = */ new Rect(),
+                        /* rightBottomBounds = */ new Rect(),
+                        /* leftTopTaskId = */ -1,
+                        /* rightBottomTaskId = */ -1,
+                        /* snapPosition = */ SplitScreenConstants.SNAP_TO_2_50_50));
         when(mSystemUiProxy.getRecentTasks(anyInt(), anyInt()))
                 .thenReturn(new ArrayList<>(Collections.singletonList(recentTaskInfos)));
 
@@ -129,7 +137,13 @@
         task1.taskDescription = new ActivityManager.TaskDescription(taskDescription);
         RecentTaskInfo task2 = new RecentTaskInfo();
         task2.taskDescription = new ActivityManager.TaskDescription();
-        GroupedTaskInfo recentTaskInfos = GroupedTaskInfo.forSplitTasks(task1, task2, null);
+        GroupedTaskInfo recentTaskInfos = GroupedTaskInfo.forSplitTasks(task1, task2,
+                new SplitBounds(
+                        /* leftTopBounds = */ new Rect(),
+                        /* rightBottomBounds = */ new Rect(),
+                        /* leftTopTaskId = */ -1,
+                        /* rightBottomTaskId = */ -1,
+                        /* snapPosition = */ SplitScreenConstants.SNAP_TO_2_50_50));
         when(mSystemUiProxy.getRecentTasks(anyInt(), anyInt()))
                 .thenReturn(new ArrayList<>(Collections.singletonList(recentTaskInfos)));
 
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/RecentsModelTest.java b/quickstep/tests/multivalentTests/src/com/android/quickstep/RecentsModelTest.java
index a5c60ce..99a1c59 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/RecentsModelTest.java
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/RecentsModelTest.java
@@ -32,6 +32,7 @@
 import android.app.ActivityManager;
 import android.content.Context;
 import android.content.res.Resources;
+import android.graphics.Rect;
 import android.platform.test.flag.junit.SetFlagsRule;
 
 import androidx.test.annotation.UiThreadTest;
@@ -42,9 +43,12 @@
 import com.android.launcher3.R;
 import com.android.launcher3.graphics.ThemeManager;
 import com.android.launcher3.icons.IconProvider;
+import com.android.launcher3.util.SplitConfigurationOptions;
 import com.android.quickstep.util.GroupTask;
+import com.android.quickstep.util.SplitTask;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.system.TaskStackChangeListeners;
+import com.android.wm.shell.shared.split.SplitScreenConstants;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -173,7 +177,13 @@
         Task.TaskKey taskKey2 = new Task.TaskKey(taskInfo2);
         Task task2 = Task.from(taskKey2, taskInfo2, false);
 
-        allTasks.add(new GroupTask(task1, task2, null));
+        allTasks.add(
+                new SplitTask(task1, task2, new SplitConfigurationOptions.SplitBounds(
+                        /* leftTopBounds = */ new Rect(),
+                        /* rightBottomBounds = */ new Rect(),
+                        /* leftTopTaskId = */ -1,
+                        /* rightBottomTaskId = */ -1,
+                        /* snapPosition = */ SplitScreenConstants.SNAP_TO_2_50_50)));
         return allTasks;
     }
 }
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt
index b6cf5bd..823f808 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt
@@ -19,13 +19,17 @@
 import android.content.ComponentName
 import android.content.Intent
 import android.graphics.Bitmap
+import android.graphics.Rect
 import android.graphics.drawable.Drawable
 import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.launcher3.util.SplitConfigurationOptions
 import com.android.launcher3.util.TestDispatcherProvider
 import com.android.quickstep.util.DesktopTask
-import com.android.quickstep.util.GroupTask
+import com.android.quickstep.util.SingleTask
+import com.android.quickstep.util.SplitTask
 import com.android.systemui.shared.recents.model.Task
 import com.android.systemui.shared.recents.model.ThumbnailData
+import com.android.wm.shell.shared.split.SplitScreenConstants.SNAP_TO_2_50_50
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.first
@@ -48,8 +52,18 @@
     private val tasks = (0..5).map(::createTaskWithId)
     private val defaultTaskList =
         listOf(
-            GroupTask(tasks[0]),
-            GroupTask(tasks[1], tasks[2], null),
+            SingleTask(tasks[0]),
+            SplitTask(
+                tasks[1],
+                tasks[2],
+                SplitConfigurationOptions.SplitBounds(
+                    /* leftTopBounds = */ Rect(),
+                    /* rightBottomBounds = */ Rect(),
+                    /* leftTopTaskId = */ -1,
+                    /* rightBottomTaskId = */ -1,
+                    /* snapPosition = */ SNAP_TO_2_50_50,
+                ),
+            ),
             DesktopTask(tasks.subList(3, 6)),
         )
     private val recentsModel = FakeRecentTasksDataSource()
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/GroupTaskTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/GroupTaskTest.kt
index 108cfb5..fa043b9 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/GroupTaskTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/GroupTaskTest.kt
@@ -21,7 +21,6 @@
 import android.graphics.Rect
 import com.android.launcher3.util.LauncherMultivalentJUnit
 import com.android.launcher3.util.SplitConfigurationOptions
-import com.android.quickstep.views.TaskViewType
 import com.android.systemui.shared.recents.model.Task
 import com.android.wm.shell.shared.split.SplitScreenConstants
 import com.google.common.truth.Truth.assertThat
@@ -33,28 +32,28 @@
 
     @Test
     fun testGroupTask_sameInstance_isEqual() {
-        val task = GroupTask(createTask(1))
+        val task = SingleTask(createTask(1))
         assertThat(task).isEqualTo(task)
     }
 
     @Test
     fun testGroupTask_identicalConstructor_isEqual() {
-        val task1 = GroupTask(createTask(1))
-        val task2 = GroupTask(createTask(1))
+        val task1 = SingleTask(createTask(1))
+        val task2 = SingleTask(createTask(1))
         assertThat(task1).isEqualTo(task2)
     }
 
     @Test
     fun testGroupTask_copy_isEqual() {
-        val task1 = GroupTask(createTask(1))
+        val task1 = SingleTask(createTask(1))
         val task2 = task1.copy()
         assertThat(task1).isEqualTo(task2)
     }
 
     @Test
     fun testGroupTask_differentId_isNotEqual() {
-        val task1 = GroupTask(createTask(1))
-        val task2 = GroupTask(createTask(2))
+        val task1 = SingleTask(createTask(1))
+        val task2 = SingleTask(createTask(2))
         assertThat(task1).isNotEqualTo(task2)
     }
 
@@ -66,10 +65,10 @@
                 Rect(),
                 1,
                 2,
-                SplitScreenConstants.SNAP_TO_2_50_50
+                SplitScreenConstants.SNAP_TO_2_50_50,
             )
-        val task1 = GroupTask(createTask(1), createTask(2), splitBounds, TaskViewType.GROUPED)
-        val task2 = GroupTask(createTask(1), createTask(2), splitBounds, TaskViewType.GROUPED)
+        val task1 = SplitTask(createTask(1), createTask(2), splitBounds)
+        val task2 = SplitTask(createTask(1), createTask(2), splitBounds)
         assertThat(task1).isEqualTo(task2)
     }
 
@@ -81,7 +80,7 @@
                 Rect(),
                 1,
                 2,
-                SplitScreenConstants.SNAP_TO_2_50_50
+                SplitScreenConstants.SNAP_TO_2_50_50,
             )
         val splitBounds2 =
             SplitConfigurationOptions.SplitBounds(
@@ -89,17 +88,17 @@
                 Rect(),
                 1,
                 2,
-                SplitScreenConstants.SNAP_TO_2_33_66
+                SplitScreenConstants.SNAP_TO_2_33_66,
             )
-        val task1 = GroupTask(createTask(1), createTask(2), splitBounds1, TaskViewType.GROUPED)
-        val task2 = GroupTask(createTask(1), createTask(2), splitBounds2, TaskViewType.GROUPED)
+        val task1 = SplitTask(createTask(1), createTask(2), splitBounds1)
+        val task2 = SplitTask(createTask(1), createTask(2), splitBounds2)
         assertThat(task1).isNotEqualTo(task2)
     }
 
     @Test
     fun testGroupTask_differentType_isNotEqual() {
-        val task1 = GroupTask(createTask(1), null, null, TaskViewType.SINGLE)
-        val task2 = GroupTask(createTask(1), null, null, TaskViewType.DESKTOP)
+        val task1 = SingleTask(createTask(1))
+        val task2 = DesktopTask(listOf(createTask(1)))
         assertThat(task1).isNotEqualTo(task2)
     }
 
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/SplitSelectStateControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/SplitSelectStateControllerTest.kt
index 708273e..0491c07 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/SplitSelectStateControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/SplitSelectStateControllerTest.kt
@@ -660,10 +660,16 @@
         intent.component = task2ComponentName
         taskInfo.baseIntent = intent
         task2.key = Task.TaskKey(taskInfo)
-        return GroupTask(
+        return SplitTask(
             task1,
             task2,
-            SplitConfigurationOptions.SplitBounds(Rect(), Rect(), -1, -1, SNAP_TO_2_50_50),
+            SplitConfigurationOptions.SplitBounds(
+                /* leftTopBounds = */ Rect(),
+                /* rightBottomBounds = */ Rect(),
+                /* leftTopTaskId = */ -1,
+                /* rightBottomTaskId = */ -1,
+                /* snapPosition = */ SNAP_TO_2_50_50,
+            ),
         )
     }
 
@@ -692,10 +698,16 @@
         intent.component = task2ComponentName
         taskInfo.baseIntent = intent
         task2.key = Task.TaskKey(taskInfo)
-        return GroupTask(
+        return SplitTask(
             task1,
             task2,
-            SplitConfigurationOptions.SplitBounds(Rect(), Rect(), -1, -1, SNAP_TO_2_50_50),
+            SplitConfigurationOptions.SplitBounds(
+                /* leftTopBounds = */ Rect(),
+                /* rightBottomBounds = */ Rect(),
+                /* leftTopTaskId = */ -1,
+                /* rightBottomTaskId = */ -1,
+                /* snapPosition = */ SNAP_TO_2_50_50,
+            ),
         )
     }
 }
diff --git a/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt b/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt
index c152ee1..52bd2ea 100644
--- a/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt
+++ b/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt
@@ -17,28 +17,33 @@
 package com.android.quickstep
 
 import android.content.ComponentName
+import android.content.Context
 import android.content.Intent
+import android.platform.test.annotations.EnableFlags
+import android.view.Display.DEFAULT_DISPLAY
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession
 import com.android.dx.mockito.inline.extended.StaticMockitoSession
+import com.android.internal.R
 import com.android.launcher3.AbstractFloatingView
 import com.android.launcher3.AbstractFloatingViewHelper
 import com.android.launcher3.Flags.enableRefactorTaskThumbnail
 import com.android.launcher3.logging.StatsLogManager
 import com.android.launcher3.logging.StatsLogManager.LauncherEvent
 import com.android.launcher3.model.data.TaskViewItemInfo
-import com.android.launcher3.uioverrides.QuickstepLauncher
 import com.android.launcher3.util.SplitConfigurationOptions
 import com.android.launcher3.util.TransformingTouchDelegate
 import com.android.quickstep.TaskOverlayFactory.TaskOverlay
 import com.android.quickstep.task.thumbnail.TaskThumbnailView
 import com.android.quickstep.views.LauncherRecentsView
+import com.android.quickstep.views.RecentsViewContainer
 import com.android.quickstep.views.TaskContainer
 import com.android.quickstep.views.TaskThumbnailViewDeprecated
 import com.android.quickstep.views.TaskView
 import com.android.quickstep.views.TaskViewIcon
 import com.android.systemui.shared.recents.model.Task
 import com.android.systemui.shared.recents.model.Task.TaskKey
+import com.android.window.flags.Flags
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
 import com.android.wm.shell.shared.desktopmode.DesktopModeTransitionSource
 import com.google.common.truth.Truth.assertThat
@@ -58,7 +63,7 @@
 /** Test for [DesktopSystemShortcut] */
 class DesktopSystemShortcutTest {
 
-    private val launcher: QuickstepLauncher = mock()
+    private val launcher: RecentsViewContainer = mock()
     private val statsLogManager: StatsLogManager = mock()
     private val statsLogger: StatsLogManager.StatsLogger = mock()
     private val recentsView: LauncherRecentsView = mock()
@@ -67,6 +72,7 @@
     private val overlayFactory: TaskOverlayFactory = mock()
     private val factory: TaskShortcutFactory =
         DesktopSystemShortcut.createFactory(abstractFloatingViewHelper)
+    private val context: Context = InstrumentationRegistry.getInstrumentation().targetContext
 
     private lateinit var mockitoSession: StaticMockitoSession
 
@@ -79,6 +85,7 @@
                 .startMocking()
         whenever(DesktopModeStatus.canEnterDesktopMode(any())).thenReturn(true)
         whenever(overlayFactory.createOverlay(any())).thenReturn(mock<TaskOverlay<*>>())
+        whenever(launcher.asContext()).thenReturn(context)
     }
 
     @After
@@ -97,6 +104,53 @@
     }
 
     @Test
+    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODALS_POLICY)
+    fun createDesktopTaskShortcutFactory_transparentTask() {
+        val baseComponent = ComponentName("", /* class */ "")
+        val taskKey =
+            TaskKey(
+                /* id */ 1,
+                /* windowingMode */ 0,
+                Intent(),
+                baseComponent,
+                /* userId */ 0,
+                /* lastActiveTime */ 2000,
+                DEFAULT_DISPLAY,
+                baseComponent,
+                /* numActivities */ 1,
+                /* isTopActivityNoDisplay */ false,
+                /* isActivityStackTransparent */ true,
+            )
+        val taskContainer = createTaskContainer(Task(taskKey))
+        val shortcuts = factory.getShortcuts(launcher, taskContainer)
+        assertThat(shortcuts).isNull()
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODALS_POLICY)
+    fun createDesktopTaskShortcutFactory_systemUiTask() {
+        val sysUiPackageName: String = context.resources.getString(R.string.config_systemUi)
+        val baseComponent = ComponentName(sysUiPackageName, /* class */ "")
+        val taskKey =
+            TaskKey(
+                /* id */ 1,
+                /* windowingMode */ 0,
+                Intent(),
+                baseComponent,
+                /* userId */ 0,
+                /* lastActiveTime */ 2000,
+                DEFAULT_DISPLAY,
+                baseComponent,
+                /* numActivities */ 1,
+                /* isTopActivityNoDisplay */ false,
+                /* isActivityStackTransparent */ false,
+            )
+        val taskContainer = createTaskContainer(Task(taskKey))
+        val shortcuts = factory.getShortcuts(launcher, taskContainer)
+        assertThat(shortcuts).isNull()
+    }
+
+    @Test
     fun createDesktopTaskShortcutFactory_undockable() {
         val unDockableTask = createTask().apply { isDockable = false }
         val taskContainer = createTaskContainer(unDockableTask)
@@ -114,8 +168,7 @@
         whenever(launcher.statsLogManager).thenReturn(statsLogManager)
         whenever(statsLogManager.logger()).thenReturn(statsLogger)
         whenever(statsLogger.withItemInfo(any())).thenReturn(statsLogger)
-        whenever(taskView.context)
-            .thenReturn(InstrumentationRegistry.getInstrumentation().targetContext)
+        whenever(taskView.context).thenReturn(context)
         whenever(recentsView.moveTaskToDesktop(any(), any(), any())).thenAnswer {
             val successCallback = it.getArgument<Runnable>(2)
             successCallback.run()
@@ -145,7 +198,22 @@
     }
 
     private fun createTask() =
-        Task(TaskKey(1, 0, Intent(), ComponentName("", ""), 0, 2000)).apply { isDockable = true }
+        Task(
+                TaskKey(
+                    /* id */ 1,
+                    /* windowingMode */ 0,
+                    Intent(),
+                    ComponentName("", ""),
+                    /* userId */ 0,
+                    /* lastActiveTime */ 2000,
+                    DEFAULT_DISPLAY,
+                    ComponentName("", ""),
+                    /* numActivities */ 1,
+                    /* isTopActivityNoDisplay */ false,
+                    /* isActivityStackTransparent */ false,
+                )
+            )
+            .apply { isDockable = true }
 
     private fun createTaskContainer(task: Task) =
         TaskContainer(
diff --git a/quickstep/tests/src/com/android/quickstep/ExternalDisplaySystemShortcutTest.kt b/quickstep/tests/src/com/android/quickstep/ExternalDisplaySystemShortcutTest.kt
index 9c2c13c..4111dec 100644
--- a/quickstep/tests/src/com/android/quickstep/ExternalDisplaySystemShortcutTest.kt
+++ b/quickstep/tests/src/com/android/quickstep/ExternalDisplaySystemShortcutTest.kt
@@ -17,23 +17,27 @@
 package com.android.quickstep
 
 import android.content.ComponentName
+import android.content.Context
 import android.content.Intent
 import android.platform.test.annotations.EnableFlags
 import android.platform.test.flag.junit.SetFlagsRule
+import android.view.Display.DEFAULT_DISPLAY
+import androidx.test.platform.app.InstrumentationRegistry
 import com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession
 import com.android.dx.mockito.inline.extended.StaticMockitoSession
+import com.android.internal.R
 import com.android.launcher3.AbstractFloatingView
 import com.android.launcher3.AbstractFloatingViewHelper
 import com.android.launcher3.Flags.enableRefactorTaskThumbnail
 import com.android.launcher3.logging.StatsLogManager
 import com.android.launcher3.logging.StatsLogManager.LauncherEvent
 import com.android.launcher3.model.data.TaskViewItemInfo
-import com.android.launcher3.uioverrides.QuickstepLauncher
 import com.android.launcher3.util.SplitConfigurationOptions
 import com.android.launcher3.util.TransformingTouchDelegate
 import com.android.quickstep.TaskOverlayFactory.TaskOverlay
 import com.android.quickstep.task.thumbnail.TaskThumbnailView
 import com.android.quickstep.views.LauncherRecentsView
+import com.android.quickstep.views.RecentsViewContainer
 import com.android.quickstep.views.TaskContainer
 import com.android.quickstep.views.TaskThumbnailViewDeprecated
 import com.android.quickstep.views.TaskView
@@ -62,7 +66,7 @@
 
     @get:Rule val setFlagsRule = SetFlagsRule(SetFlagsRule.DefaultInitValueType.DEVICE_DEFAULT)
 
-    private val launcher: QuickstepLauncher = mock()
+    private val launcher: RecentsViewContainer = mock()
     private val statsLogManager: StatsLogManager = mock()
     private val statsLogger: StatsLogManager.StatsLogger = mock()
     private val recentsView: LauncherRecentsView = mock()
@@ -71,6 +75,7 @@
     private val overlayFactory: TaskOverlayFactory = mock()
     private val factory: TaskShortcutFactory =
         ExternalDisplaySystemShortcut.createFactory(abstractFloatingViewHelper)
+    private val context: Context = InstrumentationRegistry.getInstrumentation().targetContext
 
     private lateinit var mockitoSession: StaticMockitoSession
 
@@ -83,6 +88,7 @@
                 .startMocking()
         whenever(DesktopModeStatus.canEnterDesktopMode(any())).thenReturn(true)
         whenever(overlayFactory.createOverlay(any())).thenReturn(mock<TaskOverlay<*>>())
+        whenever(launcher.asContext()).thenReturn(context)
     }
 
     @After
@@ -102,6 +108,59 @@
     }
 
     @Test
+    @EnableFlags(
+        Flags.FLAG_MOVE_TO_EXTERNAL_DISPLAY_SHORTCUT,
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODALS_POLICY,
+    )
+    fun createExternalDisplayTaskShortcut_transparentTask() {
+        val baseComponent = ComponentName("", /* class */ "")
+        val taskKey =
+            TaskKey(
+                /* id */ 1,
+                /* windowingMode */ 0,
+                Intent(),
+                baseComponent,
+                /* userId */ 0,
+                /* lastActiveTime */ 2000,
+                DEFAULT_DISPLAY,
+                baseComponent,
+                /* numActivities */ 1,
+                /* isTopActivityNoDisplay */ false,
+                /* isActivityStackTransparent */ true,
+            )
+        val taskContainer = createTaskContainer(Task(taskKey))
+        val shortcuts = factory.getShortcuts(launcher, taskContainer)
+        assertThat(shortcuts).isNull()
+    }
+
+    @Test
+    @EnableFlags(
+        Flags.FLAG_MOVE_TO_EXTERNAL_DISPLAY_SHORTCUT,
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODALS_POLICY,
+    )
+    fun createExternalDisplayTaskShortcut_systemUiTask() {
+        val sysUiPackageName: String = context.resources.getString(R.string.config_systemUi)
+        val baseComponent = ComponentName(sysUiPackageName, /* class */ "")
+        val taskKey =
+            TaskKey(
+                /* id */ 1,
+                /* windowingMode */ 0,
+                Intent(),
+                baseComponent,
+                /* userId */ 0,
+                /* lastActiveTime */ 2000,
+                DEFAULT_DISPLAY,
+                baseComponent,
+                /* numActivities */ 1,
+                /* isTopActivityNoDisplay */ false,
+                /* isActivityStackTransparent */ false,
+            )
+        val taskContainer = createTaskContainer(Task(taskKey))
+        val shortcuts = factory.getShortcuts(launcher, taskContainer)
+        assertThat(shortcuts).isNull()
+    }
+
+    @Test
     @EnableFlags(Flags.FLAG_MOVE_TO_EXTERNAL_DISPLAY_SHORTCUT)
     fun externalDisplaySystemShortcutClicked() {
         val task = createTask()
@@ -134,7 +193,22 @@
         verify(statsLogger).log(LauncherEvent.LAUNCHER_SYSTEM_SHORTCUT_EXTERNAL_DISPLAY_TAP)
     }
 
-    private fun createTask() = Task(TaskKey(1, 0, Intent(), ComponentName("", ""), 0, 2000))
+    private fun createTask() =
+        Task(
+            TaskKey(
+                /* id */ 1,
+                /* windowingMode */ 0,
+                Intent(),
+                ComponentName("", ""),
+                /* userId */ 0,
+                /* lastActiveTime */ 2000,
+                DEFAULT_DISPLAY,
+                ComponentName("", ""),
+                /* numActivities */ 1,
+                /* isTopActivityNoDisplay */ false,
+                /* isActivityStackTransparent */ false,
+            )
+        )
 
     private fun createTaskContainer(task: Task) =
         TaskContainer(
diff --git a/src/com/android/launcher3/CellLayout.java b/src/com/android/launcher3/CellLayout.java
index c38444c..988d164 100644
--- a/src/com/android/launcher3/CellLayout.java
+++ b/src/com/android/launcher3/CellLayout.java
@@ -16,6 +16,7 @@
 
 package com.android.launcher3;
 
+import static com.android.launcher3.LauncherState.NORMAL;
 import static com.android.launcher3.dragndrop.DraggableView.DRAGGABLE_ICON;
 import static com.android.launcher3.icons.IconNormalizer.ICON_VISIBLE_AREA_FACTOR;
 import static com.android.launcher3.util.MultiTranslateDelegate.INDEX_REORDER_PREVIEW_OFFSET;
@@ -47,6 +48,7 @@
 import android.view.ViewDebug;
 import android.view.ViewGroup;
 import android.view.accessibility.AccessibilityEvent;
+import android.view.accessibility.AccessibilityNodeInfo;
 
 import androidx.annotation.IntDef;
 import androidx.annotation.Nullable;
@@ -69,6 +71,7 @@
 import com.android.launcher3.folder.PreviewBackground;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
+import com.android.launcher3.statemanager.StateManager;
 import com.android.launcher3.util.CellAndSpan;
 import com.android.launcher3.util.GridOccupancy;
 import com.android.launcher3.util.MSDLPlayerWrapper;
@@ -785,6 +788,22 @@
             }
             mShortcutsAndWidgets.addView(child, index, lp);
 
+            // Whenever an app is added, if Accessibility service is enabled, focus on that app.
+            if (mActivity instanceof Launcher) {
+                Launcher.cast(mActivity).getStateManager().addStateListener(
+                        new StateManager.StateListener<LauncherState>() {
+                            @Override
+                            public void onStateTransitionComplete(LauncherState finalState) {
+                                if (finalState == NORMAL) {
+                                    child.performAccessibilityAction(
+                                            AccessibilityNodeInfo.ACTION_ACCESSIBILITY_FOCUS, null);
+                                    Launcher.cast(mActivity).getStateManager()
+                                            .removeStateListener(this);
+                                }
+                            }
+                        });
+            }
+
             if (markCells) markCellsAsOccupiedForView(child);
 
             return true;
diff --git a/src/com/android/launcher3/Utilities.java b/src/com/android/launcher3/Utilities.java
index 7563493..d93c07f 100644
--- a/src/com/android/launcher3/Utilities.java
+++ b/src/com/android/launcher3/Utilities.java
@@ -659,7 +659,8 @@
                 // Only fetch badge if the icon is on workspace
                 if (info.id != ItemInfo.NO_ID && badge == null) {
                     badge = appState.getIconCache().getShortcutInfoBadge(si)
-                            .newIcon(context, FLAG_THEMED);
+                            .newIcon(context, ThemeManager.INSTANCE.get(context)
+                                    .isMonoThemeEnabled() ? FLAG_THEMED : 0);
                 }
             }
         } else if (info.itemType == LauncherSettings.Favorites.ITEM_TYPE_FOLDER) {
diff --git a/src/com/android/launcher3/folder/FolderNameProvider.java b/src/com/android/launcher3/folder/FolderNameProvider.java
index be5f8f7..8a1f96d 100644
--- a/src/com/android/launcher3/folder/FolderNameProvider.java
+++ b/src/com/android/launcher3/folder/FolderNameProvider.java
@@ -15,6 +15,8 @@
  */
 package com.android.launcher3.folder;
 
+import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_FOLDER;
+
 import android.annotation.SuppressLint;
 import android.app.admin.DevicePolicyManager;
 import android.content.ComponentName;
@@ -37,6 +39,7 @@
 import com.android.launcher3.model.StringCache;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.CollectionInfo;
+import com.android.launcher3.model.data.FolderInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.util.IntSparseArrayMap;
 import com.android.launcher3.util.Preconditions;
@@ -197,9 +200,18 @@
         @Override
         public void execute(@NonNull ModelTaskController taskController,
                 @NonNull BgDataModel dataModel, @NonNull AllAppsList apps) {
-            mCollectionInfos = dataModel.collections.clone();
+            mCollectionInfos = getCollectionForSuggestions(dataModel);
             mAppInfos = Arrays.asList(apps.copyData());
         }
     }
 
+    public static IntSparseArrayMap<CollectionInfo> getCollectionForSuggestions(
+            BgDataModel dataModel) {
+        IntSparseArrayMap<CollectionInfo> result = new IntSparseArrayMap<>();
+        dataModel.itemsIdMap.stream()
+                .filter(item -> item.itemType == ITEM_TYPE_FOLDER)
+                .forEach(item -> result.put(item.id, (FolderInfo) item));
+        return result;
+    }
+
 }
diff --git a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
index f0e4fc4..a0b73ae 100644
--- a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
+++ b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
@@ -24,10 +24,10 @@
 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.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;
-import static com.android.launcher3.model.ModelUtils.filterCurrentWorkspaceItems;
-import static com.android.launcher3.model.ModelUtils.getMissingHotseatRanks;
+import static com.android.launcher3.model.ModelUtils.currentScreenContentFilter;
 
 import android.app.Fragment;
 import android.app.WallpaperColors;
@@ -105,7 +105,9 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 
 /**
  * Utility class for generating the preview of Launcher for a given InvariantDeviceProfile.
@@ -456,54 +458,48 @@
 
     private void populate(BgDataModel dataModel,
             Map<ComponentKey, AppWidgetProviderInfo> widgetProviderInfoMap) {
-        // Separate the items that are on the current screen, and the other remaining items.
-        ArrayList<ItemInfo> currentWorkspaceItems = new ArrayList<>();
-        ArrayList<ItemInfo> otherWorkspaceItems = new ArrayList<>();
-        ArrayList<LauncherAppWidgetInfo> currentAppWidgets = new ArrayList<>();
-        ArrayList<LauncherAppWidgetInfo> otherAppWidgets = new ArrayList<>();
+        IntSet missingHotseatRank = new IntSet();
+        IntStream.range(0, mDp.numShownHotseatIcons).forEach(missingHotseatRank::add);
 
-        IntSet currentScreenIds = IntSet.wrap(mWorkspaceScreens.keySet());
-        filterCurrentWorkspaceItems(currentScreenIds, dataModel.workspaceItems,
-                currentWorkspaceItems, otherWorkspaceItems);
-        filterCurrentWorkspaceItems(currentScreenIds, dataModel.appWidgets, currentAppWidgets,
-                otherAppWidgets);
-        for (ItemInfo itemInfo : currentWorkspaceItems) {
-            switch (itemInfo.itemType) {
-                case Favorites.ITEM_TYPE_APPLICATION:
-                case Favorites.ITEM_TYPE_DEEP_SHORTCUT:
-                    inflateAndAddIcon((WorkspaceItemInfo) itemInfo);
-                    break;
-                case Favorites.ITEM_TYPE_FOLDER:
-                case Favorites.ITEM_TYPE_APP_PAIR:
-                    inflateAndAddCollectionIcon((CollectionInfo) itemInfo);
-                    break;
-                default:
-                    break;
-            }
-        }
-        Map<ComponentKey, AppWidgetProviderInfo> widgetsMap = widgetProviderInfoMap;
-        for (ItemInfo itemInfo : currentAppWidgets) {
-            switch (itemInfo.itemType) {
-                case Favorites.ITEM_TYPE_APPWIDGET:
-                case Favorites.ITEM_TYPE_CUSTOM_APPWIDGET:
-                    if (widgetsMap == null) {
-                        widgetsMap = dataModel.widgetsModel.getWidgetsByComponentKey()
-                                .entrySet()
-                                .stream()
-                                .filter(entry -> entry.getValue().widgetInfo != null)
-                                .collect(Collectors.toMap(
-                                        Map.Entry::getKey,
-                                        entry -> entry.getValue().widgetInfo
-                                ));
+        Map<ComponentKey, AppWidgetProviderInfo>[] widgetsMap = new Map[] { widgetProviderInfoMap};
+
+        // Separate the items that are on the current screen, and the other remaining items.
+        dataModel.itemsIdMap.stream()
+                .filter(currentScreenContentFilter(IntSet.wrap(mWorkspaceScreens.keySet())))
+                .forEach(itemInfo -> {
+                    switch (itemInfo.itemType) {
+                        case Favorites.ITEM_TYPE_APPLICATION:
+                        case Favorites.ITEM_TYPE_DEEP_SHORTCUT:
+                            inflateAndAddIcon((WorkspaceItemInfo) itemInfo);
+                            break;
+                        case Favorites.ITEM_TYPE_FOLDER:
+                        case Favorites.ITEM_TYPE_APP_PAIR:
+                            inflateAndAddCollectionIcon((CollectionInfo) itemInfo);
+                            break;
+                        case Favorites.ITEM_TYPE_APPWIDGET:
+                        case Favorites.ITEM_TYPE_CUSTOM_APPWIDGET:
+                            if (widgetsMap[0] == null) {
+                                widgetsMap[0] = dataModel.widgetsModel.getWidgetsByComponentKey()
+                                        .entrySet()
+                                        .stream()
+                                        .filter(entry -> entry.getValue().widgetInfo != null)
+                                        .collect(Collectors.toMap(
+                                                Entry::getKey,
+                                                entry -> entry.getValue().widgetInfo
+                                        ));
+                            }
+                            inflateAndAddWidgets((LauncherAppWidgetInfo) itemInfo, widgetsMap[0]);
+                            break;
+                        default:
+                            break;
                     }
-                    inflateAndAddWidgets((LauncherAppWidgetInfo) itemInfo, widgetsMap);
-                    break;
-                default:
-                    break;
-            }
-        }
-        IntArray ranks = getMissingHotseatRanks(currentWorkspaceItems,
-                mDp.numShownHotseatIcons);
+
+                    if (itemInfo.container == CONTAINER_HOTSEAT) {
+                        missingHotseatRank.remove(itemInfo.screenId);
+                    }
+                });
+
+        IntArray ranks = missingHotseatRank.getArray();
         FixedContainerItems hotseatPredictions =
                 dataModel.extraItems.get(CONTAINER_HOTSEAT_PREDICTION);
         List<ItemInfo> predictions = hotseatPredictions == null
diff --git a/src/com/android/launcher3/model/BaseLauncherBinder.java b/src/com/android/launcher3/model/BaseLauncherBinder.java
index c251114..de74ae8 100644
--- a/src/com/android/launcher3/model/BaseLauncherBinder.java
+++ b/src/com/android/launcher3/model/BaseLauncherBinder.java
@@ -19,8 +19,10 @@
 import static com.android.launcher3.BuildConfig.WIDGETS_ENABLED;
 import static com.android.launcher3.Flags.enableSmartspaceRemovalToggle;
 import static com.android.launcher3.Flags.enableWorkspaceInflation;
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_DESKTOP;
 import static com.android.launcher3.model.ItemInstallQueue.FLAG_LOADER_RUNNING;
-import static com.android.launcher3.model.ModelUtils.filterCurrentWorkspaceItems;
+import static com.android.launcher3.model.ModelUtils.WIDGET_FILTER;
+import static com.android.launcher3.model.ModelUtils.currentScreenContentFilter;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
 
@@ -44,11 +46,11 @@
 import com.android.launcher3.model.BgDataModel.FixedContainerItems;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.ItemInfo;
-import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.IntSet;
+import com.android.launcher3.util.IntSparseArrayMap;
 import com.android.launcher3.util.ItemInflater;
 import com.android.launcher3.util.LooperExecutor;
 import com.android.launcher3.util.LooperIdleLock;
@@ -102,14 +104,12 @@
         Trace.beginSection("BaseLauncherBinder#bindWorkspace");
         try {
             // Save a copy of all the bg-thread collections
-            ArrayList<ItemInfo> workspaceItems = new ArrayList<>();
-            ArrayList<LauncherAppWidgetInfo> appWidgets = new ArrayList<>();
+            IntSparseArrayMap<ItemInfo> itemsIdMap;
             final IntArray orderedScreenIds = new IntArray();
             ArrayList<FixedContainerItems> extraItems = new ArrayList<>();
             final int workspaceItemCount;
             synchronized (mBgDataModel) {
-                workspaceItems.addAll(mBgDataModel.workspaceItems);
-                appWidgets.addAll(mBgDataModel.appWidgets);
+                itemsIdMap = mBgDataModel.itemsIdMap.clone();
                 orderedScreenIds.addAll(mBgDataModel.collectWorkspaceScreens());
                 mBgDataModel.extraItems.forEach(extraItems::add);
                 if (incrementBindId) {
@@ -122,7 +122,7 @@
 
             for (Callbacks cb : mCallbacksList) {
                 new UnifiedWorkspaceBinder(cb, mUiExecutor, mApp, mBgDataModel, mMyBindingId,
-                        workspaceItems, appWidgets, extraItems, orderedScreenIds)
+                        itemsIdMap, extraItems, orderedScreenIds)
                         .bind(isBindSync, workspaceItemCount);
             }
         } finally {
@@ -258,8 +258,7 @@
         private final BgDataModel mBgDataModel;
 
         private final int mMyBindingId;
-        private final ArrayList<ItemInfo> mWorkspaceItems;
-        private final ArrayList<LauncherAppWidgetInfo> mAppWidgets;
+        private final IntSparseArrayMap<ItemInfo> mItemIdMap;
         private final IntArray mOrderedScreenIds;
         private final ArrayList<FixedContainerItems> mExtraItems;
 
@@ -268,8 +267,7 @@
                 LauncherAppState app,
                 BgDataModel bgDataModel,
                 int myBindingId,
-                ArrayList<ItemInfo> workspaceItems,
-                ArrayList<LauncherAppWidgetInfo> appWidgets,
+                IntSparseArrayMap<ItemInfo> itemIdMap,
                 ArrayList<FixedContainerItems> extraItems,
                 IntArray orderedScreenIds) {
             mCallbacks = callbacks;
@@ -277,8 +275,7 @@
             mApp = app;
             mBgDataModel = bgDataModel;
             mMyBindingId = myBindingId;
-            mWorkspaceItems = workspaceItems;
-            mAppWidgets = appWidgets;
+            mItemIdMap = itemIdMap;
             mExtraItems = extraItems;
             mOrderedScreenIds = orderedScreenIds;
         }
@@ -294,10 +291,15 @@
             ArrayList<ItemInfo> currentAppWidgets = new ArrayList<>();
             ArrayList<ItemInfo> otherAppWidgets = new ArrayList<>();
 
-            filterCurrentWorkspaceItems(currentScreenIds, mWorkspaceItems, currentWorkspaceItems,
-                    otherWorkspaceItems);
-            filterCurrentWorkspaceItems(currentScreenIds, mAppWidgets, currentAppWidgets,
-                    otherAppWidgets);
+            Predicate<ItemInfo> currentScreenCheck = currentScreenContentFilter(currentScreenIds);
+            mItemIdMap.forEach(item -> {
+                if (currentScreenCheck.test(item)) {
+                    (WIDGET_FILTER.test(item) ? currentAppWidgets : currentWorkspaceItems)
+                            .add(item);
+                } else if (item.container == CONTAINER_DESKTOP) {
+                    (WIDGET_FILTER.test(item) ? otherAppWidgets : otherWorkspaceItems).add(item);
+                }
+            });
             final InvariantDeviceProfile idp = mApp.getInvariantDeviceProfile();
             sortWorkspaceItemsSpatially(idp, currentWorkspaceItems);
             sortWorkspaceItemsSpatially(idp, otherWorkspaceItems);
diff --git a/src/com/android/launcher3/model/BgDataModel.java b/src/com/android/launcher3/model/BgDataModel.java
index b9b1e98..a04cbfb 100644
--- a/src/com/android/launcher3/model/BgDataModel.java
+++ b/src/com/android/launcher3/model/BgDataModel.java
@@ -20,6 +20,11 @@
 import static com.android.launcher3.BuildConfig.QSB_ON_FIRST_SCREEN;
 import static com.android.launcher3.BuildConfig.WIDGETS_ENABLED;
 import static com.android.launcher3.Flags.enableSmartspaceRemovalToggle;
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_DESKTOP;
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT;
+import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APP_PAIR;
+import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT;
+import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_FOLDER;
 import static com.android.launcher3.Utilities.SHOULD_SHOW_FIRST_PAGE_WIDGET;
 import static com.android.launcher3.shortcuts.ShortcutRequest.PINNED;
 
@@ -31,7 +36,6 @@
 import android.content.pm.ShortcutInfo;
 import android.os.UserHandle;
 import android.text.TextUtils;
-import android.util.ArraySet;
 import android.util.Log;
 import android.util.Pair;
 import android.view.View;
@@ -39,14 +43,11 @@
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
-import com.android.launcher3.LauncherSettings;
-import com.android.launcher3.LauncherSettings.Favorites;
+import com.android.launcher3.BuildConfig;
 import com.android.launcher3.Workspace;
 import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.model.data.AppInfo;
-import com.android.launcher3.model.data.AppPairInfo;
 import com.android.launcher3.model.data.CollectionInfo;
-import com.android.launcher3.model.data.FolderInfo;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
@@ -93,22 +94,6 @@
     public final IntSparseArrayMap<ItemInfo> itemsIdMap = new IntSparseArrayMap<>();
 
     /**
-     * List of all the folders and shortcuts directly on the home screen (no widgets
-     * or shortcuts within folders).
-     */
-    public final ArrayList<ItemInfo> workspaceItems = new ArrayList<>();
-
-    /**
-     * All LauncherAppWidgetInfo created by LauncherModel.
-     */
-    public final ArrayList<LauncherAppWidgetInfo> appWidgets = new ArrayList<>();
-
-    /**
-     * Map of id to CollectionInfos of all the folders or app pairs created by LauncherModel
-     */
-    public final IntSparseArrayMap<CollectionInfo> collections = new IntSparseArrayMap<>();
-
-    /**
      * Extra container based items
      */
     public final IntSparseArrayMap<FixedContainerItems> extraItems = new IntSparseArrayMap<>();
@@ -144,9 +129,6 @@
      * Clears all the data
      */
     public synchronized void clear() {
-        workspaceItems.clear();
-        appWidgets.clear();
-        collections.clear();
         itemsIdMap.clear();
         deepShortcutMap.clear();
         extraItems.clear();
@@ -158,7 +140,7 @@
     public synchronized IntArray collectWorkspaceScreens() {
         IntSet screenSet = new IntSet();
         for (ItemInfo item: itemsIdMap) {
-            if (item.container == LauncherSettings.Favorites.CONTAINER_DESKTOP) {
+            if (item.container == CONTAINER_DESKTOP) {
                 screenSet.add(item.screenId);
             }
         }
@@ -173,26 +155,14 @@
     public synchronized void dump(String prefix, FileDescriptor fd, PrintWriter writer,
             String[] args) {
         writer.println(prefix + "Data Model:");
-        writer.println(prefix + " ---- workspace items ");
-        for (int i = 0; i < workspaceItems.size(); i++) {
-            writer.println(prefix + '\t' + workspaceItems.get(i).toString());
-        }
-        writer.println(prefix + " ---- appwidget items ");
-        for (int i = 0; i < appWidgets.size(); i++) {
-            writer.println(prefix + '\t' + appWidgets.get(i).toString());
-        }
-        writer.println(prefix + " ---- collection items ");
-        for (int i = 0; i < collections.size(); i++) {
-            writer.println(prefix + '\t' + collections.valueAt(i).toString());
+        writer.println(prefix + " ---- items id map ");
+        for (int i = 0; i < itemsIdMap.size(); i++) {
+            writer.println(prefix + '\t' + itemsIdMap.valueAt(i).toString());
         }
         writer.println(prefix + " ---- extra items ");
         for (int i = 0; i < extraItems.size(); i++) {
             writer.println(prefix + '\t' + extraItems.valueAt(i).toString());
         }
-        writer.println(prefix + " ---- items id map ");
-        for (int i = 0; i < itemsIdMap.size(); i++) {
-            writer.println(prefix + '\t' + itemsIdMap.valueAt(i).toString());
-        }
 
         if (args.length > 0 && TextUtils.equals(args[0], "--all")) {
             writer.println(prefix + "shortcut counts ");
@@ -207,94 +177,38 @@
         removeItem(context, Arrays.asList(items));
     }
 
-    public synchronized void removeItem(Context context, Iterable<? extends ItemInfo> items) {
-        ArraySet<UserHandle> updatedDeepShortcuts = new ArraySet<>();
-        for (ItemInfo item : items) {
-            switch (item.itemType) {
-                case LauncherSettings.Favorites.ITEM_TYPE_FOLDER:
-                case LauncherSettings.Favorites.ITEM_TYPE_APP_PAIR:
-                    collections.remove(item.id);
-                    if (FeatureFlags.IS_STUDIO_BUILD) {
-                        for (ItemInfo info : itemsIdMap) {
-                            if (info.container == item.id) {
-                                // We are deleting a collection which still contains items that
-                                // think they are contained by that collection.
-                                String msg = "deleting a collection (" + item + ") which still "
-                                        + "contains items (" + info + ")";
-                                Log.e(TAG, msg);
-                            }
-                        }
-                    }
-                    workspaceItems.remove(item);
-                    break;
-                case LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT: {
-                    updatedDeepShortcuts.add(item.user);
-                    // Fall through.
-                }
-                case LauncherSettings.Favorites.ITEM_TYPE_APPLICATION:
-                    workspaceItems.remove(item);
-                    break;
-                case LauncherSettings.Favorites.ITEM_TYPE_APPWIDGET:
-                case LauncherSettings.Favorites.ITEM_TYPE_CUSTOM_APPWIDGET:
-                    appWidgets.remove(item);
-                    break;
-            }
-            itemsIdMap.remove(item.id);
+    public synchronized void removeItem(Context context, List<? extends ItemInfo> items) {
+        if (BuildConfig.IS_STUDIO_BUILD) {
+            items.stream()
+                    .filter(item -> item.itemType == ITEM_TYPE_FOLDER
+                            || item.itemType == ITEM_TYPE_APP_PAIR)
+                    .forEach(item -> itemsIdMap.stream()
+                            .filter(info -> info.container == item.id)
+                            // We are deleting a collection which still contains items that
+                            // think they are contained by that collection.
+                            .forEach(info -> Log.e(TAG,
+                                    "deleting a collection (" + item + ") which still contains"
+                                            + " items (" + info + ")")));
         }
-        updatedDeepShortcuts.forEach(user -> updateShortcutPinnedState(context, user));
+
+        items.forEach(item -> itemsIdMap.remove(item.id));
+        items.stream().map(info -> info.user).distinct().forEach(
+                user -> updateShortcutPinnedState(context, user));
     }
 
     public synchronized void addItem(Context context, ItemInfo item, boolean newItem) {
-        addItem(context, item, newItem, null);
-    }
-
-    public synchronized void addItem(
-            Context context, ItemInfo item, boolean newItem, @Nullable LoaderMemoryLogger logger) {
-        if (logger != null) {
-            logger.addLog(
-                    Log.DEBUG,
-                    TAG,
-                    String.format("Adding item to ID map: %s", item.toString()),
-                    /* stackTrace= */ null);
-        }
         itemsIdMap.put(item.id, item);
-        switch (item.itemType) {
-            case LauncherSettings.Favorites.ITEM_TYPE_FOLDER:
-                collections.put(item.id, (FolderInfo) item);
-                workspaceItems.add(item);
-                break;
-            case LauncherSettings.Favorites.ITEM_TYPE_APP_PAIR:
-                collections.put(item.id, (AppPairInfo) item);
-                // Fall through here. App pairs are both containers (like folders) and containable
-                // items (can be placed in folders). So we need to add app pairs to the folders
-                // array (above) but also verify the existence of their container, like regular
-                // apps (below).
-            case LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT:
-            case LauncherSettings.Favorites.ITEM_TYPE_APPLICATION:
-                if (item.container == LauncherSettings.Favorites.CONTAINER_DESKTOP ||
-                        item.container == LauncherSettings.Favorites.CONTAINER_HOTSEAT) {
-                    workspaceItems.add(item);
-                } else {
-                    if (newItem) {
-                        if (!collections.containsKey(item.container)) {
-                            // Adding an item to a nonexistent collection.
-                            String msg = "attempted to add item: " + item + " to a nonexistent app"
-                                    + " collection";
-                            Log.e(TAG, msg);
-                        }
-                    } else {
-                        findOrMakeFolder(item.container).add(item);
-                    }
-                }
-                break;
-            case LauncherSettings.Favorites.ITEM_TYPE_APPWIDGET:
-            case LauncherSettings.Favorites.ITEM_TYPE_CUSTOM_APPWIDGET:
-                appWidgets.add((LauncherAppWidgetInfo) item);
-                break;
-        }
-        if (newItem && item.itemType == LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT) {
+        if (newItem && item.itemType == ITEM_TYPE_DEEP_SHORTCUT) {
             updateShortcutPinnedState(context, item.user);
         }
+        if (BuildConfig.IS_DEBUG_DEVICE
+                && newItem
+                && item.container != CONTAINER_DESKTOP
+                && item.container != CONTAINER_HOTSEAT
+                && !(itemsIdMap.get(item.container) instanceof CollectionInfo)) {
+            // Adding an item to a nonexistent collection.
+            Log.e(TAG, "attempted to add item: " + item + " to a nonexistent app collection");
+        }
     }
 
     /**
@@ -334,7 +248,7 @@
         Map<String, Set<String>> modelMap = Stream.concat(
                     // Model shortcuts
                     itemStream.build()
-                        .filter(wi -> wi.itemType == Favorites.ITEM_TYPE_DEEP_SHORTCUT)
+                        .filter(wi -> wi.itemType == ITEM_TYPE_DEEP_SHORTCUT)
                         .map(ShortcutKey::fromItemInfo),
                     // Pending shortcuts
                     ItemInstallQueue.INSTANCE.get(context).getPendingShortcuts(user))
@@ -375,24 +289,6 @@
     }
 
     /**
-     * Return an existing FolderInfo object if we have encountered this ID previously,
-     * or make a new one.
-     */
-    public synchronized CollectionInfo findOrMakeFolder(int id) {
-        // See if a placeholder was created for us already
-        CollectionInfo collectionInfo = collections.get(id);
-        if (collectionInfo == null) {
-            // No placeholder -- create a new blank folder instance. At this point, we don't know
-            // if the desired container is supposed to be a folder or an app pair. In the case that
-            // it is an app pair, the blank folder will be replaced by a blank app pair when the app
-            // pair is getting processed, in WorkspaceItemProcessor.processFolderOrAppPair().
-            collectionInfo = new FolderInfo();
-            collections.put(id, collectionInfo);
-        }
-        return collectionInfo;
-    }
-
-    /**
      * Clear all the deep shortcut counts for the given package, and re-add the new shortcut counts.
      */
     public synchronized void updateDeepShortcutCounts(
@@ -424,16 +320,6 @@
     }
 
     /**
-     * Returns a list containing all workspace items including widgets.
-     */
-    public synchronized ArrayList<ItemInfo> getAllWorkspaceItems() {
-        ArrayList<ItemInfo> items = new ArrayList<>(workspaceItems.size() + appWidgets.size());
-        items.addAll(workspaceItems);
-        items.addAll(appWidgets);
-        return items;
-    }
-
-    /**
      * Calls the provided {@code op} for all workspaceItems in the in-memory model (both persisted
      * items and dynamic/predicted items for the provided {@code userHandle}.
      * Note the call is not synchronized over the model, that should be handled by the called.
diff --git a/src/com/android/launcher3/model/FirstScreenBroadcastHelper.kt b/src/com/android/launcher3/model/FirstScreenBroadcastHelper.kt
index aa62c32..6ad52ea 100644
--- a/src/com/android/launcher3/model/FirstScreenBroadcastHelper.kt
+++ b/src/com/android/launcher3/model/FirstScreenBroadcastHelper.kt
@@ -30,7 +30,6 @@
 import com.android.launcher3.model.data.ItemInfo
 import com.android.launcher3.model.data.LauncherAppWidgetInfo
 import com.android.launcher3.model.data.WorkspaceItemInfo
-import com.android.launcher3.pm.InstallSessionHelper
 import com.android.launcher3.util.Executors
 import com.android.launcher3.util.PackageManagerHelper
 import com.android.launcher3.util.PackageUserKey
@@ -80,21 +79,22 @@
         packageManagerHelper: PackageManagerHelper,
         firstScreenItems: List<ItemInfo>,
         userKeyToSessionMap: Map<PackageUserKey, SessionInfo>,
-        allWidgets: List<LauncherAppWidgetInfo>
+        allWidgets: List<ItemInfo>,
     ): List<FirstScreenBroadcastModel> {
 
         // installers for installing items
-        val pendingItemInstallerMap: Map<String, MutableSet<String>> =
+        val pendingItemInstallerMap: Map<String, Set<String>> =
             createPendingItemsMap(userKeyToSessionMap)
+
         val installingPackages = pendingItemInstallerMap.values.flatten().toSet()
 
         // installers for installed items on first screen
-        val installedItemInstallerMap: Map<String, MutableSet<ItemInfo>> =
+        val installedItemInstallerMap: Map<String, List<ItemInfo>> =
             createInstalledItemsMap(firstScreenItems, installingPackages, packageManagerHelper)
 
         // installers for widgets on all screens
-        val allInstalledWidgetsMap: Map<String, MutableSet<LauncherAppWidgetInfo>> =
-            createAllInstalledWidgetsMap(allWidgets, installingPackages, packageManagerHelper)
+        val allInstalledWidgetsMap: Map<String, List<ItemInfo>> =
+            createInstalledItemsMap(allWidgets, installingPackages, packageManagerHelper)
 
         val allInstallers: Set<String> =
             pendingItemInstallerMap.keys +
@@ -131,39 +131,39 @@
                             context,
                             0 /* requestCode */,
                             Intent(),
-                            PendingIntent.FLAG_ONE_SHOT or PendingIntent.FLAG_IMMUTABLE
-                        )
+                            PendingIntent.FLAG_ONE_SHOT or PendingIntent.FLAG_IMMUTABLE,
+                        ),
                     )
                     .putStringArrayListExtra(
                         PENDING_COLLECTION_ITEM_EXTRA,
-                        ArrayList(model.pendingCollectionItems)
+                        ArrayList(model.pendingCollectionItems),
                     )
                     .putStringArrayListExtra(
                         PENDING_WORKSPACE_ITEM_EXTRA,
-                        ArrayList(model.pendingWorkspaceItems)
+                        ArrayList(model.pendingWorkspaceItems),
                     )
                     .putStringArrayListExtra(
                         PENDING_HOTSEAT_ITEM_EXTRA,
-                        ArrayList(model.pendingHotseatItems)
+                        ArrayList(model.pendingHotseatItems),
                     )
                     .putStringArrayListExtra(
                         PENDING_WIDGET_ITEM_EXTRA,
-                        ArrayList(model.pendingWidgetItems)
+                        ArrayList(model.pendingWidgetItems),
                     )
                     .putStringArrayListExtra(
                         INSTALLED_WORKSPACE_ITEMS_EXTRA,
-                        ArrayList(model.installedWorkspaceItems)
+                        ArrayList(model.installedWorkspaceItems),
                     )
                     .putStringArrayListExtra(
                         INSTALLED_HOTSEAT_ITEMS_EXTRA,
-                        ArrayList(model.installedHotseatItems)
+                        ArrayList(model.installedHotseatItems),
                     )
                     .putStringArrayListExtra(
                         ALL_INSTALLED_WIDGETS_ITEM_EXTRA,
                         ArrayList(
                             model.firstScreenInstalledWidgets +
                                 model.secondaryScreenInstalledWidgets
-                        )
+                        ),
                     )
             context.sendBroadcast(intent)
         }
@@ -172,66 +172,46 @@
     /** Maps Installer packages to Set of app packages from install sessions */
     private fun createPendingItemsMap(
         userKeyToSessionMap: Map<PackageUserKey, SessionInfo>
-    ): Map<String, MutableSet<String>> {
+    ): Map<String, Set<String>> {
         val myUser = Process.myUserHandle()
-        val result = mutableMapOf<String, MutableSet<String>>()
-        userKeyToSessionMap.forEach { entry ->
-            if (!myUser.equals(InstallSessionHelper.getUserHandle(entry.value))) return@forEach
-            val installer = entry.value.installerPackageName
-            val appPackage = entry.value.appPackageName
-            if (installer.isNullOrEmpty() || appPackage.isNullOrEmpty()) return@forEach
-            result.getOrPut(installer) { mutableSetOf() }.add(appPackage)
-        }
-        return result
-    }
-
-    /**
-     * Maps Installer packages to Set of ItemInfo from first screen. Filter out installing packages.
-     */
-    private fun createInstalledItemsMap(
-        firstScreenItems: List<ItemInfo>,
-        installingPackages: Set<String>,
-        packageManagerHelper: PackageManagerHelper
-    ): Map<String, MutableSet<ItemInfo>> {
-        val result = mutableMapOf<String, MutableSet<ItemInfo>>()
-        firstScreenItems.forEach { item ->
-            val appPackage = getPackageName(item) ?: return@forEach
-            if (installingPackages.contains(appPackage)) return@forEach
-            val installer = packageManagerHelper.getAppInstallerPackage(appPackage)
-            if (installer.isNullOrEmpty()) return@forEach
-            result.getOrPut(installer) { mutableSetOf() }.add(item)
-        }
-        return result
-    }
-
-    /**
-     * Maps Installer packages to Set of AppWidget packages installed on all screens. Filter out
-     * installing packages.
-     */
-    private fun createAllInstalledWidgetsMap(
-        allWidgets: List<LauncherAppWidgetInfo>,
-        installingPackages: Set<String>,
-        packageManagerHelper: PackageManagerHelper
-    ): Map<String, MutableSet<LauncherAppWidgetInfo>> {
-        val result = mutableMapOf<String, MutableSet<LauncherAppWidgetInfo>>()
-        allWidgets
-            .sortedBy { widget -> widget.screenId }
-            .forEach { widget ->
-                val appPackage = getPackageName(widget) ?: return@forEach
-                if (installingPackages.contains(appPackage)) return@forEach
-                val installer = packageManagerHelper.getAppInstallerPackage(appPackage)
-                if (installer.isNullOrEmpty()) return@forEach
-                result.getOrPut(installer) { mutableSetOf() }.add(widget)
+        return userKeyToSessionMap.values
+            .filter {
+                it.user == myUser &&
+                    !it.installerPackageName.isNullOrEmpty() &&
+                    !it.appPackageName.isNullOrEmpty()
             }
-        return result
+            .groupBy(
+                keySelector = { it.installerPackageName },
+                valueTransform = { it.appPackageName },
+            )
+            .mapValues { it.value.filterNotNull().toSet() } as Map<String, Set<String>>
     }
 
+    /** Maps Installer packages to Set of ItemInfos. Filter out installing packages. */
+    private fun createInstalledItemsMap(
+        allItems: Iterable<ItemInfo>,
+        installingPackages: Set<String>,
+        packageManagerHelper: PackageManagerHelper,
+    ): Map<String, List<ItemInfo>> =
+        allItems
+            .sortedBy { it.screenId }
+            .groupByTo(mutableMapOf()) {
+                getPackageName(it)?.let { pkg ->
+                    if (installingPackages.contains(pkg)) {
+                        null
+                    } else {
+                        packageManagerHelper.getAppInstallerPackage(pkg)
+                    }
+                }
+            }
+            .apply { remove(null) } as Map<String, List<ItemInfo>>
+
     /**
      * Add first screen Pending Items from Map to [FirstScreenBroadcastModel] for given installer
      */
     private fun FirstScreenBroadcastModel.addPendingItems(
         installingItems: Set<String>?,
-        firstScreenItems: List<ItemInfo>
+        firstScreenItems: List<ItemInfo>,
     ) {
         if (installingItems == null) return
         for (info in firstScreenItems) {
@@ -251,7 +231,7 @@
      */
     private fun FirstScreenBroadcastModel.addInstalledItems(
         installer: String,
-        installedItemInstallerMap: Map<String, Set<ItemInfo>>,
+        installedItemInstallerMap: Map<String, List<ItemInfo>>,
     ) {
         installedItemInstallerMap[installer]?.forEach { info ->
             val packageName: String = getPackageName(info) ?: return@forEach
@@ -265,7 +245,7 @@
     /** Add Widgets on every screen from Map to [FirstScreenBroadcastModel] for given installer */
     private fun FirstScreenBroadcastModel.addAllScreenWidgets(
         installer: String,
-        allInstalledWidgetsMap: Map<String, Set<LauncherAppWidgetInfo>>
+        allInstalledWidgetsMap: Map<String, List<ItemInfo>>,
     ) {
         allInstalledWidgetsMap[installer]?.forEach { widget ->
             val packageName: String = getPackageName(widget) ?: return@forEach
@@ -279,7 +259,7 @@
 
     private fun FirstScreenBroadcastModel.addCollectionItems(
         info: ItemInfo,
-        installingPackages: Set<String>
+        installingPackages: Set<String>,
     ) {
         if (info !is CollectionInfo) return
         pendingCollectionItems.addAll(
@@ -336,7 +316,7 @@
             Log.d(
                 TAG,
                 "Sending First Screen Broadcast for installer=$installerPackage" +
-                    ", total packages=${getTotalItemCount()}"
+                    ", total packages=${getTotalItemCount()}",
             )
             pendingCollectionItems.forEach {
                 Log.d(TAG, "$installerPackage:Pending Collection item:$it")
@@ -361,15 +341,7 @@
         }
     }
 
-    private fun getPackageName(info: ItemInfo): String? {
-        var packageName: String? = null
-        if (info is LauncherAppWidgetInfo) {
-            info.providerName?.let { packageName = info.providerName.packageName }
-        } else if (info.targetComponent != null) {
-            packageName = info.targetComponent?.packageName
-        }
-        return packageName
-    }
+    private fun getPackageName(info: ItemInfo): String? = info.targetComponent?.packageName
 
     /**
      * Clone the provided list on UI thread. This is used for [FolderInfo.getContents] which is
diff --git a/src/com/android/launcher3/model/LoaderCursor.java b/src/com/android/launcher3/model/LoaderCursor.java
index 536d4c9..1623881 100644
--- a/src/com/android/launcher3/model/LoaderCursor.java
+++ b/src/com/android/launcher3/model/LoaderCursor.java
@@ -16,6 +16,11 @@
 
 package com.android.launcher3.model;
 
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_DESKTOP;
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT;
+import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APPLICATION;
+import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APP_PAIR;
+import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT;
 import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
 import static com.android.launcher3.Utilities.SHOULD_SHOW_FIRST_PAGE_WIDGET;
 import static com.android.launcher3.icons.cache.CacheLookupFlag.DEFAULT_LOOKUP_FLAG;
@@ -48,6 +53,8 @@
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.logging.FileLog;
 import com.android.launcher3.model.data.AppInfo;
+import com.android.launcher3.model.data.CollectionInfo;
+import com.android.launcher3.model.data.FolderInfo;
 import com.android.launcher3.model.data.IconRequestInfo;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
@@ -84,6 +91,11 @@
     private final IntArray mRestoredRows = new IntArray();
     private final IntSparseArrayMap<GridOccupancy> mOccupied = new IntSparseArrayMap<>();
 
+    // CollectionInfo objects, which have not yet been loaded from the DB, but are expected to
+    // found eventually as the loading progresses
+    private final IntSparseArrayMap<CollectionInfo> mPendingCollectionInfo =
+            new IntSparseArrayMap<>();
+
     private final int mIconIndex;
     public final int mTitleIndex;
 
@@ -479,8 +491,26 @@
         info.cellY = getInt(mCellYIndex);
     }
 
-    public void checkAndAddItem(ItemInfo info, BgDataModel dataModel) {
-        checkAndAddItem(info, dataModel, null);
+    /**
+     * Return an existing FolderInfo object if we have encountered this ID previously,
+     * or make a new one.
+     */
+    public CollectionInfo findOrMakeFolder(int id, BgDataModel dataModel) {
+        // See if a placeholder was created for us already
+        ItemInfo info = dataModel.itemsIdMap.get(id);
+        if (info instanceof CollectionInfo c) return c;
+
+        CollectionInfo pending = mPendingCollectionInfo.get(id);
+        if (pending != null) return pending;
+
+        // No placeholder -- create a new blank folder instance. At this point, we don't know
+        // if the desired container is supposed to be a folder or an app pair. In the case that
+        // it is an app pair, the blank folder will be replaced by a blank app pair when the app
+        // pair is getting processed, in WorkspaceItemProcessor.processFolderOrAppPair().
+        pending = new FolderInfo();
+        pending.id = id;
+        mPendingCollectionInfo.put(id, pending);
+        return pending;
     }
 
     /**
@@ -495,7 +525,21 @@
             ShortcutKey.fromItemInfo(info);
         }
         if (checkItemPlacement(info, dataModel.isFirstPagePinnedItemEnabled)) {
-            dataModel.addItem(mContext, info, false, logger);
+            if (logger != null) {
+                logger.addLog(
+                        Log.DEBUG,
+                        TAG,
+                        String.format("Adding item to ID map: %s", info),
+                        /* stackTrace= */ null);
+            }
+            dataModel.addItem(mContext, info, false);
+            if ((info.itemType == ITEM_TYPE_APP_PAIR
+                    || info.itemType == ITEM_TYPE_DEEP_SHORTCUT
+                    || info.itemType == ITEM_TYPE_APPLICATION)
+                    && info.container != CONTAINER_DESKTOP
+                    && info.container != CONTAINER_HOTSEAT) {
+                findOrMakeFolder(info.container, dataModel).add(info);
+            }
             if (mRestoreEventLogger != null) {
                 mRestoreEventLogger.logSingleFavoritesItemRestored(itemType);
             }
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index 44b7e8b..fee9696 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -31,7 +31,8 @@
 import static com.android.launcher3.model.BgDataModel.Callbacks.FLAG_QUIET_MODE_CHANGE_PERMISSION;
 import static com.android.launcher3.model.BgDataModel.Callbacks.FLAG_QUIET_MODE_ENABLED;
 import static com.android.launcher3.model.BgDataModel.Callbacks.FLAG_WORK_PROFILE_QUIET_MODE_ENABLED;
-import static com.android.launcher3.model.ModelUtils.filterCurrentWorkspaceItems;
+import static com.android.launcher3.model.ModelUtils.WIDGET_FILTER;
+import static com.android.launcher3.model.ModelUtils.currentScreenContentFilter;
 import static com.android.launcher3.model.data.ItemInfoWithIcon.FLAG_INSTALL_SESSION_ACTIVE;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
 import static com.android.launcher3.util.PackageManagerHelper.hasShortcutsPermission;
@@ -82,7 +83,6 @@
 import com.android.launcher3.logging.FileLog;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.AppPairInfo;
-import com.android.launcher3.model.data.CollectionInfo;
 import com.android.launcher3.model.data.FolderInfo;
 import com.android.launcher3.model.data.IconRequestInfo;
 import com.android.launcher3.model.data.ItemInfo;
@@ -210,10 +210,10 @@
         final int firstScreen = allScreens.get(0);
         IntSet firstScreens = IntSet.wrap(firstScreen);
 
-        ArrayList<ItemInfo> allItems = mBgDataModel.getAllWorkspaceItems();
-        ArrayList<ItemInfo> firstScreenItems = new ArrayList<>();
-        filterCurrentWorkspaceItems(firstScreens, allItems, firstScreenItems,
-                new ArrayList<>() /* otherScreenItems are ignored */);
+        List<ItemInfo> firstScreenItems =
+                mBgDataModel.itemsIdMap.stream()
+                        .filter(currentScreenContentFilter(firstScreens))
+                        .toList();
         final int disableArchivingLauncherBroadcast = Settings.Secure.getInt(
                 mApp.getContext().getContentResolver(),
                 "disable_launcher_broadcast_installed_apps",
@@ -227,7 +227,7 @@
                             mPmHelper,
                             firstScreenItems,
                             mInstallingPkgsCached,
-                            mBgDataModel.appWidgets
+                            mBgDataModel.itemsIdMap.stream().filter(WIDGET_FILTER).toList()
                     );
             logASplit("Sending first screen broadcast with additional archiving Extras");
             FirstScreenBroadcastHelper.sendBroadcastsForModels(mApp.getContext(), broadcastModels);
@@ -523,14 +523,13 @@
      * requests high-res icons for the items that are part of an app pair.
      */
     private void processAppPairItems() {
-        for (CollectionInfo collection : mBgDataModel.collections) {
-            if (!(collection instanceof AppPairInfo appPair)) {
-                continue;
-            }
-
-            appPair.getContents().sort(Folder.ITEM_POS_COMPARATOR);
-            appPair.fetchHiResIconsIfNeeded(mIconCache);
-        }
+        mBgDataModel.itemsIdMap.stream()
+                .filter(item -> item instanceof AppPairInfo)
+                .forEach(item -> {
+                    AppPairInfo appPair = (AppPairInfo) item;
+                    appPair.getContents().sort(Folder.ITEM_POS_COMPARATOR);
+                    appPair.fetchHiResIconsIfNeeded(mIconCache);
+                });
     }
 
     /**
@@ -586,8 +585,8 @@
         // Sort the folder items, update ranks, and make sure all preview items are high res.
         List<FolderGridOrganizer> verifiers = mApp.getInvariantDeviceProfile().supportedProfiles
                 .stream().map(FolderGridOrganizer::createFolderGridOrganizer).toList();
-        for (CollectionInfo collection : mBgDataModel.collections) {
-            if (!(collection instanceof FolderInfo folder)) {
+        for (ItemInfo itemInfo : mBgDataModel.itemsIdMap) {
+            if (!(itemInfo instanceof FolderInfo folder)) {
                 continue;
             }
 
@@ -657,8 +656,6 @@
             IntArray deletedFolderIds = mApp.getModel().getModelDbController().deleteEmptyFolders();
             synchronized (mBgDataModel) {
                 for (int folderId : deletedFolderIds) {
-                    mBgDataModel.workspaceItems.remove(mBgDataModel.collections.get(folderId));
-                    mBgDataModel.collections.remove(folderId);
                     mBgDataModel.itemsIdMap.remove(folderId);
                 }
             }
@@ -676,8 +673,6 @@
 
         synchronized (mBgDataModel) {
             for (int id : deleted) {
-                mBgDataModel.workspaceItems.remove(mBgDataModel.collections.get(id));
-                mBgDataModel.collections.remove(id);
                 mBgDataModel.itemsIdMap.remove(id);
             }
         }
@@ -819,18 +814,19 @@
 
     private void loadFolderNames() {
         FolderNameProvider provider = FolderNameProvider.newInstance(mApp.getContext(),
-                mBgAllAppsList.data, mBgDataModel.collections);
+                mBgAllAppsList.data, FolderNameProvider.getCollectionForSuggestions(mBgDataModel));
 
         synchronized (mBgDataModel) {
-            for (int i = 0; i < mBgDataModel.collections.size(); i++) {
-                FolderNameInfos suggestionInfos = new FolderNameInfos();
-                CollectionInfo info = mBgDataModel.collections.valueAt(i);
-                if (info instanceof FolderInfo fi && fi.suggestedFolderNames == null) {
-                    provider.getSuggestedFolderName(mApp.getContext(), fi.getAppContents(),
-                            suggestionInfos);
-                    fi.suggestedFolderNames = suggestionInfos;
-                }
-            }
+            mBgDataModel.itemsIdMap.stream()
+                    .filter(item ->
+                            item instanceof FolderInfo fi && fi.suggestedFolderNames == null)
+                    .forEach(info -> {
+                        FolderInfo fi = (FolderInfo) info;
+                        FolderNameInfos suggestionInfos = new FolderNameInfos();
+                        provider.getSuggestedFolderName(mApp.getContext(), fi.getAppContents(),
+                                suggestionInfos);
+                        fi.suggestedFolderNames = suggestionInfos;
+                    });
         }
     }
 
diff --git a/src/com/android/launcher3/model/ModelUtils.java b/src/com/android/launcher3/model/ModelUtils.java
index 9e72e28..da79982 100644
--- a/src/com/android/launcher3/model/ModelUtils.java
+++ b/src/com/android/launcher3/model/ModelUtils.java
@@ -15,15 +15,15 @@
  */
 package com.android.launcher3.model;
 
-import com.android.launcher3.LauncherSettings;
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_DESKTOP;
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT;
+import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APPWIDGET;
+import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_CUSTOM_APPWIDGET;
+
 import com.android.launcher3.model.data.ItemInfo;
-import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.IntSet;
 
-import java.util.Collections;
-import java.util.List;
-import java.util.Objects;
-import java.util.stream.IntStream;
+import java.util.function.Predicate;
 
 /**
  * Utils class for {@link com.android.launcher3.LauncherModel}.
@@ -31,54 +31,17 @@
 public class ModelUtils {
 
     /**
-     * Filters the set of items who are directly or indirectly (via another container) on the
-     * specified screen.
+     * Returns a filter for items on hotseat or current screens
      */
-    public static <T extends ItemInfo> void filterCurrentWorkspaceItems(
-            final IntSet currentScreenIds,
-            List<? extends T> allWorkspaceItems,
-            List<T> currentScreenItems,
-            List<T> otherScreenItems) {
-        // Purge any null ItemInfos
-        allWorkspaceItems.removeIf(Objects::isNull);
-        // Order the set of items by their containers first, this allows use to walk through the
-        // list sequentially, build up a list of containers that are in the specified screen,
-        // as well as all items in those containers.
-        IntSet itemsOnScreen = new IntSet();
-        Collections.sort(allWorkspaceItems,
-                (lhs, rhs) -> Integer.compare(lhs.container, rhs.container));
-        for (T info : allWorkspaceItems) {
-            if (info.container == LauncherSettings.Favorites.CONTAINER_DESKTOP) {
-                if (currentScreenIds.contains(info.screenId)) {
-                    currentScreenItems.add(info);
-                    itemsOnScreen.add(info.id);
-                } else {
-                    otherScreenItems.add(info);
-                }
-            } else if (info.container == LauncherSettings.Favorites.CONTAINER_HOTSEAT) {
-                currentScreenItems.add(info);
-                itemsOnScreen.add(info.id);
-            } else {
-                if (itemsOnScreen.contains(info.container)) {
-                    currentScreenItems.add(info);
-                    itemsOnScreen.add(info.id);
-                } else {
-                    otherScreenItems.add(info);
-                }
-            }
-        }
+    public static Predicate<ItemInfo> currentScreenContentFilter(IntSet currentScreenIds) {
+        return item -> item.container == CONTAINER_HOTSEAT
+                || (item.container == CONTAINER_DESKTOP
+                        && currentScreenIds.contains(item.screenId));
     }
 
     /**
-     * Iterates though current workspace items and returns available hotseat ranks for prediction.
+     * Returns a filter for widget items
      */
-    public static IntArray getMissingHotseatRanks(List<ItemInfo> items, int len) {
-        IntSet seen = new IntSet();
-        items.stream().filter(
-                info -> info.container == LauncherSettings.Favorites.CONTAINER_HOTSEAT)
-                .forEach(i -> seen.add(i.screenId));
-        IntArray result = new IntArray(len);
-        IntStream.range(0, len).filter(i -> !seen.contains(i)).forEach(result::add);
-        return result;
-    }
+    public static final Predicate<ItemInfo> WIDGET_FILTER = item ->
+            item.itemType == ITEM_TYPE_APPWIDGET || item.itemType == ITEM_TYPE_CUSTOM_APPWIDGET;
 }
diff --git a/src/com/android/launcher3/model/ModelWriter.java b/src/com/android/launcher3/model/ModelWriter.java
index b477cb1..0332775 100644
--- a/src/com/android/launcher3/model/ModelWriter.java
+++ b/src/com/android/launcher3/model/ModelWriter.java
@@ -459,37 +459,14 @@
                 if (item.container != Favorites.CONTAINER_DESKTOP &&
                         item.container != Favorites.CONTAINER_HOTSEAT) {
                     // Item is in a collection, make sure this collection exists
-                    if (!mBgDataModel.collections.containsKey(item.container)) {
+                    if (!(mBgDataModel.itemsIdMap.get(item.container) instanceof CollectionInfo)) {
                         // An items container is being set to a that of an item which is not in
-                        // the list of Folders.
+                        // the list of collections.
                         String msg = "item: " + item + " container being set to: " +
                                 item.container + ", not in the list of collections";
                         Log.e(TAG, msg);
                     }
                 }
-
-                // Items are added/removed from the corresponding FolderInfo elsewhere, such
-                // as in Workspace.onDrop. Here, we just add/remove them from the list of items
-                // that are on the desktop, as appropriate
-                ItemInfo modelItem = mBgDataModel.itemsIdMap.get(itemId);
-                if (modelItem != null &&
-                        (modelItem.container == Favorites.CONTAINER_DESKTOP ||
-                                modelItem.container == Favorites.CONTAINER_HOTSEAT)) {
-                    switch (modelItem.itemType) {
-                        case Favorites.ITEM_TYPE_APPLICATION:
-                        case Favorites.ITEM_TYPE_DEEP_SHORTCUT:
-                        case Favorites.ITEM_TYPE_FOLDER:
-                        case Favorites.ITEM_TYPE_APP_PAIR:
-                            if (!mBgDataModel.workspaceItems.contains(modelItem)) {
-                                mBgDataModel.workspaceItems.add(modelItem);
-                            }
-                            break;
-                        default:
-                            break;
-                    }
-                } else {
-                    mBgDataModel.workspaceItems.remove(modelItem);
-                }
                 mVerifier.verifyModel();
             }
         }
diff --git a/src/com/android/launcher3/model/PackageInstallStateChangedTask.java b/src/com/android/launcher3/model/PackageInstallStateChangedTask.java
index d238213..4103937 100644
--- a/src/com/android/launcher3/model/PackageInstallStateChangedTask.java
+++ b/src/com/android/launcher3/model/PackageInstallStateChangedTask.java
@@ -15,6 +15,8 @@
  */
 package com.android.launcher3.model;
 
+import static com.android.launcher3.model.ModelUtils.WIDGET_FILTER;
+
 import android.content.Context;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageManager;
@@ -85,12 +87,16 @@
                 }
             });
 
-            for (LauncherAppWidgetInfo widget : dataModel.appWidgets) {
-                if (widget.providerName.getPackageName().equals(mInstallInfo.packageName)) {
-                    widget.installProgress = mInstallInfo.progress;
-                    updates.add(widget);
-                }
-            }
+            dataModel.itemsIdMap.stream()
+                    .filter(WIDGET_FILTER)
+                    .filter(item -> mInstallInfo.user.equals(item.user))
+                    .map(item -> (LauncherAppWidgetInfo) item)
+                    .filter(widget -> widget.providerName.getPackageName()
+                            .equals(mInstallInfo.packageName))
+                    .forEach(widget -> {
+                        widget.installProgress = mInstallInfo.progress;
+                        updates.add(widget);
+                    });
 
             if (!updates.isEmpty()) {
                 taskController.scheduleCallbackTask(
diff --git a/src/com/android/launcher3/model/PackageUpdatedTask.java b/src/com/android/launcher3/model/PackageUpdatedTask.java
index d619965..1153f48 100644
--- a/src/com/android/launcher3/model/PackageUpdatedTask.java
+++ b/src/com/android/launcher3/model/PackageUpdatedTask.java
@@ -18,7 +18,9 @@
 import static com.android.launcher3.model.BgDataModel.Callbacks.FLAG_PRIVATE_PROFILE_QUIET_MODE_ENABLED;
 import static com.android.launcher3.model.BgDataModel.Callbacks.FLAG_QUIET_MODE_ENABLED;
 import static com.android.launcher3.model.BgDataModel.Callbacks.FLAG_WORK_PROFILE_QUIET_MODE_ENABLED;
+import static com.android.launcher3.model.ModelUtils.WIDGET_FILTER;
 import static com.android.launcher3.model.data.ItemInfoWithIcon.FLAG_ARCHIVED;
+import static com.android.launcher3.model.data.LauncherAppWidgetInfo.FLAG_PROVIDER_NOT_READY;
 import static com.android.launcher3.model.data.WorkspaceItemInfo.FLAG_AUTOINSTALL_ICON;
 import static com.android.launcher3.model.data.WorkspaceItemInfo.FLAG_RESTORED_ICON;
 
@@ -347,24 +349,25 @@
                     }
                 });
 
-                for (LauncherAppWidgetInfo widgetInfo : dataModel.appWidgets) {
-                    if (mUser.equals(widgetInfo.user)
-                            && widgetInfo.hasRestoreFlag(
-                            LauncherAppWidgetInfo.FLAG_PROVIDER_NOT_READY)
-                            && packageSet.contains(widgetInfo.providerName.getPackageName())) {
-                        widgetInfo.restoreStatus &=
-                                ~LauncherAppWidgetInfo.FLAG_PROVIDER_NOT_READY
-                                        & ~LauncherAppWidgetInfo.FLAG_RESTORE_STARTED;
+                dataModel.itemsIdMap.stream()
+                        .filter(WIDGET_FILTER)
+                        .filter(item -> mUser.equals(item.user))
+                        .map(item -> (LauncherAppWidgetInfo) item)
+                        .filter(widget -> widget.hasRestoreFlag(FLAG_PROVIDER_NOT_READY)
+                                && packageSet.contains(widget.providerName.getPackageName()))
+                        .forEach(widgetInfo -> {
+                            widgetInfo.restoreStatus &=
+                                    ~FLAG_PROVIDER_NOT_READY
+                                            & ~LauncherAppWidgetInfo.FLAG_RESTORE_STARTED;
 
-                        // adding this flag ensures that launcher shows 'click to setup'
-                        // if the widget has a config activity. In case there is no config
-                        // activity, it will be marked as 'restored' during bind.
-                        widgetInfo.restoreStatus |= LauncherAppWidgetInfo.FLAG_UI_NOT_READY;
+                            // adding this flag ensures that launcher shows 'click to setup'
+                            // if the widget has a config activity. In case there is no config
+                            // activity, it will be marked as 'restored' during bind.
+                            widgetInfo.restoreStatus |= LauncherAppWidgetInfo.FLAG_UI_NOT_READY;
 
-                        widgets.add(widgetInfo);
-                        taskController.getModelWriter().updateItemInDatabase(widgetInfo);
-                    }
-                }
+                            widgets.add(widgetInfo);
+                            taskController.getModelWriter().updateItemInDatabase(widgetInfo);
+                        });
             }
 
             taskController.bindUpdatedWorkspaceItems(updatedWorkspaceItems);
diff --git a/src/com/android/launcher3/model/WorkspaceItemProcessor.kt b/src/com/android/launcher3/model/WorkspaceItemProcessor.kt
index dad78dd..de1df2e 100644
--- a/src/com/android/launcher3/model/WorkspaceItemProcessor.kt
+++ b/src/com/android/launcher3/model/WorkspaceItemProcessor.kt
@@ -404,18 +404,14 @@
      * stored in the BgDataModel.
      */
     private fun processFolderOrAppPair() {
-        var collection = bgDataModel.findOrMakeFolder(c.id)
+        var collection = c.findOrMakeFolder(c.id, bgDataModel)
         // If we generated a placeholder Folder before this point, it may need to be replaced with
         // an app pair.
         if (c.itemType == Favorites.ITEM_TYPE_APP_PAIR && collection is FolderInfo) {
-            val folderInfo: FolderInfo = collection
             val newAppPair = AppPairInfo()
             // Move the placeholder's contents over to the new app pair.
-            folderInfo.getContents().forEach(newAppPair::add)
+            collection.getContents().forEach(newAppPair::add)
             collection = newAppPair
-            // Remove the placeholder and add the app pair into the data model.
-            bgDataModel.collections.remove(c.id)
-            bgDataModel.collections.put(c.id, collection)
         }
 
         c.applyCommonProperties(collection)
@@ -569,7 +565,7 @@
                 logWidgetInfo(app.invariantDeviceProfile, lapi)
             }
         }
-        c.checkAndAddItem(appWidgetInfo, bgDataModel)
+        c.checkAndAddItem(appWidgetInfo, bgDataModel, memoryLogger)
     }
 
     companion object {
diff --git a/src/com/android/launcher3/model/data/ItemInfoWithIcon.java b/src/com/android/launcher3/model/data/ItemInfoWithIcon.java
index 772ea7f..7fb0152 100644
--- a/src/com/android/launcher3/model/data/ItemInfoWithIcon.java
+++ b/src/com/android/launcher3/model/data/ItemInfoWithIcon.java
@@ -16,6 +16,8 @@
 
 package com.android.launcher3.model.data;
 
+import static com.android.launcher3.icons.BitmapInfo.FLAG_THEMED;
+
 import android.content.Context;
 import android.content.Intent;
 import android.os.Process;
@@ -23,6 +25,7 @@
 import androidx.annotation.Nullable;
 
 import com.android.launcher3.Flags;
+import com.android.launcher3.graphics.ThemeManager;
 import com.android.launcher3.icons.BitmapInfo;
 import com.android.launcher3.icons.BitmapInfo.DrawableCreationFlags;
 import com.android.launcher3.icons.FastBitmapDrawable;
@@ -320,6 +323,9 @@
      * Returns a FastBitmapDrawable with the icon and context theme applied
      */
     public FastBitmapDrawable newIcon(Context context, @DrawableCreationFlags int creationFlags) {
+        if (!ThemeManager.INSTANCE.get(context).isMonoThemeEnabled()) {
+            creationFlags &= ~FLAG_THEMED;
+        }
         FastBitmapDrawable drawable = bitmap.newIcon(context, creationFlags);
         drawable.setIsDisabled(isDisabled());
         return drawable;
diff --git a/src/com/android/launcher3/util/IntSparseArrayMap.java b/src/com/android/launcher3/util/IntSparseArrayMap.java
index 9d5391b..70f74e3 100644
--- a/src/com/android/launcher3/util/IntSparseArrayMap.java
+++ b/src/com/android/launcher3/util/IntSparseArrayMap.java
@@ -19,6 +19,8 @@
 import android.util.SparseArray;
 
 import java.util.Iterator;
+import java.util.stream.Stream;
+import java.util.stream.StreamSupport;
 
 /**
  * Extension of {@link SparseArray} with some utility methods.
@@ -43,6 +45,10 @@
         return new ValueIterator();
     }
 
+    public Stream<E> stream() {
+        return StreamSupport.stream(spliterator(), false);
+    }
+
     @Thunk class ValueIterator implements Iterator<E> {
 
         private int mNextIndex = 0;
diff --git a/src/com/android/launcher3/util/LayoutImportExportHelper.kt b/src/com/android/launcher3/util/LayoutImportExportHelper.kt
index 0df9dae..8559f3b 100644
--- a/src/com/android/launcher3/util/LayoutImportExportHelper.kt
+++ b/src/com/android/launcher3/util/LayoutImportExportHelper.kt
@@ -56,7 +56,7 @@
 
         model.enqueueModelUpdateTask { _, dataModel, _ ->
             val builder = LauncherLayoutBuilder()
-            dataModel.workspaceItems.forEach { info ->
+            dataModel.itemsIdMap.forEach { info ->
                 val loc =
                     when (info.container) {
                         CONTAINER_DESKTOP ->
@@ -67,9 +67,6 @@
                     }
                 loc.addItem(context, info)
             }
-            dataModel.appWidgets.forEach { info ->
-                builder.atWorkspace(info.cellX, info.cellY, info.screenId).addItem(context, info)
-            }
 
             val layoutXml = builder.build()
             callback(layoutXml)
diff --git a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
index d850fc6..ab0f9a7 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
@@ -577,14 +577,11 @@
     public void exitSearchMode() {
         if (!mIsInSearchMode) return;
         onSearchResults(new ArrayList<>());
-        WidgetsRecyclerView searchRecyclerView = mAdapters.get(
-                AdapterHolder.SEARCH).mWidgetsRecyclerView;
         // Remove all views when exiting the search mode; this prevents animating from stale results
         // to new ones the next time we enter search mode. By the time recycler view is hidden,
         // layout may not have happened to clear up existing results. So, instead of waiting for it
         // to happen, we clear the views here.
-        searchRecyclerView.swapAdapter(
-                searchRecyclerView.getAdapter(), /*removeAndRecycleExistingViews=*/ true);
+        mAdapters.get(AdapterHolder.SEARCH).reset();
         setViewVisibilityBasedOnSearch(/*isInSearchMode=*/ false);
         if (mHasWorkProfile) {
             mViewPager.snapToPage(AdapterHolder.PRIMARY);
@@ -613,13 +610,12 @@
             mNoWidgetsView.setVisibility(GONE);
         } else {
             mAdapters.get(AdapterHolder.SEARCH).mWidgetsRecyclerView.setVisibility(GONE);
-            mAdapters.get(getCurrentAdapterHolderType()).mWidgetsRecyclerView.setVisibility(
-                    VISIBLE);
-            if (mRecommendedWidgetsCount > 0) {
-                // Display recommendations immediately, if present, so that other parts of sticky
-                // header (e.g. personal / work tabs) don't flash in interim.
-                mWidgetRecommendationsContainer.setVisibility(VISIBLE);
-            }
+            AdapterHolder currentAdapterHolder = mAdapters.get(getCurrentAdapterHolderType());
+            // Remove all views when exiting the search mode; this prevents animating / flashing old
+            // list position / state.
+            currentAdapterHolder.reset();
+            currentAdapterHolder.mWidgetsRecyclerView.setVisibility(VISIBLE);
+            post(this::onRecommendedWidgetsBound);
             // Visibility of recycler views and headers are handled in methods below.
             onWidgetsBound();
         }
@@ -1126,6 +1122,21 @@
             mWidgetsListItemAnimator = new WidgetsListItemAnimator();
         }
 
+        /**
+         * Swaps the adapter to existing adapter to prevent the recycler view from using stale view
+         * to animate in the new visibility update.
+         *
+         * <p>For instance, when clearing search text and re-entering search with new list shouldn't
+         * use stale results to animate in new results. Alternative is setting list animators to
+         * null, but, we need animations with the default item animator.
+         */
+        private void reset() {
+            mWidgetsRecyclerView.swapAdapter(
+                    mWidgetsListAdapter,
+                    /*removeAndRecycleExistingViews=*/ true
+            );
+        }
+
         private int getEmptySpaceHeight() {
             return mStickyHeaderLayout != null ? mStickyHeaderLayout.getHeaderHeight() : 0;
         }
diff --git a/tests/multivalentTests/src/com/android/launcher3/folder/PreviewItemManagerTest.kt b/tests/multivalentTests/src/com/android/launcher3/folder/PreviewItemManagerTest.kt
index 553d08c..15accbd 100644
--- a/tests/multivalentTests/src/com/android/launcher3/folder/PreviewItemManagerTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/folder/PreviewItemManagerTest.kt
@@ -17,11 +17,14 @@
 package com.android.launcher3.folder
 
 import android.R
-import android.graphics.Bitmap
 import android.os.Process
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.launcher3.LauncherAppState
+import com.android.launcher3.LauncherPrefs
+import com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_FOLDER
+import com.android.launcher3.dagger.LauncherAppComponent
+import com.android.launcher3.dagger.LauncherAppSingleton
 import com.android.launcher3.graphics.PreloadIconDrawable
 import com.android.launcher3.graphics.ThemeManager
 import com.android.launcher3.icons.BitmapInfo
@@ -30,13 +33,14 @@
 import com.android.launcher3.icons.IconCache.ItemInfoUpdateReceiver
 import com.android.launcher3.icons.PlaceHolderIconDrawable
 import com.android.launcher3.icons.UserBadgeDrawable
-import com.android.launcher3.icons.mono.MonoThemedBitmap
 import com.android.launcher3.model.data.FolderInfo
 import com.android.launcher3.model.data.ItemInfoWithIcon.FLAG_ARCHIVED
 import com.android.launcher3.model.data.ItemInfoWithIcon.FLAG_INSTALL_SESSION_ACTIVE
 import com.android.launcher3.model.data.WorkspaceItemInfo
 import com.android.launcher3.util.ActivityContextWrapper
+import com.android.launcher3.util.AllModulesForTest
 import com.android.launcher3.util.Executors
+import com.android.launcher3.util.FakePrefsModule
 import com.android.launcher3.util.FlagOp
 import com.android.launcher3.util.LauncherLayoutBuilder
 import com.android.launcher3.util.LauncherModelHelper
@@ -44,10 +48,19 @@
 import com.android.launcher3.util.TestUtil
 import com.android.launcher3.util.UserIconInfo
 import com.google.common.truth.Truth.assertThat
+import dagger.Component
+import kotlin.annotation.AnnotationRetention.RUNTIME
+import kotlin.annotation.AnnotationTarget.FUNCTION
+import kotlin.annotation.AnnotationTarget.PROPERTY_GETTER
+import kotlin.annotation.AnnotationTarget.PROPERTY_SETTER
 import org.junit.After
 import org.junit.Before
+import org.junit.Rule
 import org.junit.Test
+import org.junit.rules.TestRule
+import org.junit.runner.Description
 import org.junit.runner.RunWith
+import org.junit.runners.model.Statement
 import org.mockito.kotlin.any
 import org.mockito.kotlin.argumentCaptor
 import org.mockito.kotlin.doReturn
@@ -61,6 +74,8 @@
 @RunWith(AndroidJUnit4::class)
 class PreviewItemManagerTest {
 
+    @get:Rule val theseStateRule = ThemeStateRule()
+
     private lateinit var previewItemManager: PreviewItemManager
     private lateinit var context: SandboxModelContext
     private lateinit var folderItems: ArrayList<WorkspaceItemInfo>
@@ -68,15 +83,14 @@
     private lateinit var folderIcon: FolderIcon
     private lateinit var iconCache: IconCache
 
-    private var defaultThemedIcons = false
-
-    private val themeManager: ThemeManager
-        get() = ThemeManager.INSTANCE.get(context)
-
     @Before
     fun setup() {
         modelHelper = LauncherModelHelper()
         context = modelHelper.sandboxContext
+        context.initDaggerComponent(DaggerPreviewItemManagerTestComponent.builder())
+        theseStateRule.themeState?.let {
+            LauncherPrefs.get(context).putSync(ThemeManager.THEMED_ICONS.to(it))
+        }
         folderIcon = FolderIcon(ActivityContextWrapper(context))
 
         val app = spy(LauncherAppState.getInstance(context))
@@ -99,27 +113,16 @@
             )
             .loadModelSync()
 
+        folderIcon.mInfo =
+            modelHelper.bgDataModel.itemsIdMap.find { it.itemType == ITEM_TYPE_FOLDER }
+                as FolderInfo
         // Use getAppContents() to "cast" contents to WorkspaceItemInfo so we can set bitmaps
-        folderItems = modelHelper.bgDataModel.collections.valueAt(0).getAppContents()
-        folderIcon.mInfo = modelHelper.bgDataModel.collections.valueAt(0) as FolderInfo
-        folderIcon.mInfo.getContents().addAll(folderItems)
-
-        // Set first icon to be themed.
-        folderItems[0].bitmap.themedBitmap =
-            MonoThemedBitmap(
-                folderItems[0].bitmap.icon,
-                Bitmap.createBitmap(10, 10, Bitmap.Config.ARGB_8888),
-            )
+        folderItems = folderIcon.mInfo.getAppContents()
 
         // Set second icon to be non-themed.
         folderItems[1].bitmap.themedBitmap = null
 
         // Set third icon to be themed with badge.
-        folderItems[2].bitmap.themedBitmap =
-            MonoThemedBitmap(
-                folderItems[2].bitmap.icon,
-                Bitmap.createBitmap(10, 10, Bitmap.Config.ARGB_8888),
-            )
         folderItems[2].bitmap =
             folderItems[2].bitmap.withFlags(profileFlagOp(UserIconInfo.TYPE_WORK))
 
@@ -127,20 +130,17 @@
         folderItems[3].bitmap =
             folderItems[3].bitmap.withFlags(profileFlagOp(UserIconInfo.TYPE_WORK))
         folderItems[3].bitmap.themedBitmap = null
-
-        defaultThemedIcons = themeManager.isMonoThemeEnabled
     }
 
     @After
     @Throws(Exception::class)
     fun tearDown() {
-        themeManager.isMonoThemeEnabled = defaultThemedIcons
         modelHelper.destroy()
     }
 
     @Test
+    @MonoThemeEnabled(true)
     fun checkThemedIconWithThemingOn_iconShouldBeThemed() {
-        themeManager.isMonoThemeEnabled = true
         val drawingParams = PreviewItemDrawingParams(0f, 0f, 0f)
 
         previewItemManager.setDrawable(drawingParams, folderItems[0])
@@ -149,8 +149,8 @@
     }
 
     @Test
+    @MonoThemeEnabled(false)
     fun checkThemedIconWithThemingOff_iconShouldNotBeThemed() {
-        themeManager.isMonoThemeEnabled = false
         val drawingParams = PreviewItemDrawingParams(0f, 0f, 0f)
 
         previewItemManager.setDrawable(drawingParams, folderItems[0])
@@ -159,8 +159,8 @@
     }
 
     @Test
+    @MonoThemeEnabled(true)
     fun checkUnthemedIconWithThemingOn_iconShouldNotBeThemed() {
-        themeManager.isMonoThemeEnabled = true
         val drawingParams = PreviewItemDrawingParams(0f, 0f, 0f)
 
         previewItemManager.setDrawable(drawingParams, folderItems[1])
@@ -169,8 +169,8 @@
     }
 
     @Test
+    @MonoThemeEnabled(false)
     fun checkUnthemedIconWithThemingOff_iconShouldNotBeThemed() {
-        themeManager.isMonoThemeEnabled = false
         val drawingParams = PreviewItemDrawingParams(0f, 0f, 0f)
 
         previewItemManager.setDrawable(drawingParams, folderItems[1])
@@ -179,8 +179,8 @@
     }
 
     @Test
+    @MonoThemeEnabled(true)
     fun checkThemedIconWithBadgeWithThemingOn_iconAndBadgeShouldBeThemed() {
-        themeManager.isMonoThemeEnabled = true
         val drawingParams = PreviewItemDrawingParams(0f, 0f, 0f)
 
         previewItemManager.setDrawable(drawingParams, folderItems[2])
@@ -192,8 +192,8 @@
     }
 
     @Test
+    @MonoThemeEnabled(true)
     fun checkUnthemedIconWithBadgeWithThemingOn_badgeShouldBeThemed() {
-        themeManager.isMonoThemeEnabled = true
         val drawingParams = PreviewItemDrawingParams(0f, 0f, 0f)
 
         previewItemManager.setDrawable(drawingParams, folderItems[3])
@@ -205,8 +205,8 @@
     }
 
     @Test
+    @MonoThemeEnabled(false)
     fun checkUnthemedIconWithBadgeWithThemingOff_iconAndBadgeShouldNotBeThemed() {
-        themeManager.isMonoThemeEnabled = false
         val drawingParams = PreviewItemDrawingParams(0f, 0f, 0f)
 
         previewItemManager.setDrawable(drawingParams, folderItems[3])
@@ -278,3 +278,28 @@
     private fun profileFlagOp(type: Int) =
         UserIconInfo(Process.myUserHandle(), type).applyBitmapInfoFlags(FlagOp.NO_OP)
 }
+
+class ThemeStateRule : TestRule {
+
+    var themeState: Boolean? = null
+
+    override fun apply(base: Statement, description: Description): Statement {
+        themeState = description.getAnnotation(MonoThemeEnabled::class.java)?.value
+        return base
+    }
+}
+
+// Annotation for tests that need to be run with quickstep enabled and disabled.
+@Retention(RUNTIME)
+@Target(FUNCTION, PROPERTY_GETTER, PROPERTY_SETTER)
+annotation class MonoThemeEnabled(val value: Boolean = false)
+
+@LauncherAppSingleton
+@Component(modules = [AllModulesForTest::class, FakePrefsModule::class])
+interface PreviewItemManagerTestComponent : LauncherAppComponent {
+
+    @Component.Builder
+    interface Builder : LauncherAppComponent.Builder {
+        override fun build(): PreviewItemManagerTestComponent
+    }
+}
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java b/tests/multivalentTests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java
index 1e2431f..0ae4d00 100644
--- a/tests/multivalentTests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java
@@ -16,6 +16,8 @@
 
 package com.android.launcher3.model;
 
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_DESKTOP;
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT;
 import static com.android.launcher3.util.LauncherModelHelper.TEST_ACTIVITY;
 import static com.android.launcher3.util.LauncherModelHelper.TEST_PACKAGE;
 
@@ -42,6 +44,8 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.util.List;
+
 /**
  * Tests for layout parser for remote layout
  */
@@ -63,14 +67,23 @@
         mModelHelper.destroy();
     }
 
+    private List<ItemInfo> getWorkspaceItems() {
+        return mModelHelper
+                .getBgDataModel()
+                .itemsIdMap
+                .stream()
+                .filter(i -> i.container == CONTAINER_DESKTOP || i.container == CONTAINER_HOTSEAT)
+                .toList();
+    }
+
     @Test
     public void testCustomProfileLoaded_with_icon_on_hotseat() throws Exception {
         writeLayoutAndLoad(new LauncherLayoutBuilder().atHotseat(0)
                 .putApp(TEST_PACKAGE, TEST_ACTIVITY));
 
         // Verify one item in hotseat
-        assertEquals(1, mModelHelper.getBgDataModel().workspaceItems.size());
-        ItemInfo info = mModelHelper.getBgDataModel().workspaceItems.get(0);
+        assertEquals(1, getWorkspaceItems().size());
+        ItemInfo info = getWorkspaceItems().get(0);
         assertEquals(LauncherSettings.Favorites.CONTAINER_HOTSEAT, info.container);
         assertEquals(LauncherSettings.Favorites.ITEM_TYPE_APPLICATION, info.itemType);
     }
@@ -84,8 +97,8 @@
                 .build());
 
         // Verify folder
-        assertEquals(1, mModelHelper.getBgDataModel().workspaceItems.size());
-        ItemInfo info = mModelHelper.getBgDataModel().workspaceItems.get(0);
+        assertEquals(1, getWorkspaceItems().size());
+        ItemInfo info = getWorkspaceItems().get(0);
         assertEquals(LauncherSettings.Favorites.ITEM_TYPE_FOLDER, info.itemType);
         assertEquals(3, ((FolderInfo) info).getContents().size());
     }
@@ -99,8 +112,8 @@
                 .build());
 
         // Verify folder
-        assertEquals(1, mModelHelper.getBgDataModel().workspaceItems.size());
-        ItemInfo info = mModelHelper.getBgDataModel().workspaceItems.get(0);
+        assertEquals(1, getWorkspaceItems().size());
+        ItemInfo info = getWorkspaceItems().get(0);
         assertEquals(LauncherSettings.Favorites.ITEM_TYPE_FOLDER, info.itemType);
         assertEquals(3, ((FolderInfo) info).getContents().size());
         assertEquals("CustomFolder", info.title.toString());
@@ -124,8 +137,8 @@
                 .putWidget(pendingAppPkg, "PlaceholderWidget", 2, 2));
 
         // Verify widget
-        assertEquals(1, mModelHelper.getBgDataModel().appWidgets.size());
-        ItemInfo info = mModelHelper.getBgDataModel().appWidgets.get(0);
+        assertEquals(1, getWorkspaceItems().size());
+        ItemInfo info = getWorkspaceItems().get(0);
         assertEquals(LauncherSettings.Favorites.ITEM_TYPE_APPWIDGET, info.itemType);
         assertEquals(2, info.spanX);
         assertEquals(2, info.spanY);
@@ -138,8 +151,8 @@
                 .putShortcut(TEST_PACKAGE, "shortcut2"));
 
         // Verify one item in hotseat
-        assertEquals(1, mModelHelper.getBgDataModel().workspaceItems.size());
-        ItemInfo info = mModelHelper.getBgDataModel().workspaceItems.get(0);
+        assertEquals(1, getWorkspaceItems().size());
+        ItemInfo info = getWorkspaceItems().get(0);
         assertEquals(LauncherSettings.Favorites.CONTAINER_HOTSEAT, info.container);
         assertEquals(LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT, info.itemType);
     }
@@ -154,8 +167,8 @@
                 .build());
 
         // Verify folder
-        assertEquals(1, mModelHelper.getBgDataModel().workspaceItems.size());
-        FolderInfo info = (FolderInfo) mModelHelper.getBgDataModel().workspaceItems.get(0);
+        assertEquals(1, getWorkspaceItems().size());
+        FolderInfo info = (FolderInfo) getWorkspaceItems().get(0);
         assertEquals(3, info.getContents().size());
 
         // Verify last icon
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/FolderIconLoadTest.kt b/tests/multivalentTests/src/com/android/launcher3/model/FolderIconLoadTest.kt
index e8f778f..f357487 100644
--- a/tests/multivalentTests/src/com/android/launcher3/model/FolderIconLoadTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/model/FolderIconLoadTest.kt
@@ -18,8 +18,10 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.launcher3.LauncherAppState
+import com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_FOLDER
 import com.android.launcher3.icons.BitmapInfo
 import com.android.launcher3.icons.waitForUpdateHandlerToFinish
+import com.android.launcher3.model.data.FolderInfo
 import com.android.launcher3.model.data.WorkspaceItemInfo
 import com.android.launcher3.util.Executors
 import com.android.launcher3.util.LauncherLayoutBuilder
@@ -149,11 +151,13 @@
         // Reload again with correct icon state
         app.model.forceReload()
         modelHelper.loadModelSync()
-        val collections = modelHelper.getBgDataModel().collections
-
-        assertThat(collections.size()).isEqualTo(1)
-        assertThat(collections.valueAt(0).getAppContents().size).isEqualTo(itemCount)
-        return collections.valueAt(0).getAppContents()
+        val collections =
+            modelHelper.bgDataModel.itemsIdMap
+                .filter { it.itemType == ITEM_TYPE_FOLDER }
+                .map { it as FolderInfo }
+        assertThat(collections.size).isEqualTo(1)
+        assertThat(collections[0].getAppContents().size).isEqualTo(itemCount)
+        return collections[0].getAppContents()
     }
 
     private fun verifyHighRes(items: ArrayList<WorkspaceItemInfo>, vararg indices: Int) {
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/WorkspaceItemProcessorTest.kt b/tests/multivalentTests/src/com/android/launcher3/model/WorkspaceItemProcessorTest.kt
index d699eee..da87dfc 100644
--- a/tests/multivalentTests/src/com/android/launcher3/model/WorkspaceItemProcessorTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/model/WorkspaceItemProcessorTest.kt
@@ -494,8 +494,7 @@
     @Test
     fun `When processing Folder then create FolderInfo and mark restored`() {
         val actualFolderInfo = FolderInfo()
-        mockBgDataModel =
-            mock<BgDataModel>().apply { whenever(findOrMakeFolder(1)).thenReturn(actualFolderInfo) }
+        mockBgDataModel = mock<BgDataModel>()
         mockCursor =
             mock<LoaderCursor>().apply {
                 user = UserHandle(0)
@@ -509,6 +508,7 @@
                 whenever(getColumnIndex(Favorites.TITLE)).thenReturn(4)
                 whenever(getString(4)).thenReturn("title")
                 whenever(options).thenReturn(5)
+                whenever(findOrMakeFolder(eq(1), any())).thenReturn(actualFolderInfo)
             }
         val expectedFolderInfo =
             FolderInfo().apply {
@@ -600,7 +600,8 @@
 
         // Then
         val widgetInfoCaptor = ArgumentCaptor.forClass(LauncherAppWidgetInfo::class.java)
-        verify(mockCursor).checkAndAddItem(widgetInfoCaptor.capture(), eq(mockBgDataModel))
+        verify(mockCursor)
+            .checkAndAddItem(widgetInfoCaptor.capture(), eq(mockBgDataModel), anyOrNull())
         val actualWidgetInfo = widgetInfoCaptor.value
         with(actualWidgetInfo) {
             assertThat(providerName).isEqualTo(expectedWidgetInfo.providerName)
@@ -655,7 +656,7 @@
         itemProcessorUnderTest.processItem()
 
         // Then
-        verify(mockCursor).checkAndAddItem(any(), any())
+        verify(mockCursor).checkAndAddItem(any(), any(), anyOrNull())
     }
 
     @Test
diff --git a/tests/src/com/android/launcher3/model/LoaderTaskTest.kt b/tests/src/com/android/launcher3/model/LoaderTaskTest.kt
index d2229c4..f04688d 100644
--- a/tests/src/com/android/launcher3/model/LoaderTaskTest.kt
+++ b/tests/src/com/android/launcher3/model/LoaderTaskTest.kt
@@ -19,6 +19,10 @@
 import com.android.launcher3.LauncherPrefs
 import com.android.launcher3.LauncherPrefs.Companion.IS_FIRST_LOAD_AFTER_RESTORE
 import com.android.launcher3.LauncherPrefs.Companion.RESTORE_DEVICE
+import com.android.launcher3.LauncherSettings.Favorites.CONTAINER_DESKTOP
+import com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT
+import com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APP_PAIR
+import com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_FOLDER
 import com.android.launcher3.icons.IconCache
 import com.android.launcher3.icons.cache.CachingLogic
 import com.android.launcher3.icons.cache.IconCacheUpdateHandler
@@ -155,9 +159,24 @@
                     widgetsFilterDataProvider,
                 )
                 .runSyncOnBackgroundThread()
-            Truth.assertThat(workspaceItems.size).isAtLeast(25)
-            Truth.assertThat(appWidgets.size).isAtLeast(7)
-            Truth.assertThat(collections.size()).isAtLeast(8)
+            Truth.assertThat(
+                    itemsIdMap
+                        .filter {
+                            it.container == CONTAINER_DESKTOP || it.container == CONTAINER_HOTSEAT
+                        }
+                        .size
+                )
+                .isAtLeast(32)
+            Truth.assertThat(itemsIdMap.filter { ModelUtils.WIDGET_FILTER.test(it) }.size)
+                .isAtLeast(7)
+            Truth.assertThat(
+                    itemsIdMap
+                        .filter {
+                            it.itemType == ITEM_TYPE_FOLDER || it.itemType == ITEM_TYPE_APP_PAIR
+                        }
+                        .size
+                )
+                .isAtLeast(8)
             Truth.assertThat(itemsIdMap.size()).isAtLeast(40)
             Truth.assertThat(widgetsModel.defaultWidgetsFilter).isNotNull()
         }
diff --git a/tests/src/com/android/launcher3/model/WorkspaceItemProcessorExtraTest.kt b/tests/src/com/android/launcher3/model/WorkspaceItemProcessorExtraTest.kt
index d553f47..8db049c 100644
--- a/tests/src/com/android/launcher3/model/WorkspaceItemProcessorExtraTest.kt
+++ b/tests/src/com/android/launcher3/model/WorkspaceItemProcessorExtraTest.kt
@@ -58,6 +58,7 @@
 import org.mockito.Mockito.RETURNS_DEEP_STUBS
 import org.mockito.Mockito.verify
 import org.mockito.kotlin.any
+import org.mockito.kotlin.anyOrNull
 import org.mockito.kotlin.doAnswer
 import org.mockito.kotlin.eq
 import org.mockito.kotlin.mock
@@ -220,7 +221,8 @@
             )
             .commit()
         val widgetInfoCaptor = ArgumentCaptor.forClass(LauncherAppWidgetInfo::class.java)
-        verify(mockCursor).checkAndAddItem(widgetInfoCaptor.capture(), eq(mockBgDataModel))
+        verify(mockCursor)
+            .checkAndAddItem(widgetInfoCaptor.capture(), eq(mockBgDataModel), anyOrNull())
         val actualWidgetInfo = widgetInfoCaptor.value
         with(actualWidgetInfo) {
             assertThat(providerName).isEqualTo(expectedWidgetInfo.providerName)
@@ -271,7 +273,7 @@
         itemProcessorUnderTest.processItem()
 
         // Then
-        verify(mockCursor).checkAndAddItem(any(), any())
+        verify(mockCursor).checkAndAddItem(any(), any(), anyOrNull())
     }
 
     private fun createWorkspaceItemProcessorUnderTest(