Merge "Add two extra empty pages on two panel launcher home" into sc-v2-dev
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 94dbe00..5be93b2 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -877,11 +877,11 @@
         if (dropLayout == null) {
             // it's possible that the add screen was removed because it was
             // empty and a re-bind occurred
-            mWorkspace.addExtraEmptyScreen();
-            return mWorkspace.commitExtraEmptyScreen();
-        } else {
-            return screenId;
+            mWorkspace.addExtraEmptyScreens();
+            IntSet emptyPagesAdded = mWorkspace.commitExtraEmptyScreens();
+            return emptyPagesAdded.isEmpty() ? -1 : emptyPagesAdded.getArray().get(0);
         }
+        return screenId;
     }
 
     @Thunk
@@ -2181,7 +2181,7 @@
             orderedScreenIds.add(firstScreenPosition, Workspace.FIRST_SCREEN_ID);
         } else if (!FeatureFlags.QSB_ON_FIRST_SCREEN && orderedScreenIds.isEmpty()) {
             // If there are no screens, we need to have an empty screen
-            mWorkspace.addExtraEmptyScreen();
+            mWorkspace.addExtraEmptyScreens();
         }
         bindAddScreens(orderedScreenIds);
 
@@ -2196,17 +2196,7 @@
             // Some empty pages might have been removed while the phone was in a single panel
             // mode, so we want to add those empty pages back.
             IntSet screenIds = IntSet.wrap(orderedScreenIds);
-            for (int i = 0; i < orderedScreenIds.size(); i++) {
-                int screenId = orderedScreenIds.get(i);
-                // Don't add the page pair if the page is the last one and if the pair is on the
-                // right, because that would cause a bug when adding new pages.
-                // TODO: (b/196376162) remove this when the new screen id logic is fixed for two
-                //  panel in Workspace::commitExtraEmptyScreen
-                if (i == orderedScreenIds.size() - 1 && screenId % 2 == 0) {
-                    continue;
-                }
-                screenIds.add(mWorkspace.getPagePair(screenId));
-            }
+            orderedScreenIds.forEach(screenId -> screenIds.add(mWorkspace.getPagePair(screenId)));
             orderedScreenIds = screenIds.getArray();
         }
 
diff --git a/src/com/android/launcher3/Workspace.java b/src/com/android/launcher3/Workspace.java
index b197e1c..b1a9b76 100644
--- a/src/com/android/launcher3/Workspace.java
+++ b/src/com/android/launcher3/Workspace.java
@@ -125,6 +125,7 @@
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
+import java.util.function.Consumer;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
 
@@ -644,18 +645,35 @@
         boolean childOnFinalScreen = false;
 
         if (mDragSourceInternal != null) {
+            int dragSourceChildCount = mDragSourceInternal.getChildCount();
+
+            if (isTwoPanelEnabled()) {
+                int pagePairScreenId = getPagePair(dragObject.dragInfo.screenId);
+                CellLayout pagePair = mWorkspaceScreens.get(pagePairScreenId);
+                if (pagePair == null) {
+                    // TODO: after http://b/198820019 is fixed, remove this
+                    throw new IllegalStateException("Page pair is null, "
+                            + "dragScreenId: " + dragObject.dragInfo.screenId
+                            + ", pagePairScreenId: " + pagePairScreenId
+                            + ", mScreenOrder: " + mScreenOrder.toConcatString()
+                    );
+                }
+                dragSourceChildCount += pagePair.getShortcutsAndWidgets().getChildCount();
+            }
+
             // When the drag view content is a LauncherAppWidgetHostView, we should increment the
             // drag source child count by 1 because the widget in drag has been detached from its
             // original parent, ShortcutAndWidgetContainer, and reattached to the DragView.
-            int dragSourceChildCount =
-                    dragObject.dragView.getContentView() instanceof LauncherAppWidgetHostView
-                            ? mDragSourceInternal.getChildCount() + 1
-                            : mDragSourceInternal.getChildCount();
+            if (dragObject.dragView.getContentView() instanceof LauncherAppWidgetHostView) {
+                dragSourceChildCount++;
+            }
+
             if (dragSourceChildCount == 1) {
                 lastChildOnScreen = true;
             }
             CellLayout cl = (CellLayout) mDragSourceInternal.getParent();
-            if (indexOfChild(cl) == getChildCount() - 1) {
+            if (getLeftmostVisiblePageForIndex(indexOfChild(cl))
+                    == getLeftmostVisiblePageForIndex(getPageCount() - 1)) {
                 childOnFinalScreen = true;
             }
         }
@@ -664,40 +682,83 @@
         if (lastChildOnScreen && childOnFinalScreen) {
             return;
         }
-        if (!mWorkspaceScreens.containsKey(EXTRA_EMPTY_SCREEN_ID)) {
-            insertNewWorkspaceScreen(EXTRA_EMPTY_SCREEN_ID);
+
+        forEachExtraEmptyPageId(extraEmptyPageId -> {
+            if (!mWorkspaceScreens.containsKey(extraEmptyPageId)) {
+                insertNewWorkspaceScreen(extraEmptyPageId);
+            }
+        });
+    }
+
+    /**
+     * Inserts extra empty pages to the end of the existing workspaces.
+     * Usually we add one extra empty screen, but when two panel home is enabled we add
+     * two extra screens.
+     **/
+    public void addExtraEmptyScreens() {
+        forEachExtraEmptyPageId(extraEmptyPageId -> {
+            if (!mWorkspaceScreens.containsKey(extraEmptyPageId)) {
+                insertNewWorkspaceScreen(extraEmptyPageId);
+            }
+        });
+    }
+
+    /**
+     * Calls the consumer with all the necessary extra empty page IDs.
+     * On a normal one panel Workspace that means only EXTRA_EMPTY_SCREEN_ID,
+     * but in a two panel scenario this also includes EXTRA_EMPTY_SCREEN_SECOND_ID.
+     */
+    private void forEachExtraEmptyPageId(Consumer<Integer> callback) {
+        callback.accept(EXTRA_EMPTY_SCREEN_ID);
+        if (isTwoPanelEnabled()) {
+            callback.accept(EXTRA_EMPTY_SCREEN_SECOND_ID);
         }
     }
 
-    public boolean addExtraEmptyScreen() {
-        if (!mWorkspaceScreens.containsKey(EXTRA_EMPTY_SCREEN_ID)) {
-            insertNewWorkspaceScreen(EXTRA_EMPTY_SCREEN_ID);
-            return true;
-        }
-        return false;
-    }
-
+    /**
+     * If two panel home is enabled we convert the last two screens that are visible at the same
+     * time. In other cases we only convert the last page.
+     */
     private void convertFinalScreenToEmptyScreenIfNecessary() {
         if (mLauncher.isWorkspaceLoading()) {
             // Invalid and dangerous operation if workspace is loading
             return;
         }
 
-        if (hasExtraEmptyScreen() || mScreenOrder.size() == 0) return;
-        int finalScreenId = mScreenOrder.get(mScreenOrder.size() - 1);
+        int panelCount = getPanelCount();
+        if (hasExtraEmptyScreens() || mScreenOrder.size() < panelCount) {
+            return;
+        }
 
-        CellLayout finalScreen = mWorkspaceScreens.get(finalScreenId);
+        SparseArray<CellLayout> finalScreens = new SparseArray<>();
 
-        // If the final screen is empty, convert it to the extra empty screen
-        if (finalScreen != null
-                && finalScreen.getShortcutsAndWidgets().getChildCount() == 0
-                && !finalScreen.isDropPending()) {
-            mWorkspaceScreens.remove(finalScreenId);
-            mScreenOrder.removeValue(finalScreenId);
+        int pageCount = mScreenOrder.size();
+        // First we add the last page(s) to the finalScreens collection. The number of final pages
+        // depends on the panel count.
+        for (int pageIndex = pageCount - panelCount; pageIndex < pageCount; pageIndex++) {
+            int screenId = mScreenOrder.get(pageIndex);
+            CellLayout screen = mWorkspaceScreens.get(screenId);
+            if (screen == null || screen.getShortcutsAndWidgets().getChildCount() != 0
+                    || screen.isDropPending()) {
+                // Final screen doesn't exist or it isn't empty or there's a pending drop
+                return;
+            }
+            finalScreens.append(screenId, screen);
+        }
 
-            // if this is the last screen, convert it to the empty screen
-            mWorkspaceScreens.put(EXTRA_EMPTY_SCREEN_ID, finalScreen);
-            mScreenOrder.add(EXTRA_EMPTY_SCREEN_ID);
+        // Then we remove the final screens from the collections (but not from the view hierarchy)
+        // and we store them as extra empty screens.
+        for (int i = 0; i < finalScreens.size(); i++) {
+            int screenId = finalScreens.keyAt(i);
+            CellLayout screen = finalScreens.get(screenId);
+
+            mWorkspaceScreens.remove(screenId);
+            mScreenOrder.removeValue(screenId);
+
+            int newScreenId = mWorkspaceScreens.containsKey(EXTRA_EMPTY_SCREEN_ID)
+                    ? EXTRA_EMPTY_SCREEN_SECOND_ID : EXTRA_EMPTY_SCREEN_ID;
+            mWorkspaceScreens.put(newScreenId, screen);
+            mScreenOrder.add(newScreenId);
         }
     }
 
@@ -705,6 +766,23 @@
         removeExtraEmptyScreenDelayed(0, stripEmptyScreens, null);
     }
 
+    /**
+     * The purpose of this method is to remove empty pages from Workspace.
+     * Empty page(s) from the end of mWorkspaceScreens will always be removed. The pages with
+     * ID = Workspace.EXTRA_EMPTY_SCREEN_IDS will be removed if there are other non-empty pages.
+     * If there are no more non-empty pages left, extra empty page(s) will either stay or get added.
+     *
+     * If stripEmptyScreens is true, all empty pages (not just the ones on the end) will be removed
+     * from the Workspace, and if there are no more pages left then extra empty page(s) will be
+     * added.
+     *
+     * The number of extra empty pages is equal to what getPanelCount() returns.
+     *
+     * After the method returns the possible pages are:
+     * stripEmptyScreens = true : [non-empty pages, extra empty page(s) alone]
+     * stripEmptyScreens = false : [non-empty pages, empty pages (not in the end),
+     *                             extra empty page(s) alone]
+     */
     public void removeExtraEmptyScreenDelayed(
             int delay, boolean stripEmptyScreens, Runnable onComplete) {
         if (mLauncher.isWorkspaceLoading()) {
@@ -718,18 +796,26 @@
             return;
         }
 
+        // First we convert the last page to an extra page if the last page is empty
+        // and we don't already have an extra page.
         convertFinalScreenToEmptyScreenIfNecessary();
-        if (hasExtraEmptyScreen()) {
-            removeView(mWorkspaceScreens.get(EXTRA_EMPTY_SCREEN_ID));
+        // Then we remove the extra page(s) if they are not the only pages left in Workspace.
+        if (hasExtraEmptyScreens()) {
+            forEachExtraEmptyPageId(extraEmptyPageId -> {
+                removeView(mWorkspaceScreens.get(extraEmptyPageId));
+                mWorkspaceScreens.remove(extraEmptyPageId);
+                mScreenOrder.removeValue(extraEmptyPageId);
+            });
+
             setCurrentPage(getNextPage());
-            mWorkspaceScreens.remove(EXTRA_EMPTY_SCREEN_ID);
-            mScreenOrder.removeValue(EXTRA_EMPTY_SCREEN_ID);
 
             // Update the page indicator to reflect the removed page.
             showPageIndicatorAtCurrentScroll();
         }
 
         if (stripEmptyScreens) {
+            // This will remove all empty pages from the Workspace. If there are no more pages left,
+            // it will add extra page(s) so that users can put items on at least one page.
             stripEmptyScreens();
         }
 
@@ -738,27 +824,56 @@
         }
     }
 
-    public boolean hasExtraEmptyScreen() {
-        return mWorkspaceScreens.containsKey(EXTRA_EMPTY_SCREEN_ID) && getChildCount() > 1;
+    public boolean hasExtraEmptyScreens() {
+        return mWorkspaceScreens.containsKey(EXTRA_EMPTY_SCREEN_ID)
+                && getChildCount() > getPanelCount()
+                && (!isTwoPanelEnabled()
+                || mWorkspaceScreens.containsKey(EXTRA_EMPTY_SCREEN_SECOND_ID));
     }
 
-    public int commitExtraEmptyScreen() {
+    /**
+     *  Commits the extra empty pages then returns the screen ids of those new screens.
+     *  Usually there's only one extra empty screen, but when two panel home is enabled we commit
+     *  two extra screens.
+     *
+     *  Returns an empty IntSet in case we cannot commit any new screens.
+     */
+    public IntSet commitExtraEmptyScreens() {
         if (mLauncher.isWorkspaceLoading()) {
             // Invalid and dangerous operation if workspace is loading
-            return -1;
+            return new IntSet();
         }
 
-        CellLayout cl = mWorkspaceScreens.get(EXTRA_EMPTY_SCREEN_ID);
-        mWorkspaceScreens.remove(EXTRA_EMPTY_SCREEN_ID);
-        mScreenOrder.removeValue(EXTRA_EMPTY_SCREEN_ID);
+        IntSet extraEmptyPageIds = new IntSet();
+        forEachExtraEmptyPageId(extraEmptyPageId ->
+                extraEmptyPageIds.add(commitExtraEmptyScreen(extraEmptyPageId)));
 
-        int newId = LauncherSettings.Settings.call(getContext().getContentResolver(),
+        return extraEmptyPageIds;
+    }
+
+    private int commitExtraEmptyScreen(int emptyScreenId) {
+        CellLayout cl = mWorkspaceScreens.get(emptyScreenId);
+        mWorkspaceScreens.remove(emptyScreenId);
+        mScreenOrder.removeValue(emptyScreenId);
+
+        int newScreenId = LauncherSettings.Settings.call(getContext().getContentResolver(),
                 LauncherSettings.Settings.METHOD_NEW_SCREEN_ID)
                 .getInt(LauncherSettings.Settings.EXTRA_VALUE);
-        mWorkspaceScreens.put(newId, cl);
-        mScreenOrder.add(newId);
 
-        return newId;
+
+        // When two panel home is enabled and the last page (the page on the right) doesn't
+        // have any items, then Launcher database doesn't know about this page because it was added
+        // by Launcher::bindAddScreens but wasn't inserted into the database. LauncherSettings's
+        // generate new screen ID method will return the ID for the left page,
+        // so we need to increment it.
+        if (isTwoPanelEnabled() && emptyScreenId == EXTRA_EMPTY_SCREEN_ID && newScreenId % 2 == 1) {
+            newScreenId++;
+        }
+
+        mWorkspaceScreens.put(newScreenId, cl);
+        mScreenOrder.add(newScreenId);
+
+        return newScreenId;
     }
 
     @Override
@@ -859,9 +974,10 @@
             }
         }
 
-        // We enforce at least one page to add new items to. In the case that we remove the last
-        // such screen, we convert the last screen to the empty screen
-        int minScreens = 1;
+        // We enforce at least one page (two pages on two panel home) to add new items to.
+        // In the case that we remove the last such screen(s), we convert the last screen(s)
+        // to the empty screen(s)
+        int minScreens = getPanelCount();
 
         int pageShift = 0;
         for (int i = 0; i < removeScreens.size(); i++) {
@@ -871,14 +987,21 @@
             mScreenOrder.removeValue(id);
 
             if (getChildCount() > minScreens) {
+                // If this isn't the last page, just remove it
                 if (indexOfChild(cl) < currentPage) {
                     pageShift++;
                 }
                 removeView(cl);
             } else {
-                // if this is the last screen, convert it to the empty screen
-                mWorkspaceScreens.put(EXTRA_EMPTY_SCREEN_ID, cl);
-                mScreenOrder.add(EXTRA_EMPTY_SCREEN_ID);
+                // The last page(s) should be converted into extra empty page(s)
+                int extraScreenId = isTwoPanelEnabled() && id % 2 == 1
+                        // This is the right panel in a two panel scenario
+                        ? EXTRA_EMPTY_SCREEN_SECOND_ID
+                        // This is either the last screen in a one panel scenario, or the left panel
+                        // in a two panel scenario when there are only two empty pages left
+                        : EXTRA_EMPTY_SCREEN_ID;
+                mWorkspaceScreens.put(extraScreenId, cl);
+                mScreenOrder.add(extraScreenId);
             }
         }
 
@@ -1656,8 +1779,8 @@
         }
 
         int screenId = getIdForScreen(dropTargetLayout);
-        if (screenId == EXTRA_EMPTY_SCREEN_ID) {
-            commitExtraEmptyScreen();
+        if (Workspace.EXTRA_EMPTY_SCREEN_IDS.contains(screenId)) {
+            commitExtraEmptyScreens();
         }
 
         return true;
diff --git a/src/com/android/launcher3/WorkspaceLayoutManager.java b/src/com/android/launcher3/WorkspaceLayoutManager.java
index 44a5536..7e6e1b6 100644
--- a/src/com/android/launcher3/WorkspaceLayoutManager.java
+++ b/src/com/android/launcher3/WorkspaceLayoutManager.java
@@ -23,6 +23,7 @@
 import com.android.launcher3.folder.FolderIcon;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.touch.ItemLongClickListener;
+import com.android.launcher3.util.IntSet;
 
 public interface WorkspaceLayoutManager {
 
@@ -30,6 +31,12 @@
 
     // The screen id used for the empty screen always present at the end.
     int EXTRA_EMPTY_SCREEN_ID = -201;
+    // The screen id used for the second empty screen always present at the end for two panel home.
+    int EXTRA_EMPTY_SCREEN_SECOND_ID = -200;
+    // The screen ids used for the empty screens at the end of the workspaces.
+    IntSet EXTRA_EMPTY_SCREEN_IDS =
+            IntSet.wrap(EXTRA_EMPTY_SCREEN_ID, EXTRA_EMPTY_SCREEN_SECOND_ID);
+
     // The is the first screen. It is always present, even if its empty.
     int FIRST_SCREEN_ID = 0;
     // This is the second page. On two panel home it is always present, even if its empty.
@@ -81,9 +88,9 @@
                 return;
             }
         }
-        if (screenId == EXTRA_EMPTY_SCREEN_ID) {
+        if (EXTRA_EMPTY_SCREEN_IDS.contains(screenId)) {
             // This should never happen
-            throw new RuntimeException("Screen id should not be EXTRA_EMPTY_SCREEN_ID");
+            throw new RuntimeException("Screen id should not be extra empty screen: " + screenId);
         }
 
         final CellLayout layout;
diff --git a/src/com/android/launcher3/accessibility/LauncherAccessibilityDelegate.java b/src/com/android/launcher3/accessibility/LauncherAccessibilityDelegate.java
index 9faac5b..2032b26 100644
--- a/src/com/android/launcher3/accessibility/LauncherAccessibilityDelegate.java
+++ b/src/com/android/launcher3/accessibility/LauncherAccessibilityDelegate.java
@@ -45,6 +45,7 @@
 import com.android.launcher3.popup.PopupContainerWithArrow;
 import com.android.launcher3.touch.ItemLongClickListener;
 import com.android.launcher3.util.IntArray;
+import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.ShortcutUtil;
 import com.android.launcher3.util.Thunk;
 import com.android.launcher3.views.OptionsPopupView;
@@ -221,6 +222,9 @@
         } else if (action == ADD_TO_WORKSPACE) {
             final int[] coordinates = new int[2];
             final int screenId = findSpaceOnWorkspace(item, coordinates);
+            if (screenId == -1) {
+                return false;
+            }
             mLauncher.getStateManager().goToState(NORMAL, true, forSuccessCallback(() -> {
                 if (item instanceof AppInfo) {
                     WorkspaceItemInfo info = ((AppInfo) item).makeWorkspaceItem();
@@ -250,6 +254,9 @@
 
             final int[] coordinates = new int[2];
             final int screenId = findSpaceOnWorkspace(item, coordinates);
+            if (screenId == -1) {
+                return false;
+            }
             mLauncher.getModelWriter().moveItemInDatabase(info,
                     Favorites.CONTAINER_DESKTOP,
                     screenId, coordinates[0], coordinates[1]);
@@ -489,8 +496,14 @@
             return screenId;
         }
 
-        workspace.addExtraEmptyScreen();
-        screenId = workspace.commitExtraEmptyScreen();
+        workspace.addExtraEmptyScreens();
+        IntSet emptyScreenIds = workspace.commitExtraEmptyScreens();
+        if (emptyScreenIds.isEmpty()) {
+            // Couldn't create extra empty screens for some reason (e.g. Workspace is loading)
+            return -1;
+        }
+
+        screenId = emptyScreenIds.getArray().get(0);
         layout = workspace.getScreenWithId(screenId);
         found = layout.findCellForSpan(outCoordinates, info.spanX, info.spanY);
 
diff --git a/src/com/android/launcher3/accessibility/ShortcutMenuAccessibilityDelegate.java b/src/com/android/launcher3/accessibility/ShortcutMenuAccessibilityDelegate.java
index f96afa8..bf5a24b 100644
--- a/src/com/android/launcher3/accessibility/ShortcutMenuAccessibilityDelegate.java
+++ b/src/com/android/launcher3/accessibility/ShortcutMenuAccessibilityDelegate.java
@@ -68,6 +68,9 @@
             final WorkspaceItemInfo info = ((DeepShortcutView) host.getParent()).getFinalInfo();
             final int[] coordinates = new int[2];
             final int screenId = findSpaceOnWorkspace(item, coordinates);
+            if (screenId == -1) {
+                return false;
+            }
             mLauncher.getStateManager().goToState(NORMAL, true, forSuccessCallback(() -> {
                 mLauncher.getModelWriter().addItemToDatabase(info,
                         LauncherSettings.Favorites.CONTAINER_DESKTOP,
diff --git a/src/com/android/launcher3/model/data/ItemInfo.java b/src/com/android/launcher3/model/data/ItemInfo.java
index 7091d2b..4fdc412 100644
--- a/src/com/android/launcher3/model/data/ItemInfo.java
+++ b/src/com/android/launcher3/model/data/ItemInfo.java
@@ -232,9 +232,9 @@
      * Write the fields of this item to the DB
      */
     public void onAddToDatabase(ContentWriter writer) {
-        if (screenId == Workspace.EXTRA_EMPTY_SCREEN_ID) {
+        if (Workspace.EXTRA_EMPTY_SCREEN_IDS.contains(screenId)) {
             // We should never persist an item on the extra empty screen.
-            throw new RuntimeException("Screen id should not be EXTRA_EMPTY_SCREEN_ID");
+            throw new RuntimeException("Screen id should not be extra empty screen: " + screenId);
         }
 
         writeToValues(writer);
diff --git a/src/com/android/launcher3/util/WallpaperOffsetInterpolator.java b/src/com/android/launcher3/util/WallpaperOffsetInterpolator.java
index c51f66f..8a7cae9 100644
--- a/src/com/android/launcher3/util/WallpaperOffsetInterpolator.java
+++ b/src/com/android/launcher3/util/WallpaperOffsetInterpolator.java
@@ -150,8 +150,8 @@
      */
     private int getNumPagesExcludingEmpty() {
         int numOfPages = mWorkspace.getChildCount();
-        if (numOfPages >= MIN_PARALLAX_PAGE_SPAN && mWorkspace.hasExtraEmptyScreen()) {
-            return numOfPages - 1;
+        if (numOfPages >= MIN_PARALLAX_PAGE_SPAN && mWorkspace.hasExtraEmptyScreens()) {
+            return numOfPages - mWorkspace.getPanelCount();
         } else {
             return numOfPages;
         }