Fix visual issues with splitscreen to desktop transition.

Fixes three minor visual issues with the split to desktop transition:

1 - Modifies getFocusedDecor to always prioritize a currently dragged decor,
fixing an issue where some split -> desktop transitions would hide the
dragged task if they ended on the opposite side of the screen.

2 - Adds onSplitToDesktopAnimation to StageCoordinator to
set stages as not visible once the to desktop animation is complete,
preventing the split divider from briefly showing if we then navigate to
overview.

3- When requesting split, sets the task to WINDOWING_MODE_MULTI_WINDOW
to avoid a task flicker. Then, when the stage listener detects a change,
it will set the windowing mode to UNDEFINED to inherit the mode from
split stage, ensuring the task doesn't use freeform bounds in split
stage.

Bug: 336312668
Bug: 336311310
Bug: 323378048
Test: atest DragToDesktopTransitionHandlerTest
Test: manual
Change-Id: I55634fbe9687cd9a8f433818c8bafe983763356c
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
index 0a9e5d0..08b7c01 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
@@ -158,6 +158,10 @@
             com.android.wm.shell.R.dimen.desktop_mode_transition_area_width
         )
 
+    /** Task id of the task currently being dragged from fullscreen/split. */
+    val draggingTaskId
+        get() = dragToDesktopTransitionHandler.draggingTaskId
+
     private var recentsAnimationRunning = false
     private lateinit var splitScreenController: SplitScreenController
 
@@ -406,6 +410,7 @@
     fun onDesktopSplitSelectAnimComplete(taskInfo: RunningTaskInfo) {
         val wct = WindowContainerTransaction()
         wct.setBounds(taskInfo.token, Rect())
+        wct.setWindowingMode(taskInfo.token, WINDOWING_MODE_UNDEFINED)
         shellTaskOrganizer.applyTransaction(wct)
     }
 
@@ -447,7 +452,9 @@
         )
         val wct = WindowContainerTransaction()
         wct.setBounds(task.token, Rect())
-        addMoveToSplitChanges(wct, task)
+        // Rather than set windowing mode to multi-window at task level, set it to
+        // undefined and inherit from split stage.
+        wct.setWindowingMode(task.token, WINDOWING_MODE_UNDEFINED)
         if (Transitions.ENABLE_SHELL_TRANSITIONS) {
             transitions.startTransition(TRANSIT_CHANGE, wct, null /* handler */)
         } else {
@@ -458,10 +465,12 @@
     private fun exitSplitIfApplicable(wct: WindowContainerTransaction, taskInfo: RunningTaskInfo) {
         if (splitScreenController.isTaskInSplitScreen(taskInfo.taskId)) {
             splitScreenController.prepareExitSplitScreen(
-                    wct,
-                    splitScreenController.getStageOfTask(taskInfo.taskId),
-                    EXIT_REASON_DESKTOP_MODE
+                wct,
+                splitScreenController.getStageOfTask(taskInfo.taskId),
+                EXIT_REASON_DESKTOP_MODE
             )
+            splitScreenController.transitionHandler
+                ?.onSplitToDesktop()
         }
     }
 
@@ -1044,9 +1053,11 @@
         wct: WindowContainerTransaction,
         taskInfo: RunningTaskInfo
     ) {
-        // Explicitly setting multi-window at task level interferes with animations.
-        // Let task inherit windowing mode once transition is complete instead.
-        wct.setWindowingMode(taskInfo.token, WINDOWING_MODE_UNDEFINED)
+        // This windowing mode is to get the transition animation started; once we complete
+        // split select, we will change windowing mode to undefined and inherit from split stage.
+        // Going to undefined here causes task to flicker to the top left.
+        // Cancelling the split select flow will revert it to fullscreen.
+        wct.setWindowingMode(taskInfo.token, WINDOWING_MODE_MULTI_WINDOW)
         // The task's density may have been overridden in freeform; revert it here as we don't
         // want it overridden in multi-window.
         wct.setDensityDpi(taskInfo.token, getDefaultDensityDpi())
@@ -1237,7 +1248,7 @@
                 finalizeDragToDesktop(taskInfo, getDefaultDesktopTaskBounds(displayLayout))
             }
             DesktopModeVisualIndicator.IndicatorType.NO_INDICATOR,
-                    DesktopModeVisualIndicator.IndicatorType.TO_FULLSCREEN_INDICATOR -> {
+            DesktopModeVisualIndicator.IndicatorType.TO_FULLSCREEN_INDICATOR -> {
                 cancelDragToDesktop(taskInfo)
             }
             DesktopModeVisualIndicator.IndicatorType.TO_SPLIT_LEFT_INDICATOR -> {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandler.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandler.kt
index e341f2d..e5e435d 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandler.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandler.kt
@@ -6,6 +6,7 @@
 import android.animation.ValueAnimator
 import android.app.ActivityOptions
 import android.app.ActivityOptions.SourceInfo
+import android.app.ActivityTaskManager.INVALID_TASK_ID
 import android.app.PendingIntent
 import android.app.PendingIntent.FLAG_ALLOW_UNSAFE_IMPLICIT_INTENT
 import android.app.PendingIntent.FLAG_MUTABLE
@@ -26,6 +27,9 @@
 import android.window.WindowContainerToken
 import android.window.WindowContainerTransaction
 import com.android.wm.shell.RootTaskDisplayAreaOrganizer
+import com.android.wm.shell.common.split.SplitScreenConstants.SPLIT_POSITION_BOTTOM_OR_RIGHT
+import com.android.wm.shell.common.split.SplitScreenConstants.SPLIT_POSITION_TOP_OR_LEFT
+import com.android.wm.shell.common.split.SplitScreenConstants.SPLIT_POSITION_UNDEFINED
 import com.android.wm.shell.protolog.ShellProtoLogGroup
 import com.android.wm.shell.shared.TransitionUtil
 import com.android.wm.shell.splitscreen.SplitScreenController
@@ -68,7 +72,7 @@
             .addCategory(Intent.CATEGORY_HOME)
 
     private var dragToDesktopStateListener: DragToDesktopStateListener? = null
-    private var splitScreenController: SplitScreenController? = null
+    private lateinit var splitScreenController: SplitScreenController
     private var transitionState: TransitionState? = null
     private lateinit var onTaskResizeAnimationListener: OnTaskResizeAnimationListener
 
@@ -76,6 +80,9 @@
     val inProgress: Boolean
         get() = transitionState != null
 
+    /** The task id of the task currently being dragged from fullscreen/split. */
+    val draggingTaskId: Int
+        get() = transitionState?.draggedTaskId ?: INVALID_TASK_ID
     /** Sets a listener to receive callback about events during the transition animation. */
     fun setDragToDesktopStateListener(listener: DragToDesktopStateListener) {
         dragToDesktopStateListener = listener
@@ -130,10 +137,14 @@
                 .startTransition(TRANSIT_DESKTOP_MODE_START_DRAG_TO_DESKTOP, wct, this)
 
         transitionState = if (isSplitTask(taskId)) {
+            val otherTask = getOtherSplitTask(taskId) ?: throw IllegalStateException(
+                "Expected split task to have a counterpart."
+            )
             TransitionState.FromSplit(
                     draggedTaskId = taskId,
                     dragAnimator = dragToDesktopAnimator,
-                    startTransitionToken = startTransitionToken
+                    startTransitionToken = startTransitionToken,
+                    otherSplitTask = otherTask
             )
         } else {
             TransitionState.FromFullscreen(
@@ -347,6 +358,12 @@
                 ?: error("Start transition expected to be waiting for merge but wasn't")
         if (isEndTransition) {
             info.changes.withIndex().forEach { (i, change) ->
+                // If we're exiting split, hide the remaining split task.
+                if (state is TransitionState.FromSplit &&
+                    change.taskInfo?.taskId == state.otherSplitTask) {
+                    t.hide(change.leash)
+                    startTransactionFinishT.hide(change.leash)
+                }
                 if (change.mode == TRANSIT_CLOSE) {
                     t.hide(change.leash)
                     startTransactionFinishT.hide(change.leash)
@@ -392,7 +409,6 @@
             onTaskResizeAnimationListener.onAnimationStart(state.draggedTaskId, t,
                 unscaledStartBounds)
             finishCallback.onTransitionFinished(null /* wct */)
-
             val tx: SurfaceControl.Transaction = transactionSupplier.get()
             ValueAnimator.ofObject(rectEvaluator, unscaledStartBounds, endBounds)
                     .setDuration(DRAG_TO_DESKTOP_FINISH_ANIM_DURATION_MS)
@@ -549,7 +565,18 @@
     }
 
     private fun isSplitTask(taskId: Int): Boolean {
-        return splitScreenController?.isTaskInSplitScreen(taskId) ?: false
+        return splitScreenController.isTaskInSplitScreen(taskId)
+    }
+
+    private fun getOtherSplitTask(taskId: Int): Int? {
+        val splitPos = splitScreenController.getSplitPosition(taskId)
+        if (splitPos == SPLIT_POSITION_UNDEFINED) return null
+        val otherTaskPos = if (splitPos == SPLIT_POSITION_BOTTOM_OR_RIGHT) {
+            SPLIT_POSITION_TOP_OR_LEFT
+        } else {
+            SPLIT_POSITION_BOTTOM_OR_RIGHT
+        }
+        return splitScreenController.getTaskInfo(otherTaskPos)?.taskId
     }
 
     private fun requireTransitionState(): TransitionState {
@@ -598,6 +625,7 @@
                 override var cancelled: Boolean = false,
                 override var startAborted: Boolean = false,
                 var splitRootChange: Change? = null,
+                var otherSplitTask: Int
         ) : TransitionState()
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
index fadc970..8a488dd 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
@@ -3353,6 +3353,11 @@
                 true /* reparentLeafTaskIfRelaunch */);
     }
 
+    /** Call this when the animation from split screen to desktop is started. */
+    public void onSplitToDesktop() {
+        setSplitsVisible(false);
+    }
+
     /** Call this when the recents animation finishes by doing pair-to-pair switch. */
     public void onRecentsPairToPairAnimationFinish(WindowContainerTransaction finishWct) {
         ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "onRecentsPairToPairAnimationFinish");
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
index 6a9d17f..7d49535 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
@@ -16,6 +16,7 @@
 
 package com.android.wm.shell.windowdecor;
 
+import static android.app.ActivityTaskManager.INVALID_TASK_ID;
 import static android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD;
 import static android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM;
 import static android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN;
@@ -34,6 +35,7 @@
 import static com.android.wm.shell.common.split.SplitScreenConstants.SPLIT_POSITION_TOP_OR_LEFT;
 import static com.android.wm.shell.compatui.AppCompatUtils.isSingleTopActivityTranslucent;
 import static com.android.wm.shell.desktopmode.DesktopModeVisualIndicator.IndicatorType.TO_FULLSCREEN_INDICATOR;
+import static com.android.wm.shell.splitscreen.SplitScreen.STAGE_TYPE_UNDEFINED;
 
 import android.annotation.NonNull;
 import android.app.ActivityManager;
@@ -272,10 +274,9 @@
         mSplitScreenController.registerSplitScreenListener(new SplitScreen.SplitScreenListener() {
             @Override
             public void onTaskStageChanged(int taskId, @StageType int stage, boolean visible) {
-                if (visible) {
+                if (visible && stage != STAGE_TYPE_UNDEFINED) {
                     DesktopModeWindowDecoration decor = mWindowDecorByTaskId.get(taskId);
-                    if (decor != null && DesktopModeStatus.isEnabled()
-                            && decor.mTaskInfo.getWindowingMode() == WINDOWING_MODE_FREEFORM) {
+                    if (decor != null && DesktopModeStatus.isEnabled()) {
                         mDesktopTasksController.moveToSplit(decor.mTaskInfo);
                     }
                 }
@@ -915,6 +916,11 @@
 
     @Nullable
     private DesktopModeWindowDecoration getRelevantWindowDecor(MotionEvent ev) {
+        // If we are mid-transition, dragged task's decor is always relevant.
+        final int draggedTaskId = mDesktopTasksController.getDraggingTaskId();
+        if (draggedTaskId != INVALID_TASK_ID) {
+            return mWindowDecorByTaskId.get(draggedTaskId);
+        }
         final DesktopModeWindowDecoration focusedDecor = getFocusedDecor();
         if (focusedDecor == null) {
             return null;