From 1ef7e145d2da5b13350651990b1259d2e1dff4c3 Mon Sep 17 00:00:00 2001
From: Steven Ng <stevenckng@google.com>
Date: Tue, 25 May 2021 23:39:40 +0100
Subject: [PATCH] Fix add extra empty screen logic

The bug occurred when there are two widgets in a CellLayout. When
dragging one of the widget, ShortcutAndWidgetContainer#getChildCount
returns 1 rather than 2 because the LauncherAppWidgetHostView of the
widget in drag has been detached from its parent and reattached to
DragView (introduced in ag/14234627).

Test: Place only two widgets on the last page of the home screen. Drag
      one widget. The new page indicator on the right is shown.
Bug: 188494523
Change-Id: I8255dd1b72ad7f18f4cb35a9e14ac8e61804980c
---
 src/com/android/launcher3/Workspace.java | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/src/com/android/launcher3/Workspace.java b/src/com/android/launcher3/Workspace.java
index 98d80fedbc..05ef6ef2ef 100644
--- a/src/com/android/launcher3/Workspace.java
+++ b/src/com/android/launcher3/Workspace.java
@@ -426,10 +426,9 @@ public class Workspace extends PagedView<WorkspacePageIndicator>
         // When a accessible drag is started by the folder, we only allow rearranging withing the
         // folder.
         boolean addNewPage = !(options.isAccessibleDrag && dragObject.dragSource != this);
-
         if (addNewPage) {
             mDeferRemoveExtraEmptyScreen = false;
-            addExtraEmptyScreenOnDrag();
+            addExtraEmptyScreenOnDrag(dragObject);
 
             if (dragObject.dragInfo.itemType == LauncherSettings.Favorites.ITEM_TYPE_APPWIDGET
                     && dragObject.dragSource != this) {
@@ -636,12 +635,19 @@ public class Workspace extends PagedView<WorkspacePageIndicator>
         return newScreen;
     }
 
-    public void addExtraEmptyScreenOnDrag() {
+    private void addExtraEmptyScreenOnDrag(DragObject dragObject) {
         boolean lastChildOnScreen = false;
         boolean childOnFinalScreen = false;
 
         if (mDragSourceInternal != null) {
-            if (mDragSourceInternal.getChildCount() == 1) {
+            // 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 (dragSourceChildCount == 1) {
                 lastChildOnScreen = true;
             }
             CellLayout cl = (CellLayout) mDragSourceInternal.getParent();