Quickswitch with staged split

* UI polish/animations needed.
* One known bug (b/198310766), temp work around is to
swipe up to home.

Bug: 181704764
Test:
* Open apps in staged split and quickswitch
between GroupedTaskView and fullscreen apps.
* QS to fullscreen app and then go into overview
and re-launch split screen tasks
* QS to fullscreen app, wait 5 seconds,
swipe into overview, no GroupedTaskView shown

Change-Id: I0ce10a944d86be5c927eeaaef922559a40f39923
diff --git a/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java b/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
index b4b29aa..033cf8e 100644
--- a/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
@@ -297,7 +297,7 @@
         mActionsView = findViewById(R.id.overview_actions_view);
         RecentsView overviewPanel = (RecentsView) getOverviewPanel();
         SplitSelectStateController controller =
-                new SplitSelectStateController(mHandler, SystemUiProxy.INSTANCE.get(this));
+                new SplitSelectStateController(SystemUiProxy.INSTANCE.get(this));
         overviewPanel.init(mActionsView, controller);
         mActionsView.setDp(getDeviceProfile());
         mActionsView.updateVerticalMargin(SysUINavigationMode.getMode(this));
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index 6e90a3a..0d83510 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -531,7 +531,7 @@
         ActivityManager.RunningTaskInfo[] runningTasks;
         if (mIsSwipeForStagedSplit) {
             int[] splitTaskIds =
-                    LauncherSplitScreenListener.INSTANCE.getNoCreate().getSplitTaskIds();
+                    LauncherSplitScreenListener.INSTANCE.getNoCreate().getRunningSplitTaskIds();
             runningTasks = new ActivityManager.RunningTaskInfo[splitTaskIds.length];
             for (int i = 0; i < splitTaskIds.length; i++) {
                 int taskId = splitTaskIds[i];
diff --git a/quickstep/src/com/android/quickstep/RecentsActivity.java b/quickstep/src/com/android/quickstep/RecentsActivity.java
index 03e2395..cc6cfd7 100644
--- a/quickstep/src/com/android/quickstep/RecentsActivity.java
+++ b/quickstep/src/com/android/quickstep/RecentsActivity.java
@@ -122,7 +122,7 @@
         SYSUI_PROGRESS.set(getRootView().getSysUiScrim(), 0f);
 
         SplitSelectStateController controller =
-                new SplitSelectStateController(mHandler, SystemUiProxy.INSTANCE.get(this));
+                new SplitSelectStateController(SystemUiProxy.INSTANCE.get(this));
         mDragLayer.recreateControllers();
         mFallbackRecentsView.init(mActionsView, controller);
     }
diff --git a/quickstep/src/com/android/quickstep/SwipeUpAnimationLogic.java b/quickstep/src/com/android/quickstep/SwipeUpAnimationLogic.java
index 1f57e99..e14dbb1 100644
--- a/quickstep/src/com/android/quickstep/SwipeUpAnimationLogic.java
+++ b/quickstep/src/com/android/quickstep/SwipeUpAnimationLogic.java
@@ -90,7 +90,8 @@
         mGestureState = gestureState;
 
         mIsSwipeForStagedSplit = ENABLE_SPLIT_SELECT.get() &&
-                LauncherSplitScreenListener.INSTANCE.getNoCreate().getSplitTaskIds().length > 1;
+                LauncherSplitScreenListener.INSTANCE.getNoCreate()
+                        .getRunningSplitTaskIds().length > 1;
 
         TaskViewSimulator primaryTVS = new TaskViewSimulator(context,
                 gestureState.getActivityInterface());
diff --git a/quickstep/src/com/android/quickstep/util/LauncherSplitScreenListener.java b/quickstep/src/com/android/quickstep/util/LauncherSplitScreenListener.java
index da665d4..0f4ed01 100644
--- a/quickstep/src/com/android/quickstep/util/LauncherSplitScreenListener.java
+++ b/quickstep/src/com/android/quickstep/util/LauncherSplitScreenListener.java
@@ -11,20 +11,54 @@
 import com.android.launcher3.util.SplitConfigurationOptions.StageType;
 import com.android.launcher3.util.SplitConfigurationOptions.StagedSplitTaskPosition;
 import com.android.quickstep.SystemUiProxy;
+import com.android.systemui.shared.system.TaskStackChangeListener;
+import com.android.systemui.shared.system.TaskStackChangeListeners;
 import com.android.wm.shell.splitscreen.ISplitScreenListener;
 
 /**
  * Listeners for system wide split screen position and stage changes.
- * Use {@link #getSplitTaskIds()} to determine which tasks, if any, are in staged split.
+ *
+ * Use {@link #getRunningSplitTaskIds()} to determine which tasks, if any, are actively in
+ * staged split.
+ *
+ * Use {@link #getPersistentSplitIds()} to know if tasks were in split screen before a quickswitch
+ * gesture happened.
  */
 public class LauncherSplitScreenListener extends ISplitScreenListener.Stub {
 
     public static final MainThreadInitializedObject<LauncherSplitScreenListener> INSTANCE =
             new MainThreadInitializedObject<>(LauncherSplitScreenListener::new);
 
+    private static final int[] EMPTY_ARRAY = {};
+
     private final StagedSplitTaskPosition mMainStagePosition = new StagedSplitTaskPosition();
     private final StagedSplitTaskPosition mSideStagePosition = new StagedSplitTaskPosition();
 
+    private boolean mIsRecentsListFrozen = false;
+    private final TaskStackChangeListener mTaskStackListener = new TaskStackChangeListener() {
+        @Override
+        public void onRecentTaskListFrozenChanged(boolean frozen) {
+            super.onRecentTaskListFrozenChanged(frozen);
+            mIsRecentsListFrozen = frozen;
+
+            if (frozen) {
+                mPersistentGroupedIds = getRunningSplitTaskIds();
+            } else {
+                // TODO(b/198310766) Need to also explicitly exit split screen if
+                //  we're not currently viewing split screened apps
+                mPersistentGroupedIds = EMPTY_ARRAY;
+            }
+        }
+    };
+
+    /**
+     * Gets set to current split taskIDs whenever the task list is frozen, and set to empty array
+     * whenever task list unfreezes.
+     * When not null, this indicates that we need to load a GroupedTaskView as the most recent
+     * page, so user can quickswitch back to a grouped task.
+     */
+    private int[] mPersistentGroupedIds;
+
     public LauncherSplitScreenListener(Context context) {
         mMainStagePosition.stageType = SplitConfigurationOptions.STAGE_TYPE_MAIN;
         mSideStagePosition.stageType = SplitConfigurationOptions.STAGE_TYPE_SIDE;
@@ -33,17 +67,30 @@
     /** Also call {@link #destroy()} when done. */
     public void init() {
         SystemUiProxy.INSTANCE.getNoCreate().registerSplitScreenListener(this);
+        TaskStackChangeListeners.getInstance().registerTaskStackListener(mTaskStackListener);
     }
 
     public void destroy() {
         SystemUiProxy.INSTANCE.getNoCreate().unregisterSplitScreenListener(this);
+        TaskStackChangeListeners.getInstance().unregisterTaskStackListener(mTaskStackListener);
     }
 
     /**
+     * This method returns the active split taskIDs that were active if a user quickswitched from
+     * split screen to a fullscreen app as long as the recents task list remains frozen.
+     */
+    public int[] getPersistentSplitIds() {
+        if (mIsRecentsListFrozen) {
+            return mPersistentGroupedIds;
+        } else {
+            return getRunningSplitTaskIds();
+        }
+    }
+    /**
      * @return index 0 will be task in left/top position, index 1 in right/bottom position.
      *         Will return empty array if device is not in staged split
      */
-    public int[] getSplitTaskIds() {
+    public int[] getRunningSplitTaskIds() {
         if (mMainStagePosition.taskId == -1 || mSideStagePosition.taskId == -1) {
             return new int[]{};
         }
diff --git a/quickstep/src/com/android/quickstep/util/SplitSelectStateController.java b/quickstep/src/com/android/quickstep/util/SplitSelectStateController.java
index ac5f5d8..3069504 100644
--- a/quickstep/src/com/android/quickstep/util/SplitSelectStateController.java
+++ b/quickstep/src/com/android/quickstep/util/SplitSelectStateController.java
@@ -22,7 +22,6 @@
 
 import android.app.ActivityThread;
 import android.graphics.Rect;
-import android.os.Handler;
 import android.os.IBinder;
 import android.view.RemoteAnimationAdapter;
 import android.view.SurfaceControl;
@@ -40,6 +39,8 @@
 import com.android.systemui.shared.system.RemoteTransitionCompat;
 import com.android.systemui.shared.system.RemoteTransitionRunner;
 
+import java.util.function.Consumer;
+
 /**
  * Represent data needed for the transient state when user has selected one app for split screen
  * and is in the process of either a) selecting a second app or b) exiting intention to invoke split
@@ -52,7 +53,7 @@
     private Task mSecondTask;
     private Rect mInitialBounds;
 
-    public SplitSelectStateController(Handler handler, SystemUiProxy systemUiProxy) {
+    public SplitSelectStateController(SystemUiProxy systemUiProxy) {
         mSystemUiProxy = systemUiProxy;
     }
 
@@ -71,13 +72,14 @@
      */
     public void setSecondTaskId(Task taskView) {
         mSecondTask = taskView;
-        launchTasks(mInitialTask, mSecondTask, mStagePosition);
+        launchTasks(mInitialTask, mSecondTask, mStagePosition, null /*callback*/);
     }
 
     /**
      * @param stagePosition representing location of task1
      */
-    public void launchTasks(Task task1, Task task2, @StagePosition int stagePosition) {
+    public void launchTasks(Task task1, Task task2, @StagePosition int stagePosition,
+            Consumer<Boolean> callback) {
         // Assume initial task is for top/left part of screen
         final int[] taskIds = stagePosition == STAGE_POSITION_TOP_OR_LEFT
                 ? new int[]{task1.key.id, task2.key.id}
@@ -90,7 +92,7 @@
                     new RemoteTransitionCompat(animationRunner, MAIN_EXECUTOR));
         } else {
             RemoteSplitLaunchAnimationRunner animationRunner =
-                    new RemoteSplitLaunchAnimationRunner(task1, task2);
+                    new RemoteSplitLaunchAnimationRunner(task1, task2, callback);
             final RemoteAnimationAdapter adapter = new RemoteAnimationAdapter(
                     RemoteAnimationAdapterCompat.wrapRemoteAnimationRunner(animationRunner),
                     300, 150,
@@ -136,10 +138,13 @@
 
         private final Task mInitialTask;
         private final Task mSecondTask;
+        private final Consumer<Boolean> mSuccessCallback;
 
-        RemoteSplitLaunchAnimationRunner(Task initialTask, Task secondTask) {
+        RemoteSplitLaunchAnimationRunner(Task initialTask, Task secondTask,
+                Consumer<Boolean> successCallback) {
             mInitialTask = initialTask;
             mSecondTask = secondTask;
+            mSuccessCallback = successCallback;
         }
 
         @Override
@@ -147,13 +152,21 @@
                 RemoteAnimationTargetCompat[] wallpapers, RemoteAnimationTargetCompat[] nonApps,
                 Runnable finishedCallback) {
             TaskViewUtils.composeRecentsSplitLaunchAnimatorLegacy(mInitialTask,
-                    mSecondTask, apps, wallpapers, nonApps, finishedCallback);
+                    mSecondTask, apps, wallpapers, nonApps, () -> {
+                        finishedCallback.run();
+                        if (mSuccessCallback != null) {
+                            mSuccessCallback.accept(true);
+                        }
+                    });
             // After successful launch, call resetState
             resetState();
         }
 
         @Override
         public void onAnimationCancelled() {
+            if (mSuccessCallback != null) {
+                mSuccessCallback.accept(false);
+            }
             resetState();
         }
     }
diff --git a/quickstep/src/com/android/quickstep/views/GroupedTaskView.java b/quickstep/src/com/android/quickstep/views/GroupedTaskView.java
index e1e1c65..d3077ad 100644
--- a/quickstep/src/com/android/quickstep/views/GroupedTaskView.java
+++ b/quickstep/src/com/android/quickstep/views/GroupedTaskView.java
@@ -3,6 +3,8 @@
 import android.content.Context;
 import android.util.AttributeSet;
 
+import androidx.annotation.NonNull;
+
 import com.android.launcher3.R;
 import com.android.launcher3.util.RunnableList;
 import com.android.launcher3.util.SplitConfigurationOptions;
@@ -12,6 +14,8 @@
 import com.android.quickstep.util.RecentsOrientedState;
 import com.android.systemui.shared.recents.model.Task;
 
+import java.util.function.Consumer;
+
 /**
  * TaskView that contains and shows thumbnails for not one, BUT TWO(!!) tasks
  *
@@ -98,11 +102,17 @@
     @Override
     public RunnableList launchTaskAnimated() {
         getRecentsView().getSplitPlaceholder().launchTasks(mTask, mSecondaryTask,
-                SplitConfigurationOptions.STAGE_POSITION_TOP_OR_LEFT);
+                SplitConfigurationOptions.STAGE_POSITION_TOP_OR_LEFT, null /*callback*/);
         return null;
     }
 
     @Override
+    public void launchTask(@NonNull Consumer<Boolean> callback, boolean freezeTaskList) {
+        getRecentsView().getSplitPlaceholder().launchTasks(mTask, mSecondaryTask,
+                SplitConfigurationOptions.STAGE_POSITION_TOP_OR_LEFT, callback);
+    }
+
+    @Override
     public void onRecycle() {
         super.onRecycle();
         mSnapshotView2.setThumbnail(mSecondaryTask, null);
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 810ecce..c6e2a5d 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -1259,7 +1259,7 @@
                 mIgnoreResetTaskId == -1 ? null : getTaskViewByTaskId(mIgnoreResetTaskId);
 
         int[] splitTaskIds =
-                LauncherSplitScreenListener.INSTANCE.getNoCreate().getSplitTaskIds();
+                LauncherSplitScreenListener.INSTANCE.getNoCreate().getPersistentSplitIds();
         int requiredGroupTaskViews = splitTaskIds.length / 2;
 
         // Subtract half the number of split tasks and not total number because we've already
@@ -1957,6 +1957,9 @@
     /**
      * Called only when a swipe-up gesture from an app has completed. Only called after
      * {@link #onGestureAnimationStart} and {@link #onGestureAnimationEnd()}.
+     *
+     * TODO(b/198310766) Need to also explicitly exit split screen if
+     *  the swipe up was to home
      */
     public void onSwipeUpAnimationSuccess() {
         animateUpTaskIconScale();