Merge "Do not set isImmersive in TaskbarDelegate" into main
diff --git a/quickstep/src/com/android/launcher3/taskbar/DesktopNavbarButtonsViewController.java b/quickstep/src/com/android/launcher3/taskbar/DesktopNavbarButtonsViewController.java
index 29c5204..0a9dfff 100644
--- a/quickstep/src/com/android/launcher3/taskbar/DesktopNavbarButtonsViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/DesktopNavbarButtonsViewController.java
@@ -18,12 +18,15 @@
 import static com.android.launcher3.taskbar.TaskbarNavButtonController.BUTTON_NOTIFICATIONS;
 import static com.android.launcher3.taskbar.TaskbarNavButtonController.BUTTON_QUICK_SETTINGS;
 
+import android.content.Context;
 import android.content.pm.ActivityInfo.Config;
 import android.view.LayoutInflater;
 import android.view.View;
 import android.view.ViewGroup;
 import android.widget.FrameLayout;
 
+import androidx.annotation.Nullable;
+
 import com.android.launcher3.R;
 
 /**
@@ -40,8 +43,8 @@
     private TaskbarControllers mControllers;
 
     public DesktopNavbarButtonsViewController(TaskbarActivityContext context,
-            FrameLayout navButtonsView) {
-        super(context, navButtonsView);
+            @Nullable Context navigationBarPanelContext, FrameLayout navButtonsView) {
+        super(context, navigationBarPanelContext, navButtonsView);
         mContext = context;
         mNavButtonsView = navButtonsView;
         mNavButtonContainer = mNavButtonsView.findViewById(R.id.end_nav_buttons);
diff --git a/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java b/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
index 3514447..bed4c37 100644
--- a/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
@@ -53,6 +53,7 @@
 import android.annotation.DrawableRes;
 import android.annotation.IdRes;
 import android.annotation.LayoutRes;
+import android.content.Context;
 import android.content.pm.ActivityInfo.Config;
 import android.content.res.ColorStateList;
 import android.content.res.Resources;
@@ -80,6 +81,8 @@
 import android.widget.ImageView;
 import android.widget.LinearLayout;
 
+import androidx.annotation.Nullable;
+
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.LauncherAnimUtils;
 import com.android.launcher3.R;
@@ -146,6 +149,7 @@
     private int mState;
 
     private final TaskbarActivityContext mContext;
+    private final @Nullable Context mNavigationBarPanelContext;
     private final WindowManagerProxy mWindowManagerProxy;
     private final FrameLayout mNavButtonsView;
     private final LinearLayout mNavButtonContainer;
@@ -203,8 +207,10 @@
     private final RecentsHitboxExtender mHitboxExtender = new RecentsHitboxExtender();
     private ImageView mRecentsButton;
 
-    public NavbarButtonsViewController(TaskbarActivityContext context, FrameLayout navButtonsView) {
+    public NavbarButtonsViewController(TaskbarActivityContext context,
+            @Nullable Context navigationBarPanelContext, FrameLayout navButtonsView) {
         mContext = context;
+        mNavigationBarPanelContext = navigationBarPanelContext;
         mWindowManagerProxy = WindowManagerProxy.INSTANCE.get(mContext);
         mNavButtonsView = navButtonsView;
         mNavButtonContainer = mNavButtonsView.findViewById(R.id.end_nav_buttons);
@@ -312,7 +318,8 @@
             rotationButton.hide();
             mControllers.rotationButtonController.setRotationButton(rotationButton, null);
         } else {
-            mFloatingRotationButton = new FloatingRotationButton(mContext,
+            mFloatingRotationButton = new FloatingRotationButton(
+                    ENABLE_TASKBAR_NAVBAR_UNIFICATION ? mNavigationBarPanelContext : mContext,
                     R.string.accessibility_rotate_button,
                     R.layout.rotate_suggestion,
                     R.id.rotate_suggestion,
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index e106506..9e1c077 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -145,6 +145,8 @@
 
     private static final String WINDOW_TITLE = "Taskbar";
 
+    private final @Nullable Context mNavigationBarPanelContext;
+
     private final TaskbarDragLayer mDragLayer;
     private final TaskbarControllers mControllers;
 
@@ -180,11 +182,13 @@
 
     private DeviceProfile mPersistentTaskbarDeviceProfile;
 
-    public TaskbarActivityContext(Context windowContext, DeviceProfile launcherDp,
+    public TaskbarActivityContext(Context windowContext,
+            @Nullable Context navigationBarPanelContext, DeviceProfile launcherDp,
             TaskbarNavButtonController buttonController, ScopedUnfoldTransitionProgressProvider
             unfoldTransitionProgressProvider) {
         super(windowContext);
 
+        mNavigationBarPanelContext = navigationBarPanelContext;
         applyDeviceProfile(launcherDp);
         final Resources resources = getResources();
 
@@ -258,8 +262,10 @@
                 new TaskbarDragController(this),
                 buttonController,
                 isDesktopMode
-                        ? new DesktopNavbarButtonsViewController(this, navButtonsView)
-                        : new NavbarButtonsViewController(this, navButtonsView),
+                        ? new DesktopNavbarButtonsViewController(this, mNavigationBarPanelContext,
+                                navButtonsView)
+                        : new NavbarButtonsViewController(this, mNavigationBarPanelContext,
+                                navButtonsView),
                 rotationButtonController,
                 new TaskbarDragLayerController(this, mDragLayer),
                 new TaskbarViewController(this, taskbarView),
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
index c0b07e7..bbac116 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
@@ -106,6 +106,7 @@
             Settings.Secure.NAV_BAR_KIDS_MODE);
 
     private final Context mContext;
+    private final @Nullable Context mNavigationBarPanelContext;
     private WindowManager mWindowManager;
     private FrameLayout mTaskbarRootLayout;
     private boolean mAddedWindow;
@@ -198,6 +199,9 @@
         mContext = service.createWindowContext(display,
                 ENABLE_TASKBAR_NAVBAR_UNIFICATION ? TYPE_NAVIGATION_BAR : TYPE_NAVIGATION_BAR_PANEL,
                 null);
+        mNavigationBarPanelContext = ENABLE_TASKBAR_NAVBAR_UNIFICATION
+                ? service.createWindowContext(display, TYPE_NAVIGATION_BAR_PANEL, null)
+                : null;
         if (enableTaskbarNoRecreate()) {
             mWindowManager = mContext.getSystemService(WindowManager.class);
             mTaskbarRootLayout = new FrameLayout(mContext) {
@@ -435,8 +439,9 @@
             }
 
             if (enableTaskbarNoRecreate() || mTaskbarActivityContext == null) {
-                mTaskbarActivityContext = new TaskbarActivityContext(mContext, dp,
-                        mNavButtonController, mUnfoldProgressProvider);
+                mTaskbarActivityContext = new TaskbarActivityContext(mContext,
+                        mNavigationBarPanelContext, dp, mNavButtonController,
+                        mUnfoldProgressProvider);
             } else {
                 mTaskbarActivityContext.updateDeviceProfile(dp);
             }
diff --git a/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java b/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
index 0bcdb19..6cbe171 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
+++ b/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
@@ -355,6 +355,8 @@
     public void testPressBack() throws Exception {
         InstrumentationRegistry.getInstrumentation().getUiAutomation().adoptShellPermissionIdentity(
                 READ_DEVICE_CONFIG_PERMISSION);
+        // Debug if we need to goHome to prevent wrong previous state b/315525621
+        mLauncher.goHome();
         assumeFalse(FeatureFlags.ENABLE_BACK_SWIPE_LAUNCHER_ANIMATION.get());
         mLauncher.getWorkspace().switchToAllApps();
         mLauncher.pressBack();
diff --git a/res/values/attrs.xml b/res/values/attrs.xml
index 20e7089..8d84c90 100644
--- a/res/values/attrs.xml
+++ b/res/values/attrs.xml
@@ -164,7 +164,19 @@
 
         <!-- numFolderRows & numFolderColumns defaults to numRows & numColumns, if not specified -->
         <attr name="numFolderRows" format="integer" />
+        <!-- defaults to numFolderRows, if not specified -->
+        <attr name="numFolderRowsLandscape" format="integer" />
+        <!-- defaults to numFolderRows, if not specified -->
+        <attr name="numFolderRowsTwoPanelLandscape" format="integer" />
+        <!-- defaults to numFolderRows, if not specified -->
+        <attr name="numFolderRowsTwoPanelPortrait" format="integer" />
         <attr name="numFolderColumns" format="integer" />
+        <!-- defaults to numFolderColumns, if not specified -->
+        <attr name="numFolderColumnsLandscape" format="integer" />
+        <!-- defaults to numFolderColumns, if not specified -->
+        <attr name="numFolderColumnsTwoPanelLandscape" format="integer" />
+        <!-- defaults to numFolderColumns, if not specified -->
+        <attr name="numFolderColumnsTwoPanelPortrait" format="integer" />
         <!-- Support attributes in FolderStyle -->
         <attr name="folderStyle" format="reference" />
 
diff --git a/src/com/android/launcher3/DeviceProfile.java b/src/com/android/launcher3/DeviceProfile.java
index 8aa8c1e..b6e8ec3 100644
--- a/src/com/android/launcher3/DeviceProfile.java
+++ b/src/com/android/launcher3/DeviceProfile.java
@@ -182,6 +182,8 @@
     public int cellYPaddingPx = -1;
 
     // Folder
+    public final int numFolderRows;
+    public final int numFolderColumns;
     public final float folderLabelTextScale;
     public int folderLabelTextSizePx;
     public int folderFooterHeightPx;
@@ -439,6 +441,8 @@
         }
 
         folderLabelTextScale = res.getFloat(R.dimen.folder_label_text_scale);
+        numFolderRows = inv.numFolderRows[mTypeIndex];
+        numFolderColumns = inv.numFolderColumns[mTypeIndex];
 
         if (mIsScalableGrid && inv.folderStyle != INVALID_RESOURCE_HANDLE) {
             TypedArray folderStyle = context.obtainStyledAttributes(inv.folderStyle,
@@ -645,11 +649,11 @@
                             isTwoPanels ? inv.folderSpecsTwoPanelId : inv.folderSpecsId),
                     ResponsiveSpecType.Folder);
             mResponsiveFolderWidthSpec = folderSpecs.getCalculatedSpec(responsiveAspectRatio,
-                    DimensionType.WIDTH, inv.numFolderColumns,
+                    DimensionType.WIDTH, numFolderColumns,
                     mResponsiveWorkspaceWidthSpec.getAvailableSpace(),
                     mResponsiveWorkspaceWidthSpec);
             mResponsiveFolderHeightSpec = folderSpecs.getCalculatedSpec(responsiveAspectRatio,
-                    DimensionType.HEIGHT, inv.numFolderRows,
+                    DimensionType.HEIGHT, numFolderRows,
                     mResponsiveWorkspaceHeightSpec.getAvailableSpace(),
                     mResponsiveWorkspaceHeightSpec);
 
@@ -1406,16 +1410,16 @@
         Point totalWorkspacePadding = getTotalWorkspacePadding();
 
         // Check if the folder fit within the available height.
-        float contentUsedHeight = folderCellHeightPx * inv.numFolderRows
-                + ((inv.numFolderRows - 1) * folderCellLayoutBorderSpacePx.y)
+        float contentUsedHeight = folderCellHeightPx * numFolderRows
+                + ((numFolderRows - 1) * folderCellLayoutBorderSpacePx.y)
                 + folderFooterHeightPx
                 + folderContentPaddingTop;
         int contentMaxHeight = availableHeightPx - totalWorkspacePadding.y;
         float scaleY = contentMaxHeight / contentUsedHeight;
 
         // Check if the folder fit within the available width.
-        float contentUsedWidth = folderCellWidthPx * inv.numFolderColumns
-                + ((inv.numFolderColumns - 1) * folderCellLayoutBorderSpacePx.x)
+        float contentUsedWidth = folderCellWidthPx * numFolderColumns
+                + ((numFolderColumns - 1) * folderCellLayoutBorderSpacePx.x)
                 + folderContentPaddingLeftRight * 2;
         int contentMaxWidth = availableWidthPx - totalWorkspacePadding.x;
         float scaleX = contentMaxWidth / contentUsedWidth;
@@ -2045,8 +2049,8 @@
         writer.println(prefix + pxToDpStr("iconTextSizePx", iconTextSizePx));
         writer.println(prefix + pxToDpStr("iconDrawablePaddingPx", iconDrawablePaddingPx));
 
-        writer.println(prefix + "\tinv.numFolderRows: " + inv.numFolderRows);
-        writer.println(prefix + "\tinv.numFolderColumns: " + inv.numFolderColumns);
+        writer.println(prefix + "\tnumFolderRows: " + numFolderRows);
+        writer.println(prefix + "\tnumFolderColumns: " + numFolderColumns);
         writer.println(prefix + pxToDpStr("folderCellWidthPx", folderCellWidthPx));
         writer.println(prefix + pxToDpStr("folderCellHeightPx", folderCellHeightPx));
         writer.println(prefix + pxToDpStr("folderChildIconSizePx", folderChildIconSizePx));
diff --git a/src/com/android/launcher3/InvariantDeviceProfile.java b/src/com/android/launcher3/InvariantDeviceProfile.java
index 567d0c5..dfbbcaa 100644
--- a/src/com/android/launcher3/InvariantDeviceProfile.java
+++ b/src/com/android/launcher3/InvariantDeviceProfile.java
@@ -122,8 +122,8 @@
     /**
      * Number of icons per row and column in the folder.
      */
-    public int numFolderRows;
-    public int numFolderColumns;
+    public int[] numFolderRows;
+    public int[] numFolderColumns;
     public float[] iconSize;
     public float[] iconTextSize;
     public int iconBitmapSize;
@@ -810,8 +810,8 @@
         public final int numSearchContainerColumns;
         public final int deviceCategory;
 
-        private final int numFolderRows;
-        private final int numFolderColumns;
+        private final int[] numFolderRows = new int[COUNT_SIZES];
+        private final int[] numFolderColumns = new int[COUNT_SIZES];
         private final @StyleRes int folderStyle;
         private final @StyleRes int cellStyle;
 
@@ -888,11 +888,39 @@
                     a.getResourceId(R.styleable.GridDisplayOption_inlineNavButtonsEndSpacing,
                             R.dimen.taskbar_button_margin_default);
 
-            numFolderRows = a.getInt(
+            numFolderRows[INDEX_DEFAULT] = a.getInt(
                     R.styleable.GridDisplayOption_numFolderRows, numRows);
-            numFolderColumns = a.getInt(
+            numFolderColumns[INDEX_DEFAULT] = a.getInt(
                     R.styleable.GridDisplayOption_numFolderColumns, numColumns);
 
+            if (FeatureFlags.enableResponsiveWorkspace()) {
+                numFolderRows[INDEX_LANDSCAPE] = a.getInt(
+                        R.styleable.GridDisplayOption_numFolderRowsLandscape,
+                        numFolderRows[INDEX_DEFAULT]);
+                numFolderColumns[INDEX_LANDSCAPE] = a.getInt(
+                        R.styleable.GridDisplayOption_numFolderColumnsLandscape,
+                        numFolderColumns[INDEX_DEFAULT]);
+                numFolderRows[INDEX_TWO_PANEL_PORTRAIT] = a.getInt(
+                        R.styleable.GridDisplayOption_numFolderRowsTwoPanelPortrait,
+                        numFolderRows[INDEX_DEFAULT]);
+                numFolderColumns[INDEX_TWO_PANEL_PORTRAIT] = a.getInt(
+                        R.styleable.GridDisplayOption_numFolderColumnsTwoPanelPortrait,
+                        numFolderColumns[INDEX_DEFAULT]);
+                numFolderRows[INDEX_TWO_PANEL_LANDSCAPE] = a.getInt(
+                        R.styleable.GridDisplayOption_numFolderRowsTwoPanelLandscape,
+                        numFolderRows[INDEX_DEFAULT]);
+                numFolderColumns[INDEX_TWO_PANEL_LANDSCAPE] = a.getInt(
+                        R.styleable.GridDisplayOption_numFolderColumnsTwoPanelLandscape,
+                        numFolderColumns[INDEX_DEFAULT]);
+            } else {
+                numFolderRows[INDEX_LANDSCAPE] = numFolderRows[INDEX_DEFAULT];
+                numFolderColumns[INDEX_LANDSCAPE] = numFolderColumns[INDEX_DEFAULT];
+                numFolderRows[INDEX_TWO_PANEL_PORTRAIT] = numFolderRows[INDEX_DEFAULT];
+                numFolderColumns[INDEX_TWO_PANEL_PORTRAIT] = numFolderColumns[INDEX_DEFAULT];
+                numFolderRows[INDEX_TWO_PANEL_LANDSCAPE] = numFolderRows[INDEX_DEFAULT];
+                numFolderColumns[INDEX_TWO_PANEL_LANDSCAPE] = numFolderColumns[INDEX_DEFAULT];
+            }
+
             folderStyle = a.getResourceId(R.styleable.GridDisplayOption_folderStyle,
                     INVALID_RESOURCE_HANDLE);
 
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index e301bdb..17bcfa4 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -80,7 +80,6 @@
 import static com.android.launcher3.logging.StatsLogManager.LauncherLatencyEvent.LAUNCHER_LATENCY_STARTUP_ACTIVITY_ON_CREATE;
 import static com.android.launcher3.logging.StatsLogManager.LauncherLatencyEvent.LAUNCHER_LATENCY_STARTUP_TOTAL_DURATION;
 import static com.android.launcher3.logging.StatsLogManager.LauncherLatencyEvent.LAUNCHER_LATENCY_STARTUP_VIEW_INFLATION;
-import static com.android.launcher3.logging.StatsLogManager.LauncherLatencyEvent.LAUNCHER_LATENCY_STARTUP_WORKSPACE_LOADER_ASYNC;
 import static com.android.launcher3.logging.StatsLogManager.StatsLatencyLogger.LatencyType.COLD;
 import static com.android.launcher3.logging.StatsLogManager.StatsLatencyLogger.LatencyType.COLD_DEVICE_REBOOTING;
 import static com.android.launcher3.logging.StatsLogManager.StatsLatencyLogger.LatencyType.WARM;
@@ -139,7 +138,6 @@
 import android.view.MotionEvent;
 import android.view.View;
 import android.view.ViewGroup;
-import android.view.ViewTreeObserver;
 import android.view.ViewTreeObserver.OnPreDrawListener;
 import android.view.WindowManager.LayoutParams;
 import android.view.accessibility.AccessibilityEvent;
@@ -159,7 +157,6 @@
 import com.android.launcher3.accessibility.LauncherAccessibilityDelegate;
 import com.android.launcher3.allapps.ActivityAllAppsContainerView;
 import com.android.launcher3.allapps.AllAppsRecyclerView;
-import com.android.launcher3.allapps.AllAppsStore;
 import com.android.launcher3.allapps.AllAppsTransitionController;
 import com.android.launcher3.allapps.DiscoveryBounce;
 import com.android.launcher3.anim.AnimationSuccessListener;
@@ -189,6 +186,7 @@
 import com.android.launcher3.logging.InstanceIdSequence;
 import com.android.launcher3.logging.StartupLatencyLogger;
 import com.android.launcher3.logging.StatsLogManager;
+import com.android.launcher3.logging.StatsLogManager.LauncherLatencyEvent;
 import com.android.launcher3.model.BgDataModel.Callbacks;
 import com.android.launcher3.model.ItemInstallQueue;
 import com.android.launcher3.model.ModelWriter;
@@ -791,7 +789,7 @@
         if (info.container >= 0) {
             View folderIcon = getWorkspace().getHomescreenIconByItemId(info.container);
             if (folderIcon instanceof FolderIcon && folderIcon.getTag() instanceof FolderInfo) {
-                if (new FolderGridOrganizer(getDeviceProfile().inv)
+                if (new FolderGridOrganizer(getDeviceProfile())
                         .setFolderInfo((FolderInfo) folderIcon.getTag())
                         .isItemInPreview(info.rank)) {
                     folderIcon.invalidate();
@@ -2475,39 +2473,20 @@
         }
     }
 
-    @Override
+    /**
+     * Call back when ModelCallbacks finish binding the Launcher data.
+     */
     @TargetApi(Build.VERSION_CODES.S)
-    public void onInitialBindComplete(IntSet boundPages, RunnableList pendingTasks,
-            int workspaceItemCount, boolean isBindSync) {
-        mModelCallbacks.setSynchronouslyBoundPages(boundPages);
-        mModelCallbacks.setPagesToBindSynchronously(new IntSet());
-
-        mModelCallbacks.clearPendingBinds();
-        ViewOnDrawExecutor executor = new ViewOnDrawExecutor(pendingTasks);
-        mModelCallbacks.setPendingExecutor(executor);
-        if (!isInState(ALL_APPS)) {
-            mAppsView.getAppsStore().enableDeferUpdates(AllAppsStore.DEFER_UPDATES_NEXT_DRAW);
-            pendingTasks.add(() -> mAppsView.getAppsStore().disableDeferUpdates(
-                    AllAppsStore.DEFER_UPDATES_NEXT_DRAW));
-        }
-
+    public void bindComplete(int workspaceItemCount, boolean isBindSync) {
         if (mOnInitialBindListener != null) {
             getRootView().getViewTreeObserver().removeOnPreDrawListener(mOnInitialBindListener);
             mOnInitialBindListener = null;
         }
-
-        executor.onLoadAnimationCompleted();
-        executor.attachTo(this);
-        if (Utilities.ATLEAST_S) {
-            Trace.endAsyncSection(DISPLAY_WORKSPACE_TRACE_METHOD_NAME,
-                    DISPLAY_WORKSPACE_TRACE_COOKIE);
-        }
         if (!isBindSync) {
             mStartupLatencyLogger
                     .logCardinality(workspaceItemCount)
-                    .logEnd(LAUNCHER_LATENCY_STARTUP_WORKSPACE_LOADER_ASYNC);
+                    .logEnd(LauncherLatencyEvent.LAUNCHER_LATENCY_STARTUP_WORKSPACE_LOADER_ASYNC);
         }
-
         MAIN_EXECUTOR.getHandler().postAtFrontOfQueue(() -> {
             mStartupLatencyLogger
                     .logEnd(LAUNCHER_LATENCY_STARTUP_TOTAL_DURATION)
@@ -2518,15 +2497,13 @@
                         COLD_STARTUP_TRACE_COOKIE);
             }
         });
-        getRootView().getViewTreeObserver().addOnDrawListener(
-                new ViewTreeObserver.OnDrawListener() {
-                    @Override
-                    public void onDraw() {
-                        MAIN_EXECUTOR.getHandler().postAtFrontOfQueue(
-                                () -> getRootView().getViewTreeObserver()
-                                        .removeOnDrawListener(this));
-                    }
-                });
+    }
+
+    @Override
+    public void onInitialBindComplete(IntSet boundPages, RunnableList pendingTasks,
+            int workspaceItemCount, boolean isBindSync) {
+        mModelCallbacks.onInitialBindComplete(boundPages, pendingTasks, workspaceItemCount,
+                isBindSync);
     }
 
     /**
diff --git a/src/com/android/launcher3/ModelCallbacks.kt b/src/com/android/launcher3/ModelCallbacks.kt
index c05158b..5172999 100644
--- a/src/com/android/launcher3/ModelCallbacks.kt
+++ b/src/com/android/launcher3/ModelCallbacks.kt
@@ -1,6 +1,11 @@
 package com.android.launcher3
 
+import android.annotation.TargetApi
+import android.os.Build
+import android.os.Trace
+import android.view.ViewTreeObserver.OnDrawListener
 import androidx.annotation.UiThread
+import com.android.launcher3.LauncherConstants.TraceEvents
 import com.android.launcher3.WorkspaceLayoutManager.FIRST_SCREEN_ID
 import com.android.launcher3.allapps.AllAppsStore
 import com.android.launcher3.config.FeatureFlags
@@ -13,11 +18,12 @@
 import com.android.launcher3.model.data.WorkspaceItemInfo
 import com.android.launcher3.popup.PopupContainerWithArrow
 import com.android.launcher3.util.ComponentKey
+import com.android.launcher3.util.Executors
 import com.android.launcher3.util.IntArray as LIntArray
 import com.android.launcher3.util.IntSet as LIntSet
-import com.android.launcher3.util.IntSet
 import com.android.launcher3.util.PackageUserKey
 import com.android.launcher3.util.Preconditions
+import com.android.launcher3.util.RunnableList
 import com.android.launcher3.util.TraceHelper
 import com.android.launcher3.util.ViewOnDrawExecutor
 import com.android.launcher3.widget.PendingAddWidgetInfo
@@ -64,6 +70,46 @@
         TraceHelper.INSTANCE.endSection()
     }
 
+    @TargetApi(Build.VERSION_CODES.S)
+    override fun onInitialBindComplete(
+        boundPages: LIntSet,
+        pendingTasks: RunnableList,
+        workspaceItemCount: Int,
+        isBindSync: Boolean
+    ) {
+        synchronouslyBoundPages = boundPages
+        pagesToBindSynchronously = LIntSet()
+        clearPendingBinds()
+        val executor = ViewOnDrawExecutor(pendingTasks)
+        pendingExecutor = executor
+        if (!launcher.isInState(LauncherState.ALL_APPS)) {
+            launcher.appsView.appsStore.enableDeferUpdates(AllAppsStore.DEFER_UPDATES_NEXT_DRAW)
+            pendingTasks.add {
+                launcher.appsView.appsStore.disableDeferUpdates(
+                    AllAppsStore.DEFER_UPDATES_NEXT_DRAW
+                )
+            }
+        }
+        executor.onLoadAnimationCompleted()
+        executor.attachTo(launcher)
+        if (Utilities.ATLEAST_S) {
+            Trace.endAsyncSection(
+                TraceEvents.DISPLAY_WORKSPACE_TRACE_METHOD_NAME,
+                TraceEvents.DISPLAY_WORKSPACE_TRACE_COOKIE
+            )
+        }
+        launcher.bindComplete(workspaceItemCount, isBindSync)
+        launcher.rootView.viewTreeObserver.addOnDrawListener(
+            object : OnDrawListener {
+                override fun onDraw() {
+                    Executors.MAIN_EXECUTOR.handler.postAtFrontOfQueue {
+                        launcher.rootView.getViewTreeObserver().removeOnDrawListener(this)
+                    }
+                }
+            }
+        )
+    }
+
     /**
      * Callback saying that there aren't any more items to bind.
      *
@@ -83,12 +129,12 @@
         // Since we are just resetting the current page without user interaction,
         // override the previous page so we don't log the page switch.
         launcher.workspace.setCurrentPage(currentPage, currentPage /* overridePrevPage */)
-        pagesToBindSynchronously = IntSet()
+        pagesToBindSynchronously = LIntSet()
 
         // Cache one page worth of icons
         launcher.viewCache.setCacheSize(
             R.layout.folder_application,
-            deviceProfile.inv.numFolderColumns * deviceProfile.inv.numFolderRows
+            deviceProfile.numFolderColumns * deviceProfile.numFolderRows
         )
         launcher.viewCache.setCacheSize(R.layout.folder_page, 2)
         TraceHelper.INSTANCE.endSection()
@@ -319,7 +365,7 @@
             } else {
                 // Some empty pages might have been removed while the phone was in a single panel
                 // mode, so we want to add those empty pages back.
-                val screenIds = IntSet.wrap(orderedScreenIds)
+                val screenIds = LIntSet.wrap(orderedScreenIds)
                 orderedScreenIds.forEach { screenId: Int ->
                     screenIds.add(launcher.workspace.getScreenPair(screenId))
                 }
@@ -343,7 +389,7 @@
      * if not present.
      */
     private fun filterTwoPanelScreenIds(orderedScreenIds: LIntArray): LIntArray {
-        val screenIds = IntSet.wrap(orderedScreenIds)
+        val screenIds = LIntSet.wrap(orderedScreenIds)
         orderedScreenIds
             .filter { screenId -> screenId % 2 == 1 }
             .forEach { screenId ->
diff --git a/src/com/android/launcher3/celllayout/ReorderAlgorithm.java b/src/com/android/launcher3/celllayout/ReorderAlgorithm.java
index 42b6991..8754b74 100644
--- a/src/com/android/launcher3/celllayout/ReorderAlgorithm.java
+++ b/src/com/android/launcher3/celllayout/ReorderAlgorithm.java
@@ -285,6 +285,11 @@
         return foundSolution;
     }
 
+    private void revertDir(int[] direction) {
+        direction[0] *= -1;
+        direction[1] *= -1;
+    }
+
     // This method tries to find a reordering solution which satisfies the push mechanic by trying
     // to push items in each of the cardinal directions, in an order based on the direction vector
     // passed.
@@ -293,91 +298,36 @@
         if ((Math.abs(direction[0]) + Math.abs(direction[1])) > 1) {
             // If the direction vector has two non-zero components, we try pushing
             // separately in each of the components.
-            int temp = direction[1];
-            direction[1] = 0;
-
-            if (pushViewsToTempLocation(intersectingViews, occupied, direction, ignoreView,
-                    solution)) {
-                return true;
+            int temp;
+            for (int j = 0; j < 2; j++) {
+                for (int i = 1; i >= 0; i--) {
+                    temp = direction[i];
+                    direction[i] = 0;
+                    if (pushViewsToTempLocation(intersectingViews, occupied, direction, ignoreView,
+                            solution)) {
+                        return true;
+                    }
+                    direction[i] = temp;
+                }
+                revertDir(direction);
             }
-            direction[1] = temp;
-            temp = direction[0];
-            direction[0] = 0;
-
-            if (pushViewsToTempLocation(intersectingViews, occupied, direction, ignoreView,
-                    solution)) {
-                return true;
-            }
-            // Revert the direction
-            direction[0] = temp;
-
-            // Now we try pushing in each component of the opposite direction
-            direction[0] *= -1;
-            direction[1] *= -1;
-            temp = direction[1];
-            direction[1] = 0;
-            if (pushViewsToTempLocation(intersectingViews, occupied, direction, ignoreView,
-                    solution)) {
-                return true;
-            }
-
-            direction[1] = temp;
-            temp = direction[0];
-            direction[0] = 0;
-            if (pushViewsToTempLocation(intersectingViews, occupied, direction, ignoreView,
-                    solution)) {
-                return true;
-            }
-            // revert the direction
-            direction[0] = temp;
-            direction[0] *= -1;
-            direction[1] *= -1;
-
         } else {
             // If the direction vector has a single non-zero component, we push first in the
             // direction of the vector
-            if (pushViewsToTempLocation(intersectingViews, occupied, direction, ignoreView,
-                    solution)) {
-                return true;
+            int temp;
+            for (int j = 0; j < 2; j++) {
+                for (int i = 0; i < 2; i++) {
+                    if (pushViewsToTempLocation(intersectingViews, occupied, direction, ignoreView,
+                            solution)) {
+                        return true;
+                    }
+                    revertDir(direction);
+                }
+                // Swap the components
+                temp = direction[1];
+                direction[1] = direction[0];
+                direction[0] = temp;
             }
-            // Then we try the opposite direction
-            direction[0] *= -1;
-            direction[1] *= -1;
-            if (pushViewsToTempLocation(intersectingViews, occupied, direction, ignoreView,
-                    solution)) {
-                return true;
-            }
-            // Switch the direction back
-            direction[0] *= -1;
-            direction[1] *= -1;
-
-            // If we have failed to find a push solution with the above, then we try
-            // to find a solution by pushing along the perpendicular axis.
-
-            // Swap the components
-            int temp = direction[1];
-            direction[1] = direction[0];
-            direction[0] = temp;
-            if (pushViewsToTempLocation(intersectingViews, occupied, direction, ignoreView,
-                    solution)) {
-                return true;
-            }
-
-            // Then we try the opposite direction
-            direction[0] *= -1;
-            direction[1] *= -1;
-            if (pushViewsToTempLocation(intersectingViews, occupied, direction, ignoreView,
-                    solution)) {
-                return true;
-            }
-            // Switch the direction back
-            direction[0] *= -1;
-            direction[1] *= -1;
-
-            // Swap the components back
-            temp = direction[1];
-            direction[1] = direction[0];
-            direction[0] = temp;
         }
         return false;
     }
diff --git a/src/com/android/launcher3/folder/Folder.java b/src/com/android/launcher3/folder/Folder.java
index 76e74c0..084f829 100644
--- a/src/com/android/launcher3/folder/Folder.java
+++ b/src/com/android/launcher3/folder/Folder.java
@@ -1084,7 +1084,7 @@
 
     private void updateItemLocationsInDatabaseBatch(boolean isBind) {
         FolderGridOrganizer verifier = new FolderGridOrganizer(
-                mActivityContext.getDeviceProfile().inv).setFolderInfo(mInfo);
+                mActivityContext.getDeviceProfile()).setFolderInfo(mInfo);
 
         ArrayList<ItemInfo> items = new ArrayList<>();
         int total = mInfo.contents.size();
@@ -1383,7 +1383,7 @@
     @Override
     public void onAdd(WorkspaceItemInfo item, int rank) {
         FolderGridOrganizer verifier = new FolderGridOrganizer(
-                mActivityContext.getDeviceProfile().inv).setFolderInfo(mInfo);
+                mActivityContext.getDeviceProfile()).setFolderInfo(mInfo);
         verifier.updateRankAndPos(item, rank);
         mLauncherDelegate.getModelWriter().addOrMoveItemInDatabase(item, mInfo.id, 0, item.cellX,
                 item.cellY);
diff --git a/src/com/android/launcher3/folder/FolderAnimationManager.java b/src/com/android/launcher3/folder/FolderAnimationManager.java
index 9e2e2bf..a91373b 100644
--- a/src/com/android/launcher3/folder/FolderAnimationManager.java
+++ b/src/com/android/launcher3/folder/FolderAnimationManager.java
@@ -96,7 +96,7 @@
 
         mContext = folder.getContext();
         mDeviceProfile = folder.mActivityContext.getDeviceProfile();
-        mPreviewVerifier = new FolderGridOrganizer(mDeviceProfile.inv);
+        mPreviewVerifier = new FolderGridOrganizer(mDeviceProfile);
 
         mIsOpening = isOpening;
 
diff --git a/src/com/android/launcher3/folder/FolderGridOrganizer.java b/src/com/android/launcher3/folder/FolderGridOrganizer.java
index 4be82ed..cc24761 100644
--- a/src/com/android/launcher3/folder/FolderGridOrganizer.java
+++ b/src/com/android/launcher3/folder/FolderGridOrganizer.java
@@ -20,7 +20,7 @@
 
 import android.graphics.Point;
 
-import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.model.data.FolderInfo;
 import com.android.launcher3.model.data.ItemInfo;
 
@@ -41,11 +41,13 @@
     private int mCountX;
     private int mCountY;
     private boolean mDisplayingUpperLeftQuadrant = false;
+    private static final int PREVIEW_MAX_ROWS = 2;
+    private static final int PREVIEW_MAX_COLUMNS = 2;
 
     /**
      * Note: must call {@link #setFolderInfo(FolderInfo)} manually for verifier to work.
      */
-    public FolderGridOrganizer(InvariantDeviceProfile profile) {
+    public FolderGridOrganizer(DeviceProfile profile) {
         mMaxCountX = profile.numFolderColumns;
         mMaxCountY = profile.numFolderRows;
         mMaxItemsPerPage = mMaxCountX * mMaxCountY;
@@ -127,6 +129,7 @@
 
     /**
      * Updates the item's cellX, cellY and rank corresponding to the provided rank.
+     *
      * @return true if there was any change
      */
     public boolean updateRankAndPos(ItemInfo item, int rank) {
@@ -189,7 +192,7 @@
         if (page > 0 || mDisplayingUpperLeftQuadrant) {
             int col = rank % mCountX;
             int row = rank / mCountX;
-            return col < 2 && row < 2;
+            return col < PREVIEW_MAX_COLUMNS && row < PREVIEW_MAX_ROWS;
         }
         return rank < MAX_NUM_ITEMS_IN_PREVIEW;
     }
diff --git a/src/com/android/launcher3/folder/FolderIcon.java b/src/com/android/launcher3/folder/FolderIcon.java
index 2cea974..f058ae4 100644
--- a/src/com/android/launcher3/folder/FolderIcon.java
+++ b/src/com/android/launcher3/folder/FolderIcon.java
@@ -220,7 +220,7 @@
 
         icon.setAccessibilityDelegate(activity.getAccessibilityDelegate());
 
-        icon.mPreviewVerifier = new FolderGridOrganizer(activity.getDeviceProfile().inv);
+        icon.mPreviewVerifier = new FolderGridOrganizer(activity.getDeviceProfile());
         icon.mPreviewVerifier.setFolderInfo(folderInfo);
         icon.updatePreviewItems(false);
 
diff --git a/src/com/android/launcher3/folder/FolderPagedView.java b/src/com/android/launcher3/folder/FolderPagedView.java
index 36e5e1b..f2bed92 100644
--- a/src/com/android/launcher3/folder/FolderPagedView.java
+++ b/src/com/android/launcher3/folder/FolderPagedView.java
@@ -37,8 +37,6 @@
 import com.android.launcher3.BubbleTextView;
 import com.android.launcher3.CellLayout;
 import com.android.launcher3.DeviceProfile;
-import com.android.launcher3.InvariantDeviceProfile;
-import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.PagedView;
 import com.android.launcher3.R;
 import com.android.launcher3.ShortcutAndWidgetContainer;
@@ -101,14 +99,15 @@
 
     public FolderPagedView(Context context, AttributeSet attrs) {
         super(context, attrs);
-        InvariantDeviceProfile profile = LauncherAppState.getIDP(context);
+        ActivityContext activityContext = ActivityContext.lookupContext(context);
+        DeviceProfile profile = activityContext.getDeviceProfile();
         mOrganizer = new FolderGridOrganizer(profile);
 
         mIsRtl = Utilities.isRtl(getResources());
         setImportantForAccessibility(IMPORTANT_FOR_ACCESSIBILITY_YES);
 
         mFocusIndicatorHelper = new ViewGroupFocusHelper(this);
-        mViewCache = ActivityContext.lookupContext(context).getViewCache();
+        mViewCache = activityContext.getViewCache();
     }
 
     public void setFolder(Folder folder) {
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index f4ce360..6ea3e8a 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -481,16 +481,17 @@
             mItemsDeleted = c.commitDeleted();
 
             // Sort the folder items, update ranks, and make sure all preview items are high res.
-            FolderGridOrganizer verifier =
-                    new FolderGridOrganizer(mApp.getInvariantDeviceProfile());
+            List<FolderGridOrganizer> verifiers =
+                    mApp.getInvariantDeviceProfile().supportedProfiles.stream().map(
+                            FolderGridOrganizer::new).toList();
             for (FolderInfo folder : mBgDataModel.folders) {
                 Collections.sort(folder.contents, Folder.ITEM_POS_COMPARATOR);
-                verifier.setFolderInfo(folder);
+                verifiers.forEach(verifier -> verifier.setFolderInfo(folder));
                 int size = folder.contents.size();
 
                 // Update ranks here to ensure there are no gaps caused by removed folder items.
-                // Ranks are the source of truth for folder items, so cellX and cellY can be ignored
-                // for now. Database will be updated once user manually modifies folder.
+                // Ranks are the source of truth for folder items, so cellX and cellY can be
+                // ignored for now. Database will be updated once user manually modifies folder.
                 for (int rank = 0; rank < size; ++rank) {
                     WorkspaceItemInfo info = folder.contents.get(rank);
                     // rank is used differently in app pairs, so don't reset
@@ -498,9 +499,9 @@
                         info.rank = rank;
                     }
 
-                    if (info.usingLowResIcon()
-                            && info.itemType == Favorites.ITEM_TYPE_APPLICATION
-                            && verifier.isItemInPreview(info.rank)) {
+                    if (info.usingLowResIcon() && info.itemType == Favorites.ITEM_TYPE_APPLICATION
+                            && verifiers.stream().anyMatch(
+                                verifier -> verifier.isItemInPreview(info.rank))) {
                         mIconCache.getTitleAndIcon(info, false);
                     }
                 }
diff --git a/tests/assets/dumpTests/DeviceProfileDumpTest/phonePortrait.txt b/tests/assets/dumpTests/DeviceProfileDumpTest/phonePortrait.txt
index 5f5cf5e..18752e9 100644
--- a/tests/assets/dumpTests/DeviceProfileDumpTest/phonePortrait.txt
+++ b/tests/assets/dumpTests/DeviceProfileDumpTest/phonePortrait.txt
@@ -38,8 +38,8 @@
 	iconSizePx: 147.0px (56.0dp)
 	iconTextSizePx: 38.0px (14.476191dp)
 	iconDrawablePaddingPx: 12.0px (4.571429dp)
-	inv.numFolderRows: 4
-	inv.numFolderColumns: 4
+	numFolderRows: 4
+	numFolderColumns: 4
 	folderCellWidthPx: 195.0px (74.28571dp)
 	folderCellHeightPx: 230.0px (87.61905dp)
 	folderChildIconSizePx: 147.0px (56.0dp)
diff --git a/tests/assets/dumpTests/DeviceProfileDumpTest/phonePortrait3Button.txt b/tests/assets/dumpTests/DeviceProfileDumpTest/phonePortrait3Button.txt
index 02bab0e..c0de8d8 100644
--- a/tests/assets/dumpTests/DeviceProfileDumpTest/phonePortrait3Button.txt
+++ b/tests/assets/dumpTests/DeviceProfileDumpTest/phonePortrait3Button.txt
@@ -38,8 +38,8 @@
 	iconSizePx: 147.0px (56.0dp)
 	iconTextSizePx: 38.0px (14.476191dp)
 	iconDrawablePaddingPx: 12.0px (4.571429dp)
-	inv.numFolderRows: 4
-	inv.numFolderColumns: 4
+	numFolderRows: 4
+	numFolderColumns: 4
 	folderCellWidthPx: 195.0px (74.28571dp)
 	folderCellHeightPx: 230.0px (87.61905dp)
 	folderChildIconSizePx: 147.0px (56.0dp)
diff --git a/tests/assets/dumpTests/DeviceProfileDumpTest/phoneVerticalBar.txt b/tests/assets/dumpTests/DeviceProfileDumpTest/phoneVerticalBar.txt
index 1ade779..920ba6f 100644
--- a/tests/assets/dumpTests/DeviceProfileDumpTest/phoneVerticalBar.txt
+++ b/tests/assets/dumpTests/DeviceProfileDumpTest/phoneVerticalBar.txt
@@ -38,8 +38,8 @@
 	iconSizePx: 147.0px (56.0dp)
 	iconTextSizePx: 0.0px (0.0dp)
 	iconDrawablePaddingPx: 0.0px (0.0dp)
-	inv.numFolderRows: 4
-	inv.numFolderColumns: 4
+	numFolderRows: 4
+	numFolderColumns: 4
 	folderCellWidthPx: 163.0px (62.095238dp)
 	folderCellHeightPx: 192.0px (73.14286dp)
 	folderChildIconSizePx: 123.0px (46.857143dp)
diff --git a/tests/assets/dumpTests/DeviceProfileDumpTest/phoneVerticalBar3Button.txt b/tests/assets/dumpTests/DeviceProfileDumpTest/phoneVerticalBar3Button.txt
index b0b745d..65460ec 100644
--- a/tests/assets/dumpTests/DeviceProfileDumpTest/phoneVerticalBar3Button.txt
+++ b/tests/assets/dumpTests/DeviceProfileDumpTest/phoneVerticalBar3Button.txt
@@ -38,8 +38,8 @@
 	iconSizePx: 147.0px (56.0dp)
 	iconTextSizePx: 0.0px (0.0dp)
 	iconDrawablePaddingPx: 0.0px (0.0dp)
-	inv.numFolderRows: 4
-	inv.numFolderColumns: 4
+	numFolderRows: 4
+	numFolderColumns: 4
 	folderCellWidthPx: 173.0px (65.90476dp)
 	folderCellHeightPx: 205.0px (78.09524dp)
 	folderChildIconSizePx: 131.0px (49.904762dp)
diff --git a/tests/assets/dumpTests/DeviceProfileDumpTest/tabletLandscape.txt b/tests/assets/dumpTests/DeviceProfileDumpTest/tabletLandscape.txt
index d7f3c1a..1781673 100644
--- a/tests/assets/dumpTests/DeviceProfileDumpTest/tabletLandscape.txt
+++ b/tests/assets/dumpTests/DeviceProfileDumpTest/tabletLandscape.txt
@@ -38,8 +38,8 @@
 	iconSizePx: 120.0px (60.0dp)
 	iconTextSizePx: 28.0px (14.0dp)
 	iconDrawablePaddingPx: 9.0px (4.5dp)
-	inv.numFolderRows: 3
-	inv.numFolderColumns: 3
+	numFolderRows: 3
+	numFolderColumns: 3
 	folderCellWidthPx: 240.0px (120.0dp)
 	folderCellHeightPx: 208.0px (104.0dp)
 	folderChildIconSizePx: 120.0px (60.0dp)
diff --git a/tests/assets/dumpTests/DeviceProfileDumpTest/tabletLandscape3Button.txt b/tests/assets/dumpTests/DeviceProfileDumpTest/tabletLandscape3Button.txt
index 20a2a99..bd9e267 100644
--- a/tests/assets/dumpTests/DeviceProfileDumpTest/tabletLandscape3Button.txt
+++ b/tests/assets/dumpTests/DeviceProfileDumpTest/tabletLandscape3Button.txt
@@ -38,8 +38,8 @@
 	iconSizePx: 120.0px (60.0dp)
 	iconTextSizePx: 28.0px (14.0dp)
 	iconDrawablePaddingPx: 9.0px (4.5dp)
-	inv.numFolderRows: 3
-	inv.numFolderColumns: 3
+	numFolderRows: 3
+	numFolderColumns: 3
 	folderCellWidthPx: 240.0px (120.0dp)
 	folderCellHeightPx: 208.0px (104.0dp)
 	folderChildIconSizePx: 120.0px (60.0dp)
diff --git a/tests/assets/dumpTests/DeviceProfileDumpTest/tabletPortrait.txt b/tests/assets/dumpTests/DeviceProfileDumpTest/tabletPortrait.txt
index 94022e4..e983ef7 100644
--- a/tests/assets/dumpTests/DeviceProfileDumpTest/tabletPortrait.txt
+++ b/tests/assets/dumpTests/DeviceProfileDumpTest/tabletPortrait.txt
@@ -38,8 +38,8 @@
 	iconSizePx: 120.0px (60.0dp)
 	iconTextSizePx: 28.0px (14.0dp)
 	iconDrawablePaddingPx: 9.0px (4.5dp)
-	inv.numFolderRows: 3
-	inv.numFolderColumns: 3
+	numFolderRows: 3
+	numFolderColumns: 3
 	folderCellWidthPx: 204.0px (102.0dp)
 	folderCellHeightPx: 240.0px (120.0dp)
 	folderChildIconSizePx: 120.0px (60.0dp)
diff --git a/tests/assets/dumpTests/DeviceProfileDumpTest/tabletPortrait3Button.txt b/tests/assets/dumpTests/DeviceProfileDumpTest/tabletPortrait3Button.txt
index 7977204..aa92838 100644
--- a/tests/assets/dumpTests/DeviceProfileDumpTest/tabletPortrait3Button.txt
+++ b/tests/assets/dumpTests/DeviceProfileDumpTest/tabletPortrait3Button.txt
@@ -38,8 +38,8 @@
 	iconSizePx: 120.0px (60.0dp)
 	iconTextSizePx: 28.0px (14.0dp)
 	iconDrawablePaddingPx: 9.0px (4.5dp)
-	inv.numFolderRows: 3
-	inv.numFolderColumns: 3
+	numFolderRows: 3
+	numFolderColumns: 3
 	folderCellWidthPx: 204.0px (102.0dp)
 	folderCellHeightPx: 240.0px (120.0dp)
 	folderChildIconSizePx: 120.0px (60.0dp)
diff --git a/tests/assets/dumpTests/DeviceProfileDumpTest/twoPanelLandscape.txt b/tests/assets/dumpTests/DeviceProfileDumpTest/twoPanelLandscape.txt
index 0b17996..43e4a60 100644
--- a/tests/assets/dumpTests/DeviceProfileDumpTest/twoPanelLandscape.txt
+++ b/tests/assets/dumpTests/DeviceProfileDumpTest/twoPanelLandscape.txt
@@ -38,8 +38,8 @@
 	iconSizePx: 141.0px (53.714287dp)
 	iconTextSizePx: 34.0px (12.952381dp)
 	iconDrawablePaddingPx: 13.0px (4.952381dp)
-	inv.numFolderRows: 3
-	inv.numFolderColumns: 4
+	numFolderRows: 3
+	numFolderColumns: 4
 	folderCellWidthPx: 189.0px (72.0dp)
 	folderCellHeightPx: 219.0px (83.42857dp)
 	folderChildIconSizePx: 141.0px (53.714287dp)
diff --git a/tests/assets/dumpTests/DeviceProfileDumpTest/twoPanelLandscape3Button.txt b/tests/assets/dumpTests/DeviceProfileDumpTest/twoPanelLandscape3Button.txt
index 71fffe8..e7ea839 100644
--- a/tests/assets/dumpTests/DeviceProfileDumpTest/twoPanelLandscape3Button.txt
+++ b/tests/assets/dumpTests/DeviceProfileDumpTest/twoPanelLandscape3Button.txt
@@ -38,8 +38,8 @@
 	iconSizePx: 141.0px (53.714287dp)
 	iconTextSizePx: 34.0px (12.952381dp)
 	iconDrawablePaddingPx: 13.0px (4.952381dp)
-	inv.numFolderRows: 3
-	inv.numFolderColumns: 4
+	numFolderRows: 3
+	numFolderColumns: 4
 	folderCellWidthPx: 189.0px (72.0dp)
 	folderCellHeightPx: 219.0px (83.42857dp)
 	folderChildIconSizePx: 141.0px (53.714287dp)
diff --git a/tests/assets/dumpTests/DeviceProfileDumpTest/twoPanelPortrait.txt b/tests/assets/dumpTests/DeviceProfileDumpTest/twoPanelPortrait.txt
index 5da4ed0..043380c 100644
--- a/tests/assets/dumpTests/DeviceProfileDumpTest/twoPanelPortrait.txt
+++ b/tests/assets/dumpTests/DeviceProfileDumpTest/twoPanelPortrait.txt
@@ -38,8 +38,8 @@
 	iconSizePx: 141.0px (53.714287dp)
 	iconTextSizePx: 34.0px (12.952381dp)
 	iconDrawablePaddingPx: 13.0px (4.952381dp)
-	inv.numFolderRows: 3
-	inv.numFolderColumns: 4
+	numFolderRows: 3
+	numFolderColumns: 4
 	folderCellWidthPx: 189.0px (72.0dp)
 	folderCellHeightPx: 219.0px (83.42857dp)
 	folderChildIconSizePx: 141.0px (53.714287dp)
diff --git a/tests/assets/dumpTests/DeviceProfileDumpTest/twoPanelPortrait3Button.txt b/tests/assets/dumpTests/DeviceProfileDumpTest/twoPanelPortrait3Button.txt
index 359e530..a1b3e95 100644
--- a/tests/assets/dumpTests/DeviceProfileDumpTest/twoPanelPortrait3Button.txt
+++ b/tests/assets/dumpTests/DeviceProfileDumpTest/twoPanelPortrait3Button.txt
@@ -38,8 +38,8 @@
 	iconSizePx: 141.0px (53.714287dp)
 	iconTextSizePx: 34.0px (12.952381dp)
 	iconDrawablePaddingPx: 13.0px (4.952381dp)
-	inv.numFolderRows: 3
-	inv.numFolderColumns: 4
+	numFolderRows: 3
+	numFolderColumns: 4
 	folderCellWidthPx: 189.0px (72.0dp)
 	folderCellHeightPx: 219.0px (83.42857dp)
 	folderChildIconSizePx: 141.0px (53.714287dp)
diff --git a/tests/src/com/android/launcher3/FakeInvariantDeviceProfileTest.kt b/tests/src/com/android/launcher3/FakeInvariantDeviceProfileTest.kt
index a421006..30b5663 100644
--- a/tests/src/com/android/launcher3/FakeInvariantDeviceProfileTest.kt
+++ b/tests/src/com/android/launcher3/FakeInvariantDeviceProfileTest.kt
@@ -121,8 +121,8 @@
                     listOf(PointF(16f, 16f), PointF(16f, 16f), PointF(16f, 16f), PointF(16f, 16f))
                         .toTypedArray()
 
-                numFolderRows = 3
-                numFolderColumns = 3
+                numFolderRows = intArrayOf(3, 3, 3, 3)
+                numFolderColumns = intArrayOf(3, 3, 3, 3)
                 folderStyle = R.style.FolderStyleDefault
 
                 inlineNavButtonsEndSpacing = R.dimen.taskbar_button_margin_split
@@ -204,8 +204,8 @@
                     listOf(PointF(16f, 64f), PointF(64f, 16f), PointF(16f, 64f), PointF(16f, 64f))
                         .toTypedArray()
 
-                numFolderRows = 3
-                numFolderColumns = 3
+                numFolderRows = intArrayOf(3, 3, 3, 3)
+                numFolderColumns = intArrayOf(3, 3, 3, 3)
                 folderStyle = R.style.FolderStyleDefault
 
                 inlineNavButtonsEndSpacing = R.dimen.taskbar_button_margin_6_5
@@ -288,8 +288,8 @@
                     listOf(PointF(16f, 16f), PointF(16f, 16f), PointF(16f, 20f), PointF(20f, 20f))
                         .toTypedArray()
 
-                numFolderRows = 3
-                numFolderColumns = 3
+                numFolderRows = intArrayOf(3, 3, 3, 3)
+                numFolderColumns = intArrayOf(3, 3, 3, 3)
                 folderStyle = R.style.FolderStyleDefault
 
                 inlineNavButtonsEndSpacing = R.dimen.taskbar_button_margin_split
diff --git a/tests/src/com/android/launcher3/model/FolderIconLoadTest.kt b/tests/src/com/android/launcher3/model/FolderIconLoadTest.kt
index 21bd900..60a4d2d 100644
--- a/tests/src/com/android/launcher3/model/FolderIconLoadTest.kt
+++ b/tests/src/com/android/launcher3/model/FolderIconLoadTest.kt
@@ -24,9 +24,10 @@
 import com.android.launcher3.util.LauncherModelHelper
 import com.android.launcher3.util.LauncherModelHelper.*
 import com.android.launcher3.util.TestUtil
+import com.google.common.truth.Truth.assertThat
+import com.google.common.truth.Truth.assertWithMessage
 import java.util.concurrent.CountDownLatch
 import org.junit.After
-import org.junit.Assert
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -71,7 +72,7 @@
     @Throws(Exception::class)
     fun folderLoadedWithHighRes_2x2() {
         val items = setupAndLoadFolder(4)
-        Assert.assertEquals(4, items.size.toLong())
+        assertThat(items.size).isEqualTo(4)
         verifyHighRes(items, 0, 1, 2, 3)
     }
 
@@ -79,7 +80,7 @@
     @Throws(Exception::class)
     fun folderLoadedWithHighRes_3x2() {
         val items = setupAndLoadFolder(6)
-        Assert.assertEquals(6, items.size.toLong())
+        assertThat(items.size).isEqualTo(6)
         verifyHighRes(items, 0, 1, 3, 4)
         verifyLowRes(items, 2, 5)
     }
@@ -88,8 +89,10 @@
     @Throws(Exception::class)
     fun folderLoadedWithHighRes_max_3x3() {
         val idp = LauncherAppState.getIDP(modelHelper.sandboxContext)
-        idp.numFolderColumns = 3
-        idp.numFolderRows = 3
+        idp.numFolderColumns = intArrayOf(3, 3, 3, 3)
+        idp.numFolderRows = intArrayOf(3, 3, 3, 3)
+        recreateSupportedDeviceProfiles()
+
         val items = setupAndLoadFolder(14)
         verifyHighRes(items, 0, 1, 3, 4)
         verifyLowRes(items, 2, 5, 6, 7, 8, 9, 10, 11, 12, 13)
@@ -99,13 +102,28 @@
     @Throws(Exception::class)
     fun folderLoadedWithHighRes_max_4x4() {
         val idp = LauncherAppState.getIDP(modelHelper.sandboxContext)
-        idp.numFolderColumns = 4
-        idp.numFolderRows = 4
+        idp.numFolderColumns = intArrayOf(4, 4, 4, 4)
+        idp.numFolderRows = intArrayOf(4, 4, 4, 4)
+        recreateSupportedDeviceProfiles()
+
         val items = setupAndLoadFolder(14)
         verifyHighRes(items, 0, 1, 4, 5)
         verifyLowRes(items, 2, 3, 6, 7, 8, 9, 10, 11, 12, 13)
     }
 
+    @Test
+    @Throws(Exception::class)
+    fun folderLoadedWithHighRes_differentFolderConfigurations() {
+        val idp = LauncherAppState.getIDP(modelHelper.sandboxContext)
+        idp.numFolderColumns = intArrayOf(4, 3, 4, 4)
+        idp.numFolderRows = intArrayOf(4, 3, 4, 4)
+        recreateSupportedDeviceProfiles()
+
+        val items = setupAndLoadFolder(14)
+        verifyHighRes(items, 0, 1, 3, 4, 5)
+        verifyLowRes(items, 2, 6, 7, 8, 9, 10, 11, 12, 13)
+    }
+
     @Throws(Exception::class)
     private fun setupAndLoadFolder(itemCount: Int): ArrayList<WorkspaceItemInfo> {
         val builder =
@@ -113,9 +131,7 @@
                 .atWorkspace(0, 0, 1)
                 .putFolder("Sample")
                 .apply {
-                    for (i in 0..itemCount - 1) {
-                        this.addApp(TEST_PACKAGE, uniqueActivities[i])
-                    }
+                    for (i in 0..itemCount - 1) this.addApp(TEST_PACKAGE, uniqueActivities[i])
                 }
                 .build()
 
@@ -136,20 +152,33 @@
         app.model.forceReload()
         modelHelper.loadModelSync()
         val folders = modelHelper.getBgDataModel().folders
-        Assert.assertEquals(1, folders.size())
-        Assert.assertEquals(itemCount, folders.valueAt(0).contents.size)
+
+        assertThat(folders.size()).isEqualTo(1)
+        assertThat(folders.valueAt(0).contents.size).isEqualTo(itemCount)
         return folders.valueAt(0).contents
     }
 
     private fun verifyHighRes(items: ArrayList<WorkspaceItemInfo>, vararg indices: Int) {
         for (index in indices) {
-            Assert.assertFalse("Index $index was not highRes", items[index].bitmap.isNullOrLowRes)
+            assertWithMessage("Index $index was not highRes")
+                .that(items[index].bitmap.isNullOrLowRes)
+                .isFalse()
         }
     }
 
     private fun verifyLowRes(items: ArrayList<WorkspaceItemInfo>, vararg indices: Int) {
         for (index in indices) {
-            Assert.assertTrue("Index $index was not lowRes", items[index].bitmap.isNullOrLowRes)
+            assertWithMessage("Index $index was not lowRes")
+                .that(items[index].bitmap.isNullOrLowRes)
+                .isTrue()
         }
     }
+
+    /** Recreate DeviceProfiles after changing InvariantDeviceProfile */
+    private fun recreateSupportedDeviceProfiles() {
+        LauncherAppState.getIDP(modelHelper.sandboxContext).supportedProfiles =
+            LauncherAppState.getIDP(modelHelper.sandboxContext).supportedProfiles.map {
+                it.copy(modelHelper.sandboxContext)
+            }
+    }
 }
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index 8ad2249..95ed401 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -131,6 +131,8 @@
     /** Detects activity leaks and throws an exception if a leak is found. */
     public static void checkDetectedLeaks(LauncherInstrumentation launcher,
             boolean requireOneActiveActivityUnused) {
+        if (TestStabilityRule.isPresubmit()) return; // b/313501215
+
         final boolean requireOneActiveActivity =
                 false; // workaround for leaks when there is an unexpected Recents activity
 
diff --git a/tests/src/com/android/launcher3/ui/widget/TaplAddConfigWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/TaplAddConfigWidgetTest.java
index 3263d60..d96287f 100644
--- a/tests/src/com/android/launcher3/ui/widget/TaplAddConfigWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/TaplAddConfigWidgetTest.java
@@ -40,7 +40,6 @@
 import com.android.launcher3.ui.TestViewHelpers;
 import com.android.launcher3.util.Wait;
 import com.android.launcher3.util.rule.ShellCommandRule;
-import com.android.launcher3.util.rule.TestStabilityRule;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 
 import org.junit.Before;
@@ -120,8 +119,7 @@
     }
 
     private void setResultAndWaitForAnimation(boolean success) {
-        if (mLauncher.isLauncher3()
-                || TestStabilityRule.isPresubmit() /* b/313926097 */) {
+        if (mLauncher.isLauncher3()) {
             setResult(success);
         } else {
             mLauncher.executeAndWaitForWallpaperAnimation(
diff --git a/tests/src/com/android/launcher3/ui/widget/TaplWidgetPickerTest.java b/tests/src/com/android/launcher3/ui/widget/TaplWidgetPickerTest.java
index 465e9b4..bc73683 100644
--- a/tests/src/com/android/launcher3/ui/widget/TaplWidgetPickerTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/TaplWidgetPickerTest.java
@@ -15,10 +15,6 @@
  */
 package com.android.launcher3.ui.widget;
 
-import static com.android.launcher3.ui.AbstractLauncherUiTest.initialize;
-import static com.android.launcher3.util.rule.TestStabilityRule.LOCAL;
-import static com.android.launcher3.util.rule.TestStabilityRule.PLATFORM_POSTSUBMIT;
-
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
@@ -31,7 +27,6 @@
 import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.ui.PortraitLandscapeRunner.PortraitLandscape;
 import com.android.launcher3.util.rule.ScreenRecordRule.ScreenRecord;
-import com.android.launcher3.util.rule.TestStabilityRule.Stability;
 import com.android.launcher3.widget.picker.WidgetsFullSheet;
 import com.android.launcher3.widget.picker.WidgetsRecyclerView;
 
@@ -65,11 +60,9 @@
      * Open Widget picker, make sure the widget picker can scroll and then go to home screen.
      */
     @Test
-    @Stability(flavors = LOCAL | PLATFORM_POSTSUBMIT) // b/303263644
     @ScreenRecord
     @PortraitLandscape
     public void testWidgets() {
-        // Testing if this will fix b/303263644
         mLauncher.goHome();
         // Test opening widgets.
         executeOnLauncher(launcher ->
diff --git a/tests/src/com/android/launcher3/util/viewcapture_analysis/ViewCaptureAnalyzer.java b/tests/src/com/android/launcher3/util/viewcapture_analysis/ViewCaptureAnalyzer.java
index ba02473..b27ccbf 100644
--- a/tests/src/com/android/launcher3/util/viewcapture_analysis/ViewCaptureAnalyzer.java
+++ b/tests/src/com/android/launcher3/util/viewcapture_analysis/ViewCaptureAnalyzer.java
@@ -35,7 +35,7 @@
 
     // All detectors. They will be invoked in the order listed here.
     private static final AnomalyDetector[] ANOMALY_DETECTORS = {
-            new AlphaJumpDetector(),
+//            new AlphaJumpDetector(), // b/309014345
 //            new FlashDetector(), // b/309014345
             new PositionJumpDetector()
     };