Merge "Merging unnecessary subclass of ActivityAllAppsContainerView" into tm-qpr-dev
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarDragLayer.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarDragLayer.java
index d0059f7..7114849 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarDragLayer.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarDragLayer.java
@@ -194,7 +194,8 @@
     public boolean dispatchKeyEvent(KeyEvent event) {
         if (event.getAction() == ACTION_UP && event.getKeyCode() == KEYCODE_BACK) {
             AbstractFloatingView topView = AbstractFloatingView.getTopOpenView(mActivity);
-            if (topView != null && topView.onBackPressed()) {
+            if (topView != null && topView.canHandleBack()) {
+                topView.onBackInvoked();
                 // Handled by the floating view.
                 return true;
             }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
index 98c45d5..a58906f 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
@@ -156,10 +156,14 @@
                 } else {
                     // Config change might be handled without re-creating the taskbar
                     if (mTaskbarActivityContext != null) {
-                        if (dp != null && isTaskbarPresent(dp)) {
-                            mTaskbarActivityContext.updateDeviceProfile(dp, mNavMode);
+                        if (dp != null && !isTaskbarPresent(dp)) {
+                            destroyExistingTaskbar();
+                        } else {
+                            if (dp != null && isTaskbarPresent(dp)) {
+                                mTaskbarActivityContext.updateDeviceProfile(dp, mNavMode);
+                            }
+                            mTaskbarActivityContext.onConfigurationChanged(configDiff);
                         }
-                        mTaskbarActivityContext.onConfigurationChanged(configDiff);
                     }
                 }
                 mOldConfig = newConfig;
diff --git a/quickstep/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayDragLayer.java b/quickstep/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayDragLayer.java
index 044afd6..d91b650 100644
--- a/quickstep/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayDragLayer.java
+++ b/quickstep/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayDragLayer.java
@@ -71,7 +71,8 @@
     public boolean dispatchKeyEvent(KeyEvent event) {
         if (event.getAction() == ACTION_UP && event.getKeyCode() == KEYCODE_BACK) {
             AbstractFloatingView topView = AbstractFloatingView.getTopOpenView(mActivity);
-            if (topView != null && topView.onBackPressed()) {
+            if (topView != null && topView.canHandleBack()) {
+                topView.onBackInvoked();
                 return true;
             }
         }
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
index 0673dc6..30850b9 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
@@ -84,6 +84,7 @@
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.LauncherState;
+import com.android.launcher3.OnBackPressedHandler;
 import com.android.launcher3.QuickstepAccessibilityDelegate;
 import com.android.launcher3.QuickstepTransitionManager;
 import com.android.launcher3.R;
@@ -638,20 +639,49 @@
         getOnBackInvokedDispatcher().registerOnBackInvokedCallback(
                 OnBackInvokedDispatcher.PRIORITY_DEFAULT,
                 new OnBackAnimationCallback() {
+
+                    @Nullable OnBackPressedHandler mActiveOnBackPressedHandler;
+
+                    @Override
+                    public void onBackStarted(@NonNull BackEvent backEvent) {
+                        if (mActiveOnBackPressedHandler != null) {
+                            mActiveOnBackPressedHandler.onBackCancelled();
+                        }
+                        mActiveOnBackPressedHandler = getOnBackPressedHandler();
+                        mActiveOnBackPressedHandler.onBackStarted();
+                    }
+
                     @Override
                     public void onBackInvoked() {
-                        onBackPressed();
+                        // Recreate mActiveOnBackPressedHandler if necessary to avoid NPE because:
+                        // 1. b/260636433: In 3-button-navigation mode, onBackStarted() is not
+                        // called on ACTION_DOWN before onBackInvoked() is called in ACTION_UP.
+                        // 2. Launcher#onBackPressed() will call onBackInvoked() without calling
+                        // onBackInvoked() beforehand.
+                        if (mActiveOnBackPressedHandler == null) {
+                            mActiveOnBackPressedHandler = getOnBackPressedHandler();
+                        }
+                        mActiveOnBackPressedHandler.onBackInvoked();
+                        mActiveOnBackPressedHandler = null;
                         TestLogging.recordEvent(TestProtocol.SEQUENCE_MAIN, "onBackInvoked");
                     }
 
                     @Override
                     public void onBackProgressed(@NonNull BackEvent backEvent) {
-                        QuickstepLauncher.this.onBackProgressed(backEvent.getProgress());
+                        if (!FeatureFlags.IS_STUDIO_BUILD && mActiveOnBackPressedHandler == null) {
+                            return;
+                        }
+                        mActiveOnBackPressedHandler
+                                .onBackProgressed(backEvent.getProgress());
                     }
 
                     @Override
                     public void onBackCancelled() {
-                        QuickstepLauncher.this.onBackCancelled();
+                        if (!FeatureFlags.IS_STUDIO_BUILD && mActiveOnBackPressedHandler == null) {
+                            return;
+                        }
+                        mActiveOnBackPressedHandler.onBackCancelled();
+                        mActiveOnBackPressedHandler = null;
                     }
                 });
     }
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index e3427b7..deeb027 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -51,6 +51,7 @@
 import static com.android.quickstep.MultiStateCallback.DEBUG_STATES;
 import static com.android.quickstep.util.ActiveGestureErrorDetector.GestureEvent.CANCEL_RECENTS_ANIMATION;
 import static com.android.quickstep.util.ActiveGestureErrorDetector.GestureEvent.EXPECTING_TASK_APPEARED;
+import static com.android.quickstep.util.ActiveGestureErrorDetector.GestureEvent.LAUNCHER_DESTROYED;
 import static com.android.quickstep.util.ActiveGestureErrorDetector.GestureEvent.ON_SETTLED_ON_END_TARGET;
 import static com.android.quickstep.views.RecentsView.UPDATE_SYSUI_FLAGS_THRESHOLD;
 import static com.android.systemui.shared.system.ActivityManagerWrapper.CLOSE_SYSTEM_WINDOWS_REASON_RECENTS;
@@ -182,6 +183,7 @@
                     if (mActivity != activity) {
                         return;
                     }
+                    handleActivityDestroyed();
                     mRecentsView = null;
                     mActivity = null;
                 }
@@ -462,6 +464,7 @@
         if (mStateCallback.hasStates(STATE_HANDLER_INVALIDATED)) {
             return false;
         }
+        mStateCallback.resumeCallbacks();
 
         T createdActivity = mActivityInterface.getCreatedActivity();
         if (createdActivity != null) {
@@ -531,6 +534,15 @@
         return true;
     }
 
+    private void handleActivityDestroyed() {
+        ActiveGestureLog.INSTANCE.addLog("Launcher activity destroyed", LAUNCHER_DESTROYED);
+        if (mActivityInterface.shouldCancelGestureOnDestroy()) {
+            onGestureCancelled();
+        } else {
+            mStateCallback.pauseCallbacks();
+        }
+    }
+
     /**
      * Return true if the window should be translated horizontally if the recents view scrolls
      */
diff --git a/quickstep/src/com/android/quickstep/BaseActivityInterface.java b/quickstep/src/com/android/quickstep/BaseActivityInterface.java
index 274b686..22eaa97 100644
--- a/quickstep/src/com/android/quickstep/BaseActivityInterface.java
+++ b/quickstep/src/com/android/quickstep/BaseActivityInterface.java
@@ -123,6 +123,14 @@
     public abstract AnimationFactory prepareRecentsUI(RecentsAnimationDeviceState deviceState,
             boolean activityVisible, Consumer<AnimatorControllerWithResistance> callback);
 
+    /**
+     * Returns {@code true} iff an ongoing navigational gesture should be cancelled on activity
+     * destroy. Otherwise, the MultiStateCallbacks will be paused until the activity is recreated.
+     */
+    public boolean shouldCancelGestureOnDestroy() {
+        return true;
+    }
+
     public abstract ActivityInitListener createActivityInitListener(
             Predicate<Boolean> onInitListener);
 
diff --git a/quickstep/src/com/android/quickstep/FallbackActivityInterface.java b/quickstep/src/com/android/quickstep/FallbackActivityInterface.java
index ae9fb0b..4cb4665 100644
--- a/quickstep/src/com/android/quickstep/FallbackActivityInterface.java
+++ b/quickstep/src/com/android/quickstep/FallbackActivityInterface.java
@@ -88,6 +88,11 @@
     }
 
     @Override
+    public boolean shouldCancelGestureOnDestroy() {
+        return false;
+    }
+
+    @Override
     public ActivityInitListener createActivityInitListener(
             Predicate<Boolean> onInitListener) {
         return new ActivityInitListener<>((activity, alreadyOnHome) ->
diff --git a/quickstep/src/com/android/quickstep/MultiStateCallback.java b/quickstep/src/com/android/quickstep/MultiStateCallback.java
index a68bea2..6d767ed 100644
--- a/quickstep/src/com/android/quickstep/MultiStateCallback.java
+++ b/quickstep/src/com/android/quickstep/MultiStateCallback.java
@@ -31,6 +31,7 @@
 
 import java.util.ArrayList;
 import java.util.LinkedList;
+import java.util.List;
 import java.util.StringJoiner;
 import java.util.function.Consumer;
 
@@ -52,6 +53,9 @@
 
     private int mState = 0;
 
+    private boolean mCallbacksPaused = false;
+    private final List<Runnable> mPendingCallbacks = new ArrayList<>();
+
     public MultiStateCallback(String[] stateNames) {
         this(stateNames, stateFlag -> null);
     }
@@ -78,6 +82,24 @@
         }
     }
 
+    /** Pauses callbacks. */
+    public void pauseCallbacks() {
+        mCallbacksPaused = true;
+    }
+
+    /** Immediately queues any callbacks that were pending paused. */
+    public void resumeCallbacks() {
+        if (!mCallbacksPaused) {
+            return;
+        }
+        mCallbacksPaused = false;
+        List<Runnable> queuedCallbacks = new ArrayList<>(mPendingCallbacks);
+        mPendingCallbacks.clear();
+        for (Runnable runnable : queuedCallbacks) {
+            runnable.run();
+        }
+    }
+
     /**
      * Adds the provided state flags to the global state and executes any callbacks as a result.
      */
@@ -99,7 +121,12 @@
             if ((mState & state) == state) {
                 LinkedList<Runnable> callbacks = mCallbacks.valueAt(i);
                 while (!callbacks.isEmpty()) {
-                    callbacks.pollFirst().run();
+                    Runnable cb = callbacks.pollFirst();
+                    if (mCallbacksPaused) {
+                        mPendingCallbacks.add(cb);
+                    } else {
+                        cb.run();
+                    }
                 }
             }
         }
@@ -151,7 +178,11 @@
             if (wasOn != isOn) {
                 ArrayList<Consumer<Boolean>> listeners = mStateChangeListeners.valueAt(i);
                 for (Consumer<Boolean> listener : listeners) {
-                    listener.accept(isOn);
+                    if (mCallbacksPaused) {
+                        mPendingCallbacks.add(() -> listener.accept(isOn));
+                    } else {
+                        listener.accept(isOn);
+                    }
                 }
             }
         }
diff --git a/quickstep/src/com/android/quickstep/util/ActiveGestureErrorDetector.java b/quickstep/src/com/android/quickstep/util/ActiveGestureErrorDetector.java
index 60065fb..0fdd8b5 100644
--- a/quickstep/src/com/android/quickstep/util/ActiveGestureErrorDetector.java
+++ b/quickstep/src/com/android/quickstep/util/ActiveGestureErrorDetector.java
@@ -37,7 +37,7 @@
         ON_SETTLED_ON_END_TARGET, START_RECENTS_ANIMATION, FINISH_RECENTS_ANIMATION,
         CANCEL_RECENTS_ANIMATION, SET_ON_PAGE_TRANSITION_END_CALLBACK, CANCEL_CURRENT_ANIMATION,
         CLEANUP_SCREENSHOT, SCROLLER_ANIMATION_ABORTED, TASK_APPEARED, EXPECTING_TASK_APPEARED,
-        FLAG_USING_OTHER_ACTIVITY_INPUT_CONSUMER,
+        FLAG_USING_OTHER_ACTIVITY_INPUT_CONSUMER, LAUNCHER_DESTROYED,
 
         /**
          * These GestureEvents are specifically associated to state flags that get set in
@@ -162,6 +162,13 @@
                                         + "before/without setting end target to new task",
                                 writer);
                         break;
+                    case LAUNCHER_DESTROYED:
+                        errorDetected |= printErrorIfTrue(
+                                true,
+                                prefix,
+                                /* errorMessage= */ "Launcher destroyed mid-gesture",
+                                writer);
+                        break;
                     case STATE_GESTURE_COMPLETED:
                         errorDetected |= printErrorIfTrue(
                                 !encounteredEvents.contains(GestureEvent.MOTION_UP),
diff --git a/quickstep/src/com/android/quickstep/views/AllAppsEduView.java b/quickstep/src/com/android/quickstep/views/AllAppsEduView.java
index d79b318..716d389 100644
--- a/quickstep/src/com/android/quickstep/views/AllAppsEduView.java
+++ b/quickstep/src/com/android/quickstep/views/AllAppsEduView.java
@@ -112,11 +112,6 @@
     }
 
     @Override
-    public boolean onBackPressed() {
-        return true;
-    }
-
-    @Override
     public boolean canInterceptEventsInSystemGestureRegion() {
         return true;
     }
diff --git a/quickstep/src/com/android/quickstep/views/DesktopTaskView.java b/quickstep/src/com/android/quickstep/views/DesktopTaskView.java
index 858f6ab..44aa41a 100644
--- a/quickstep/src/com/android/quickstep/views/DesktopTaskView.java
+++ b/quickstep/src/com/android/quickstep/views/DesktopTaskView.java
@@ -18,6 +18,7 @@
 
 import static android.view.ViewGroup.LayoutParams.WRAP_CONTENT;
 
+import static com.android.launcher3.LauncherState.NORMAL;
 import static com.android.launcher3.util.SplitConfigurationOptions.STAGE_POSITION_UNDEFINED;
 
 import android.content.Context;
@@ -35,6 +36,7 @@
 import androidx.annotation.Nullable;
 
 import com.android.launcher3.DeviceProfile;
+import com.android.launcher3.Launcher;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.touch.PagedOrientationHandler;
 import com.android.launcher3.util.RunnableList;
@@ -295,7 +297,7 @@
     @Override
     public RunnableList launchTasks() {
         SystemUiProxy.INSTANCE.get(getContext()).showDesktopApps();
-        getRecentsView().startHome();
+        Launcher.getLauncher(mActivity).getStateManager().goToState(NORMAL, false /* animated */);
         return null;
     }
 
diff --git a/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java b/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
index 47bef7b..a4dbf6a 100644
--- a/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
+++ b/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
@@ -168,7 +168,6 @@
 
     // b/143488140
     //@NavigationModeSwitch
-    @Ignore
     @Test
     public void goToOverviewFromHome() {
         mDevice.pressHome();
@@ -216,7 +215,6 @@
 
     // b/143488140
     //@NavigationModeSwitch
-    @Ignore
     @Test
     public void testOverview() {
         startAppFast(getAppPackageName());
diff --git a/src/com/android/launcher3/AbstractFloatingView.java b/src/com/android/launcher3/AbstractFloatingView.java
index 49afc1f..796fa80 100644
--- a/src/com/android/launcher3/AbstractFloatingView.java
+++ b/src/com/android/launcher3/AbstractFloatingView.java
@@ -46,7 +46,8 @@
 /**
  * Base class for a View which shows a floating UI on top of the launcher UI.
  */
-public abstract class AbstractFloatingView extends LinearLayout implements TouchController {
+public abstract class AbstractFloatingView extends LinearLayout implements TouchController,
+        OnBackPressedHandler {
 
     @IntDef(flag = true, value = {
             TYPE_FOLDER,
@@ -165,13 +166,17 @@
 
     protected abstract boolean isOfType(@FloatingViewType int type);
 
-    /** @return Whether the back is consumed. If false, Launcher will handle the back as well. */
-    public boolean onBackPressed() {
-        close(true);
+    /** Return true if this view can consume back press. */
+    public boolean canHandleBack() {
         return true;
     }
 
     @Override
+    public void onBackInvoked() {
+        close(true);
+    }
+
+    @Override
     public boolean onControllerTouchEvent(MotionEvent ev) {
         return false;
     }
diff --git a/src/com/android/launcher3/BaseDraggingActivity.java b/src/com/android/launcher3/BaseDraggingActivity.java
index ca1fe40..6f3e948 100644
--- a/src/com/android/launcher3/BaseDraggingActivity.java
+++ b/src/com/android/launcher3/BaseDraggingActivity.java
@@ -125,9 +125,13 @@
         mCurrentActionMode = null;
     }
 
+    protected boolean isInAutoCancelActionMode() {
+        return mCurrentActionMode != null && AUTO_CANCEL_ACTION_MODE == mCurrentActionMode.getTag();
+    }
+
     @Override
     public boolean finishAutoCancelActionMode() {
-        if (mCurrentActionMode != null && AUTO_CANCEL_ACTION_MODE == mCurrentActionMode.getTag()) {
+        if (isInAutoCancelActionMode()) {
             mCurrentActionMode.finish();
             return true;
         }
diff --git a/src/com/android/launcher3/InvariantDeviceProfile.java b/src/com/android/launcher3/InvariantDeviceProfile.java
index a4020f9..ffe81ad 100644
--- a/src/com/android/launcher3/InvariantDeviceProfile.java
+++ b/src/com/android/launcher3/InvariantDeviceProfile.java
@@ -206,8 +206,6 @@
         if (!newGridName.equals(gridName)) {
             LauncherPrefs.getPrefs(context).edit().putString(KEY_IDP_GRID_NAME, newGridName)
                     .apply();
-            Log.d("b/258560494", "InvariantDeviceProfile - setting newGridName: " + newGridName
-                    + ", gridName: " + gridName);
         }
         new DeviceGridState(this).writeToPrefs(context);
 
@@ -455,7 +453,6 @@
     public void setCurrentGrid(Context context, String gridName) {
         Context appContext = context.getApplicationContext();
         LauncherPrefs.getPrefs(appContext).edit().putString(KEY_IDP_GRID_NAME, gridName).apply();
-        Log.d("b/258560494", "setCurrentGrid: " + gridName);
         MAIN_EXECUTOR.execute(() -> onConfigChanged(appContext));
     }
 
@@ -520,10 +517,6 @@
             }
         }
         if (filteredProfiles.isEmpty()) {
-            if (gridName != null) {
-                Log.d("b/258560494", "No matching grid from for gridName: " + gridName
-                        + ", deviceType: " + deviceType);
-            }
             // No grid found, use the default options
             for (DisplayOption option : profiles) {
                 if (option.canBeDefault) {
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 5b21fb8..e9b9f31 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -559,6 +559,61 @@
         setTitle(R.string.home_screen);
     }
 
+    /**
+     * Provide {@link OnBackPressedHandler} in below order:
+     * <ol>
+     *  <li> auto cancel action mode handler
+     *  <li> drag handler
+     *  <li> view handler
+     *  <li> state handler
+     * </ol>
+     *
+     * A back gesture (a single click on back button, or a swipe back gesture that contains a series
+     * of swipe events) should be handled by the same handler from above list. For a new back
+     * gesture, a new handler should be regenerated.
+     *
+     * Note that state handler will always be handling the back press event if the previous 3 don't.
+     */
+    @NonNull
+    protected OnBackPressedHandler getOnBackPressedHandler() {
+        // #1 auto cancel action mode handler
+        if (isInAutoCancelActionMode()) {
+            return this::finishAutoCancelActionMode;
+        }
+
+        // #2 drag handler
+        if (mDragController.isDragging()) {
+            return mDragController::cancelDrag;
+        }
+
+        // #3 view handler
+        AbstractFloatingView topView =
+                AbstractFloatingView.getTopOpenView(Launcher.this);
+        if (topView != null && topView.canHandleBack()) {
+            return topView;
+        }
+
+        // #4 state handler
+        return new OnBackPressedHandler() {
+            @Override
+            public void onBackInvoked() {
+                onStateBack();
+            }
+
+            @Override
+            public void onBackProgressed(
+                    @FloatRange(from = 0.0, to = 1.0) float backProgress) {
+                mStateManager.getState().onBackProgressed(
+                        Launcher.this, backProgress);
+            }
+
+            @Override
+            public void onBackCancelled() {
+                mStateManager.getState().onBackCancelled(Launcher.this);
+            }
+        };
+    }
+
     protected LauncherOverlayManager getDefaultOverlay() {
         return new LauncherOverlayManager() { };
     }
@@ -2046,36 +2101,13 @@
 
     @Override
     public void onBackPressed() {
-        if (finishAutoCancelActionMode()) {
-            return;
-        }
-
-        if (mDragController.isDragging()) {
-            mDragController.cancelDrag();
-            return;
-        }
-
-        // Note: There should be at most one log per method call. This is enforced implicitly
-        // by using if-else statements.
-        AbstractFloatingView topView = AbstractFloatingView.getTopOpenView(this);
-        if (topView == null || !topView.onBackPressed()) {
-            // Not handled by the floating view.
-            onStateBack();
-        }
+        getOnBackPressedHandler().onBackInvoked();
     }
 
     protected void onStateBack() {
         mStateManager.getState().onBackPressed(this);
     }
 
-    protected void onBackProgressed(@FloatRange(from = 0.0, to = 1.0) float backProgress) {
-        mStateManager.getState().onBackProgressed(this, backProgress);
-    }
-
-    protected void onBackCancelled() {
-        mStateManager.getState().onBackCancelled(this);
-    }
-
     protected void onScreenOnChanged(boolean isOn) {
         // Reset AllApps to its initial state only if we are not in the middle of
         // processing a multi-step drop
diff --git a/src/com/android/launcher3/OnBackPressedHandler.java b/src/com/android/launcher3/OnBackPressedHandler.java
new file mode 100644
index 0000000..485aa0a
--- /dev/null
+++ b/src/com/android/launcher3/OnBackPressedHandler.java
@@ -0,0 +1,45 @@
+/*
+ * Copyright (C) 2023 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 androidx.annotation.FloatRange;
+
+/**
+ * Interface that mimics {@link android.window.OnBackInvokedCallback} without dependencies on U's
+ * API such as {@link android.window.BackEvent}.
+ *
+ * <p> Impl can assume below order during a back gesture:
+ * <ol>
+ *  <li> [optional] one {@link #onBackStarted()} will be called to start the gesture
+ *  <li> zero or multiple {@link #onBackProgressed(float)} will be called during swipe gesture
+ *  <li> either one of {@link #onBackInvoked()} or {@link #onBackCancelled()} will be called to end
+ *  the gesture
+ */
+public interface OnBackPressedHandler {
+
+    /** Called when back has started. */
+    default void onBackStarted() {}
+
+    /** Called when back is committed. */
+    void onBackInvoked();
+
+    /** Called with back gesture's progress. */
+    default void onBackProgressed(@FloatRange(from = 0.0, to = 1.0) float backProgress) {}
+
+    /** Called when user drops the back gesture. */
+    default void onBackCancelled() {}
+}
diff --git a/src/com/android/launcher3/allapps/DiscoveryBounce.java b/src/com/android/launcher3/allapps/DiscoveryBounce.java
index be261f7..0188a47 100644
--- a/src/com/android/launcher3/allapps/DiscoveryBounce.java
+++ b/src/com/android/launcher3/allapps/DiscoveryBounce.java
@@ -81,10 +81,10 @@
     }
 
     @Override
-    public boolean onBackPressed() {
-        super.onBackPressed();
-        // Go back to the previous state (from a user's perspective this floating view isn't
-        // something to go back from).
+    public boolean canHandleBack() {
+        // Since DiscoveryBounce doesn't handle back, onBackInvoked() won't be called and we should
+        // close it without animation.
+        close(false);
         return false;
     }
 
diff --git a/src/com/android/launcher3/folder/Folder.java b/src/com/android/launcher3/folder/Folder.java
index 94eea35..9a5d77e 100644
--- a/src/com/android/launcher3/folder/Folder.java
+++ b/src/com/android/launcher3/folder/Folder.java
@@ -1579,17 +1579,14 @@
         return getOpenView(activityContext, TYPE_FOLDER);
     }
 
-    /**
-     * Navigation bar back key or hardware input back key has been issued.
-     */
+    /** Navigation bar back key or hardware input back key has been issued. */
     @Override
-    public boolean onBackPressed() {
+    public void onBackInvoked() {
         if (isEditingName()) {
             mFolderName.dispatchBackKey();
         } else {
-            super.onBackPressed();
+            super.onBackInvoked();
         }
-        return true;
     }
 
     @Override
diff --git a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
index e7b0446..4810b15 100644
--- a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
+++ b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
@@ -68,6 +68,7 @@
 import com.android.launcher3.InsettableFrameLayout;
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
@@ -137,7 +138,7 @@
                 new ConcurrentLinkedQueue<>();
 
         public PreviewContext(Context base, InvariantDeviceProfile idp) {
-            super(base, UserCache.INSTANCE, InstallSessionHelper.INSTANCE,
+            super(base, UserCache.INSTANCE, InstallSessionHelper.INSTANCE, LauncherPrefs.INSTANCE,
                     LauncherAppState.INSTANCE, InvariantDeviceProfile.INSTANCE,
                     CustomWidgetManager.INSTANCE, PluginManagerWrapper.INSTANCE,
                     WindowManagerProxy.INSTANCE, DisplayController.INSTANCE);
diff --git a/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncher.java b/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncher.java
index a2353d8..56dffa9 100644
--- a/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncher.java
+++ b/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncher.java
@@ -175,8 +175,9 @@
         // Note: There should be at most one log per method call. This is enforced implicitly
         // by using if-else statements.
         AbstractFloatingView topView = AbstractFloatingView.getTopOpenView(this);
-        if (topView != null && topView.onBackPressed()) {
+        if (topView != null && topView.canHandleBack()) {
             // Handled by the floating view.
+            topView.onBackInvoked();
         } else {
             showAppDrawer(false);
         }
diff --git a/src/com/android/launcher3/testing/TestInformationHandler.java b/src/com/android/launcher3/testing/TestInformationHandler.java
index acb7eb3..9a34478 100644
--- a/src/com/android/launcher3/testing/TestInformationHandler.java
+++ b/src/com/android/launcher3/testing/TestInformationHandler.java
@@ -201,6 +201,13 @@
                 });
             }
 
+            case TestProtocol.REQUEST_WORKSPACE_COLUMNS_ROWS: {
+                return getLauncherUIProperty(Bundle::putParcelable, launcher -> new Point(
+                        InvariantDeviceProfile.INSTANCE.get(mContext).numColumns,
+                        InvariantDeviceProfile.INSTANCE.get(mContext).numRows)
+                );
+            }
+
             case TestProtocol.REQUEST_HOTSEAT_CELL_CENTER: {
                 final HotseatCellCenterRequest request = extra.getParcelable(
                         TestProtocol.TEST_INFO_REQUEST_FIELD);
diff --git a/src/com/android/launcher3/testing/shared/TestProtocol.java b/src/com/android/launcher3/testing/shared/TestProtocol.java
index f5ee91b..9b2ce9a 100644
--- a/src/com/android/launcher3/testing/shared/TestProtocol.java
+++ b/src/com/android/launcher3/testing/shared/TestProtocol.java
@@ -123,6 +123,7 @@
 
     public static final String REQUEST_WORKSPACE_CELL_LAYOUT_SIZE = "workspace-cell-layout-size";
     public static final String REQUEST_WORKSPACE_CELL_CENTER = "workspace-cell-center";
+    public static final String REQUEST_WORKSPACE_COLUMNS_ROWS = "workspace-columns-rows";
 
     public static final String REQUEST_HOTSEAT_CELL_CENTER = "hotseat-cell-center";
 
diff --git a/src/com/android/launcher3/testing/shared/WorkspaceCellCenterRequest.java b/src/com/android/launcher3/testing/shared/WorkspaceCellCenterRequest.java
index 80dbef8..e2cd8ea 100644
--- a/src/com/android/launcher3/testing/shared/WorkspaceCellCenterRequest.java
+++ b/src/com/android/launcher3/testing/shared/WorkspaceCellCenterRequest.java
@@ -124,7 +124,7 @@
          * Set span Height in cells
          */
         public WorkspaceCellCenterRequest.Builder setSpanY(int y) {
-            this.mCellY = y;
+            this.mSpanY = y;
             return this;
         }
 
diff --git a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
index 38cdb4b..c78ecf5 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
@@ -17,7 +17,9 @@
 
 import static android.view.View.MeasureSpec.makeMeasureSpec;
 
+import static com.android.launcher3.LauncherAnimUtils.SCALE_PROPERTY;
 import static com.android.launcher3.LauncherAnimUtils.VIEW_TRANSLATE_Y;
+import static com.android.launcher3.allapps.AllAppsTransitionController.SWIPE_ALL_APPS_TO_HOME_MIN_SCALE;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_WIDGETSTRAY_SEARCHED;
 import static com.android.launcher3.testing.shared.TestProtocol.NORMAL_STATE_ORDINAL;
 
@@ -45,6 +47,7 @@
 import android.widget.Button;
 import android.widget.TextView;
 
+import androidx.annotation.FloatRange;
 import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 import androidx.recyclerview.widget.DefaultItemAnimator;
@@ -54,6 +57,7 @@
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
+import com.android.launcher3.anim.Interpolators;
 import com.android.launcher3.anim.PendingAnimation;
 import com.android.launcher3.compat.AccessibilityManagerCompat;
 import com.android.launcher3.model.UserManagerState;
@@ -264,6 +268,15 @@
         attachScrollbarToRecyclerView(currentRecyclerView);
     }
 
+    @Override
+    public void onBackProgressed(@FloatRange(from = 0.0, to = 1.0) float progress) {
+        float deceleratedProgress =
+                Interpolators.PREDICTIVE_BACK_DECELERATED_EASE.getInterpolation(progress);
+        float scaleProgress = SWIPE_ALL_APPS_TO_HOME_MIN_SCALE
+                + (1 - SWIPE_ALL_APPS_TO_HOME_MIN_SCALE) * (1 - deceleratedProgress);
+        SCALE_PROPERTY.set(this, scaleProgress);
+    }
+
     private void attachScrollbarToRecyclerView(WidgetsRecyclerView recyclerView) {
         recyclerView.bindFastScrollbar();
         if (mCurrentWidgetsRecyclerView != recyclerView) {
@@ -736,12 +749,12 @@
     }
 
     @Override
-    public boolean onBackPressed() {
+    public void onBackInvoked() {
         if (mIsInSearchMode) {
             mSearchBar.reset();
-            return true;
+        } else {
+            super.onBackInvoked();
         }
-        return super.onBackPressed();
     }
 
     @Override
diff --git a/tests/src/com/android/launcher3/celllayout/CellLayoutBoard.java b/tests/src/com/android/launcher3/celllayout/CellLayoutBoard.java
index 469d79f..6d75180 100644
--- a/tests/src/com/android/launcher3/celllayout/CellLayoutBoard.java
+++ b/tests/src/com/android/launcher3/celllayout/CellLayoutBoard.java
@@ -29,7 +29,73 @@
 import java.util.stream.Collectors;
 
 
-public class CellLayoutBoard {
+public class CellLayoutBoard implements Comparable<CellLayoutBoard> {
+
+    private boolean intersects(Rect r1, Rect r2) {
+        // If one rectangle is on left side of other
+        if (r1.left > r2.right || r2.left > r1.right) {
+            return false;
+        }
+
+        // If one rectangle is above other
+        if (r1.bottom > r2.top || r2.bottom > r1.top) {
+            return false;
+        }
+
+        return true;
+    }
+
+    private boolean overlapsWithIgnored(Set<Rect> ignoredRectangles, Rect rect) {
+        for (Rect ignoredRect : ignoredRectangles) {
+            // Using the built in intersects doesn't work because it doesn't account for area 0
+            if (intersects(ignoredRect, rect)) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    @Override
+    public int compareTo(CellLayoutBoard cellLayoutBoard) {
+        // to be equal they need to have the same number of widgets and the same dimensions
+        // their order can be different
+        Set<Rect> widgetsSet = new HashSet<>();
+        Set<Rect> ignoredRectangles = new HashSet<>();
+        for (WidgetRect rect : mWidgetsRects) {
+            if (rect.shouldIgnore()) {
+                ignoredRectangles.add(rect.mBounds);
+            } else {
+                widgetsSet.add(rect.mBounds);
+            }
+        }
+        for (WidgetRect rect : cellLayoutBoard.mWidgetsRects) {
+            // ignore rectangles overlapping with the area marked by x
+            if (overlapsWithIgnored(ignoredRectangles, rect.mBounds)) {
+                continue;
+            }
+            if (!widgetsSet.contains(rect.mBounds)) {
+                return -1;
+            }
+            widgetsSet.remove(rect.mBounds);
+        }
+        if (!widgetsSet.isEmpty()) {
+            return 1;
+        }
+
+        // to be equal they need to have the same number of icons their order can be different
+        Set<Point> iconsSet = new HashSet<>();
+        mIconPoints.forEach(icon -> iconsSet.add(icon.getCoord()));
+        for (IconPoint icon : cellLayoutBoard.mIconPoints) {
+            if (!iconsSet.contains(icon.getCoord())) {
+                return -1;
+            }
+            iconsSet.remove(icon.getCoord());
+        }
+        if (!iconsSet.isEmpty()) {
+            return 1;
+        }
+        return 0;
+    }
 
     public static class CellType {
         // The cells marked by this will be filled by 1x1 widgets and will be ignored when
@@ -115,7 +181,7 @@
     List<IconPoint> mIconPoints = new ArrayList<>();
     Map<Character, IconPoint> mIconsMap = new HashMap<>();
 
-    Point mMain = new Point();
+    WidgetRect mMain = null;
 
     CellLayoutBoard() {
         for (int x = 0; x < mWidget.length; x++) {
@@ -133,7 +199,7 @@
         return mIconPoints;
     }
 
-    public Point getMain() {
+    public WidgetRect getMain() {
         return mMain;
     }
 
@@ -273,6 +339,16 @@
         return iconPoints;
     }
 
+    public static WidgetRect getMainFromList(List<CellLayoutBoard> boards) {
+        for (CellLayoutBoard board : boards) {
+            WidgetRect main = board.getMain();
+            if (main != null) {
+                return main;
+            }
+        }
+        return null;
+    }
+
     public static CellLayoutBoard boardFromString(String boardStr) {
         String[] lines = boardStr.split("\n");
         CellLayoutBoard board = new CellLayoutBoard();
@@ -281,17 +357,18 @@
             String line = lines[y];
             for (int x = 0; x < line.length(); x++) {
                 char c = line.charAt(x);
-                if (c == CellType.MAIN_WIDGET) {
-                    board.mMain = new Point(x, y);
-                }
                 if (c != CellType.EMPTY) {
                     board.mWidget[x][y] = line.charAt(x);
                 }
             }
         }
         board.mWidgetsRects = getRects(board.mWidget);
-        board.mWidgetsRects.forEach(
-                widgetRect -> board.mWidgetsMap.put(widgetRect.mType, widgetRect));
+        board.mWidgetsRects.forEach(widgetRect -> {
+            if (widgetRect.mType == CellType.MAIN_WIDGET) {
+                board.mMain = widgetRect;
+            }
+            board.mWidgetsMap.put(widgetRect.mType, widgetRect);
+        });
         board.mIconPoints = getIconPoints(board.mWidget);
         return board;
     }
@@ -308,4 +385,24 @@
         }
         return s.toString();
     }
+
+    public static List<CellLayoutBoard> boardListFromString(String boardsStr) {
+        String[] lines = boardsStr.split("\n");
+        ArrayList<String> individualBoards = new ArrayList<>();
+        ArrayList<CellLayoutBoard> boards = new ArrayList<>();
+        for (String line : lines) {
+            String[] boardSegment = line.split("\\|");
+            for (int i = 0; i < boardSegment.length; i++) {
+                if (i >= individualBoards.size()) {
+                    individualBoards.add(boardSegment[i]);
+                } else {
+                    individualBoards.set(i, individualBoards.get(i) + "\n" + boardSegment[i]);
+                }
+            }
+        }
+        for (String board : individualBoards) {
+            boards.add(CellLayoutBoard.boardFromString(board));
+        }
+        return boards;
+    }
 }
diff --git a/tests/src/com/android/launcher3/celllayout/MultipleCellLayoutsSimpleReorder.java b/tests/src/com/android/launcher3/celllayout/MultipleCellLayoutsSimpleReorder.java
new file mode 100644
index 0000000..9b52fe8
--- /dev/null
+++ b/tests/src/com/android/launcher3/celllayout/MultipleCellLayoutsSimpleReorder.java
@@ -0,0 +1,88 @@
+/*
+ * Copyright (C) 2022 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.celllayout.testcases;
+
+import android.graphics.Point;
+
+import java.util.Map;
+
+/**
+ * The grids represent the workspace to be build by TestWorkspaceBuilder, to see what each character
+ * in the board mean refer to {@code CellType}
+ */
+public class MultipleCellLayoutsSimpleReorder {
+
+    /** 5x5 Test
+     **/
+    private static final String START_BOARD_STR_5x5 = ""
+            + "xxxxx|-----\n"
+            + "--mm-|-----\n"
+            + "--mm-|-----\n"
+            + "-----|-----\n"
+            + "-----|-----";
+    private static final Point MOVE_TO_5x5 = new Point(8, 3);
+    private static final String END_BOARD_STR_5x5 = ""
+            + "xxxxx|-----\n"
+            + "-----|-----\n"
+            + "-----|-----\n"
+            + "-----|---mm\n"
+            + "-----|---mm";
+    private static final ReorderTestCase TEST_CASE_5x5 = new ReorderTestCase(START_BOARD_STR_5x5,
+            MOVE_TO_5x5,
+            END_BOARD_STR_5x5);
+
+    /** 4x4 Test
+     **/
+    private static final String START_BOARD_STR_4x4 = ""
+            + "xxxx|----\n"
+            + "--mm|----\n"
+            + "--mm|----\n"
+            + "----|----";
+    private static final Point MOVE_TO_4x4 = new Point(5, 3);
+    private static final String END_BOARD_STR_4x4 = ""
+            + "xxxx|----\n"
+            + "----|----\n"
+            + "----|-mm-\n"
+            + "----|-mm-";
+    private static final ReorderTestCase TEST_CASE_4x4 = new ReorderTestCase(START_BOARD_STR_4x4,
+            MOVE_TO_4x4,
+            END_BOARD_STR_4x4);
+
+
+    /** 6x5 Test
+     **/
+    private static final String START_BOARD_STR_6x5 = ""
+            + "xxxxxx|------\n"
+            + "--m---|------\n"
+            + "------|------\n"
+            + "------|------\n"
+            + "------|------";
+    private static final Point MOVE_TO_6x5 = new Point(10, 4);
+    private static final String END_BOARD_STR_6x5 = ""
+            + "xxxxxx|------\n"
+            + "------|------\n"
+            + "------|------\n"
+            + "------|------\n"
+            + "------|----m-";
+    private static final ReorderTestCase TEST_CASE_6x5 = new ReorderTestCase(START_BOARD_STR_6x5,
+            MOVE_TO_6x5,
+            END_BOARD_STR_6x5);
+
+    public static final Map<Point, ReorderTestCase> TEST_BY_GRID_SIZE =
+            Map.of(new Point(5, 5), TEST_CASE_5x5,
+                    new Point(4, 4), TEST_CASE_4x4,
+                    new Point(6, 5), TEST_CASE_6x5);
+}
diff --git a/tests/src/com/android/launcher3/celllayout/ReorderWidgets.java b/tests/src/com/android/launcher3/celllayout/ReorderWidgets.java
index 7b38ed6..35f6cbc 100644
--- a/tests/src/com/android/launcher3/celllayout/ReorderWidgets.java
+++ b/tests/src/com/android/launcher3/celllayout/ReorderWidgets.java
@@ -29,6 +29,7 @@
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.celllayout.testcases.FullReorderCase;
 import com.android.launcher3.celllayout.testcases.MoveOutReorderCase;
+import com.android.launcher3.celllayout.testcases.MultipleCellLayoutsSimpleReorder;
 import com.android.launcher3.celllayout.testcases.PushReorderCase;
 import com.android.launcher3.celllayout.testcases.ReorderTestCase;
 import com.android.launcher3.celllayout.testcases.SimpleReorderCase;
@@ -38,7 +39,6 @@
 import com.android.launcher3.ui.TaplTestsLauncher3;
 import com.android.launcher3.util.rule.ShellCommandRule;
 import com.android.launcher3.views.DoubleShadowBubbleTextView;
-import com.android.launcher3.widget.LauncherAppWidgetHostView;
 
 import org.junit.Assume;
 import org.junit.Before;
@@ -48,6 +48,7 @@
 import org.junit.runner.RunWith;
 
 import java.util.ArrayList;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ExecutionException;
 
@@ -62,20 +63,6 @@
 
     TestWorkspaceBuilder mWorkspaceBuilder;
 
-    private View getViewAt(int cellX, int cellY) {
-        return getFromLauncher(l -> l.getWorkspace().getScreenWithId(
-                l.getWorkspace().getScreenIdForPageIndex(0)).getChildAt(cellX, cellY));
-    }
-
-    private Point getCellDimensions() {
-        return getFromLauncher(l -> {
-            CellLayout cellLayout = l.getWorkspace().getScreenWithId(
-                    l.getWorkspace().getScreenIdForPageIndex(0));
-            return new Point(cellLayout.getWidth() / cellLayout.getCountX(),
-                    cellLayout.getHeight() / cellLayout.getCountY());
-        });
-    }
-
     @Before
     public void setup() throws Throwable {
         mWorkspaceBuilder = new TestWorkspaceBuilder(this, mTargetContext);
@@ -86,28 +73,26 @@
     /**
      * Validate if the given board represent the current CellLayout
      **/
-    private boolean validateBoard(CellLayoutBoard board) {
-        boolean match = true;
-        Point cellDimensions = getCellDimensions();
-        for (CellLayoutBoard.WidgetRect widgetRect : board.getWidgets()) {
-            if (widgetRect.shouldIgnore()) {
-                continue;
+    private boolean validateBoard(List<CellLayoutBoard> testBoards) {
+        ArrayList<CellLayoutBoard> workspaceBoards = workspaceToBoards();
+        if (workspaceBoards.size() < testBoards.size()) {
+            return false;
+        }
+        for (int i = 0; i < testBoards.size(); i++) {
+            if (testBoards.get(i).compareTo(workspaceBoards.get(i)) != 0) {
+                return false;
             }
-            View widget = getViewAt(widgetRect.getCellX(), widgetRect.getCellY());
-            assertTrue("The view selected at " + board + " is not a widget",
-                    widget instanceof LauncherAppWidgetHostView);
-            match &= widgetRect.getSpanX()
-                    == Math.round(widget.getWidth() / (float) cellDimensions.x);
-            match &= widgetRect.getSpanY()
-                    == Math.round(widget.getHeight() / (float) cellDimensions.y);
-            if (!match) return match;
         }
-        for (CellLayoutBoard.IconPoint iconPoint : board.getIcons()) {
-            View icon = getViewAt(iconPoint.getCoord().x, iconPoint.getCoord().y);
-            assertTrue("The view selected at " + iconPoint.coord + " is not an Icon",
-                    icon instanceof DoubleShadowBubbleTextView);
+        return true;
+    }
+
+    private FavoriteItemsTransaction buildWorkspaceFromBoards(List<CellLayoutBoard> boards,
+            FavoriteItemsTransaction transaction) {
+        for (int i = 0; i < boards.size(); i++) {
+            CellLayoutBoard board = boards.get(i);
+            mWorkspaceBuilder.buildFromBoard(board, transaction, i);
         }
-        return match;
+        return transaction;
     }
 
     private void printCurrentWorkspace() {
@@ -148,22 +133,25 @@
 
     private void runTestCase(ReorderTestCase testCase)
             throws ExecutionException, InterruptedException {
-        Point mainWidgetCellPos = testCase.mStart.getMain();
+        CellLayoutBoard.WidgetRect mainWidgetCellPos = CellLayoutBoard.getMainFromList(
+                testCase.mStart);
 
         FavoriteItemsTransaction transaction =
                 new FavoriteItemsTransaction(mTargetContext, this);
-        mWorkspaceBuilder.buildFromBoard(testCase.mStart, transaction).commit();
+        transaction = buildWorkspaceFromBoards(testCase.mStart, transaction);
+        transaction.commit();
         waitForLauncherCondition("Workspace didn't finish loading", l -> !l.isWorkspaceLoading());
-        Widget widget = mLauncher.getWorkspace().getWidgetAtCell(mainWidgetCellPos.x,
-                mainWidgetCellPos.y);
+        Widget widget = mLauncher.getWorkspace().getWidgetAtCell(mainWidgetCellPos.getCellX(),
+                mainWidgetCellPos.getCellY());
         assertNotNull(widget);
         WidgetResizeFrame resizeFrame = widget.dragWidgetToWorkspace(testCase.moveMainTo.x,
-                testCase.moveMainTo.y);
+                testCase.moveMainTo.y, mainWidgetCellPos.getSpanX(), mainWidgetCellPos.getSpanY());
         resizeFrame.dismiss();
 
         boolean isValid = false;
-        for (CellLayoutBoard board : testCase.mEnd) {
-            isValid |= validateBoard(board);
+        for (List<CellLayoutBoard> boards : testCase.mEnd) {
+            isValid |= validateBoard(boards);
+            if (isValid) break;
         }
         printCurrentWorkspace();
         assertTrue("Non of the valid boards match with the current state", isValid);
@@ -207,4 +195,11 @@
         runTestCaseMap(MoveOutReorderCase.TEST_BY_GRID_SIZE,
                 MoveOutReorderCase.class.getSimpleName());
     }
+
+    @Test
+    public void multipleCellLayoutsSimpleReorder() throws ExecutionException, InterruptedException {
+        Assume.assumeTrue("Test doesn't support foldables", !mLauncher.isTwoPanels());
+        runTestCaseMap(MultipleCellLayoutsSimpleReorder.TEST_BY_GRID_SIZE,
+                MultipleCellLayoutsSimpleReorder.class.getSimpleName());
+    }
 }
diff --git a/tests/src/com/android/launcher3/celllayout/TestWorkspaceBuilder.java b/tests/src/com/android/launcher3/celllayout/TestWorkspaceBuilder.java
index 7e3588b..5945605 100644
--- a/tests/src/com/android/launcher3/celllayout/TestWorkspaceBuilder.java
+++ b/tests/src/com/android/launcher3/celllayout/TestWorkspaceBuilder.java
@@ -15,7 +15,6 @@
  */
 package com.android.launcher3.celllayout;
 
-import static com.android.launcher3.WorkspaceLayoutManager.FIRST_SCREEN_ID;
 import static com.android.launcher3.util.WidgetUtils.createWidgetInfo;
 
 import android.content.ComponentName;
@@ -62,7 +61,7 @@
      * Fills the given rect in WidgetRect with 1x1 widgets. This is useful to equalize cases.
      */
     private FavoriteItemsTransaction fillWithWidgets(CellLayoutBoard.WidgetRect widgetRect,
-            FavoriteItemsTransaction transaction) {
+            FavoriteItemsTransaction transaction, int screenId) {
         int initX = widgetRect.getCellX();
         int initY = widgetRect.getCellY();
         for (int x = initX; x < initX + widgetRect.getSpanX(); x++) {
@@ -71,8 +70,7 @@
                     // this widgets are filling, we don't care if we can't place them
                     ItemInfo item = createWidgetInCell(
                             new CellLayoutBoard.WidgetRect(CellLayoutBoard.CellType.IGNORE,
-                                    new Rect(x, y, x, y))
-                    );
+                                    new Rect(x, y, x, y)), screenId);
                     transaction.addItem(item);
                 } catch (Exception e) {
                     Log.d(TAG, "Unable to place filling widget at " + x + "," + y);
@@ -94,11 +92,11 @@
     }
 
     private void addCorrespondingWidgetRect(CellLayoutBoard.WidgetRect widgetRect,
-            FavoriteItemsTransaction transaction) {
+            FavoriteItemsTransaction transaction, int screenId) {
         if (widgetRect.mType == 'x') {
-            fillWithWidgets(widgetRect, transaction);
+            fillWithWidgets(widgetRect, transaction, screenId);
         } else {
-            transaction.addItem(createWidgetInCell(widgetRect));
+            transaction.addItem(createWidgetInCell(widgetRect, screenId));
         }
     }
 
@@ -106,11 +104,11 @@
      * Builds the given board into the transaction
      */
     public FavoriteItemsTransaction buildFromBoard(CellLayoutBoard board,
-            FavoriteItemsTransaction transaction) {
+            FavoriteItemsTransaction transaction, final int screenId) {
         board.getWidgets().forEach(
-                (widgetRect) -> addCorrespondingWidgetRect(widgetRect, transaction));
+                (widgetRect) -> addCorrespondingWidgetRect(widgetRect, transaction, screenId));
         board.getIcons().forEach((iconPoint) ->
-                transaction.addItem(createIconInCell(iconPoint))
+                transaction.addItem(createIconInCell(iconPoint, screenId))
         );
         return transaction;
     }
@@ -127,7 +125,7 @@
         return transaction;
     }
 
-    private ItemInfo createWidgetInCell(CellLayoutBoard.WidgetRect widgetRect) {
+    private ItemInfo createWidgetInCell(CellLayoutBoard.WidgetRect widgetRect, int screenId) {
         LauncherAppWidgetProviderInfo info = TestViewHelpers.findWidgetProvider(mTest, false);
         LauncherAppWidgetInfo item = createWidgetInfo(info,
                 ApplicationProvider.getApplicationContext(), true);
@@ -136,14 +134,14 @@
         item.cellY = widgetRect.getCellY();
         item.spanX = widgetRect.getSpanX();
         item.spanY = widgetRect.getSpanY();
-        item.screenId = FIRST_SCREEN_ID;
+        item.screenId = screenId;
         return item;
     }
 
-    private ItemInfo createIconInCell(CellLayoutBoard.IconPoint iconPoint) {
+    private ItemInfo createIconInCell(CellLayoutBoard.IconPoint iconPoint, int screenId) {
         WorkspaceItemInfo item = new WorkspaceItemInfo(getApp());
         item.id = getID();
-        item.screenId = FIRST_SCREEN_ID;
+        item.screenId = screenId;
         item.cellX = iconPoint.getCoord().x;
         item.cellY = iconPoint.getCoord().y;
         item.minSpanY = item.minSpanX = item.spanX = item.spanY = 1;
diff --git a/tests/src/com/android/launcher3/celllayout/testcases/ReorderTestCase.java b/tests/src/com/android/launcher3/celllayout/testcases/ReorderTestCase.java
index 0a28668..1689e47 100644
--- a/tests/src/com/android/launcher3/celllayout/testcases/ReorderTestCase.java
+++ b/tests/src/com/android/launcher3/celllayout/testcases/ReorderTestCase.java
@@ -24,23 +24,23 @@
 import java.util.stream.Collectors;
 
 public class ReorderTestCase {
-    public CellLayoutBoard mStart;
+    public List<CellLayoutBoard> mStart;
     public Point moveMainTo;
-    public List<CellLayoutBoard> mEnd;
+    public List<List<CellLayoutBoard>> mEnd;
 
-    ReorderTestCase(CellLayoutBoard start, Point moveMainTo, CellLayoutBoard ... end) {
+    ReorderTestCase(List<CellLayoutBoard> start, Point moveMainTo, List<CellLayoutBoard> ... end) {
         mStart = start;
         this.moveMainTo = moveMainTo;
         mEnd = Arrays.asList(end);
     }
 
     ReorderTestCase(String start, Point moveMainTo, String ... end) {
-        mStart = CellLayoutBoard.boardFromString(start);
+        mStart = CellLayoutBoard.boardListFromString(start);
         this.moveMainTo = moveMainTo;
         mEnd = Arrays
                 .asList(end)
                 .stream()
-                .map(CellLayoutBoard::boardFromString)
+                .map(CellLayoutBoard::boardListFromString)
                 .collect(Collectors.toList());
     }
 }
diff --git a/tests/src/com/android/launcher3/celllayout/testcases/SimpleReorderCase.java b/tests/src/com/android/launcher3/celllayout/testcases/SimpleReorderCase.java
index 546c48b..c9f6849 100644
--- a/tests/src/com/android/launcher3/celllayout/testcases/SimpleReorderCase.java
+++ b/tests/src/com/android/launcher3/celllayout/testcases/SimpleReorderCase.java
@@ -33,13 +33,13 @@
             + "--mm-\n"
             + "-----\n"
             + "-----";
-    private static final Point MOVE_TO_5x5 = new Point(4, 4);
+    private static final Point MOVE_TO_5x5 = new Point(0, 4);
     private static final String END_BOARD_STR_5x5 = ""
             + "xxxxx\n"
             + "-----\n"
             + "-----\n"
-            + "---mm\n"
-            + "---mm";
+            + "mm---\n"
+            + "mm---";
     private static final ReorderTestCase TEST_CASE_5x5 = new ReorderTestCase(START_BOARD_STR_5x5,
             MOVE_TO_5x5,
             END_BOARD_STR_5x5);
@@ -61,7 +61,27 @@
             MOVE_TO_4x4,
             END_BOARD_STR_4x4);
 
+    /** 6x5 Test
+     **/
+    private static final String START_BOARD_STR_6x5 = ""
+            + "xxxxxx\n"
+            + "-mm---\n"
+            + "-mm---\n"
+            + "------\n"
+            + "------";
+    private static final Point MOVE_TO_6x5 = new Point(4, 3);
+    private static final String END_BOARD_STR_6x5 = ""
+            + "xxxxxx\n"
+            + "------\n"
+            + "------\n"
+            + "----mm\n"
+            + "----mm";
+    private static final ReorderTestCase TEST_CASE_6x5 = new ReorderTestCase(START_BOARD_STR_6x5,
+            MOVE_TO_6x5,
+            END_BOARD_STR_6x5);
+
     public static final Map<Point, ReorderTestCase> TEST_BY_GRID_SIZE =
             Map.of(new Point(5, 5), TEST_CASE_5x5,
-                    new Point(4, 4), TEST_CASE_4x4);
+                    new Point(4, 4), TEST_CASE_4x4,
+                    new Point(6, 5), TEST_CASE_6x5);
 }
diff --git a/tests/tapl/com/android/launcher3/tapl/Widget.java b/tests/tapl/com/android/launcher3/tapl/Widget.java
index 046308b..d440903 100644
--- a/tests/tapl/com/android/launcher3/tapl/Widget.java
+++ b/tests/tapl/com/android/launcher3/tapl/Widget.java
@@ -71,7 +71,7 @@
     public WidgetResizeFrame dragWidgetToWorkspace() {
         try (LauncherInstrumentation.Closable e = mLauncher.eventsCheck()) {
             return dragWidgetToWorkspace(/* configurable= */ false, /* acceptsConfig= */ false, -1,
-                    -1);
+                    -1, 1, 1);
         }
     }
 
@@ -80,12 +80,12 @@
      * cellY and returns the resize frame that is shown after the widget is added.
      */
     @NonNull
-    public WidgetResizeFrame dragWidgetToWorkspace(int cellX, int cellY) {
+    public WidgetResizeFrame dragWidgetToWorkspace(int cellX, int cellY, int spanX, int spanY) {
         try (LauncherInstrumentation.Closable e = mLauncher.eventsCheck();
              LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
                      "Dragging widget to workspace cell " + cellX + "," + cellY)) {
             return dragWidgetToWorkspace(/* configurable= */ false, /* acceptsConfig= */ false,
-                    cellX, cellY);
+                    cellX, cellY, spanX, spanY);
         }
     }
 
@@ -98,7 +98,7 @@
     public WidgetResizeFrame dragConfigWidgetToWorkspace(boolean acceptsConfig) {
         // TODO(b/239438337, fransebas) add correct event checking for this case
         //try (LauncherInstrumentation.Closable e = mLauncher.eventsCheck()) {
-        return dragWidgetToWorkspace(/* configurable= */ true, acceptsConfig, -1, -1);
+        return dragWidgetToWorkspace(/* configurable= */ true, acceptsConfig, -1, -1, 1, 1);
         //}
     }
 
@@ -110,18 +110,17 @@
      * @param cellX            X position in the CellLayout
      * @param cellY            Y position in the CellLayout
      */
-    private void dragToWorkspace(boolean startsActivity, boolean isWidgetShortcut, int cellX,
-            int cellY) {
+    private void dragToWorkspaceCellPosition(boolean startsActivity, boolean isWidgetShortcut,
+            int cellX, int cellY, int spanX, int spanY) {
         Launchable launchable = getLaunchable();
         LauncherInstrumentation launcher = launchable.mLauncher;
-        Workspace.dragIconToWorkspace(
+        Workspace.dragIconToWorkspaceCellPosition(
                 launcher,
                 launchable,
-                () -> Workspace.getCellCenter(launchable.mLauncher, cellX, cellY),
+                cellX, cellY, spanX, spanY,
                 startsActivity,
                 isWidgetShortcut,
                 launchable::addExpectedEventsForLongClick);
-
     }
 
     /**
@@ -144,13 +143,13 @@
      */
     @Nullable
     private WidgetResizeFrame dragWidgetToWorkspace(boolean configurable, boolean acceptsConfig,
-            int cellX, int cellY) {
+            int cellX, int cellY, int spanX, int spanY) {
         if (cellX == -1 || cellY == -1) {
             internalDragToWorkspace(/* startsActivity= */ configurable, /* isWidgetShortcut= */
                     false);
         } else {
-            dragToWorkspace(/* startsActivity= */ configurable, /* isWidgetShortcut= */ false,
-                    cellX, cellY);
+            dragToWorkspaceCellPosition(/* startsActivity= */ configurable, /* isWidgetShortcut= */
+                    false, cellX, cellY, spanX, spanY);
         }
 
         if (configurable) {
diff --git a/tests/tapl/com/android/launcher3/tapl/Workspace.java b/tests/tapl/com/android/launcher3/tapl/Workspace.java
index 388955c..62665de 100644
--- a/tests/tapl/com/android/launcher3/tapl/Workspace.java
+++ b/tests/tapl/com/android/launcher3/tapl/Workspace.java
@@ -409,9 +409,15 @@
     }
 
     static Point getCellCenter(LauncherInstrumentation launcher, int cellX, int cellY) {
-        return launcher.getTestInfo(WorkspaceCellCenterRequest.builder().setCellX(
-                cellX).setCellY(cellY).build()).getParcelable(
-                TestProtocol.TEST_INFO_RESPONSE_FIELD);
+        return launcher.getTestInfo(WorkspaceCellCenterRequest.builder().setCellX(cellX).setCellY(
+                cellY).build()).getParcelable(TestProtocol.TEST_INFO_RESPONSE_FIELD);
+    }
+
+    static Point getCellCenter(LauncherInstrumentation launcher, int cellX, int cellY, int spanX,
+            int spanY) {
+        return launcher.getTestInfo(WorkspaceCellCenterRequest.builder().setCellX(cellX)
+                .setCellY(cellY).setSpanX(spanX).setSpanY(spanY).build())
+                .getParcelable(TestProtocol.TEST_INFO_RESPONSE_FIELD);
     }
 
     static Point getHotseatCellCenter(LauncherInstrumentation launcher, int cellInd) {
@@ -419,6 +425,12 @@
                 .setCellInd(cellInd).build()).getParcelable(TestProtocol.TEST_INFO_RESPONSE_FIELD);
     }
 
+    /** Returns the number of rows and columns in the workspace */
+    public Point getRowsAndCols() {
+        return mLauncher.getTestInfo(TestProtocol.REQUEST_WORKSPACE_COLUMNS_ROWS).getParcelable(
+                TestProtocol.TEST_INFO_RESPONSE_FIELD);
+    }
+
     /**
      * Finds folder icons in the current workspace.
      *
@@ -457,6 +469,19 @@
                 launcher, launchable, dest, expectLongClickEvents, expectDropEvents);
     }
 
+    static void dragIconToWorkspaceCellPosition(LauncherInstrumentation launcher,
+            Launchable launchable, int cellX, int cellY, int spanX, int spanY,
+            boolean startsActivity, boolean isWidgetShortcut, Runnable expectLongClickEvents) {
+        Runnable expectDropEvents = null;
+        if (startsActivity || isWidgetShortcut) {
+            expectDropEvents = () -> launcher.expectEvent(TestProtocol.SEQUENCE_MAIN,
+                    LauncherInstrumentation.EVENT_START);
+        }
+        dragIconToWorkspaceCellPosition(
+                launcher, launchable, cellX, cellY, spanX, spanY, true, expectLongClickEvents,
+                expectDropEvents);
+    }
+
     /**
      * Drag icon in workspace to else where and drop it immediately.
      * (There is no slow down time before drop event)
@@ -526,6 +551,51 @@
         }
     }
 
+    static void dragIconToWorkspaceCellPosition(
+            LauncherInstrumentation launcher,
+            Launchable launchable,
+            int cellX, int cellY, int spanX, int spanY,
+            boolean isDecelerating,
+            Runnable expectLongClickEvents,
+            @Nullable Runnable expectDropEvents) {
+        try (LauncherInstrumentation.Closable ignored = launcher.addContextLayer(
+                "want to drag icon to workspace")) {
+            Point rowsAndCols = launcher.getWorkspace().getRowsAndCols();
+            int destinationWorkspace = cellX / rowsAndCols.x;
+            cellX = cellX % rowsAndCols.x;
+
+            final long downTime = SystemClock.uptimeMillis();
+            Point dragStart = launchable.startDrag(
+                    downTime,
+                    expectLongClickEvents,
+                    /* runToSpringLoadedState= */ true);
+            Point targetDest = getCellCenter(launcher, cellX, cellY, spanX, spanY);
+            int displayX = launcher.getRealDisplaySize().x;
+
+            // Since the destination can be on another page, we need to drag to the edge first
+            // until we reach the target page
+            for (int i = 0; i < destinationWorkspace; i++) {
+                // Don't drag all the way to the edge to prevent touch events from getting out of
+                //screen bounds.
+                Point screenEdge = new Point(displayX - 1, targetDest.y);
+                Point finalDragStart = dragStart;
+                executeAndWaitForPageScroll(launcher,
+                        () -> launcher.movePointer(finalDragStart, screenEdge, DEFAULT_DRAG_STEPS,
+                                true, downTime, downTime, true,
+                                LauncherInstrumentation.GestureScope.INSIDE));
+                dragStart = screenEdge;
+            }
+
+            // targetDest.x is now between 0 and displayX so we found the target page,
+            // we just have to put move the icon to the destination and drop it
+            launcher.movePointer(dragStart, targetDest, DEFAULT_DRAG_STEPS, isDecelerating,
+                    downTime, SystemClock.uptimeMillis(), false,
+                    LauncherInstrumentation.GestureScope.INSIDE);
+            launcher.runCallbackIfActive(CALLBACK_HOLD_BEFORE_DROP);
+            dropDraggedIcon(launcher, targetDest, downTime, expectDropEvents);
+        }
+    }
+
     private static void executeAndWaitForPageScroll(LauncherInstrumentation launcher,
             Runnable command) {
         launcher.executeAndWaitForEvent(command,