Merge changes from topic "splitscreen_singletop" into tm-dev

* changes:
  Consolidate split screen rotation in shell transition
  Add a single-top root task for split screen
diff --git a/core/java/android/app/TaskInfo.java b/core/java/android/app/TaskInfo.java
index 1a38fcf..1b87945 100644
--- a/core/java/android/app/TaskInfo.java
+++ b/core/java/android/app/TaskInfo.java
@@ -472,7 +472,8 @@
                 && isFocused == that.isFocused
                 && isVisible == that.isVisible
                 && isSleeping == that.isSleeping
-                && Objects.equals(mTopActivityLocusId, that.mTopActivityLocusId);
+                && Objects.equals(mTopActivityLocusId, that.mTopActivityLocusId)
+                && parentTaskId == that.parentTaskId;
     }
 
     /**
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/ShellCommandHandlerImpl.java b/libs/WindowManager/Shell/src/com/android/wm/shell/ShellCommandHandlerImpl.java
index 8483f07..06f4367 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/ShellCommandHandlerImpl.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/ShellCommandHandlerImpl.java
@@ -119,8 +119,6 @@
                 return runRemoveFromSideStage(args, pw);
             case "setSideStagePosition":
                 return runSetSideStagePosition(args, pw);
-            case "setSideStageVisibility":
-                return runSetSideStageVisibility(args, pw);
             case "help":
                 return runHelp(pw);
             default:
@@ -186,18 +184,6 @@
         return true;
     }
 
-    private boolean runSetSideStageVisibility(String[] args, PrintWriter pw) {
-        if (args.length < 3) {
-            // First arguments are "WMShell" and command name.
-            pw.println("Error: side stage visibility should be provided as arguments");
-            return false;
-        }
-        final Boolean visible = new Boolean(args[2]);
-
-        mSplitScreenOptional.ifPresent(split -> split.setSideStageVisibility(visible));
-        return true;
-    }
-
     private boolean runHelp(PrintWriter pw) {
         pw.println("Window Manager Shell commands:");
         pw.println("  help");
@@ -215,8 +201,6 @@
         pw.println("    Enable/Disable outline on the side-stage.");
         pw.println("  setSideStagePosition <SideStagePosition>");
         pw.println("    Sets the position of the side-stage.");
-        pw.println("  setSideStageVisibility <true/false>");
-        pw.println("    Show/hide side-stage.");
         return true;
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/apppairs/AppPair.java b/libs/WindowManager/Shell/src/com/android/wm/shell/apppairs/AppPair.java
index 33eec33..e528df8 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/apppairs/AppPair.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/apppairs/AppPair.java
@@ -19,7 +19,6 @@
 import static android.app.ActivityTaskManager.INVALID_TASK_ID;
 import static android.app.WindowConfiguration.WINDOWING_MODE_MULTI_WINDOW;
 import static android.app.WindowConfiguration.WINDOWING_MODE_UNDEFINED;
-import static android.view.WindowManagerPolicyConstants.SPLIT_DIVIDER_LAYER;
 
 import static com.android.wm.shell.common.split.SplitScreenConstants.SPLIT_POSITION_BOTTOM_OR_RIGHT;
 import static com.android.wm.shell.common.split.SplitScreenConstants.SPLIT_POSITION_TOP_OR_LEFT;
@@ -83,7 +82,7 @@
         public void onLeashReady(SurfaceControl leash) {
             mSyncQueue.runInSync(t -> t
                     .show(leash)
-                    .setLayer(leash, SPLIT_DIVIDER_LAYER)
+                    .setLayer(leash, Integer.MAX_VALUE)
                     .setPosition(leash,
                             mSplitLayout.getDividerBounds().left,
                             mSplitLayout.getDividerBounds().top));
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitDecorManager.java b/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitDecorManager.java
index 3ac7cbf..5dc6bd1 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitDecorManager.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitDecorManager.java
@@ -21,7 +21,6 @@
 import static android.view.WindowManager.LayoutParams.PRIVATE_FLAG_NO_MOVE_ANIMATION;
 import static android.view.WindowManager.LayoutParams.PRIVATE_FLAG_TRUSTED_OVERLAY;
 import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_OVERLAY;
-import static android.view.WindowManagerPolicyConstants.SPLIT_DIVIDER_LAYER;
 
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
@@ -162,7 +161,7 @@
             mBackgroundLeash = SurfaceUtils.makeColorLayer(mHostLeash,
                     RESIZING_BACKGROUND_SURFACE_NAME, mSurfaceSession);
             t.setColor(mBackgroundLeash, getResizingBackgroundColor(resizingTask))
-                    .setLayer(mBackgroundLeash, SPLIT_DIVIDER_LAYER - 1);
+                    .setLayer(mBackgroundLeash, Integer.MAX_VALUE - 1);
         }
 
         if (mIcon == null && resizingTask.topActivityInfo != null) {
@@ -175,7 +174,7 @@
             lp.width = mIcon.getIntrinsicWidth();
             lp.height = mIcon.getIntrinsicHeight();
             mViewHost.relayout(lp);
-            t.setLayer(mIconLeash, SPLIT_DIVIDER_LAYER);
+            t.setLayer(mIconLeash, Integer.MAX_VALUE);
         }
         t.setPosition(mIconLeash,
                 newBounds.width() / 2 - mIcon.getIntrinsicWidth() / 2,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitLayout.java b/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitLayout.java
index d05a4df..116d352 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitLayout.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitLayout.java
@@ -23,7 +23,6 @@
 import static android.view.WindowManager.DOCKED_LEFT;
 import static android.view.WindowManager.DOCKED_RIGHT;
 import static android.view.WindowManager.DOCKED_TOP;
-import static android.view.WindowManagerPolicyConstants.SPLIT_DIVIDER_LAYER;
 
 import static com.android.internal.policy.DividerSnapAlgorithm.SnapTarget.FLAG_DISMISS_END;
 import static com.android.internal.policy.DividerSnapAlgorithm.SnapTarget.FLAG_DISMISS_START;
@@ -202,24 +201,24 @@
 
     /** Applies new configuration, returns {@code false} if there's no effect to the layout. */
     public boolean updateConfiguration(Configuration configuration) {
+        // Always update configuration after orientation changed to make sure to render divider bar
+        // with proper resources that matching screen orientation.
+        final int orientation = configuration.orientation;
+        if (mOrientation != orientation) {
+            mContext = mContext.createConfigurationContext(configuration);
+            mSplitWindowManager.setConfiguration(configuration);
+            mOrientation = orientation;
+        }
+
         // Update the split bounds when necessary. Besides root bounds changed, split bounds need to
         // be updated when the rotation changed to cover the case that users rotated the screen 180
         // degrees.
-        // Make sure to render the divider bar with proper resources that matching the screen
-        // orientation.
         final int rotation = configuration.windowConfiguration.getRotation();
         final Rect rootBounds = configuration.windowConfiguration.getBounds();
-        final int orientation = configuration.orientation;
-
-        if (mOrientation == orientation
-                && rotation == mRotation
-                && mRootBounds.equals(rootBounds)) {
+        if (mRotation == rotation && mRootBounds.equals(rootBounds)) {
             return false;
         }
 
-        mContext = mContext.createConfigurationContext(configuration);
-        mSplitWindowManager.setConfiguration(configuration);
-        mOrientation = orientation;
         mTempRect.set(mRootBounds);
         mRootBounds.set(rootBounds);
         mRotation = rotation;
@@ -493,7 +492,7 @@
             mTempRect.set(getRefDividerBounds());
             t.setPosition(dividerLeash, mTempRect.left, mTempRect.top);
             // Resets layer of divider bar to make sure it is always on top.
-            t.setLayer(dividerLeash, SPLIT_DIVIDER_LAYER);
+            t.setLayer(dividerLeash, Integer.MAX_VALUE);
         }
         mTempRect.set(getRefBounds1());
         t.setPosition(leash1, mTempRect.left, mTempRect.top)
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/ISplitScreen.aidl b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/ISplitScreen.aidl
index 13137fa..9adf196 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/ISplitScreen.aidl
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/ISplitScreen.aidl
@@ -42,11 +42,6 @@
     oneway void unregisterSplitScreenListener(in ISplitScreenListener listener) = 2;
 
     /**
-     * Hides the side-stage if it is currently visible.
-     */
-    oneway void setSideStageVisibility(boolean visible) = 3;
-
-    /**
      * Removes a task from the side stage.
      */
     oneway void removeFromSideStage(int taskId) = 4;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/MainStage.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/MainStage.java
index 22dd9b9..ae5e075 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/MainStage.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/MainStage.java
@@ -52,9 +52,6 @@
         if (mIsActive) return;
 
         final WindowContainerToken rootToken = mRootTaskInfo.token;
-        // Moving the root task to top after the child tasks were re-parented , or the root
-        // task cannot be visible and focused.
-        wct.reorder(rootToken, true /* onTop */);
         if (includingTopTask) {
             wct.reparentTasks(
                     null /* currentParent */,
@@ -83,9 +80,6 @@
                         null /* newParent */,
                         CONTROLLED_WINDOWING_MODES_WHEN_ACTIVE,
                         CONTROLLED_ACTIVITY_TYPES,
-                        toTop)
-                // We want this re-order to the bottom regardless since we are re-parenting
-                // all its tasks.
-                .reorder(rootToken, false /* onTop */);
+                        toTop);
     }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SideStage.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SideStage.java
index 122fc9f..d55619f 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SideStage.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SideStage.java
@@ -45,9 +45,6 @@
     }
 
     boolean removeAllTasks(WindowContainerTransaction wct, boolean toTop) {
-        // No matter if the root task is empty or not, moving the root to bottom because it no
-        // longer preserves visible child task.
-        wct.reorder(mRootTaskInfo.token, false /* onTop */);
         if (mChildrenTaskInfo.size() == 0) return false;
         wct.reparentTasks(
                 mRootTaskInfo.token,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitDecorManager.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitDecorManager.java
deleted file mode 100644
index 8e5cc6d..0000000
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitDecorManager.java
+++ /dev/null
@@ -1,182 +0,0 @@
-/*
- * Copyright (C) 2021 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.wm.shell.splitscreen;
-
-import static android.view.WindowManager.LayoutParams.FLAG_NOT_FOCUSABLE;
-import static android.view.WindowManager.LayoutParams.FLAG_NOT_TOUCHABLE;
-import static android.view.WindowManager.LayoutParams.PRIVATE_FLAG_NO_MOVE_ANIMATION;
-import static android.view.WindowManager.LayoutParams.PRIVATE_FLAG_TRUSTED_OVERLAY;
-import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_OVERLAY;
-import static android.view.WindowManagerPolicyConstants.SPLIT_DIVIDER_LAYER;
-
-import android.app.ActivityManager;
-import android.content.Context;
-import android.content.res.Configuration;
-import android.graphics.Color;
-import android.graphics.PixelFormat;
-import android.graphics.Rect;
-import android.graphics.drawable.Drawable;
-import android.os.Binder;
-import android.view.IWindow;
-import android.view.LayoutInflater;
-import android.view.SurfaceControl;
-import android.view.SurfaceControlViewHost;
-import android.view.SurfaceSession;
-import android.view.View;
-import android.view.WindowManager;
-import android.view.WindowlessWindowManager;
-import android.widget.FrameLayout;
-import android.widget.ImageView;
-
-import com.android.launcher3.icons.IconProvider;
-import com.android.wm.shell.R;
-import com.android.wm.shell.common.SurfaceUtils;
-
-/**
- * Handles split decor like showing resizing hint for a specific split.
- */
-class SplitDecorManager extends WindowlessWindowManager {
-    private static final String TAG = SplitDecorManager.class.getSimpleName();
-    private static final String RESIZING_BACKGROUND_SURFACE_NAME = "ResizingBackground";
-
-    private final IconProvider mIconProvider;
-    private final SurfaceSession mSurfaceSession;
-
-    private Drawable mIcon;
-    private ImageView mResizingIconView;
-    private SurfaceControlViewHost mViewHost;
-    private SurfaceControl mHostLeash;
-    private SurfaceControl mIconLeash;
-    private SurfaceControl mBackgroundLeash;
-
-    SplitDecorManager(Configuration configuration, IconProvider iconProvider,
-            SurfaceSession surfaceSession) {
-        super(configuration, null /* rootSurface */, null /* hostInputToken */);
-        mIconProvider = iconProvider;
-        mSurfaceSession = surfaceSession;
-    }
-
-    @Override
-    protected void attachToParentSurface(IWindow window, SurfaceControl.Builder b) {
-        // Can't set position for the ViewRootImpl SC directly. Create a leash to manipulate later.
-        final SurfaceControl.Builder builder = new SurfaceControl.Builder(new SurfaceSession())
-                .setContainerLayer()
-                .setName(TAG)
-                .setHidden(true)
-                .setParent(mHostLeash)
-                .setCallsite("SplitDecorManager#attachToParentSurface");
-        mIconLeash = builder.build();
-        b.setParent(mIconLeash);
-    }
-
-    void inflate(Context context, SurfaceControl rootLeash, Rect rootBounds) {
-        if (mIconLeash != null && mViewHost != null) {
-            return;
-        }
-
-        context = context.createWindowContext(context.getDisplay(), TYPE_APPLICATION_OVERLAY,
-                null /* options */);
-        mHostLeash = rootLeash;
-        mViewHost = new SurfaceControlViewHost(context, context.getDisplay(), this);
-
-        final FrameLayout rootLayout = (FrameLayout) LayoutInflater.from(context)
-                .inflate(R.layout.split_decor, null);
-        mResizingIconView = rootLayout.findViewById(R.id.split_resizing_icon);
-
-        final WindowManager.LayoutParams lp = new WindowManager.LayoutParams(
-                0 /* width */, 0 /* height */, TYPE_APPLICATION_OVERLAY,
-                FLAG_NOT_FOCUSABLE | FLAG_NOT_TOUCHABLE, PixelFormat.TRANSLUCENT);
-        lp.width = rootBounds.width();
-        lp.height = rootBounds.height();
-        lp.token = new Binder();
-        lp.setTitle(TAG);
-        lp.privateFlags |= PRIVATE_FLAG_NO_MOVE_ANIMATION | PRIVATE_FLAG_TRUSTED_OVERLAY;
-        // TODO(b/189839391): Set INPUT_FEATURE_NO_INPUT_CHANNEL after WM supports
-        //  TRUSTED_OVERLAY for windowless window without input channel.
-        mViewHost.setView(rootLayout, lp);
-    }
-
-    void release(SurfaceControl.Transaction t) {
-        if (mViewHost != null) {
-            mViewHost.release();
-            mViewHost = null;
-        }
-        if (mIconLeash != null) {
-            t.remove(mIconLeash);
-            mIconLeash = null;
-        }
-        if (mBackgroundLeash != null) {
-            t.remove(mBackgroundLeash);
-            mBackgroundLeash = null;
-        }
-        mHostLeash = null;
-        mIcon = null;
-        mResizingIconView = null;
-    }
-
-    /** Showing resizing hint. */
-    void onResizing(ActivityManager.RunningTaskInfo resizingTask, Rect newBounds,
-            SurfaceControl.Transaction t) {
-        if (mResizingIconView == null) {
-            return;
-        }
-
-        if (mIcon == null) {
-            // TODO: add fade-in animation.
-            mBackgroundLeash = SurfaceUtils.makeColorLayer(mHostLeash,
-                    RESIZING_BACKGROUND_SURFACE_NAME, mSurfaceSession);
-            t.setColor(mBackgroundLeash, getResizingBackgroundColor(resizingTask))
-                    .setLayer(mBackgroundLeash, SPLIT_DIVIDER_LAYER - 1)
-                    .show(mBackgroundLeash);
-
-            mIcon = mIconProvider.getIcon(resizingTask.topActivityInfo);
-            mResizingIconView.setImageDrawable(mIcon);
-            mResizingIconView.setVisibility(View.VISIBLE);
-
-            WindowManager.LayoutParams lp =
-                    (WindowManager.LayoutParams) mViewHost.getView().getLayoutParams();
-            lp.width = mIcon.getIntrinsicWidth();
-            lp.height = mIcon.getIntrinsicHeight();
-            mViewHost.relayout(lp);
-            t.show(mIconLeash).setLayer(mIconLeash, SPLIT_DIVIDER_LAYER);
-        }
-
-        t.setPosition(mIconLeash,
-                newBounds.width() / 2 - mIcon.getIntrinsicWidth() / 2,
-                newBounds.height() / 2 - mIcon.getIntrinsicWidth() / 2);
-    }
-
-    /** Stops showing resizing hint. */
-    void onResized(Rect newBounds, SurfaceControl.Transaction t) {
-        if (mResizingIconView == null) {
-            return;
-        }
-
-        if (mIcon != null) {
-            mResizingIconView.setVisibility(View.GONE);
-            mResizingIconView.setImageDrawable(null);
-            t.remove(mBackgroundLeash).hide(mIconLeash);
-            mIcon = null;
-            mBackgroundLeash = null;
-        }
-    }
-
-    private static float[] getResizingBackgroundColor(ActivityManager.RunningTaskInfo taskInfo) {
-        final int taskBgColor = taskInfo.taskDescription.getBackgroundColor();
-        return Color.valueOf(taskBgColor == -1 ? Color.WHITE : taskBgColor).getComponents();
-    }
-}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
index 2da5bec..f20870f 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
@@ -188,7 +188,7 @@
         if (mStageCoordinator == null) {
             // TODO: Multi-display
             mStageCoordinator = new StageCoordinator(mContext, DEFAULT_DISPLAY, mSyncQueue,
-                    mRootTDAOrganizer, mTaskOrganizer, mDisplayController, mDisplayImeController,
+                    mTaskOrganizer, mDisplayController, mDisplayImeController,
                     mDisplayInsetsController, mTransitions, mTransactionPool, mLogger,
                     mIconProvider, mRecentTasksOptional, mUnfoldControllerProvider);
         }
@@ -239,10 +239,6 @@
         mStageCoordinator.setSideStagePosition(sideStagePosition, null /* wct */);
     }
 
-    public void setSideStageVisibility(boolean visible) {
-        mStageCoordinator.setSideStageVisibility(visible);
-    }
-
     public void enterSplitScreen(int taskId, boolean leftOrTop) {
         enterSplitScreen(taskId, leftOrTop, new WindowContainerTransaction());
     }
@@ -645,14 +641,6 @@
         }
 
         @Override
-        public void setSideStageVisibility(boolean visible) {
-            executeRemoteCallWithTaskPermission(mController, "setSideStageVisibility",
-                    (controller) -> {
-                        controller.setSideStageVisibility(visible);
-                    });
-        }
-
-        @Override
         public void removeFromSideStage(int taskId) {
             executeRemoteCallWithTaskPermission(mController, "removeFromSideStage",
                     (controller) -> {
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 880693d..b10d50d 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
@@ -28,7 +28,6 @@
 import static android.view.WindowManager.TRANSIT_TO_BACK;
 import static android.view.WindowManager.TRANSIT_TO_FRONT;
 import static android.view.WindowManager.transitTypeToString;
-import static android.view.WindowManagerPolicyConstants.SPLIT_DIVIDER_LAYER;
 import static android.window.TransitionInfo.FLAG_IS_DISPLAY;
 
 import static com.android.wm.shell.common.split.SplitScreenConstants.SPLIT_POSITION_BOTTOM_OR_RIGHT;
@@ -43,6 +42,7 @@
 import static com.android.wm.shell.splitscreen.SplitScreenController.EXIT_REASON_DEVICE_FOLDED;
 import static com.android.wm.shell.splitscreen.SplitScreenController.EXIT_REASON_DRAG_DIVIDER;
 import static com.android.wm.shell.splitscreen.SplitScreenController.EXIT_REASON_RETURN_HOME;
+import static com.android.wm.shell.splitscreen.SplitScreenController.EXIT_REASON_ROOT_TASK_VANISHED;
 import static com.android.wm.shell.splitscreen.SplitScreenController.EXIT_REASON_SCREEN_LOCKED_SHOW_ON_TOP;
 import static com.android.wm.shell.splitscreen.SplitScreenController.EXIT_REASON_UNKNOWN;
 import static com.android.wm.shell.splitscreen.SplitScreenController.exitReasonToString;
@@ -53,6 +53,7 @@
 import static com.android.wm.shell.transition.Transitions.isClosingType;
 import static com.android.wm.shell.transition.Transitions.isOpeningType;
 
+import android.annotation.CallSuper;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.app.ActivityManager;
@@ -77,7 +78,6 @@
 import android.view.SurfaceControl;
 import android.view.SurfaceSession;
 import android.view.WindowManager;
-import android.window.DisplayAreaInfo;
 import android.window.RemoteTransition;
 import android.window.TransitionInfo;
 import android.window.TransitionRequestInfo;
@@ -88,7 +88,6 @@
 import com.android.internal.logging.InstanceId;
 import com.android.internal.protolog.common.ProtoLog;
 import com.android.launcher3.icons.IconProvider;
-import com.android.wm.shell.RootTaskDisplayAreaOrganizer;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.common.DisplayController;
 import com.android.wm.shell.common.DisplayImeController;
@@ -122,13 +121,13 @@
  * - The {@link MainStage} should only have children if the coordinator is active.
  * - The {@link SplitLayout} divider is only visible if both the {@link MainStage}
  * and {@link SideStage} are visible.
- * - The {@link MainStage} configuration is fullscreen when the {@link SideStage} isn't visible.
+ * - Both stages are put under a single-top root task.
  * This rules are mostly implemented in {@link #onStageVisibilityChanged(StageListenerImpl)} and
  * {@link #onStageHasChildrenChanged(StageListenerImpl).}
  */
 class StageCoordinator implements SplitLayout.SplitLayoutHandler,
-        RootTaskDisplayAreaOrganizer.RootTaskDisplayAreaListener,
-        DisplayController.OnDisplaysChangedListener, Transitions.TransitionHandler {
+        DisplayController.OnDisplaysChangedListener, Transitions.TransitionHandler,
+        ShellTaskOrganizer.TaskListener {
 
     private static final String TAG = StageCoordinator.class.getSimpleName();
 
@@ -148,9 +147,7 @@
     private SplitLayout mSplitLayout;
     private boolean mDividerVisible;
     private final SyncTransactionQueue mSyncQueue;
-    private final RootTaskDisplayAreaOrganizer mRootTDAOrganizer;
     private final ShellTaskOrganizer mTaskOrganizer;
-    private DisplayAreaInfo mDisplayAreaInfo;
     private final Context mContext;
     private final List<SplitScreen.SplitScreenListener> mListeners = new ArrayList<>();
     private final DisplayController mDisplayController;
@@ -160,6 +157,15 @@
     private final SplitScreenTransitions mSplitTransitions;
     private final SplitscreenEventLogger mLogger;
     private final Optional<RecentTasksController> mRecentTasks;
+
+    /**
+     * A single-top root task which the split divider attached to.
+     */
+    @VisibleForTesting
+    ActivityManager.RunningTaskInfo mRootTaskInfo;
+
+    private SurfaceControl mRootTaskLeash;
+
     // Tracks whether we should update the recent tasks.  Only allow this to happen in between enter
     // and exit, since exit itself can trigger a number of changes that update the stages.
     private boolean mShouldUpdateRecents;
@@ -174,7 +180,7 @@
             new SplitWindowManager.ParentContainerCallbacks() {
                 @Override
                 public void attachToParentSurface(SurfaceControl.Builder b) {
-                    mRootTDAOrganizer.attachToDisplayArea(mDisplayId, b);
+                    b.setParent(mRootTaskLeash);
                 }
 
                 @Override
@@ -184,23 +190,21 @@
             };
 
     StageCoordinator(Context context, int displayId, SyncTransactionQueue syncQueue,
-            RootTaskDisplayAreaOrganizer rootTDAOrganizer, ShellTaskOrganizer taskOrganizer,
-            DisplayController displayController,
+            ShellTaskOrganizer taskOrganizer, DisplayController displayController,
             DisplayImeController displayImeController,
             DisplayInsetsController displayInsetsController, Transitions transitions,
             TransactionPool transactionPool, SplitscreenEventLogger logger,
-            IconProvider iconProvider,
-            Optional<RecentTasksController> recentTasks,
+            IconProvider iconProvider, Optional<RecentTasksController> recentTasks,
             Provider<Optional<StageTaskUnfoldController>> unfoldControllerProvider) {
         mContext = context;
         mDisplayId = displayId;
         mSyncQueue = syncQueue;
-        mRootTDAOrganizer = rootTDAOrganizer;
         mTaskOrganizer = taskOrganizer;
         mLogger = logger;
         mRecentTasks = recentTasks;
         mMainUnfoldController = unfoldControllerProvider.get().orElse(null);
         mSideUnfoldController = unfoldControllerProvider.get().orElse(null);
+        taskOrganizer.createRootTask(displayId, WINDOWING_MODE_FULLSCREEN, this /* listener */);
 
         mMainStage = new MainStage(
                 mContext,
@@ -224,7 +228,6 @@
         mDisplayImeController = displayImeController;
         mDisplayInsetsController = displayInsetsController;
         mTransactionPool = transactionPool;
-        mRootTDAOrganizer.registerListener(displayId, this);
         final DeviceStateManager deviceStateManager =
                 mContext.getSystemService(DeviceStateManager.class);
         deviceStateManager.registerCallback(taskOrganizer.getExecutor(),
@@ -238,9 +241,8 @@
 
     @VisibleForTesting
     StageCoordinator(Context context, int displayId, SyncTransactionQueue syncQueue,
-            RootTaskDisplayAreaOrganizer rootTDAOrganizer, ShellTaskOrganizer taskOrganizer,
-            MainStage mainStage, SideStage sideStage, DisplayController displayController,
-            DisplayImeController displayImeController,
+            ShellTaskOrganizer taskOrganizer, MainStage mainStage, SideStage sideStage,
+            DisplayController displayController, DisplayImeController displayImeController,
             DisplayInsetsController displayInsetsController, SplitLayout splitLayout,
             Transitions transitions, TransactionPool transactionPool,
             SplitscreenEventLogger logger,
@@ -249,7 +251,6 @@
         mContext = context;
         mDisplayId = displayId;
         mSyncQueue = syncQueue;
-        mRootTDAOrganizer = rootTDAOrganizer;
         mTaskOrganizer = taskOrganizer;
         mMainStage = mainStage;
         mSideStage = sideStage;
@@ -257,7 +258,6 @@
         mDisplayImeController = displayImeController;
         mDisplayInsetsController = displayInsetsController;
         mTransactionPool = transactionPool;
-        mRootTDAOrganizer.registerListener(displayId, this);
         mSplitLayout = splitLayout;
         mSplitTransitions = new SplitScreenTransitions(transactionPool, transitions,
                 this::onTransitionAnimationComplete, this);
@@ -357,6 +357,7 @@
         // while task 1 is on the side stage.
         mMainStage.activate(wct, false /* reparent */);
         updateWindowBounds(mSplitLayout, wct);
+        wct.reorder(mRootTaskInfo.token, true);
 
         // Make sure the launch options will put tasks in the corresponding split roots
         addActivityOptions(mainOptions, mMainStage);
@@ -470,15 +471,13 @@
         setSideStagePosition(sidePosition, wct);
 
         mSplitLayout.setDivideRatio(splitRatio);
-        if (mMainStage.isActive()) {
-            mMainStage.moveToTop(wct);
-        } else {
+        if (!mMainStage.isActive()) {
             // Build a request WCT that will launch both apps such that task 0 is on the main stage
             // while task 1 is on the side stage.
             mMainStage.activate(wct, false /* reparent */);
         }
-        mSideStage.moveToTop(wct);
         updateWindowBounds(mSplitLayout, wct);
+        wct.reorder(mRootTaskInfo.token, true);
 
         // Make sure the launch options will put tasks in the corresponding split roots
         addActivityOptions(mainOptions, mMainStage);
@@ -603,14 +602,6 @@
         }
     }
 
-    void setSideStageVisibility(boolean visible) {
-        if (mSideStageListener.mVisible == visible) return;
-
-        final WindowContainerTransaction wct = new WindowContainerTransaction();
-        mSideStage.setVisibility(visible, wct);
-        mTaskOrganizer.applyTransaction(wct);
-    }
-
     void onKeyguardVisibilityChanged(boolean showing) {
         if (!mMainStage.isActive()) {
             return;
@@ -682,15 +673,18 @@
         applyExitSplitScreen(childrenToTop, wct, exitReason);
     }
 
-    private void exitSplitScreen(StageTaskListener childrenToTop, @ExitReason int exitReason) {
+    private void exitSplitScreen(@Nullable StageTaskListener childrenToTop,
+            @ExitReason int exitReason) {
         if (!mMainStage.isActive()) return;
 
         final WindowContainerTransaction wct = new WindowContainerTransaction();
         applyExitSplitScreen(childrenToTop, wct, exitReason);
     }
 
-    private void applyExitSplitScreen(StageTaskListener childrenToTop,
+    private void applyExitSplitScreen(@Nullable StageTaskListener childrenToTop,
             WindowContainerTransaction wct, @ExitReason int exitReason) {
+        if (!mMainStage.isActive()) return;
+
         mRecentTasks.ifPresent(recentTasks -> {
             // Notify recents if we are exiting in a way that breaks the pair, and disable further
             // updates to splits in the recents until we enter split again
@@ -705,8 +699,9 @@
         // we want the tasks to be put to bottom instead of top, otherwise it will end up
         // a fullscreen plus a pinned task instead of pinned only at the end of the transition.
         final boolean fromEnteringPip = exitReason == EXIT_REASON_CHILD_TASK_ENTER_PIP;
-        mSideStage.removeAllTasks(wct, !fromEnteringPip && childrenToTop == mSideStage);
-        mMainStage.deactivate(wct, !fromEnteringPip && childrenToTop == mMainStage);
+        mSideStage.removeAllTasks(wct, !fromEnteringPip && mSideStage == childrenToTop);
+        mMainStage.deactivate(wct, !fromEnteringPip && mMainStage == childrenToTop);
+        wct.reorder(mRootTaskInfo.token, false /* onTop */);
         mTaskOrganizer.applyTransaction(wct);
         mSyncQueue.runInSync(t -> t
                 .setWindowCrop(mMainStage.mRootLeash, null)
@@ -777,8 +772,8 @@
             mSideStage.addTask(taskInfo, wct);
         }
         mMainStage.activate(wct, true /* includingTopTask */);
-        mSideStage.moveToTop(wct);
         updateWindowBounds(mSplitLayout, wct);
+        wct.reorder(mRootTaskInfo.token, true);
     }
 
     void finishEnterSplitScreen(SurfaceControl.Transaction t) {
@@ -916,25 +911,97 @@
         }
     }
 
-    private void onStageRootTaskAppeared(StageListenerImpl stageListener) {
-        if (mMainStageListener.mHasRootTask && mSideStageListener.mHasRootTask) {
-            final WindowContainerTransaction wct = new WindowContainerTransaction();
-            // Make the stages adjacent to each other so they occlude what's behind them.
-            wct.setAdjacentRoots(mMainStage.mRootTaskInfo.token, mSideStage.mRootTaskInfo.token,
-                    true /* moveTogether */);
-            wct.setLaunchAdjacentFlagRoot(mSideStage.mRootTaskInfo.token);
-            mTaskOrganizer.applyTransaction(wct);
+    @Override
+    @CallSuper
+    public void onTaskAppeared(ActivityManager.RunningTaskInfo taskInfo, SurfaceControl leash) {
+        if (mRootTaskInfo != null || taskInfo.hasParentTask()) {
+            throw new IllegalArgumentException(this + "\n Unknown task appeared: " + taskInfo);
+        }
+
+        mRootTaskInfo = taskInfo;
+        mRootTaskLeash = leash;
+
+        if (mSplitLayout == null) {
+            mSplitLayout = new SplitLayout(TAG + "SplitDivider", mContext,
+                    mRootTaskInfo.configuration, this, mParentContainerCallbacks,
+                    mDisplayImeController, mTaskOrganizer, false /* applyDismissingParallax */);
+            mDisplayInsetsController.addInsetsChangedListener(mDisplayId, mSplitLayout);
+        }
+
+        if (mMainUnfoldController != null && mSideUnfoldController != null) {
+            mMainUnfoldController.init();
+            mSideUnfoldController.init();
+        }
+
+        onRootTaskAppeared();
+    }
+
+    @Override
+    @CallSuper
+    public void onTaskInfoChanged(ActivityManager.RunningTaskInfo taskInfo) {
+        if (mRootTaskInfo == null || mRootTaskInfo.taskId != taskInfo.taskId) {
+            throw new IllegalArgumentException(this + "\n Unknown task info changed: " + taskInfo);
+        }
+
+        mRootTaskInfo = taskInfo;
+        if (mSplitLayout != null
+                && mSplitLayout.updateConfiguration(mRootTaskInfo.configuration)
+                && mMainStage.isActive()) {
+            // TODO(b/204925795): With Shell transition, We are handling split bounds rotation at
+            //  onRotateDisplay. But still need to handle unfold case.
+            if (ENABLE_SHELL_TRANSITIONS) {
+                updateUnfoldBounds();
+                return;
+            }
+            mSplitLayout.update(null /* t */);
+            onLayoutSizeChanged(mSplitLayout);
         }
     }
 
-    private void onStageRootTaskVanished(StageListenerImpl stageListener) {
-        if (stageListener == mMainStageListener || stageListener == mSideStageListener) {
-            final WindowContainerTransaction wct = new WindowContainerTransaction();
-            wct.clearLaunchAdjacentFlagRoot(mSideStage.mRootTaskInfo.token);
-            // Deactivate the main stage if it no longer has a root task.
-            mMainStage.deactivate(wct);
-            mTaskOrganizer.applyTransaction(wct);
+    @Override
+    @CallSuper
+    public void onTaskVanished(ActivityManager.RunningTaskInfo taskInfo) {
+        if (mRootTaskInfo == null) {
+            throw new IllegalArgumentException(this + "\n Unknown task vanished: " + taskInfo);
         }
+
+        onRootTaskVanished();
+
+        if (mSplitLayout != null) {
+            mSplitLayout.release();
+            mSplitLayout = null;
+        }
+
+        mRootTaskInfo = null;
+    }
+
+
+    @VisibleForTesting
+    void onRootTaskAppeared() {
+        // Wait unit all root tasks appeared.
+        if (mRootTaskInfo == null
+                || !mMainStageListener.mHasRootTask
+                || !mSideStageListener.mHasRootTask) {
+            return;
+        }
+
+        final WindowContainerTransaction wct = new WindowContainerTransaction();
+        wct.reparent(mMainStage.mRootTaskInfo.token, mRootTaskInfo.token, true);
+        wct.reparent(mSideStage.mRootTaskInfo.token, mRootTaskInfo.token, true);
+        // Make the stages adjacent to each other so they occlude what's behind them.
+        wct.setAdjacentRoots(mMainStage.mRootTaskInfo.token, mSideStage.mRootTaskInfo.token,
+                true /* moveTogether */);
+        wct.setLaunchAdjacentFlagRoot(mSideStage.mRootTaskInfo.token);
+        mTaskOrganizer.applyTransaction(wct);
+    }
+
+    private void onRootTaskVanished() {
+        final WindowContainerTransaction wct = new WindowContainerTransaction();
+        if (mRootTaskInfo != null) {
+            wct.clearLaunchAdjacentFlagRoot(mRootTaskInfo.token);
+        }
+        applyExitSplitScreen(null /* childrenToTop */, wct, EXIT_REASON_ROOT_TASK_VANISHED);
+        mDisplayInsetsController.removeInsetsChangedListener(mDisplayId, mSplitLayout);
     }
 
     private void onStageVisibilityChanged(StageListenerImpl stageListener) {
@@ -981,7 +1048,7 @@
         if (mDividerVisible) {
             t.show(dividerLeash);
             t.setAlpha(dividerLeash, 1);
-            t.setLayer(dividerLeash, SPLIT_DIVIDER_LAYER);
+            t.setLayer(dividerLeash, Integer.MAX_VALUE);
             t.setPosition(dividerLeash,
                     mSplitLayout.getRefDividerBounds().left,
                     mSplitLayout.getRefDividerBounds().top);
@@ -1140,45 +1207,6 @@
     }
 
     @Override
-    public void onDisplayAreaAppeared(DisplayAreaInfo displayAreaInfo) {
-        mDisplayAreaInfo = displayAreaInfo;
-        if (mSplitLayout == null) {
-            mSplitLayout = new SplitLayout(TAG + "SplitDivider", mContext,
-                    mDisplayAreaInfo.configuration, this, mParentContainerCallbacks,
-                    mDisplayImeController, mTaskOrganizer, false /* applyDismissingParallax */);
-            mDisplayInsetsController.addInsetsChangedListener(mDisplayId, mSplitLayout);
-
-            if (mMainUnfoldController != null && mSideUnfoldController != null) {
-                mMainUnfoldController.init();
-                mSideUnfoldController.init();
-            }
-        }
-    }
-
-    @Override
-    public void onDisplayAreaVanished(DisplayAreaInfo displayAreaInfo) {
-        throw new IllegalStateException("Well that was unexpected...");
-    }
-
-    @Override
-    public void onDisplayAreaInfoChanged(DisplayAreaInfo displayAreaInfo) {
-        mDisplayAreaInfo = displayAreaInfo;
-        if (mSplitLayout != null
-                && mSplitLayout.updateConfiguration(mDisplayAreaInfo.configuration)
-                && mMainStage.isActive()) {
-            // TODO(b/204925795): With Shell transition, We are handle roation case for apply split
-            //  bounds at onRotateDisplay. But still need to handle unfold case.
-            if (ENABLE_SHELL_TRANSITIONS) {
-                updateUnfoldBounds();
-                return;
-            }
-
-            mSplitLayout.update(null /* t */);
-            onLayoutSizeChanged(mSplitLayout);
-        }
-    }
-
-    @Override
     public void onDisplayAdded(int displayId) {
         if (displayId != DEFAULT_DISPLAY) {
             return;
@@ -1200,13 +1228,10 @@
         // Only do this when shell transition
         if (!ENABLE_SHELL_TRANSITIONS) return;
 
-        final SurfaceControl.Transaction t = mTransactionPool.acquire();
         mDisplayLayout.rotateTo(mContext.getResources(), toRotation);
         mSplitLayout.rotateTo(toRotation, mDisplayLayout.stableInsets());
         updateWindowBounds(mSplitLayout, wct);
         updateUnfoldBounds();
-        t.apply();
-        mTransactionPool.release(t);
     }
 
     private void onFoldedStateChanged(boolean folded) {
@@ -1258,7 +1283,10 @@
         final ActivityManager.RunningTaskInfo triggerTask = request.getTriggerTask();
         if (triggerTask == null) {
             if (mMainStage.isActive()) {
-                if (request.getType() == TRANSIT_CHANGE && request.getDisplayChange() != null) {
+                final TransitionRequestInfo.DisplayChange displayChange =
+                        request.getDisplayChange();
+                if (request.getType() == TRANSIT_CHANGE && displayChange != null
+                        && displayChange.getStartRotation() != displayChange.getEndRotation()) {
                     mSplitLayout.setFreezeDividerWindow(true);
                 }
                 // Still want to monitor everything while in split-screen, so return non-null.
@@ -1602,7 +1630,7 @@
         // Be default, make it visible. The remote animator can adjust alpha if it plans to animate.
         if (show) {
             t.setAlpha(leash, 1.f);
-            t.setLayer(leash, SPLIT_DIVIDER_LAYER);
+            t.setLayer(leash, Integer.MAX_VALUE);
             t.setPosition(leash, bounds.left, bounds.top);
             t.show(leash);
         }
@@ -1685,7 +1713,7 @@
         @Override
         public void onRootTaskAppeared() {
             mHasRootTask = true;
-            StageCoordinator.this.onStageRootTaskAppeared(this);
+            StageCoordinator.this.onRootTaskAppeared();
         }
 
         @Override
@@ -1715,7 +1743,7 @@
         @Override
         public void onRootTaskVanished() {
             reset();
-            StageCoordinator.this.onStageRootTaskVanished(this);
+            StageCoordinator.this.onRootTaskVanished();
         }
 
         @Override
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java
index fb0eebd..9fd5d20 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java
@@ -39,7 +39,6 @@
 
 import androidx.annotation.NonNull;
 
-import com.android.internal.R;
 import com.android.launcher3.icons.IconProvider;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.common.SurfaceUtils;
@@ -108,12 +107,7 @@
         mSurfaceSession = surfaceSession;
         mIconProvider = iconProvider;
         mStageTaskUnfoldController = stageTaskUnfoldController;
-
-        // No need to create root task if the device is using legacy split screen.
-        // TODO(b/199236198): Remove this check after totally deprecated legacy split.
-        if (!context.getResources().getBoolean(R.bool.config_useLegacySplit)) {
-            taskOrganizer.createRootTask(displayId, WINDOWING_MODE_MULTI_WINDOW, this);
-        }
+        taskOrganizer.createRootTask(displayId, WINDOWING_MODE_MULTI_WINDOW, this);
     }
 
     int getChildCount() {
@@ -187,7 +181,7 @@
     @Override
     @CallSuper
     public void onTaskAppeared(ActivityManager.RunningTaskInfo taskInfo, SurfaceControl leash) {
-        if (mRootTaskInfo == null && !taskInfo.hasParentTask()) {
+        if (mRootTaskInfo == null) {
             mRootLeash = leash;
             mRootTaskInfo = taskInfo;
             mSplitDecorManager = new SplitDecorManager(
@@ -334,11 +328,6 @@
         wct.reparent(task.token, mRootTaskInfo.token, true /* onTop*/);
     }
 
-    void moveToTop(WindowContainerTransaction wct) {
-        final WindowContainerToken rootToken = mRootTaskInfo.token;
-        wct.reorder(rootToken, true /* onTop */);
-    }
-
     void reorderChild(int taskId, boolean onTop, WindowContainerTransaction wct) {
         if (!containsTask(taskId)) {
             return;
@@ -354,10 +343,6 @@
         }
     }
 
-    void setVisibility(boolean visible, WindowContainerTransaction wct) {
-        wct.reorder(mRootTaskInfo.token, visible /* onTop */);
-    }
-
     void onSplitScreenListenerRegistered(SplitScreen.SplitScreenListener listener,
             @StageType int stage) {
         for (int i = mChildrenTaskInfo.size() - 1; i >= 0; --i) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/stagesplit/StageCoordinator.java b/libs/WindowManager/Shell/src/com/android/wm/shell/stagesplit/StageCoordinator.java
index a17942f..6ef20e3 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/stagesplit/StageCoordinator.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/stagesplit/StageCoordinator.java
@@ -23,7 +23,6 @@
 import static android.view.WindowManager.TRANSIT_TO_BACK;
 import static android.view.WindowManager.TRANSIT_TO_FRONT;
 import static android.view.WindowManager.transitTypeToString;
-import static android.view.WindowManagerPolicyConstants.SPLIT_DIVIDER_LAYER;
 
 import static com.android.internal.util.FrameworkStatsLog.SPLITSCREEN_UICHANGED__EXIT_REASON__APP_DOES_NOT_SUPPORT_MULTIWINDOW;
 import static com.android.internal.util.FrameworkStatsLog.SPLITSCREEN_UICHANGED__EXIT_REASON__APP_FINISHED;
@@ -722,7 +721,7 @@
 
         if (mDividerVisible) {
             t.show(dividerLeash)
-                    .setLayer(dividerLeash, SPLIT_DIVIDER_LAYER)
+                    .setLayer(dividerLeash, Integer.MAX_VALUE)
                     .setPosition(dividerLeash,
                             mSplitLayout.getDividerBounds().left,
                             mSplitLayout.getDividerBounds().top);
@@ -738,7 +737,7 @@
         }
 
         if (mDividerVisible) {
-            t.show(outlineLeash).setLayer(outlineLeash, SPLIT_DIVIDER_LAYER);
+            t.show(outlineLeash).setLayer(outlineLeash, Integer.MAX_VALUE);
         } else {
             t.hide(outlineLeash);
         }
@@ -1190,7 +1189,7 @@
         // Be default, make it visible. The remote animator can adjust alpha if it plans to animate.
         if (show) {
             t.setAlpha(leash, 1.f);
-            t.setLayer(leash, SPLIT_DIVIDER_LAYER);
+            t.setLayer(leash, Integer.MAX_VALUE);
             t.setPosition(leash, bounds.left, bounds.top);
             t.show(leash);
         }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultTransitionHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultTransitionHandler.java
index b0e44a1..542ddee 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultTransitionHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultTransitionHandler.java
@@ -369,12 +369,20 @@
             }
 
             if (change.getMode() == TRANSIT_CHANGE) {
-                // If task is child task, only set position in parent.
+                // If task is child task, only set position in parent and update crop when needed.
                 if (isTask && change.getParent() != null
                         && info.getChange(change.getParent()).getTaskInfo() != null) {
                     final Point positionInParent = change.getTaskInfo().positionInParent;
                     startTransaction.setPosition(change.getLeash(),
                             positionInParent.x, positionInParent.y);
+
+                    if (!change.getEndAbsBounds().equals(
+                            info.getChange(change.getParent()).getEndAbsBounds())) {
+                        startTransaction.setWindowCrop(change.getLeash(),
+                                change.getEndAbsBounds().width(),
+                                change.getEndAbsBounds().height());
+                    }
+
                     continue;
                 }
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/split/SplitLayoutTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/split/SplitLayoutTests.java
index daec336..8b4e1f8 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/split/SplitLayoutTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/split/SplitLayoutTests.java
@@ -16,7 +16,6 @@
 
 package com.android.wm.shell.common.split;
 
-import static android.content.res.Configuration.ORIENTATION_LANDSCAPE;
 import static android.content.res.Configuration.ORIENTATION_PORTRAIT;
 
 import static com.google.common.truth.Truth.assertThat;
@@ -85,10 +84,6 @@
         // Verify it returns true if new config won't affect split layout.
         assertThat(mSplitLayout.updateConfiguration(config)).isFalse();
 
-        // Verify updateConfiguration returns true if the orientation changed.
-        config.orientation = ORIENTATION_LANDSCAPE;
-        assertThat(mSplitLayout.updateConfiguration(config)).isTrue();
-
         // Verify updateConfiguration returns true if it rotated.
         config.windowConfiguration.setRotation(1);
         assertThat(mSplitLayout.updateConfiguration(config)).isTrue();
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTestUtils.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTestUtils.java
index 85527c8..49f36a4 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTestUtils.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTestUtils.java
@@ -16,22 +16,18 @@
 
 package com.android.wm.shell.splitscreen;
 
-import static android.view.Display.DEFAULT_DISPLAY;
-import static android.window.DisplayAreaOrganizer.FEATURE_DEFAULT_TASK_CONTAINER;
-
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 
+import android.app.ActivityManager;
 import android.content.Context;
 import android.graphics.Rect;
 import android.view.SurfaceControl;
-import android.window.DisplayAreaInfo;
-import android.window.IWindowContainerToken;
-import android.window.WindowContainerToken;
+import android.view.SurfaceSession;
 
 import com.android.dx.mockito.inline.extended.ExtendedMockito;
-import com.android.wm.shell.RootTaskDisplayAreaOrganizer;
 import com.android.wm.shell.ShellTaskOrganizer;
+import com.android.wm.shell.TestRunningTaskInfoBuilder;
 import com.android.wm.shell.common.DisplayController;
 import com.android.wm.shell.common.DisplayImeController;
 import com.android.wm.shell.common.DisplayInsetsController;
@@ -68,26 +64,24 @@
     }
 
     static class TestStageCoordinator extends StageCoordinator {
-        final DisplayAreaInfo mDisplayAreaInfo;
+        final ActivityManager.RunningTaskInfo mRootTask;
+        final SurfaceControl mRootLeash;
 
         TestStageCoordinator(Context context, int displayId, SyncTransactionQueue syncQueue,
-                RootTaskDisplayAreaOrganizer rootTDAOrganizer, ShellTaskOrganizer taskOrganizer,
-                MainStage mainStage, SideStage sideStage, DisplayController displayController,
-                DisplayImeController imeController, DisplayInsetsController insetsController,
-                SplitLayout splitLayout, Transitions transitions, TransactionPool transactionPool,
-                SplitscreenEventLogger logger,
-                Optional<RecentTasksController> recentTasks,
+                ShellTaskOrganizer taskOrganizer, MainStage mainStage, SideStage sideStage,
+                DisplayController displayController, DisplayImeController imeController,
+                DisplayInsetsController insetsController, SplitLayout splitLayout,
+                Transitions transitions, TransactionPool transactionPool,
+                SplitscreenEventLogger logger, Optional<RecentTasksController> recentTasks,
                 Provider<Optional<StageTaskUnfoldController>> unfoldController) {
-            super(context, displayId, syncQueue, rootTDAOrganizer, taskOrganizer, mainStage,
+            super(context, displayId, syncQueue, taskOrganizer, mainStage,
                     sideStage, displayController, imeController, insetsController, splitLayout,
                     transitions, transactionPool, logger, recentTasks, unfoldController);
 
-            // Prepare default TaskDisplayArea for testing.
-            mDisplayAreaInfo = new DisplayAreaInfo(
-                    new WindowContainerToken(new IWindowContainerToken.Default()),
-                    DEFAULT_DISPLAY,
-                    FEATURE_DEFAULT_TASK_CONTAINER);
-            this.onDisplayAreaAppeared(mDisplayAreaInfo);
+            // Prepare root task for testing.
+            mRootTask = new TestRunningTaskInfoBuilder().build();
+            mRootLeash = new SurfaceControl.Builder(new SurfaceSession()).setName("test").build();
+            onTaskAppeared(mRootTask, mRootLeash);
         }
     }
 }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTransitionTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTransitionTests.java
index 19d2a7e..f847e6e 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTransitionTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTransitionTests.java
@@ -124,9 +124,9 @@
                 mIconProvider, null);
         mSideStage.onTaskAppeared(new TestRunningTaskInfoBuilder().build(), createMockSurface());
         mStageCoordinator = new SplitTestUtils.TestStageCoordinator(mContext, DEFAULT_DISPLAY,
-                mSyncQueue, mRootTDAOrganizer, mTaskOrganizer, mMainStage, mSideStage,
-                mDisplayController, mDisplayImeController, mDisplayInsetsController, mSplitLayout,
-                mTransitions, mTransactionPool, mLogger, Optional.empty(), Optional::empty);
+                mSyncQueue, mTaskOrganizer, mMainStage, mSideStage, mDisplayController,
+                mDisplayImeController, mDisplayInsetsController, mSplitLayout, mTransitions,
+                mTransactionPool, mLogger, Optional.empty(), Optional::empty);
         mSplitScreenTransitions = mStageCoordinator.getSplitTransitions();
         doAnswer((Answer<IBinder>) invocation -> mock(IBinder.class))
                 .when(mTransitions).startTransition(anyInt(), any(), any());
@@ -424,19 +424,14 @@
     }
 
     private boolean containsSplitEnter(@NonNull WindowContainerTransaction wct) {
-        boolean movedMainToFront = false;
-        boolean movedSideToFront = false;
         for (int i = 0; i < wct.getHierarchyOps().size(); ++i) {
             WindowContainerTransaction.HierarchyOp op = wct.getHierarchyOps().get(i);
-            if (op.getType() == HIERARCHY_OP_TYPE_REORDER) {
-                if (op.getContainer() == mMainStage.mRootTaskInfo.token.asBinder()) {
-                    movedMainToFront = true;
-                } else if (op.getContainer() == mSideStage.mRootTaskInfo.token.asBinder()) {
-                    movedSideToFront = true;
-                }
+            if (op.getType() == HIERARCHY_OP_TYPE_REORDER
+                    && op.getContainer() == mStageCoordinator.mRootTaskInfo.token.asBinder()) {
+                return true;
             }
         }
-        return movedMainToFront && movedSideToFront;
+        return false;
     }
 
     private boolean containsSplitExit(@NonNull WindowContainerTransaction wct) {
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageCoordinatorTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageCoordinatorTests.java
index a6ee8ff..061136c6 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageCoordinatorTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageCoordinatorTests.java
@@ -34,21 +34,20 @@
 import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.doReturn;
-import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.app.ActivityManager;
+import android.content.res.Configuration;
 import android.graphics.Rect;
 import android.view.SurfaceControl;
-import android.window.DisplayAreaInfo;
+import android.view.SurfaceSession;
 import android.window.WindowContainerTransaction;
 
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
 
-import com.android.wm.shell.RootTaskDisplayAreaOrganizer;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.ShellTestCase;
 import com.android.wm.shell.TestRunningTaskInfoBuilder;
@@ -81,8 +80,6 @@
     @Mock
     private SyncTransactionQueue mSyncQueue;
     @Mock
-    private RootTaskDisplayAreaOrganizer mRootTDAOrganizer;
-    @Mock
     private MainStage mMainStage;
     @Mock
     private SideStage mSideStage;
@@ -108,17 +105,27 @@
     private final Rect mBounds1 = new Rect(10, 20, 30, 40);
     private final Rect mBounds2 = new Rect(5, 10, 15, 20);
 
+    private SurfaceSession mSurfaceSession = new SurfaceSession();
+    private SurfaceControl mRootLeash;
+    private ActivityManager.RunningTaskInfo mRootTask;
     private StageCoordinator mStageCoordinator;
 
     @Before
     public void setup() {
         MockitoAnnotations.initMocks(this);
-        mStageCoordinator = spy(createStageCoordinator(mSplitLayout));
+        mStageCoordinator = spy(new StageCoordinator(mContext, DEFAULT_DISPLAY, mSyncQueue,
+                mTaskOrganizer, mMainStage, mSideStage, mDisplayController, mDisplayImeController,
+                mDisplayInsetsController, mSplitLayout, mTransitions, mTransactionPool, mLogger,
+                Optional.empty(), new UnfoldControllerProvider()));
         doNothing().when(mStageCoordinator).updateActivityOptions(any(), anyInt());
 
         when(mSplitLayout.getBounds1()).thenReturn(mBounds1);
         when(mSplitLayout.getBounds2()).thenReturn(mBounds2);
         when(mSplitLayout.isLandscape()).thenReturn(false);
+
+        mRootTask = new TestRunningTaskInfoBuilder().build();
+        mRootLeash = new SurfaceControl.Builder(mSurfaceSession).setName("test").build();
+        mStageCoordinator.onTaskAppeared(mRootTask, mRootLeash);
     }
 
     @Test
@@ -158,14 +165,17 @@
     }
 
     @Test
-    public void testDisplayAreaAppeared_initializesUnfoldControllers() {
-        // Create a stage coordinator with null split layout to test layout init flow.
-        mStageCoordinator = createStageCoordinator(null /* splitLayout */);
-
-        mStageCoordinator.onDisplayAreaAppeared(mock(DisplayAreaInfo.class));
-
+    public void testRootTaskAppeared_initializesUnfoldControllers() {
         verify(mMainUnfoldController).init();
         verify(mSideUnfoldController).init();
+        verify(mStageCoordinator).onRootTaskAppeared();
+    }
+
+    @Test
+    public void testRootTaskInfoChanged_updatesSplitLayout() {
+        mStageCoordinator.onTaskInfoChanged(mRootTask);
+
+        verify(mSplitLayout).updateConfiguration(any(Configuration.class));
     }
 
     @Test
@@ -302,14 +312,6 @@
         verify(mSplitLayout).applySurfaceChanges(any(), any(), any(), any(), any());
     }
 
-    private StageCoordinator createStageCoordinator(SplitLayout splitLayout) {
-        return new SplitTestUtils.TestStageCoordinator(mContext, DEFAULT_DISPLAY,
-                mSyncQueue, mRootTDAOrganizer, mTaskOrganizer, mMainStage, mSideStage,
-                mDisplayController, mDisplayImeController, mDisplayInsetsController, splitLayout,
-                mTransitions, mTransactionPool, mLogger, Optional.empty(),
-                new UnfoldControllerProvider());
-    }
-
     private class UnfoldControllerProvider implements
             Provider<Optional<StageTaskUnfoldController>> {
 
diff --git a/services/core/java/com/android/server/wm/Task.java b/services/core/java/com/android/server/wm/Task.java
index 731a045..0542c0b 100644
--- a/services/core/java/com/android/server/wm/Task.java
+++ b/services/core/java/com/android/server/wm/Task.java
@@ -2701,7 +2701,7 @@
     @Override
     void onDisplayChanged(DisplayContent dc) {
         final boolean isRootTask = isRootTask();
-        if (!isRootTask) {
+        if (!isRootTask && !mCreatedByOrganizer) {
             adjustBoundsForDisplayChangeIfNeeded(dc);
         }
         super.onDisplayChanged(dc);
@@ -3468,9 +3468,9 @@
         forAllActivities(r -> {
             info.addLaunchCookie(r.mLaunchCookie);
         });
-        final Task rootTask = getRootTask();
-        info.parentTaskId = rootTask == getParent() && rootTask.mCreatedByOrganizer
-                ? rootTask.mTaskId
+        final Task parentTask = getParent() != null ? getParent().asTask() : null;
+        info.parentTaskId = parentTask != null && parentTask.mCreatedByOrganizer
+                ? parentTask.mTaskId
                 : INVALID_TASK_ID;
         info.isFocused = isFocused();
         info.isVisible = hasVisibleChildren();
@@ -4146,13 +4146,13 @@
 
     private boolean canBeOrganized() {
         // All root tasks can be organized
-        if (isRootTask()) {
+        if (isRootTask() || mCreatedByOrganizer) {
             return true;
         }
 
-        // Task could be organized if it's the direct child of the root created by organizer.
-        final Task rootTask = getRootTask();
-        return rootTask == getParent() && rootTask.mCreatedByOrganizer;
+        // Task could be organized if it's the direct child of a task created by organizer.
+        final Task parentTask = getParent().asTask();
+        return parentTask != null && parentTask.mCreatedByOrganizer;
     }
 
     @Override