diff --git a/quickstep/res/values/dimens.xml b/quickstep/res/values/dimens.xml
index 50453ac..a9e0f7d 100644
--- a/quickstep/res/values/dimens.xml
+++ b/quickstep/res/values/dimens.xml
@@ -16,7 +16,7 @@
 
 <resources>
     <dimen name="task_thumbnail_icon_size">48dp</dimen>
-    <dimen name="task_thumbnail_icon_size_grid">32dp</dimen>
+    <dimen name="task_thumbnail_icon_size_grid">40dp</dimen>
     <!-- For screens without rounded corners -->
     <dimen name="task_corner_radius_small">2dp</dimen>
     <!-- For Launchers that want to override the default dialog corner radius -->
@@ -35,10 +35,10 @@
     <dimen name="overview_actions_horizontal_margin">16dp</dimen>
 
     <dimen name="overview_grid_top_margin">77dp</dimen>
-    <dimen name="overview_grid_bottom_margin">90dp</dimen>
+    <dimen name="overview_grid_bottom_margin">70dp</dimen>
     <dimen name="overview_grid_side_margin">54dp</dimen>
     <dimen name="overview_grid_row_spacing">42dp</dimen>
-    <dimen name="overview_grid_focus_vertical_margin">90dp</dimen>
+    <dimen name="overview_grid_focus_vertical_margin">40dp</dimen>
     <dimen name="split_placeholder_size">110dp</dimen>
 
     <!-- These speeds are in dp/s -->
diff --git a/quickstep/robolectric_tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java b/quickstep/robolectric_tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
index 7c97b93..5471e49 100644
--- a/quickstep/robolectric_tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
+++ b/quickstep/robolectric_tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
@@ -48,6 +48,7 @@
 import com.android.launcher3.shadows.ShadowDeviceFlag;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.IntArray;
+import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.ItemInfoMatcher;
 import com.android.launcher3.util.LauncherModelHelper;
 import com.android.launcher3.util.ViewOnDrawExecutor;
@@ -239,8 +240,8 @@
         }
 
         @Override
-        public int getPageToBindSynchronously() {
-            return 0;
+        public IntSet getPagesToBindSynchronously() {
+            return IntSet.wrap(0);
         }
 
         @Override
@@ -259,7 +260,7 @@
         public void finishFirstPageBind(ViewOnDrawExecutor executor) { }
 
         @Override
-        public void finishBindingItems(int pageBoundFirst) { }
+        public void finishBindingItems(IntSet pagesBoundFirst) { }
 
         @Override
         public void preAddApps() { }
@@ -287,7 +288,7 @@
         public void bindAllWidgets(List<WidgetsListBaseEntry> widgets) { }
 
         @Override
-        public void onPageBoundSynchronously(int page) { }
+        public void onPagesBoundSynchronously(IntSet pages) { }
 
         @Override
         public void executeOnNextDraw(ViewOnDrawExecutor executor) { }
diff --git a/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java b/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
index 0b41f15..67a4225 100644
--- a/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
@@ -55,6 +55,7 @@
 import com.android.launcher3.taskbar.TaskbarStateHandler;
 import com.android.launcher3.uioverrides.RecentsViewStateController;
 import com.android.launcher3.util.ActivityOptionsWrapper;
+import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.ObjectWrapper;
 import com.android.launcher3.util.UiThreadHelper;
 import com.android.quickstep.RecentsModel;
@@ -410,8 +411,8 @@
     }
 
     @Override
-    public void finishBindingItems(int pageBoundFirst) {
-        super.finishBindingItems(pageBoundFirst);
+    public void finishBindingItems(IntSet pagesBoundFirst) {
+        super.finishBindingItems(pagesBoundFirst);
         // Instantiate and initialize WellbeingModel now that its loading won't interfere with
         // populating workspace.
         // TODO: Find a better place for this
diff --git a/quickstep/src/com/android/launcher3/uioverrides/RecentsViewStateController.java b/quickstep/src/com/android/launcher3/uioverrides/RecentsViewStateController.java
index 996d36a..6cad3dd 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/RecentsViewStateController.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/RecentsViewStateController.java
@@ -106,8 +106,7 @@
         float clearAllButtonAlpha = state.areElementsVisible(mLauncher, CLEAR_ALL_BUTTON) ? 1 : 0;
         propertySetter.setFloat(mRecentsView.getClearAllButton(), ClearAllButton.VISIBILITY_ALPHA,
                 clearAllButtonAlpha, LINEAR);
-        float overviewButtonAlpha = state.areElementsVisible(mLauncher, OVERVIEW_ACTIONS)
-                && mRecentsView.shouldShowOverviewActionsForState(state) ? 1 : 0;
+        float overviewButtonAlpha = state.areElementsVisible(mLauncher, OVERVIEW_ACTIONS) ? 1 : 0;
         propertySetter.setFloat(mLauncher.getActionsView().getVisibilityAlpha(),
                 MultiValueAlpha.VALUE, overviewButtonAlpha, config.getInterpolator(
                         ANIM_OVERVIEW_ACTIONS_FADE, LINEAR));
diff --git a/quickstep/src/com/android/quickstep/BaseActivityInterface.java b/quickstep/src/com/android/quickstep/BaseActivityInterface.java
index 86bf119..270a790 100644
--- a/quickstep/src/com/android/quickstep/BaseActivityInterface.java
+++ b/quickstep/src/com/android/quickstep/BaseActivityInterface.java
@@ -197,17 +197,18 @@
             Rect gridRect = new Rect();
             calculateGridSize(context, dp, gridRect);
 
-            int verticalMargin = res.getDimensionPixelSize(
-                    R.dimen.overview_grid_focus_vertical_margin);
-            float taskHeight = gridRect.height() - verticalMargin * 2;
+            int verticalMargin = Math.max(
+                    res.getDimensionPixelSize(R.dimen.overview_grid_focus_vertical_margin),
+                    res.getDimensionPixelSize(R.dimen.overview_actions_height));
+            float taskHeight =
+                    gridRect.height() - verticalMargin * 2 - dp.overviewTaskThumbnailTopMarginPx;
 
             PointF taskDimension = getTaskDimension(context, dp);
-            float scale = taskHeight / Math.max(taskDimension.x, taskDimension.y);
+            float scale = taskHeight / taskDimension.y;
             int outWidth = Math.round(scale * taskDimension.x);
             int outHeight = Math.round(scale * taskDimension.y);
 
-            int gravity = Gravity.CENTER_VERTICAL;
-            gravity |= orientedState.getRecentsRtlSetting(res) ? Gravity.RIGHT : Gravity.LEFT;
+            int gravity = Gravity.CENTER;
             Gravity.apply(gravity, outWidth, outHeight, gridRect, outRect);
         } else {
             int taskMargin = dp.overviewTaskMarginPx;
@@ -296,8 +297,7 @@
         float rowHeight = (gridRect.height() - rowSpacing) / 2f;
 
         PointF taskDimension = getTaskDimension(context, dp);
-        float scale = (rowHeight - dp.overviewTaskThumbnailTopMarginPx) / Math.max(
-                taskDimension.x, taskDimension.y);
+        float scale = (rowHeight - dp.overviewTaskThumbnailTopMarginPx) / taskDimension.y;
         int outWidth = Math.round(scale * taskDimension.x);
         int outHeight = Math.round(scale * taskDimension.y);
 
diff --git a/quickstep/src/com/android/quickstep/fallback/FallbackRecentsStateController.java b/quickstep/src/com/android/quickstep/fallback/FallbackRecentsStateController.java
index f0364eb..52083bb 100644
--- a/quickstep/src/com/android/quickstep/fallback/FallbackRecentsStateController.java
+++ b/quickstep/src/com/android/quickstep/fallback/FallbackRecentsStateController.java
@@ -77,9 +77,7 @@
         float clearAllButtonAlpha = state.hasClearAllButton() ? 1 : 0;
         setter.setFloat(mRecentsView.getClearAllButton(), ClearAllButton.VISIBILITY_ALPHA,
                 clearAllButtonAlpha, LINEAR);
-        float overviewButtonAlpha =
-                state.hasOverviewActions() && mRecentsView.shouldShowOverviewActionsForState(state)
-                        ? 1 : 0;
+        float overviewButtonAlpha = state.hasOverviewActions() ? 1 : 0;
         setter.setFloat(mActivity.getActionsView().getVisibilityAlpha(),
                 MultiValueAlpha.VALUE, overviewButtonAlpha, LINEAR);
 
diff --git a/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java b/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java
index ac3fb27..efce650 100644
--- a/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java
+++ b/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java
@@ -26,14 +26,12 @@
 import android.content.Context;
 import android.os.Build;
 import android.util.AttributeSet;
-import android.util.Log;
 
 import androidx.annotation.Nullable;
 
 import com.android.launcher3.anim.AnimatorPlaybackController;
 import com.android.launcher3.anim.PendingAnimation;
 import com.android.launcher3.statemanager.StateManager.StateListener;
-import com.android.launcher3.testing.TestProtocol;
 import com.android.quickstep.FallbackActivityInterface;
 import com.android.quickstep.GestureState;
 import com.android.quickstep.RecentsActivity;
@@ -122,6 +120,12 @@
         }
     }
 
+    @Nullable
+    @Override
+    protected TaskView getHomeTaskView() {
+        return mHomeTaskInfo != null ? getTaskView(mHomeTaskInfo.taskId) : null;
+    }
+
     @Override
     protected boolean shouldAddStubTaskView(RunningTaskInfo runningTaskInfo) {
         if (mHomeTaskInfo != null && runningTaskInfo != null &&
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 376cb31..fda40e7 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -32,7 +32,6 @@
 import static com.android.launcher3.Utilities.squaredHypot;
 import static com.android.launcher3.Utilities.squaredTouchSlop;
 import static com.android.launcher3.anim.Interpolators.ACCEL;
-import static com.android.launcher3.anim.Interpolators.ACCEL_0_5;
 import static com.android.launcher3.anim.Interpolators.ACCEL_0_75;
 import static com.android.launcher3.anim.Interpolators.ACCEL_DEACCEL;
 import static com.android.launcher3.anim.Interpolators.FAST_OUT_SLOW_IN;
@@ -341,6 +340,7 @@
     private static final int ADDITION_TASK_DURATION = 200;
     private static final float INITIAL_DISMISS_TRANSLATION_INTERPOLATION_OFFSET = 0.55f;
     private static final float ADDITIONAL_DISMISS_TRANSLATION_INTERPOLATION_OFFSET = 0.05f;
+    private static final float ANIMATION_DISMISS_PROGRESS_MIDPOINT = 0.5f;
 
     protected final RecentsOrientedState mOrientationState;
     protected final BaseActivityInterface<STATE_TYPE, ACTIVITY_TYPE> mSizeStrategy;
@@ -348,6 +348,13 @@
     protected SurfaceTransactionApplier mSyncTransactionApplier;
     protected int mTaskWidth;
     protected int mTaskHeight;
+    // Used to position the top of a task in the top row of the grid
+    private float mTaskGridVerticalDiff;
+    // The vertical space one grid task takes + space between top and bottom row.
+    private float mTopBottomRowHeightDiff;
+    // mTaskGridVerticalDiff and mTopBottomRowHeightDiff summed together provides the top
+    // position for bottom row of grid tasks.
+
     protected final TransformParams mLiveTileParams = new TransformParams();
     protected final TaskViewSimulator mLiveTileTaskViewSimulator;
     protected final Rect mLastComputedTaskSize = new Rect();
@@ -487,7 +494,6 @@
     protected boolean mRunningTaskTileHidden;
     private Task mTmpRunningTask;
     protected int mFocusedTaskId = -1;
-    private float mFocusedTaskRatio;
 
     private boolean mRunningTaskIconScaledDown = false;
 
@@ -812,7 +818,7 @@
         // taskview for entering split screen, we only pretend to dismiss the task
         if (child instanceof TaskView && child != mSplitHiddenTaskView) {
             TaskView taskView = (TaskView) child;
-            mHasVisibleTaskData.delete(taskView.getTask().key.id);
+            mHasVisibleTaskData.delete(taskView.getTaskId());
             mTaskViewPool.recycle(taskView);
             mActionsView.updateHiddenFlags(HIDDEN_NO_TASKS, getTaskViewCount() == 0);
         }
@@ -919,7 +925,7 @@
     public TaskView getTaskView(int taskId) {
         for (int i = 0; i < getTaskViewCount(); i++) {
             TaskView taskView = getTaskViewAt(i);
-            if (taskView.hasTaskId(taskId)) {
+            if (taskView.getTaskId() == taskId) {
                 return taskView;
             }
         }
@@ -1111,6 +1117,9 @@
             final TaskView taskView = (TaskView) getChildAt(pageIndex);
             taskView.bind(task, mOrientationState);
         }
+        if (mFocusedTaskId == -1 && getTaskViewCount() > 0) {
+            mFocusedTaskId = getTaskViewAt(0).getTaskId();
+        }
         updateTaskSize();
 
         if (mNextPage == INVALID_PAGE) {
@@ -1167,7 +1176,7 @@
     public void resetTaskVisuals() {
         for (int i = getTaskViewCount() - 1; i >= 0; i--) {
             TaskView taskView = getTaskViewAt(i);
-            if (mIgnoreResetTaskId != taskView.getTask().key.id) {
+            if (mIgnoreResetTaskId != taskView.getTaskId()) {
                 taskView.resetViewTransforms();
                 taskView.setStableAlpha(mContentAlpha);
                 taskView.setFullscreenProgress(mFullscreenProgress);
@@ -1286,6 +1295,11 @@
         mSizeStrategy.calculateGridTaskSize(mActivity, mActivity.getDeviceProfile(),
                 mLastComputedGridTaskSize, mOrientationHandler);
 
+        mTaskGridVerticalDiff = mLastComputedGridTaskSize.top - mLastComputedTaskSize.top;
+        mTopBottomRowHeightDiff =
+                mLastComputedGridTaskSize.height() + dp.overviewTaskThumbnailTopMarginPx
+                        + mRowSpacing;
+
         // Force TaskView to update size from thumbnail
         updateTaskSize();
 
@@ -1316,6 +1330,15 @@
      * Updates TaskView scaling and translation required to support variable width.
      */
     private void updateTaskSize() {
+        updateTaskSize(false);
+    }
+
+    /**
+     * Updates TaskView scaling and translation required to support variable width.
+     *
+     * @param isTaskDismissal indicates if update was called due to task dismissal
+     */
+    private void updateTaskSize(boolean isTaskDismissal) {
         final int taskCount = getTaskViewCount();
         if (taskCount == 0) {
             return;
@@ -1336,7 +1359,7 @@
 
         mClearAllButton.setFullscreenTranslationPrimary(accumulatedTranslationX);
 
-        updateGridProperties();
+        updateGridProperties(isTaskDismissal);
     }
 
     public void getTaskSize(Rect outRect) {
@@ -1351,19 +1374,7 @@
     public Point getSelectedTaskSize() {
         mSizeStrategy.calculateTaskSize(mActivity, mActivity.getDeviceProfile(), mTempRect,
                 mOrientationHandler);
-        int taskWidth = mTempRect.width();
-        int taskHeight = mTempRect.height();
-        if (mRunningTaskId != -1) {
-            int boxLength = Math.max(taskWidth, taskHeight);
-            if (mFocusedTaskRatio > 1) {
-                taskWidth = boxLength;
-                taskHeight = (int) (boxLength / mFocusedTaskRatio);
-            } else {
-                taskWidth = (int) (boxLength * mFocusedTaskRatio);
-                taskHeight = boxLength;
-            }
-        }
-        return new Point(taskWidth, taskHeight);
+        return new Point(mTempRect.width(), mTempRect.height());
     }
 
     /** Gets the last computed task size */
@@ -1395,13 +1406,7 @@
             loadVisibleTaskData(TaskView.FLAG_UPDATE_ALL);
 
             // After scrolling, update ActionsView's visibility.
-            TaskView focusedTaskView = getFocusedTaskView();
-            if (focusedTaskView != null) {
-                float scrollDiff = Math.abs(getScrollForPage(indexOfChild(focusedTaskView))
-                        - mOrientationHandler.getPrimaryScroll(this));
-                float delta = (mGridSideMargin - scrollDiff) / (float) mGridSideMargin;
-                mActionsView.getScrollAlpha().setValue(Utilities.boundToRange(delta, 0, 1));
-            }
+            updateActionsViewScrollAlpha();
         }
 
         // Update the high res thumbnail loader state
@@ -1409,6 +1414,20 @@
         return scrolling;
     }
 
+    private void updateActionsViewScrollAlpha() {
+        float scrollAlpha = 1f;
+        if (showAsGrid()) {
+            TaskView focusedTaskView = getFocusedTaskView();
+            if (focusedTaskView != null) {
+                float scrollDiff = Math.abs(getScrollForPage(indexOfChild(focusedTaskView))
+                        - mOrientationHandler.getPrimaryScroll(this));
+                float delta = (mGridSideMargin - scrollDiff) / (float) mGridSideMargin;
+                scrollAlpha = Utilities.boundToRange(delta, 0, 1);
+            }
+        }
+        mActionsView.getScrollAlpha().setValue(scrollAlpha);
+    }
+
     /**
      * Scales and adjusts translation of adjacent pages as if on a curved carousel.
      */
@@ -1551,7 +1570,7 @@
         setCurrentTask(-1);
         mIgnoreResetTaskId = -1;
         mTaskListChangeId = -1;
-        mFocusedTaskId = -1;
+        mFocusedTaskId = getTaskViewCount() > 0 ? getTaskViewAt(0).getTaskId() : -1;
 
         if (mRecentsAnimationController != null) {
             if (LIVE_TILE.get() && mEnableDrawingLiveTile) {
@@ -1592,11 +1611,8 @@
         return getTaskView(mFocusedTaskId);
     }
 
-    /**
-     * Returns the width to height ratio of the focused {@link TaskView}.
-     */
-    public float getFocusedTaskRatio() {
-        return mFocusedTaskRatio;
+    protected @Nullable TaskView getHomeTaskView() {
+        return null;
     }
 
     /**
@@ -1731,11 +1747,7 @@
         }
         setRunningTaskHidden(false);
         animateUpRunningTaskIconScale();
-
-        if (mCurrentGestureEndTarget == GestureState.GestureEndTarget.RECENTS
-                && (!showAsGrid() || getFocusedTaskView() != null)) {
-            animateActionsViewIn();
-        }
+        animateActionsViewIn();
 
         mCurrentGestureEndTarget = null;
     }
@@ -1777,9 +1789,7 @@
         boolean runningTaskTileHidden = mRunningTaskTileHidden;
         int runningTaskId = runningTaskInfo == null ? -1 : runningTaskInfo.taskId;
         setCurrentTask(runningTaskId);
-        if (mActivity.getDeviceProfile().isTablet && FeatureFlags.ENABLE_OVERVIEW_GRID.get()) {
-            setFocusedTask(runningTaskId);
-        }
+        mFocusedTaskId = runningTaskId;
         setCurrentPage(getRunningTaskIndex());
         setRunningTaskViewShowScreenshot(false);
         setRunningTaskHidden(runningTaskTileHidden);
@@ -1808,15 +1818,6 @@
     }
 
     /**
-     * Sets the focused task id and store the width to height ratio of the focused task.
-     */
-    protected void setFocusedTask(int focusedTaskId) {
-        mFocusedTaskId = focusedTaskId;
-        mFocusedTaskRatio =
-                mLastComputedTaskSize.width() / (float) mLastComputedTaskSize.height();
-    }
-
-    /**
      * Hides the tile associated with {@link #mRunningTaskId}
      */
     public void setRunningTaskHidden(boolean isHidden) {
@@ -1865,13 +1866,6 @@
         anim.start();
     }
 
-    private void animateActionsViewOut() {
-        ObjectAnimator anim = ObjectAnimator.ofFloat(
-                mActionsView.getVisibilityAlpha(), MultiValueAlpha.VALUE, 1, 0);
-        anim.setDuration(TaskView.SCALE_ICON_DURATION);
-        anim.start();
-    }
-
     public void animateUpRunningTaskIconScale() {
         mRunningTaskIconScaledDown = false;
         TaskView firstTask = getRunningTaskView();
@@ -1905,20 +1899,8 @@
             return;
         }
 
-        final int boxLength = Math.max(mLastComputedGridTaskSize.width(),
-                mLastComputedGridTaskSize.height());
         int taskTopMargin = mActivity.getDeviceProfile().overviewTaskThumbnailTopMarginPx;
 
-        /*
-         * taskGridVerticalDiff is used to position the top of a task in the top row of the grid
-         * heightOffset is the vertical space one grid task takes + space between top and
-         *   bottom row
-         * Summed together they provide the top position for bottom row of grid tasks
-         */
-        final float taskGridVerticalDiff =
-                mLastComputedGridTaskSize.top - mLastComputedTaskSize.top;
-        final float heightOffset = (boxLength + taskTopMargin) + mRowSpacing;
-
         int topRowWidth = 0;
         int bottomRowWidth = 0;
         float topAccumulatedTranslationX = 0;
@@ -1938,6 +1920,8 @@
         int snappedTaskRowWidth = 0;
         int snappedPage = getNextPage();
         TaskView snappedTaskView = getTaskViewAtByAbsoluteIndex(snappedPage);
+        TaskView homeTaskView = getHomeTaskView();
+        TaskView nextFocusedTaskView = null;
 
         if (!isTaskDismissal) {
             mTopRowIdSet.clear();
@@ -1975,15 +1959,20 @@
                     // calculate the distance focused task need to shift.
                     focusedTaskShift += mIsRtl ? taskWidthAndSpacing : -taskWidthAndSpacing;
                 }
-                int taskId = taskView.getTask().key.id;
+                int taskId = taskView.getTaskId();
                 boolean isTopRow = isTaskDismissal ? mTopRowIdSet.contains(taskId)
                         : topRowWidth <= bottomRowWidth;
                 if (isTopRow) {
-                    topRowWidth += taskWidthAndSpacing;
+                    if (homeTaskView != null && nextFocusedTaskView == null) {
+                        // TaskView will be focused when swipe up, don't count towards row width.
+                        nextFocusedTaskView = taskView;
+                    } else {
+                        topRowWidth += taskWidthAndSpacing;
+                    }
                     topSet.add(i);
                     mTopRowIdSet.add(taskId);
 
-                    taskView.setGridTranslationY(taskGridVerticalDiff);
+                    taskView.setGridTranslationY(mTaskGridVerticalDiff);
 
                     // Move horizontally into empty space.
                     float widthOffset = 0;
@@ -2002,7 +1991,7 @@
                     bottomSet.add(i);
 
                     // Move into bottom row.
-                    taskView.setGridTranslationY(heightOffset + taskGridVerticalDiff);
+                    taskView.setGridTranslationY(mTopBottomRowHeightDiff + mTaskGridVerticalDiff);
 
                     // Move horizontally into empty space.
                     float widthOffset = 0;
@@ -2085,8 +2074,10 @@
         // of swiping up after quick switch.
         if (snappedTaskView != null) {
             int distanceFromClearAll = longRowWidth - snappedTaskRowWidth;
+            // ClearAllButton should be off screen when snapped task is in its snapped position.
             int minimumDistance =
-                    mLastComputedGridSize.width() - snappedTaskView.getLayoutParams().width;
+                    mTaskWidth - snappedTaskView.getLayoutParams().width
+                            + (mLastComputedGridSize.width() - mTaskWidth) / 2;
             if (distanceFromClearAll < minimumDistance) {
                 int distanceDifference = minimumDistance - distanceFromClearAll;
                 clearAllTotalTranslationX += mIsRtl ? -distanceDifference : distanceDifference;
@@ -2106,8 +2097,8 @@
         if (taskView1 == null || taskView2 == null) {
             return false;
         }
-        int taskId1 = taskView1.getTask().key.id;
-        int taskId2 = taskView2.getTask().key.id;
+        int taskId1 = taskView1.getTaskId();
+        int taskId2 = taskView2.getTaskId();
         if (taskId1 == mFocusedTaskId || taskId2 == mFocusedTaskId) {
             return false;
         }
@@ -2236,8 +2227,16 @@
         }
     }
 
-    public PendingAnimation createTaskDismissAnimation(TaskView taskView, boolean animateTaskView,
-            boolean shouldRemoveTask, long duration) {
+    /**
+     * Creates a {@link PendingAnimation} for dismissing the specified {@link TaskView}.
+     * @param dismissedTaskView the {@link TaskView} to be dismissed
+     * @param animateTaskView whether the {@link TaskView} to be dismissed should be animated
+     * @param shouldRemoveTask whether the associated {@link Task} should be removed from
+     *                         ActivityManager after dismissal
+     * @param duration duration of the animation
+     */
+    public PendingAnimation createTaskDismissAnimation(TaskView dismissedTaskView,
+            boolean animateTaskView, boolean shouldRemoveTask, long duration) {
         if (mPendingAnimation != null) {
             mPendingAnimation.createPlaybackController().dispatchOnCancel();
         }
@@ -2248,30 +2247,65 @@
             return anim;
         }
 
+        boolean showAsGrid = showAsGrid();
+        int taskCount = getTaskViewCount();
+        int dismissedIndex = indexOfChild(dismissedTaskView);
+        int dismissedTaskId = dismissedTaskView.getTaskId();
+
+        // Grid specific properties.
+        boolean isFocusedTaskDismissed = false;
+        TaskView nextFocusedTaskView = null;
+        boolean nextFocusedTaskFromTop = false;
+        float dismissedTaskWidth = 0;
+        float nextFocusedTaskWidth = 0;
+
+        // Non-grid specific properties.
         int[] oldScroll = new int[count];
         int[] newScroll = new int[count];
-        getPageScrolls(oldScroll, false, SIMPLE_SCROLL_LOGIC);
-        getPageScrolls(newScroll, false, (v) -> v.getVisibility() != GONE && v != taskView);
-        int taskCount = getTaskViewCount();
         int scrollDiffPerPage = 0;
-        if (count > 1) {
-            scrollDiffPerPage = Math.abs(oldScroll[1] - oldScroll[0]);
-        }
-        int draggedIndex = indexOfChild(taskView);
-
-        boolean isFocusedTaskDismissed = taskView.getTask().key.id == mFocusedTaskId;
-        if (isFocusedTaskDismissed && showAsGrid()) {
-            anim.setFloat(mActionsView, VIEW_ALPHA, 0, clampToProgress(ACCEL_0_5, 0, 0.5f));
-        }
-        float dismissedTaskWidth = taskView.getLayoutParams().width + mPageSpacing;
         boolean needsCurveUpdates = false;
+
+        if (showAsGrid) {
+            dismissedTaskWidth = dismissedTaskView.getLayoutParams().width + mPageSpacing;
+            isFocusedTaskDismissed = dismissedTaskId == mFocusedTaskId;
+            if (isFocusedTaskDismissed) {
+                nextFocusedTaskFromTop =
+                        mTopRowIdSet.size() > 0 && mTopRowIdSet.size() >= (taskCount - 1) / 2f;
+                // Pick the next focused task from the preferred row.
+                for (int i = 0; i < taskCount; i++) {
+                    TaskView taskView = getTaskViewAt(i);
+                    if (taskView == dismissedTaskView) {
+                        continue;
+                    }
+                    boolean isTopRow = mTopRowIdSet.contains(taskView.getTaskId());
+                    if ((nextFocusedTaskFromTop && isTopRow
+                            || (!nextFocusedTaskFromTop && !isTopRow))) {
+                        nextFocusedTaskView = taskView;
+                        break;
+                    }
+                }
+                if (nextFocusedTaskView != null) {
+                    nextFocusedTaskWidth =
+                            nextFocusedTaskView.getLayoutParams().width + mPageSpacing;
+                }
+            }
+        } else {
+            getPageScrolls(oldScroll, false, SIMPLE_SCROLL_LOGIC);
+            getPageScrolls(newScroll, false,
+                    v -> v.getVisibility() != GONE && v != dismissedTaskView);
+            if (count > 1) {
+                scrollDiffPerPage = Math.abs(oldScroll[1] - oldScroll[0]);
+            }
+        }
+
+        int distanceFromDismissedTask = 0;
         for (int i = 0; i < count; i++) {
             View child = getChildAt(i);
-            if (child == taskView) {
+            if (child == dismissedTaskView) {
                 if (animateTaskView) {
-                    addDismissedTaskAnimations(taskView, duration, anim);
+                    addDismissedTaskAnimations(dismissedTaskView, duration, anim);
                 }
-            } else if (!showAsGrid()) {
+            } else if (!showAsGrid) {
                 // Compute scroll offsets from task dismissal for animation.
                 // If we just take newScroll - oldScroll, everything to the right of dragged task
                 // translates to the left. We need to offset this in some cases:
@@ -2280,15 +2314,15 @@
                 // - Current page is rightmost page (leftmost for RTL)
                 // - Dragging an adjacent page on the left side (right side for RTL)
                 int offset = mIsRtl ? scrollDiffPerPage : 0;
-                if (mCurrentPage == draggedIndex) {
+                if (mCurrentPage == dismissedIndex) {
                     int lastPage = taskCount - 1;
                     if (mCurrentPage == lastPage) {
                         offset += mIsRtl ? -scrollDiffPerPage : scrollDiffPerPage;
                     }
                 } else {
-                    // Dragging an adjacent page.
+                    // Dismissing an adjacent page.
                     int negativeAdjacent = mCurrentPage - 1; // (Right in RTL, left in LTR)
-                    if (draggedIndex == negativeAdjacent) {
+                    if (dismissedIndex == negativeAdjacent) {
                         offset += mIsRtl ? -scrollDiffPerPage : scrollDiffPerPage;
                     }
                 }
@@ -2301,28 +2335,53 @@
 
                     float additionalDismissDuration =
                             ADDITIONAL_DISMISS_TRANSLATION_INTERPOLATION_OFFSET * Math.abs(
-                                    i - draggedIndex);
+                                    i - dismissedIndex);
                     anim.setFloat(child, translationProperty, scrollDiff, clampToProgress(LINEAR,
                             Utilities.boundToRange(INITIAL_DISMISS_TRANSLATION_INTERPOLATION_OFFSET
                                     + additionalDismissDuration, 0f, 1f), 1));
                     needsCurveUpdates = true;
                 }
             } else if (child instanceof TaskView) {
+                TaskView taskView = (TaskView) child;
+                if (isFocusedTaskDismissed) {
+                    if (!isSameGridRow(taskView, nextFocusedTaskView)) {
+                        continue;
+                    }
+                } else {
+                    if (i < dismissedIndex || !isSameGridRow(taskView, dismissedTaskView)) {
+                        continue;
+                    }
+                }
                 // Animate task with index >= dismissed index and in the same row as the
-                // dismissed index, or if the dismissed task was the focused task. Offset
-                // successive task dismissal durations for a staggered effect.
-                if (isFocusedTaskDismissed || (i >= draggedIndex && isSameGridRow((TaskView) child,
-                        taskView))) {
-                    FloatProperty translationProperty =
-                            ((TaskView) child).getPrimaryDismissTranslationProperty();
-                    float additionalDismissDuration =
-                            ADDITIONAL_DISMISS_TRANSLATION_INTERPOLATION_OFFSET * Math.abs(
-                                    i - draggedIndex);
-                    anim.setFloat(child, translationProperty,
-                            !mIsRtl ? -dismissedTaskWidth : dismissedTaskWidth,
-                            clampToProgress(LINEAR, Utilities.boundToRange(
-                                    INITIAL_DISMISS_TRANSLATION_INTERPOLATION_OFFSET
-                                            + additionalDismissDuration, 0f, 1f), 1));
+                // dismissed index or next focused index. Offset successive task dismissal
+                // durations for a staggered effect.
+                float animationStartProgress = Utilities.boundToRange(
+                        INITIAL_DISMISS_TRANSLATION_INTERPOLATION_OFFSET
+                                + ADDITIONAL_DISMISS_TRANSLATION_INTERPOLATION_OFFSET
+                                * ++distanceFromDismissedTask, 0f, 1f);
+                if (taskView == nextFocusedTaskView) {
+                    // Enlarge the task to be focused next, and translate into focus position.
+                    float scale = mTaskWidth / (float) mLastComputedGridTaskSize.width();
+                    anim.setFloat(taskView, TaskView.SNAPSHOT_SCALE, scale,
+                            clampToProgress(LINEAR, animationStartProgress, 1f));
+                    anim.setFloat(taskView, taskView.getPrimaryDismissTranslationProperty(),
+                            mIsRtl ? dismissedTaskWidth : -dismissedTaskWidth,
+                            clampToProgress(LINEAR, animationStartProgress, 1f));
+                    float secondaryTranslation = -mTaskGridVerticalDiff;
+                    if (!nextFocusedTaskFromTop) {
+                        secondaryTranslation -= mTopBottomRowHeightDiff;
+                    }
+                    anim.setFloat(taskView, taskView.getSecondaryDissmissTranslationProperty(),
+                            secondaryTranslation,
+                            clampToProgress(LINEAR, animationStartProgress, 1f));
+                    anim.setFloat(taskView, TaskView.FOCUS_TRANSITION, 0f,
+                            clampToProgress(LINEAR, 0f, ANIMATION_DISMISS_PROGRESS_MIDPOINT));
+                } else {
+                    float primaryTranslation =
+                            isFocusedTaskDismissed ? nextFocusedTaskWidth : dismissedTaskWidth;
+                    anim.setFloat(taskView, taskView.getPrimaryDismissTranslationProperty(),
+                            mIsRtl ? primaryTranslation : -primaryTranslation,
+                            clampToProgress(LINEAR, animationStartProgress, 1f));
                 }
             }
         }
@@ -2333,14 +2392,15 @@
 
         // Add a tiny bit of translation Z, so that it draws on top of other views
         if (animateTaskView) {
-            taskView.setTranslationZ(0.1f);
+            dismissedTaskView.setTranslationZ(0.1f);
         }
 
         mPendingAnimation = anim;
+        final TaskView finalNextFocusedTaskView = nextFocusedTaskView;
         mPendingAnimation.addEndListener(new Consumer<Boolean>() {
             @Override
             public void accept(Boolean success) {
-                if (LIVE_TILE.get() && mEnableDrawingLiveTile && taskView.isRunningTask()
+                if (LIVE_TILE.get() && mEnableDrawingLiveTile && dismissedTaskView.isRunningTask()
                         && success) {
                     finishRecentsAnimation(true /* toHome */, () -> onEnd(success));
                 } else {
@@ -2352,11 +2412,11 @@
             private void onEnd(boolean success) {
                 if (success) {
                     if (shouldRemoveTask) {
-                        if (taskView.getTask() != null) {
+                        if (dismissedTaskView.getTask() != null) {
                             UI_HELPER_EXECUTOR.execute(() -> ActivityManagerWrapper.getInstance()
-                                    .removeTask(taskView.getTask().key.id));
+                                    .removeTask(dismissedTaskId));
                             mActivity.getStatsLogManager().logger()
-                                    .withItemInfo(taskView.getItemInfo())
+                                    .withItemInfo(dismissedTaskView.getItemInfo())
                                     .log(LAUNCHER_TASK_DISMISS_SWIPE_UP);
                         }
                     }
@@ -2366,31 +2426,31 @@
                     resetTaskVisuals();
 
                     int pageToSnapTo = mCurrentPage;
-                    // Snap to start if focused task was dismissed, as after quick switch it could
-                    // be at any page but the focused task always displays at the start.
-                    if (taskView.getTask().key.id == mFocusedTaskId) {
-                        pageToSnapTo = mTaskViewStartIndex;
-                    } else if (draggedIndex < pageToSnapTo || pageToSnapTo == (getTaskViewCount()
-                            - 1)) {
+                    if (finalNextFocusedTaskView != null) {
+                        pageToSnapTo = indexOfChild(finalNextFocusedTaskView);
+                    }
+                    if (dismissedIndex < pageToSnapTo || pageToSnapTo == (taskCount - 1)) {
                         pageToSnapTo -= 1;
                     }
-                    removeViewInLayout(taskView);
+                    removeViewInLayout(dismissedTaskView);
+                    mTopRowIdSet.remove(dismissedTaskId);
 
-                    if (getTaskViewCount() == 0) {
+                    if (taskCount == 1) {
                         removeViewInLayout(mClearAllButton);
                         startHome();
                     } else {
-                        snapToPageImmediately(pageToSnapTo);
-                        // Grid got messed up, reapply.
-                        updateGridProperties(true);
-                        if (showAsGrid() && getFocusedTaskView() == null
-                                && mActionsView.getVisibilityAlpha().getValue() == 1) {
-                            animateActionsViewOut();
+                        // Update focus task and its size.
+                        if (finalNextFocusedTaskView != null) {
+                            mFocusedTaskId = finalNextFocusedTaskView.getTaskId();
+                            mTopRowIdSet.remove(mFocusedTaskId);
+                            finalNextFocusedTaskView.animateIconScaleAndDimIntoView();
+                            setCurrentPage(pageToSnapTo);
                         }
+                        updateTaskSize(true);
+                        // Update scroll and snap to page.
+                        updateScrollSynchronously();
+                        snapToPageImmediately(pageToSnapTo);
                     }
-                    // Update the layout synchronously so that the position of next view is
-                    // immediately available.
-                    onLayout(false /*  changed */, getLeft(), getTop(), getRight(), getBottom());
                 }
                 onDismissAnimationEnds();
                 mPendingAnimation = null;
@@ -2426,7 +2486,8 @@
         int taskCount = getTaskViewCount();
         for (int i = 0; i < taskCount; i++) {
             TaskView taskView = getTaskViewAt(i);
-            if (taskView == mSplitHiddenTaskView && taskView != getFocusedTaskView()) {
+            if (taskView == mSplitHiddenTaskView
+                    && !(showAsGrid() && taskView == getFocusedTaskView())) {
                 // Case where the hidden task view would have overlapped w/ placeholder,
                 // but because it's going to hide we don't care
                 // TODO (b/187312247) edge case for thumbnails that are off screen but scroll on
@@ -2571,7 +2632,7 @@
         mContentAlpha = alpha;
         for (int i = getTaskViewCount() - 1; i >= 0; i--) {
             TaskView child = getTaskViewAt(i);
-            if (!mRunningTaskTileHidden || child.getTask().key.id != mRunningTaskId) {
+            if (!mRunningTaskTileHidden || child.getTaskId() != mRunningTaskId) {
                 child.setStableAlpha(alpha);
             }
         }
@@ -3396,9 +3457,13 @@
     }
 
     /**
-     * Updates page scroll synchronously and layout child views.
+     * Updates page scroll synchronously after measure and layout child views.
      */
     public void updateScrollSynchronously() {
+        // onMeasure is needed to update child's measured width which is used in scroll calculation,
+        // in case TaskView sizes has changed when being focused/unfocused.
+        onMeasure(makeMeasureSpec(getMeasuredWidth(), EXACTLY),
+                makeMeasureSpec(getMeasuredHeight(), EXACTLY));
         onLayout(false /*  changed */, getLeft(), getTop(), getRight(), getBottom());
         updateMinAndMaxScrollX();
     }
@@ -3596,6 +3661,7 @@
     public void setOverviewGridEnabled(boolean overviewGridEnabled) {
         if (mOverviewGridEnabled != overviewGridEnabled) {
             mOverviewGridEnabled = overviewGridEnabled;
+            updateActionsViewScrollAlpha();
             // Request layout to ensure scroll position is recalculated with updated mGridProgress.
             requestLayout();
         }
@@ -3733,11 +3799,6 @@
                 && mCurrentGestureEndTarget != GestureState.GestureEndTarget.RECENTS;
     }
 
-    public boolean shouldShowOverviewActionsForState(STATE_TYPE state) {
-        return !state.displayOverviewTasksAsGrid(mActivity.getDeviceProfile())
-                || getFocusedTaskView() != null;
-    }
-
     /**
      * Used to register callbacks for when our empty message state changes.
      *
diff --git a/quickstep/src/com/android/quickstep/views/TaskThumbnailView.java b/quickstep/src/com/android/quickstep/views/TaskThumbnailView.java
index 02888a1..1bfd430 100644
--- a/quickstep/src/com/android/quickstep/views/TaskThumbnailView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskThumbnailView.java
@@ -411,7 +411,9 @@
      */
     public static class PreviewPositionHelper {
 
-        // Contains the portion of the thumbnail that is clipped when fullscreen progress = 0.
+        private static final RectF EMPTY_RECT_F = new RectF();
+
+        // Contains the portion of the thumbnail that is unclipped when fullscreen progress = 1.
         private final RectF mClippedInsets = new RectF();
         private final Matrix mMatrix = new Matrix();
         private boolean mIsOrientationChanged;
@@ -616,15 +618,17 @@
                     break;
             }
             mClippedInsets.offsetTo(newLeftInset * scale, newTopInset * scale);
-            mMatrix.postTranslate(translateX - mClippedInsets.left,
-                    translateY - mClippedInsets.top);
+            mMatrix.postTranslate(translateX, translateY);
+            if (TaskView.FULL_THUMBNAIL) {
+                mMatrix.postTranslate(-mClippedInsets.left, -mClippedInsets.top);
+            }
         }
 
         /**
          * Insets to used for clipping the thumbnail (in case it is drawing outside its own space)
          */
         public RectF getInsetsToDrawInFullscreen() {
-            return mClippedInsets;
+            return TaskView.FULL_THUMBNAIL ? mClippedInsets : EMPTY_RECT_F;
         }
     }
 }
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.java b/quickstep/src/com/android/quickstep/views/TaskView.java
index cfac32c..46729ec 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskView.java
@@ -141,6 +141,11 @@
      */
     public static final boolean CLIP_STATUS_AND_NAV_BARS = false;
 
+    /**
+     * Should the TaskView scale down to fit whole thumbnail in fullscreen.
+     */
+    public static final boolean FULL_THUMBNAIL = false;
+
     private static final float EDGE_SCALE_DOWN_FACTOR_CAROUSEL = 0.03f;
     private static final float EDGE_SCALE_DOWN_FACTOR_GRID = 0.00f;
 
@@ -161,7 +166,7 @@
     private static final List<Rect> SYSTEM_GESTURE_EXCLUSION_RECT =
             Collections.singletonList(new Rect());
 
-    private static final FloatProperty<TaskView> FOCUS_TRANSITION =
+    public static final FloatProperty<TaskView> FOCUS_TRANSITION =
             new FloatProperty<TaskView>("focusTransition") {
                 @Override
                 public void setValue(TaskView taskView, float v) {
@@ -330,6 +335,19 @@
                 }
             };
 
+    public static final FloatProperty<TaskView> SNAPSHOT_SCALE =
+            new FloatProperty<TaskView>("snapshotScale") {
+                @Override
+                public void setValue(TaskView taskView, float v) {
+                    taskView.setSnapshotScale(v);
+                }
+
+                @Override
+                public Float get(TaskView taskView) {
+                    return taskView.mSnapshotView.getScaleX();
+                }
+            };
+
     private final TaskOutlineProvider mOutlineProvider;
 
     private Task mTask;
@@ -518,8 +536,8 @@
         return mTask;
     }
 
-    public boolean hasTaskId(int taskId) {
-        return mTask != null && mTask.key != null && mTask.key.id == taskId;
+    public int getTaskId() {
+        return mTask != null && mTask.key != null ? mTask.key.id : -1;
     }
 
     public TaskThumbnailView getThumbnail() {
@@ -841,6 +859,7 @@
                 mSplitSelectTranslationX = 0f;
         mDismissTranslationY = mTaskOffsetTranslationY = mTaskResistanceTranslationY =
                 mSplitSelectTranslationY = 0f;
+        setSnapshotScale(1f);
         applyTranslationX();
         applyTranslationY();
         setTranslationZ(0);
@@ -920,6 +939,9 @@
         if (mActivity.getDeviceProfile().isTablet && FeatureFlags.ENABLE_OVERVIEW_GRID.get()) {
             setPivotX(getLayoutDirection() == LAYOUT_DIRECTION_RTL ? 0 : right - left);
             setPivotY(mSnapshotView.getTop());
+            mSnapshotView.setPivotX(
+                    getLayoutDirection() == LAYOUT_DIRECTION_RTL ? 0 : right - left);
+            mSnapshotView.setPivotY(0);
         } else {
             setPivotX((right - left) * 0.5f);
             setPivotY(mSnapshotView.getTop() + mSnapshotView.getHeight() * 0.5f);
@@ -950,6 +972,11 @@
         return mFullscreenScale;
     }
 
+    private void setSnapshotScale(float dismissScale) {
+        mSnapshotView.setScaleX(dismissScale);
+        mSnapshotView.setScaleY(dismissScale);
+    }
+
     /**
      * Moves TaskView between carousel and 2 row grid.
      *
@@ -1343,46 +1370,26 @@
 
             int boxWidth;
             int boxHeight;
-            float thumbnailRatio;
             boolean isFocusedTask = isFocusedTask();
             if (isFocusedTask) {
                 // Task will be focused and should use focused task size. Use focusTaskRatio
                 // that is associated with the original orientation of the focused task.
                 boxWidth = taskWidth;
                 boxHeight = taskHeight;
-                thumbnailRatio = getRecentsView().getFocusedTaskRatio();
             } else {
                 // Otherwise task is in grid, and should use lastComputedGridTaskSize.
                 Rect lastComputedGridTaskSize = getRecentsView().getLastComputedGridTaskSize();
                 boxWidth = lastComputedGridTaskSize.width();
                 boxHeight = lastComputedGridTaskSize.height();
-                thumbnailRatio = mTask != null ? mTask.getVisibleThumbnailRatio(
-                        TaskView.CLIP_STATUS_AND_NAV_BARS) : 0f;
             }
-            int boxLength = Math.max(boxWidth, boxHeight);
 
             // Bound width/height to the box size.
-            if (thumbnailRatio == 0f) {
-                expectedWidth = boxWidth;
-                expectedHeight = boxHeight + thumbnailPadding;
-            } else if (thumbnailRatio > 1) {
-                expectedWidth = boxLength;
-                expectedHeight = (int) (boxLength / thumbnailRatio) + thumbnailPadding;
-            } else {
-                expectedWidth = (int) (boxLength * thumbnailRatio);
-                expectedHeight = boxLength + thumbnailPadding;
-            }
+            expectedWidth = boxWidth;
+            expectedHeight = boxHeight + thumbnailPadding;
 
             // Scale to to fit task Rect.
             fullscreenScale = taskWidth / (float) boxWidth;
 
-            // In full screen, scale back TaskView to original size.
-            if (expectedWidth > boxWidth) {
-                fullscreenScale *= boxWidth / (float) expectedWidth;
-            } else if (expectedHeight - thumbnailPadding > boxHeight) {
-                fullscreenScale *= boxHeight / (float) (expectedHeight - thumbnailPadding);
-            }
-
             // Align to top of task Rect.
             boxTranslationY = (expectedHeight - thumbnailPadding - taskHeight) / 2.0f;
         } else {
diff --git a/res/drawable/widgets_list_bottom_ripple.xml b/res/drawable/widgets_list_bottom_ripple.xml
index 72262d4..971d6f3 100644
--- a/res/drawable/widgets_list_bottom_ripple.xml
+++ b/res/drawable/widgets_list_bottom_ripple.xml
@@ -21,6 +21,7 @@
     android:color="?android:attr/colorControlHighlight">
     <item android:id="@android:id/mask">
         <shape android:shape="rectangle">
+            <solid android:color="@color/widgets_picker_surface" />
             <corners
                 android:topLeftRadius="@dimen/widget_list_content_corner_radius"
                 android:topRightRadius="@dimen/widget_list_content_corner_radius"
diff --git a/res/drawable/widgets_list_middle_ripple.xml b/res/drawable/widgets_list_middle_ripple.xml
index 1136bea..2b77d4d 100644
--- a/res/drawable/widgets_list_middle_ripple.xml
+++ b/res/drawable/widgets_list_middle_ripple.xml
@@ -21,6 +21,7 @@
     android:color="?android:attr/colorControlHighlight">
     <item android:id="@android:id/mask">
         <shape android:shape="rectangle">
+            <solid android:color="@color/widgets_picker_surface" />
             <corners
                 android:topLeftRadius="@dimen/widget_list_content_corner_radius"
                 android:topRightRadius="@dimen/widget_list_content_corner_radius"
diff --git a/res/drawable/widgets_list_single_item_ripple.xml b/res/drawable/widgets_list_single_item_ripple.xml
index a82918e..c09944d 100644
--- a/res/drawable/widgets_list_single_item_ripple.xml
+++ b/res/drawable/widgets_list_single_item_ripple.xml
@@ -21,6 +21,7 @@
     android:color="?android:attr/colorControlHighlight">
     <item android:id="@android:id/mask">
         <shape android:shape="rectangle">
+            <solid android:color="@color/widgets_picker_surface" />
             <corners
                 android:topLeftRadius="@dimen/widget_list_top_bottom_corner_radius"
                 android:topRightRadius="@dimen/widget_list_top_bottom_corner_radius"
diff --git a/res/drawable/widgets_list_top_ripple.xml b/res/drawable/widgets_list_top_ripple.xml
index 4ad185c..f79ab72 100644
--- a/res/drawable/widgets_list_top_ripple.xml
+++ b/res/drawable/widgets_list_top_ripple.xml
@@ -21,6 +21,7 @@
     android:color="?android:attr/colorControlHighlight">
     <item android:id="@android:id/mask">
         <shape android:shape="rectangle">
+            <solid android:color="@color/widgets_picker_surface" />
             <corners
                 android:topLeftRadius="@dimen/widget_list_top_bottom_corner_radius"
                 android:topRightRadius="@dimen/widget_list_top_bottom_corner_radius"
diff --git a/res/layout/widgets_full_sheet_paged_view.xml b/res/layout/widgets_full_sheet_paged_view.xml
index 580ca49..f0ddc2b 100644
--- a/res/layout/widgets_full_sheet_paged_view.xml
+++ b/res/layout/widgets_full_sheet_paged_view.xml
@@ -21,6 +21,7 @@
         android:layout_width="match_parent"
         android:layout_height="match_parent"
         android:clipToPadding="false"
+        android:paddingTop="@dimen/widget_picker_view_pager_top_padding"
         android:descendantFocusability="afterDescendants"
         launcher:pageIndicator="@+id/tabs">
 
diff --git a/res/layout/widgets_list_row_header.xml b/res/layout/widgets_list_row_header.xml
index a0a0456..278d3dc 100644
--- a/res/layout/widgets_list_row_header.xml
+++ b/res/layout/widgets_list_row_header.xml
@@ -74,6 +74,7 @@
         android:layout_height="wrap_content"
         android:layout_gravity="center_vertical"
         android:layout_alignParentEnd="true"
+        android:enabled="false"
         android:clickable="false"
         android:importantForAccessibility="no"
         android:button="@drawable/widgets_tray_expand_button"/>
diff --git a/res/layout/widgets_personal_work_tabs.xml b/res/layout/widgets_personal_work_tabs.xml
index 72d83e8..15275a6 100644
--- a/res/layout/widgets_personal_work_tabs.xml
+++ b/res/layout/widgets_personal_work_tabs.xml
@@ -20,10 +20,9 @@
     android:id="@+id/tabs"
     android:layout_width="match_parent"
     android:layout_height="@dimen/all_apps_header_pill_height"
-    android:layout_marginHorizontal="16dp"
+    android:layout_marginHorizontal="32dp"
     android:orientation="horizontal"
     android:background="@drawable/all_apps_tabs_background"
-    android:elevation="2dp"
     style="@style/TextHeadline">
 
     <Button
diff --git a/res/values/dimens.xml b/res/values/dimens.xml
index bfa02c0..2882b1f 100644
--- a/res/values/dimens.xml
+++ b/res/values/dimens.xml
@@ -146,6 +146,8 @@
     <dimen name="widget_picker_education_tip_width">120dp</dimen>
     <dimen name="widget_picker_education_tip_min_margin">4dp</dimen>
 
+    <dimen name="widget_picker_view_pager_top_padding">16dp</dimen>
+
     <!-- Padding applied to shortcut previews -->
     <dimen name="shortcut_preview_padding_left">0dp</dimen>
     <dimen name="shortcut_preview_padding_right">0dp</dimen>
diff --git a/robolectric_tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java b/robolectric_tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java
index a2abfd5..275cf81 100644
--- a/robolectric_tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/ModelMultiCallbacksTest.java
@@ -26,12 +26,12 @@
 
 import android.os.Process;
 
-import com.android.launcher3.PagedView;
 import com.android.launcher3.model.BgDataModel.Callbacks;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.shadows.ShadowLooperExecutor;
 import com.android.launcher3.util.Executors;
+import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.LauncherLayoutBuilder;
 import com.android.launcher3.util.LauncherModelHelper;
 import com.android.launcher3.util.LooperExecutor;
@@ -92,7 +92,7 @@
         // Add a new callback
         cb1.reset();
         MyCallbacks cb2 = spy(MyCallbacks.class);
-        cb2.mPageToBindSync = 2;
+        cb2.mPageToBindSync = IntSet.wrap(2);
         mModelHelper.getModel().addCallbacksAndLoad(cb2);
 
         waitForLoaderAndTempMainThread();
@@ -178,16 +178,16 @@
     private abstract static class MyCallbacks implements Callbacks {
 
         final List<ItemInfo> mItems = new ArrayList<>();
-        int mPageToBindSync = 0;
-        int mPageBoundSync = PagedView.INVALID_PAGE;
+        IntSet mPageToBindSync = IntSet.wrap(0);
+        IntSet mPageBoundSync = new IntSet();
         ViewOnDrawExecutor mDeferredExecutor;
         AppInfo[] mAppInfos;
 
         MyCallbacks() { }
 
         @Override
-        public void onPageBoundSynchronously(int page) {
-            mPageBoundSync = page;
+        public void onPagesBoundSynchronously(IntSet pages) {
+            mPageBoundSync = pages;
         }
 
         @Override
@@ -206,13 +206,13 @@
         }
 
         @Override
-        public int getPageToBindSynchronously() {
+        public IntSet getPagesToBindSynchronously() {
             return mPageToBindSync;
         }
 
         public void reset() {
             mItems.clear();
-            mPageBoundSync = PagedView.INVALID_PAGE;
+            mPageBoundSync = new IntSet();
             mDeferredExecutor = null;
             mAppInfos = null;
         }
diff --git a/robolectric_tests/src/com/android/launcher3/util/LauncherPageRestoreHelperTest.java b/robolectric_tests/src/com/android/launcher3/util/LauncherPageRestoreHelperTest.java
new file mode 100644
index 0000000..51f5851
--- /dev/null
+++ b/robolectric_tests/src/com/android/launcher3/util/LauncherPageRestoreHelperTest.java
@@ -0,0 +1,224 @@
+/**
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.launcher3.util;
+
+import android.os.Bundle;
+
+import com.android.launcher3.LauncherPageRestoreHelper;
+import com.android.launcher3.Workspace;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.RobolectricTestRunner;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.when;
+
+@RunWith(RobolectricTestRunner.class)
+public class LauncherPageRestoreHelperTest {
+
+    // Type: int
+    private static final String RUNTIME_STATE_CURRENT_SCREEN = "launcher.current_screen";
+    // Type: int
+    private static final String RUNTIME_STATE_CURRENT_SCREEN_COUNT =
+            "launcher.current_screen_count";
+
+    private LauncherPageRestoreHelper mPageRestoreHelper;
+    private Bundle mState;
+
+    @Mock
+    private Workspace mWorkspace;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mPageRestoreHelper = new LauncherPageRestoreHelper(mWorkspace);
+        mState = new Bundle();
+    }
+
+    @Test
+    public void givenNoChildrenInWorkspace_whenSavePages_thenNothingSaved() {
+        when(mWorkspace.getChildCount()).thenReturn(0);
+
+        mPageRestoreHelper.savePagesToRestore(mState);
+
+        assertFalse(mState.containsKey(RUNTIME_STATE_CURRENT_SCREEN_COUNT));
+        assertFalse(mState.containsKey(RUNTIME_STATE_CURRENT_SCREEN));
+    }
+
+    @Test
+    public void givenMultipleCurrentPages_whenSavePages_thenSavedCorrectly() {
+        when(mWorkspace.getChildCount()).thenReturn(5);
+        when(mWorkspace.getCurrentPage()).thenReturn(2);
+        givenPanelCount(2);
+
+        mPageRestoreHelper.savePagesToRestore(mState);
+
+        assertEquals(5, mState.getInt(RUNTIME_STATE_CURRENT_SCREEN_COUNT));
+        assertEquals(2, mState.getInt(RUNTIME_STATE_CURRENT_SCREEN));
+    }
+
+    @Test
+    public void givenNullSavedState_whenRestorePages_thenReturnEmptyIntSet() {
+        IntSet result = mPageRestoreHelper.getPagesToRestore(null);
+
+        assertTrue(result.isEmpty());
+    }
+
+    @Test
+    public void givenTotalPageCountMissing_whenRestorePages_thenReturnEmptyIntSet() {
+        givenSavedCurrentPage(1);
+        givenPanelCount(1);
+
+        IntSet result = mPageRestoreHelper.getPagesToRestore(mState);
+
+        assertTrue(result.isEmpty());
+    }
+
+    @Test
+    public void givenCurrentPageMissing_whenRestorePages_thenReturnEmptyIntSet() {
+        givenSavedPageCount(3);
+        givenPanelCount(2);
+
+        IntSet result = mPageRestoreHelper.getPagesToRestore(mState);
+
+        assertTrue(result.isEmpty());
+    }
+
+    @Test
+    public void givenOnePanel_whenRestorePages_thenReturnThatPage() {
+        givenSavedCurrentPage(2);
+        givenSavedPageCount(5);
+        givenPanelCount(1);
+
+        IntSet result = mPageRestoreHelper.getPagesToRestore(mState);
+
+        assertEquals(1, result.size());
+        assertEquals(2, result.getArray().get(0));
+    }
+
+    @Test
+    public void givenTwoPanelOnFirstPages_whenRestorePages_thenReturnThosePages() {
+        givenSavedCurrentPage(0, 1);
+        givenSavedPageCount(2);
+        givenPanelCount(2);
+
+        IntSet result = mPageRestoreHelper.getPagesToRestore(mState);
+
+        assertEquals(IntSet.wrap(0, 1), result);
+    }
+
+    @Test
+    public void givenTwoPanelOnMiddlePages_whenRestorePages_thenReturnThosePages() {
+        givenSavedCurrentPage(2, 3);
+        givenSavedPageCount(5);
+        givenPanelCount(2);
+
+        IntSet result = mPageRestoreHelper.getPagesToRestore(mState);
+
+        assertEquals(IntSet.wrap(2, 3), result);
+    }
+
+    @Test
+    public void givenTwoPanelOnLastPage_whenRestorePages_thenReturnOnlyLastPage() {
+        // The device has two panel home but the current page is the last page, so we don't have
+        // a right panel, only the left one.
+        givenSavedCurrentPage(2);
+        givenSavedPageCount(3);
+        givenPanelCount(2);
+
+        IntSet result = mPageRestoreHelper.getPagesToRestore(mState);
+
+        assertEquals(IntSet.wrap(2), result);
+    }
+
+    @Test
+    public void givenOnlyOnePageAndPhoneFolding_whenRestorePages_thenReturnOnlyOnePage() {
+        givenSavedCurrentPage(0);
+        givenSavedPageCount(1);
+        givenPanelCount(1);
+
+        IntSet result = mPageRestoreHelper.getPagesToRestore(mState);
+
+        assertEquals(IntSet.wrap(0), result);
+    }
+
+    @Test
+    public void givenPhoneFolding_whenRestorePages_thenReturnOnlyTheFirstCurrentPage() {
+        givenSavedCurrentPage(2, 3);
+        givenSavedPageCount(4);
+        givenPanelCount(1);
+
+        IntSet result = mPageRestoreHelper.getPagesToRestore(mState);
+
+        assertEquals(IntSet.wrap(2), result);
+    }
+
+    @Test
+    public void givenPhoneUnfolding_whenRestorePages_thenReturnCurrentPagePlusTheNextOne() {
+        givenSavedCurrentPage(2);
+        givenSavedPageCount(4);
+        givenPanelCount(2);
+
+        IntSet result = mPageRestoreHelper.getPagesToRestore(mState);
+
+        assertEquals(IntSet.wrap(2, 3), result);
+    }
+
+    @Test
+    public void givenPhoneUnfoldingOnLastPage_whenRestorePages_thenReturnOnlyLastPage() {
+        givenSavedCurrentPage(4);
+        givenSavedPageCount(5);
+        givenPanelCount(2);
+
+        IntSet result = mPageRestoreHelper.getPagesToRestore(mState);
+
+        assertEquals(IntSet.wrap(4), result);
+    }
+
+    @Test
+    public void givenOnlyOnePageAndPhoneUnfolding_whenRestorePages_thenReturnOnlyOnePage() {
+        givenSavedCurrentPage(0);
+        givenSavedPageCount(1);
+        givenPanelCount(2);
+
+        IntSet result = mPageRestoreHelper.getPagesToRestore(mState);
+
+        assertEquals(IntSet.wrap(0), result);
+    }
+
+    private void givenPanelCount(int panelCount) {
+        when(mWorkspace.getPanelCount()).thenReturn(panelCount);
+        when(mWorkspace.getLeftmostVisiblePageForIndex(anyInt())).thenAnswer(invocation -> {
+            int pageIndex = invocation.getArgument(0);
+            return pageIndex * panelCount / panelCount;
+        });
+    }
+
+    private void givenSavedPageCount(int pageCount) {
+        mState.putInt(RUNTIME_STATE_CURRENT_SCREEN_COUNT, pageCount);
+    }
+
+    private void givenSavedCurrentPage(int... pages) {
+        mState.putInt(RUNTIME_STATE_CURRENT_SCREEN, pages[0]);
+    }
+}
diff --git a/src/com/android/launcher3/DeleteDropTarget.java b/src/com/android/launcher3/DeleteDropTarget.java
index 80ec192..ba55834 100644
--- a/src/com/android/launcher3/DeleteDropTarget.java
+++ b/src/com/android/launcher3/DeleteDropTarget.java
@@ -33,6 +33,7 @@
 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.util.IntSet;
 import com.android.launcher3.views.Snackbar;
 
 public class DeleteDropTarget extends ButtonDropTarget {
@@ -131,7 +132,7 @@
             onAccessibilityDrop(null, item);
             ModelWriter modelWriter = mLauncher.getModelWriter();
             Runnable onUndoClicked = () -> {
-                mLauncher.setPageToBindSynchronously(itemPage);
+                mLauncher.setPagesToBindSynchronously(IntSet.wrap(itemPage));
                 modelWriter.abortDelete();
                 mLauncher.getStatsLogManager().logger().log(LAUNCHER_UNDO);
             };
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 8889e60..4a7937b 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -105,6 +105,7 @@
 import android.widget.Toast;
 
 import androidx.annotation.CallSuper;
+import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 import androidx.annotation.StringRes;
 import androidx.annotation.VisibleForTesting;
@@ -164,6 +165,7 @@
 import com.android.launcher3.util.ActivityTracker;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.IntArray;
+import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.ItemInfoMatcher;
 import com.android.launcher3.util.MultiValueAlpha;
 import com.android.launcher3.util.MultiValueAlpha.AlphaProperty;
@@ -245,8 +247,6 @@
     protected static final int REQUEST_LAST = 100;
 
     // Type: int
-    private static final String RUNTIME_STATE_CURRENT_SCREEN = "launcher.current_screen";
-    // Type: int
     private static final String RUNTIME_STATE = "launcher.state";
     // Type: PendingRequestArgs
     private static final String RUNTIME_STATE_PENDING_REQUEST_ARGS = "launcher.request_args";
@@ -284,6 +284,8 @@
     private WidgetManagerHelper mAppWidgetManager;
     private LauncherAppWidgetHost mAppWidgetHost;
 
+    private LauncherPageRestoreHelper mPageRestoreHelper;
+
     private final int[] mTmpAddItemCellCoordinates = new int[2];
 
     @Thunk
@@ -319,8 +321,8 @@
 
     private PopupDataProvider mPopupDataProvider;
 
-    private int mSynchronouslyBoundPage = PagedView.INVALID_PAGE;
-    private int mPageToBindSynchronously = PagedView.INVALID_PAGE;
+    private IntSet mSynchronouslyBoundPages = new IntSet();
+    private IntSet mPagesToBindSynchronously = new IntSet();
 
     // We only want to get the SharedPreferences once since it does an FS stat each time we get
     // it from the context.
@@ -455,13 +457,10 @@
         restoreState(savedInstanceState);
         mStateManager.reapplyState();
 
-        // We only load the page synchronously if the user rotates (or triggers a
-        // configuration change) while launcher is in the foreground
-        int currentScreen = PagedView.INVALID_PAGE;
+        mPageRestoreHelper = new LauncherPageRestoreHelper(mWorkspace);
         if (savedInstanceState != null) {
-            currentScreen = savedInstanceState.getInt(RUNTIME_STATE_CURRENT_SCREEN, currentScreen);
+            mPagesToBindSynchronously = mPageRestoreHelper.getPagesToRestore(savedInstanceState);
         }
-        mPageToBindSynchronously = currentScreen;
 
         if (!mModel.addCallbacksAndLoad(this)) {
             if (!internalStateHandled) {
@@ -1525,18 +1524,17 @@
     @Override
     public void onRestoreInstanceState(Bundle state) {
         super.onRestoreInstanceState(state);
-        mWorkspace.restoreInstanceStateForChild(mSynchronouslyBoundPage);
+        if (mSynchronouslyBoundPages != null) {
+            mSynchronouslyBoundPages.forEach(page -> mWorkspace.restoreInstanceStateForChild(page));
+        }
     }
 
     @Override
     protected void onSaveInstanceState(Bundle outState) {
-        if (mWorkspace.getChildCount() > 0) {
-            outState.putInt(RUNTIME_STATE_CURRENT_SCREEN, mWorkspace.getNextPage());
+        mPageRestoreHelper.savePagesToRestore(outState);
 
-        }
         outState.putInt(RUNTIME_STATE, mStateManager.getState().ordinal);
 
-
         AbstractFloatingView widgets = AbstractFloatingView
                 .getOpenView(this, AbstractFloatingView.TYPE_WIDGETS_FULL_SHEET);
         if (widgets != null) {
@@ -2015,24 +2013,24 @@
     }
 
     /**
-     * Sets the next page to bind synchronously on next bind.
-     * @param page
+     * Sets the next pages to bind synchronously on next bind.
+     * @param pages should not be null.
      */
-    public void setPageToBindSynchronously(int page) {
-        mPageToBindSynchronously = page;
+    public void setPagesToBindSynchronously(@NonNull IntSet pages) {
+        mPagesToBindSynchronously = pages;
     }
 
     /**
      * Implementation of the method from LauncherModel.Callbacks.
      */
     @Override
-    public int getPageToBindSynchronously() {
-        if (mPageToBindSynchronously != PagedView.INVALID_PAGE) {
-            return mPageToBindSynchronously;
-        } else  if (mWorkspace != null) {
-            return mWorkspace.getCurrentPage();
+    public IntSet getPagesToBindSynchronously() {
+        if (mPagesToBindSynchronously != null && !mPagesToBindSynchronously.isEmpty()) {
+            return mPagesToBindSynchronously;
+        } else if (mWorkspace != null) {
+            return mWorkspace.getVisiblePageIndices();
         } else {
-            return 0;
+            return new IntSet();
         }
     }
 
@@ -2448,10 +2446,10 @@
         return info;
     }
 
-    public void onPageBoundSynchronously(int page) {
-        mSynchronouslyBoundPage = page;
-        mWorkspace.setCurrentPage(page);
-        mPageToBindSynchronously = PagedView.INVALID_PAGE;
+    public void onPagesBoundSynchronously(IntSet pages) {
+        mSynchronouslyBoundPages = pages;
+        mWorkspace.setCurrentPage(pages.getArray().get(0));
+        mPagesToBindSynchronously = new IntSet();
     }
 
     @Override
@@ -2497,7 +2495,7 @@
      *
      * Implementation of the method from LauncherModel.Callbacks.
      */
-    public void finishBindingItems(int pageBoundFirst) {
+    public void finishBindingItems(IntSet pagesBoundFirst) {
         Object traceToken = TraceHelper.INSTANCE.beginSection("finishBindingItems");
         mWorkspace.restoreInstanceStateForRemainingPages();
 
@@ -2512,11 +2510,13 @@
         ItemInstallQueue.INSTANCE.get(this)
                 .resumeModelPush(FLAG_LOADER_RUNNING);
 
+        int currentPage = pagesBoundFirst != null && !pagesBoundFirst.isEmpty()
+                ? pagesBoundFirst.getArray().get(0) : PagedView.INVALID_PAGE;
         // When undoing the removal of the last item on a page, return to that page.
         // Since we are just resetting the current page without user interaction,
         // override the previous page so we don't log the page switch.
-        mWorkspace.setCurrentPage(pageBoundFirst, pageBoundFirst /* overridePrevPage */);
-        mPageToBindSynchronously = PagedView.INVALID_PAGE;
+        mWorkspace.setCurrentPage(currentPage, currentPage /* overridePrevPage */);
+        mPagesToBindSynchronously = new IntSet();
 
         // Cache one page worth of icons
         getViewCache().setCacheSize(R.layout.folder_application,
diff --git a/src/com/android/launcher3/LauncherPageRestoreHelper.java b/src/com/android/launcher3/LauncherPageRestoreHelper.java
new file mode 100644
index 0000000..e679a12
--- /dev/null
+++ b/src/com/android/launcher3/LauncherPageRestoreHelper.java
@@ -0,0 +1,92 @@
+/**
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.launcher3;
+
+import android.os.Bundle;
+import android.util.Log;
+
+import androidx.annotation.VisibleForTesting;
+
+import com.android.launcher3.util.IntSet;
+
+import static androidx.annotation.VisibleForTesting.PACKAGE_PRIVATE;
+
+/**
+ * There's a logic which prioritizes the binding for the current page and defers the other pages'
+ * binding. If two panel home is enabled, we want to bind both pages together.
+ * LauncherPageRestoreHelper's purpose is to contain the logic for persisting, restoring and
+ * calculating which pages to load immediately.
+ */
+public class LauncherPageRestoreHelper {
+
+    public static final String TAG = "LauncherPageRestoreHelper";
+
+    // Type: int
+    private static final String RUNTIME_STATE_CURRENT_SCREEN = "launcher.current_screen";
+    // Type: int
+    private static final String RUNTIME_STATE_CURRENT_SCREEN_COUNT =
+            "launcher.current_screen_count";
+
+    private Workspace mWorkspace;
+
+    public LauncherPageRestoreHelper(Workspace workspace) {
+        this.mWorkspace = workspace;
+    }
+
+    /**
+     * Some configuration changes trigger Launcher to recreate itself, and we want to give more
+     * priority to the currently active pages in the restoration process.
+     */
+    @VisibleForTesting(otherwise = PACKAGE_PRIVATE)
+    public IntSet getPagesToRestore(Bundle savedInstanceState) {
+        IntSet pagesToRestore = new IntSet();
+
+        if (savedInstanceState == null) {
+            return pagesToRestore;
+        }
+
+        int currentPage = savedInstanceState.getInt(RUNTIME_STATE_CURRENT_SCREEN, -1);
+        int totalPageCount = savedInstanceState.getInt(RUNTIME_STATE_CURRENT_SCREEN_COUNT, -1);
+        int panelCount = mWorkspace.getPanelCount();
+
+        if (totalPageCount <= 0 || currentPage < 0) {
+            Log.e(TAG, "getPagesToRestore: Invalid input: " + totalPageCount + ", " + currentPage);
+            return pagesToRestore;
+        }
+
+        int newCurrentPage = mWorkspace.getLeftmostVisiblePageForIndex(currentPage);
+        for (int page = newCurrentPage; page < newCurrentPage + panelCount
+                && page < totalPageCount; page++) {
+            pagesToRestore.add(page);
+        }
+
+        return pagesToRestore;
+    }
+
+    /**
+     * This should be called from Launcher's onSaveInstanceState method to persist everything that
+     * is necessary to calculate later which pages need to be initialized first after a
+     * configuration change.
+     */
+    @VisibleForTesting(otherwise = PACKAGE_PRIVATE)
+    public void savePagesToRestore(Bundle outState) {
+        int pageCount = mWorkspace.getChildCount();
+        if (pageCount > 0) {
+            outState.putInt(RUNTIME_STATE_CURRENT_SCREEN, mWorkspace.getCurrentPage());
+            outState.putInt(RUNTIME_STATE_CURRENT_SCREEN_COUNT, pageCount);
+        }
+    }
+}
diff --git a/src/com/android/launcher3/PagedView.java b/src/com/android/launcher3/PagedView.java
index b26a7ea..123ae6c 100644
--- a/src/com/android/launcher3/PagedView.java
+++ b/src/com/android/launcher3/PagedView.java
@@ -16,6 +16,7 @@
 
 package com.android.launcher3;
 
+import static androidx.annotation.VisibleForTesting.PACKAGE_PRIVATE;
 import static com.android.launcher3.anim.Interpolators.SCROLL;
 import static com.android.launcher3.compat.AccessibilityManagerCompat.isAccessibilityEnabled;
 import static com.android.launcher3.compat.AccessibilityManagerCompat.isObservedEventType;
@@ -48,6 +49,7 @@
 import android.widget.ScrollView;
 
 import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
 
 import com.android.launcher3.compat.AccessibilityManagerCompat;
 import com.android.launcher3.config.FeatureFlags;
@@ -55,6 +57,7 @@
 import com.android.launcher3.touch.PagedOrientationHandler;
 import com.android.launcher3.touch.PagedOrientationHandler.ChildBounds;
 import com.android.launcher3.util.EdgeEffectCompat;
+import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.Thunk;
 import com.android.launcher3.views.ActivityContext;
 
@@ -282,9 +285,15 @@
         return newPage;
     }
 
-    private int getLeftmostVisiblePageForIndex(int pageIndex) {
+    /**
+     * In most cases where panelCount is 1, this method will just return the page index that was
+     * passed in.
+     * But for example when two panel home is enabled we might need the leftmost visible page index
+     * because that page is the current page.
+     */
+    public int getLeftmostVisiblePageForIndex(int pageIndex) {
         int panelCount = getPanelCount();
-        return (pageIndex / panelCount) * panelCount;
+        return pageIndex - pageIndex % panelCount;
     }
 
     /**
@@ -295,16 +304,34 @@
     }
 
     /**
+     * Returns an IntSet with the indices of the currently visible pages
+     */
+    @VisibleForTesting(otherwise = PACKAGE_PRIVATE)
+    public IntSet getVisiblePageIndices() {
+        IntSet visiblePageIndices = new IntSet();
+        int panelCount = getPanelCount();
+        int pageCount = getPageCount();
+
+        // If a device goes from one panel to two panel (i.e. unfolding a foldable device) while
+        // an odd indexed page is the current page, then the new leftmost visible page will be
+        // different from the old mCurrentPage.
+        int currentPage = getLeftmostVisiblePageForIndex(mCurrentPage);
+        for (int page = currentPage; page < currentPage + panelCount && page < pageCount; page++) {
+            visiblePageIndices.add(page);
+        }
+        return visiblePageIndices;
+    }
+
+    /**
      * Executes the callback against each visible page
      */
     public void forEachVisiblePage(Consumer<View> callback) {
-        int panelCount = getPanelCount();
-        for (int i = mCurrentPage; i < mCurrentPage + panelCount; i++) {
-            View page = getPageAt(i);
+        getVisiblePageIndices().forEach(pageIndex -> {
+            View page = getPageAt(pageIndex);
             if (page != null) {
                 callback.accept(page);
             }
-        }
+        });
     }
 
     /**
diff --git a/src/com/android/launcher3/Workspace.java b/src/com/android/launcher3/Workspace.java
index 98d80fe..78e8048 100644
--- a/src/com/android/launcher3/Workspace.java
+++ b/src/com/android/launcher3/Workspace.java
@@ -16,6 +16,7 @@
 
 package com.android.launcher3;
 
+import static androidx.annotation.VisibleForTesting.PROTECTED;
 import static com.android.launcher3.LauncherAnimUtils.SPRING_LOADED_EXIT_DELAY;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APPLICATION;
 import static com.android.launcher3.LauncherState.ALL_APPS;
@@ -63,6 +64,8 @@
 import android.view.accessibility.AccessibilityNodeInfo;
 import android.widget.Toast;
 
+import androidx.annotation.VisibleForTesting;
+
 import com.android.launcher3.accessibility.AccessibleDragListenerAdapter;
 import com.android.launcher3.accessibility.WorkspaceAccessibilityHelper;
 import com.android.launcher3.anim.Interpolators;
@@ -461,7 +464,8 @@
     }
 
     @Override
-    protected int getPanelCount() {
+    @VisibleForTesting(otherwise = PROTECTED)
+    public int getPanelCount() {
         return isTwoPanelEnabled() ? 2 : super.getPanelCount();
     }
 
diff --git a/src/com/android/launcher3/config/FeatureFlags.java b/src/com/android/launcher3/config/FeatureFlags.java
index 0c0c92e..1715161 100644
--- a/src/com/android/launcher3/config/FeatureFlags.java
+++ b/src/com/android/launcher3/config/FeatureFlags.java
@@ -214,14 +214,14 @@
             + "predictions to be updated while they are visible to the user.");
 
     public static final BooleanFlag ENABLE_TASKBAR = getDebugFlag(
-            "ENABLE_TASKBAR", false, "Allows a system Taskbar to be shown on larger devices.");
+            "ENABLE_TASKBAR", true, "Allows a system Taskbar to be shown on larger devices.");
 
     public static final BooleanFlag ENABLE_OVERVIEW_GRID = getDebugFlag(
-            "ENABLE_OVERVIEW_GRID", false, "Uses grid overview layout. "
+            "ENABLE_OVERVIEW_GRID", true, "Uses grid overview layout. "
             + "Only applicable on large screen devices.");
 
     public static final BooleanFlag ENABLE_TWO_PANEL_HOME = getDebugFlag(
-            "ENABLE_TWO_PANEL_HOME", false,
+            "ENABLE_TWO_PANEL_HOME", true,
             "Uses two panel on home screen. Only applicable on large screen devices.");
 
     public static final BooleanFlag ENABLE_SPLIT_SELECT = getDebugFlag(
diff --git a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
index 2a1aec8..952b850 100644
--- a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
+++ b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
@@ -79,6 +79,7 @@
 import com.android.launcher3.uioverrides.plugins.PluginManagerWrapper;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.IntArray;
+import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.views.BaseDragLayer;
@@ -391,11 +392,14 @@
         ArrayList<ItemInfo> otherWorkspaceItems = new ArrayList<>();
         ArrayList<LauncherAppWidgetInfo> currentAppWidgets = new ArrayList<>();
         ArrayList<LauncherAppWidgetInfo> otherAppWidgets = new ArrayList<>();
-        filterCurrentWorkspaceItems(0 /* currentScreenId */,
-                dataModel.workspaceItems, currentWorkspaceItems,
-                otherWorkspaceItems);
-        filterCurrentWorkspaceItems(0 /* currentScreenId */, dataModel.appWidgets,
-                currentAppWidgets, otherAppWidgets);
+
+        IntSet currentScreenIds = IntSet.wrap(0);
+        // TODO(b/185508060): support two panel preview.
+        filterCurrentWorkspaceItems(currentScreenIds, dataModel.workspaceItems,
+                currentWorkspaceItems, otherWorkspaceItems);
+        filterCurrentWorkspaceItems(currentScreenIds, dataModel.appWidgets, currentAppWidgets,
+                otherAppWidgets);
+
         sortWorkspaceItemsSpatially(mIdp, currentWorkspaceItems);
         for (ItemInfo itemInfo : currentWorkspaceItems) {
             switch (itemInfo.itemType) {
diff --git a/src/com/android/launcher3/model/BaseLoaderResults.java b/src/com/android/launcher3/model/BaseLoaderResults.java
index 5c85bab..12ee676 100644
--- a/src/com/android/launcher3/model/BaseLoaderResults.java
+++ b/src/com/android/launcher3/model/BaseLoaderResults.java
@@ -24,13 +24,13 @@
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherModel.CallbackTask;
-import com.android.launcher3.PagedView;
 import com.android.launcher3.model.BgDataModel.Callbacks;
 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.util.IntArray;
+import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.LooperExecutor;
 import com.android.launcher3.util.LooperIdleLock;
 import com.android.launcher3.util.ViewOnDrawExecutor;
@@ -160,20 +160,26 @@
         }
 
         private void bind() {
-            final int currentScreen;
+            IntSet currentScreenIndices;
             {
                 // Create an anonymous scope to calculate currentScreen as it has to be a
                 // final variable.
-                int currScreen = mCallbacks.getPageToBindSynchronously();
-                if (currScreen >= mOrderedScreenIds.size()) {
-                    // There may be no workspace screens (just hotseat items and an empty page).
-                    currScreen = PagedView.INVALID_PAGE;
+                IntSet screenIndices = mCallbacks.getPagesToBindSynchronously();
+                if (screenIndices == null || screenIndices.isEmpty()
+                        || screenIndices.getArray().get(screenIndices.size() - 1)
+                        >= mOrderedScreenIds.size()) {
+                    // There maybe no workspace screens (just hotseat items and an empty page).
+                    // Also we want to prevent IndexOutOfBoundsExceptions.
+                    screenIndices = new IntSet();
                 }
-                currentScreen = currScreen;
+                currentScreenIndices = screenIndices;
             }
-            final boolean validFirstPage = currentScreen >= 0;
-            final int currentScreenId =
-                    validFirstPage ? mOrderedScreenIds.get(currentScreen) : INVALID_SCREEN_ID;
+
+            final boolean validFirstPage = !currentScreenIndices.isEmpty();
+
+            IntSet currentScreenIds  = new IntSet();
+            currentScreenIndices.forEach(
+                    index -> currentScreenIds.add(mOrderedScreenIds.get(index)));
 
             // Separate the items that are on the current screen, and all the other remaining items
             ArrayList<ItemInfo> currentWorkspaceItems = new ArrayList<>();
@@ -181,9 +187,9 @@
             ArrayList<LauncherAppWidgetInfo> currentAppWidgets = new ArrayList<>();
             ArrayList<LauncherAppWidgetInfo> otherAppWidgets = new ArrayList<>();
 
-            filterCurrentWorkspaceItems(currentScreenId, mWorkspaceItems, currentWorkspaceItems,
+            filterCurrentWorkspaceItems(currentScreenIds, mWorkspaceItems, currentWorkspaceItems,
                     otherWorkspaceItems);
-            filterCurrentWorkspaceItems(currentScreenId, mAppWidgets, currentAppWidgets,
+            filterCurrentWorkspaceItems(currentScreenIds, mAppWidgets, currentAppWidgets,
                     otherAppWidgets);
             final InvariantDeviceProfile idp = mApp.getInvariantDeviceProfile();
             sortWorkspaceItemsSpatially(idp, currentWorkspaceItems);
@@ -220,14 +226,14 @@
             bindWorkspaceItems(otherWorkspaceItems, deferredExecutor);
             bindAppWidgets(otherAppWidgets, deferredExecutor);
             // Tell the workspace that we're done binding items
-            executeCallbacksTask(c -> c.finishBindingItems(currentScreen), deferredExecutor);
+            executeCallbacksTask(c -> c.finishBindingItems(currentScreenIndices), deferredExecutor);
 
             if (validFirstPage) {
                 executeCallbacksTask(c -> {
                     // We are loading synchronously, which means, some of the pages will be
                     // bound after first draw. Inform the mCallbacks that page binding is
                     // not complete, and schedule the remaining pages.
-                    c.onPageBoundSynchronously(currentScreen);
+                    c.onPagesBoundSynchronously(currentScreenIndices);
                     c.executeOnNextDraw((ViewOnDrawExecutor) deferredExecutor);
 
                 }, mUiExecutor);
diff --git a/src/com/android/launcher3/model/BgDataModel.java b/src/com/android/launcher3/model/BgDataModel.java
index 1d7d1a2..037f408 100644
--- a/src/com/android/launcher3/model/BgDataModel.java
+++ b/src/com/android/launcher3/model/BgDataModel.java
@@ -446,15 +446,16 @@
         int FLAG_QUIET_MODE_CHANGE_PERMISSION = 1 << 2;
 
         /**
-         * Returns the page number to bind first, synchronously if possible or -1
+         * Returns an IntSet of page numbers to bind first, synchronously if possible
+         * or an empty IntSet
          */
-        int getPageToBindSynchronously();
+        IntSet getPagesToBindSynchronously();
         void clearPendingBinds();
         void startBinding();
         void bindItems(List<ItemInfo> shortcuts, boolean forceAnimateIcons);
         void bindScreens(IntArray orderedScreenIds);
         void finishFirstPageBind(ViewOnDrawExecutor executor);
-        void finishBindingItems(int pageBoundFirst);
+        void finishBindingItems(IntSet pagesBoundFirst);
         void preAddApps();
         void bindAppsAdded(IntArray newScreens,
                 ArrayList<ItemInfo> addNotAnimated, ArrayList<ItemInfo> addAnimated);
@@ -468,7 +469,7 @@
         void bindRestoreItemsChange(HashSet<ItemInfo> updates);
         void bindWorkspaceComponentsRemoved(ItemInfoMatcher matcher);
         void bindAllWidgets(List<WidgetsListBaseEntry> widgets);
-        void onPageBoundSynchronously(int page);
+        void onPagesBoundSynchronously(IntSet pages);
         void executeOnNextDraw(ViewOnDrawExecutor executor);
         void bindDeepShortcutMap(HashMap<ComponentKey, Integer> deepShortcutMap);
 
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index 79396b1..34a21fe 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -83,6 +83,7 @@
 import com.android.launcher3.shortcuts.ShortcutRequest.QueryResult;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.IOUtils;
+import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.LooperIdleLock;
 import com.android.launcher3.util.PackageManagerHelper;
 import com.android.launcher3.util.PackageUserKey;
@@ -173,8 +174,9 @@
         ArrayList<ItemInfo> allItems = mBgDataModel.getAllWorkspaceItems();
         // Screen set is never empty
         final int firstScreen = mBgDataModel.collectWorkspaceScreens().get(0);
+        // TODO(b/185515153): support two panel home.
 
-        filterCurrentWorkspaceItems(firstScreen, allItems, firstScreenItems,
+        filterCurrentWorkspaceItems(IntSet.wrap(firstScreen), allItems, firstScreenItems,
                 new ArrayList<>() /* otherScreenItems are ignored */);
         mFirstScreenBroadcast.sendBroadcasts(mApp.getContext(), firstScreenItems);
     }
diff --git a/src/com/android/launcher3/model/ModelUtils.java b/src/com/android/launcher3/model/ModelUtils.java
index 9b5fac8..58aa9e5 100644
--- a/src/com/android/launcher3/model/ModelUtils.java
+++ b/src/com/android/launcher3/model/ModelUtils.java
@@ -51,7 +51,8 @@
      * Filters the set of items who are directly or indirectly (via another container) on the
      * specified screen.
      */
-    public static <T extends ItemInfo> void filterCurrentWorkspaceItems(int currentScreenId,
+    public static <T extends ItemInfo> void filterCurrentWorkspaceItems(
+            IntSet currentScreenIds,
             ArrayList<T> allWorkspaceItems,
             ArrayList<T> currentScreenItems,
             ArrayList<T> otherScreenItems) {
@@ -65,7 +66,7 @@
                 (lhs, rhs) -> Integer.compare(lhs.container, rhs.container));
         for (T info : allWorkspaceItems) {
             if (info.container == LauncherSettings.Favorites.CONTAINER_DESKTOP) {
-                if (info.screenId == currentScreenId) {
+                if (currentScreenIds.contains(info.screenId)) {
                     currentScreenItems.add(info);
                     itemsOnScreen.add(info.id);
                 } else {
diff --git a/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncher.java b/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncher.java
index 5999091..b271a6a 100644
--- a/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncher.java
+++ b/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncher.java
@@ -41,6 +41,7 @@
 import com.android.launcher3.popup.PopupDataProvider;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.IntArray;
+import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.ItemInfoMatcher;
 import com.android.launcher3.util.Themes;
 import com.android.launcher3.util.ViewOnDrawExecutor;
@@ -175,8 +176,8 @@
     }
 
     @Override
-    public int getPageToBindSynchronously() {
-        return 0;
+    public IntSet getPagesToBindSynchronously() {
+        return new IntSet();
     }
 
     @Override
@@ -199,7 +200,7 @@
     }
 
     @Override
-    public void finishBindingItems(int pageBoundFirst) { }
+    public void finishBindingItems(IntSet pagesBoundFirst) { }
 
     @Override
     public void preAddApps() { }
@@ -229,7 +230,7 @@
     public void bindAllWidgets(List<WidgetsListBaseEntry> widgets) { }
 
     @Override
-    public void onPageBoundSynchronously(int page) { }
+    public void onPagesBoundSynchronously(IntSet pages) { }
 
     @Override
     public void executeOnNextDraw(ViewOnDrawExecutor executor) {
diff --git a/src/com/android/launcher3/util/IntArray.java b/src/com/android/launcher3/util/IntArray.java
index 7252f7a..e7235e7 100644
--- a/src/com/android/launcher3/util/IntArray.java
+++ b/src/com/android/launcher3/util/IntArray.java
@@ -17,13 +17,14 @@
 package com.android.launcher3.util;
 
 import java.util.Arrays;
+import java.util.Iterator;
 import java.util.StringTokenizer;
 
 /**
  * Copy of the platform hidden implementation of android.util.IntArray.
  * Implements a growing array of int primitives.
  */
-public class IntArray implements Cloneable {
+public class IntArray implements Cloneable, Iterable<Integer> {
     private static final int MIN_CAPACITY_INCREMENT = 12;
 
     private static final int[] EMPTY_INT = new int[0];
@@ -272,4 +273,30 @@
             throw new ArrayIndexOutOfBoundsException("length=" + len + "; index=" + index);
         }
     }
+
+    @Override
+    public Iterator<Integer> iterator() {
+        return new ValueIterator();
+    }
+
+    @Thunk
+    class ValueIterator implements Iterator<Integer> {
+
+        private int mNextIndex = 0;
+
+        @Override
+        public boolean hasNext() {
+            return mNextIndex < size();
+        }
+
+        @Override
+        public Integer next() {
+            return get(mNextIndex++);
+        }
+
+        @Override
+        public void remove() {
+            throw new UnsupportedOperationException();
+        }
+    }
 }
\ No newline at end of file
diff --git a/src/com/android/launcher3/util/IntSet.java b/src/com/android/launcher3/util/IntSet.java
index 851f129..0f4df62 100644
--- a/src/com/android/launcher3/util/IntSet.java
+++ b/src/com/android/launcher3/util/IntSet.java
@@ -16,11 +16,13 @@
 package com.android.launcher3.util;
 
 import java.util.Arrays;
+import java.util.Iterator;
 
 /**
  * A wrapper over IntArray implementing a growing set of int primitives.
+ * The elements in the array are sorted in ascending order.
  */
-public class IntSet {
+public class IntSet implements Iterable<Integer> {
 
     final IntArray mArray = new IntArray();
 
@@ -34,6 +36,16 @@
         }
     }
 
+    /**
+     * Removes the specified value from the set if it exist.
+     */
+    public void remove(int value) {
+        int index = Arrays.binarySearch(mArray.mValues, 0, mArray.mSize, value);
+        if (index >= 0) {
+            mArray.removeIndex(index);
+        }
+    }
+
     public boolean contains(int value) {
         return Arrays.binarySearch(mArray.mValues, 0, mArray.mSize, value) >= 0;
     }
@@ -61,6 +73,9 @@
         return (obj instanceof IntSet) && ((IntSet) obj).mArray.equals(mArray);
     }
 
+    /**
+     * Returns the wrapped IntArray. The elements in the array are sorted in ascending order.
+     */
     public IntArray getArray() {
         return mArray;
     }
@@ -78,4 +93,21 @@
         Arrays.sort(set.mArray.mValues, 0, set.mArray.mSize);
         return set;
     }
+
+    /**
+     * Returns an IntSet with the given values.
+     */
+    public static IntSet wrap(int... array) {
+        return wrap(IntArray.wrap(array));
+    }
+
+    @Override
+    public Iterator<Integer> iterator() {
+        return mArray.iterator();
+    }
+
+    @Override
+    public String toString() {
+        return "IntSet{" + mArray.toConcatString() + '}';
+    }
 }
diff --git a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
index 039cad8..ee5afd9 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
@@ -126,6 +126,7 @@
         }
     };
     private final int mTabsHeight;
+    private final int mViewPagerTopPadding;
     private final int mWidgetCellHorizontalPadding;
 
     @Nullable private WidgetsRecyclerView mCurrentWidgetsRecyclerView;
@@ -148,6 +149,10 @@
                 ? getContext().getResources()
                         .getDimensionPixelSize(R.dimen.all_apps_header_pill_height)
                 : 0;
+        mViewPagerTopPadding = mHasWorkProfile
+                ? getContext().getResources()
+                    .getDimensionPixelSize(R.dimen.widget_picker_view_pager_top_padding)
+                : 0;
         mWidgetCellHorizontalPadding = 2 * getResources().getDimensionPixelOffset(
                 R.dimen.widget_cell_horizontal_padding);
     }
@@ -499,8 +504,8 @@
                 noWidgetsViewHeight = noWidgetsViewTextBounds.height();
             }
             float maxTableHeight = (mActivityContext.getDeviceProfile().availableHeightPx
-                                        - mTabsHeight - getHeaderViewHeight() - noWidgetsViewHeight)
-                                                * RECOMMENDATION_TABLE_HEIGHT_RATIO;
+                    - mTabsHeight - mViewPagerTopPadding - getHeaderViewHeight()
+                    - noWidgetsViewHeight) * RECOMMENDATION_TABLE_HEIGHT_RATIO;
 
             List<ArrayList<WidgetItem>> recommendedWidgetsInTable =
                     WidgetsTableUtils.groupWidgetItemsIntoTable(recommendedWidgets,
@@ -605,7 +610,7 @@
         return measureHeightWithVerticalMargins(mSearchAndRecommendationViewHolder.mCollapseHandle)
                 + measureHeightWithVerticalMargins(mSearchAndRecommendationViewHolder.mHeaderTitle)
                 + measureHeightWithVerticalMargins(
-                (View) mSearchAndRecommendationViewHolder.mSearchBar);
+                        (View) mSearchAndRecommendationViewHolder.mSearchBarContainer);
     }
 
     /** private the height, in pixel, + the vertical margins of a given view. */