Merge "Add two panel home support for page binding logic" into sc-v2-dev
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/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/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 a8dc93b..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();
 
@@ -71,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;
     }
@@ -88,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() + '}';
+    }
 }