Revert "Migrate unfold animation to Shell transitions [Part 3]"

Revert "Update TaskListener constructors to match the upstream"

Revert submission 18628270-unfold-shell-p3

Reason for revert: b/235090315
Reverted Changes:
Ibe684471a:Migrate unfold animation to Shell transitions [Par...
I7bd398c96:Update TaskListener constructors to match the upst...
I7bd398c96:Migrate unfold animation to Shell transitions [Par...

Bug: 235090315
Change-Id: I8639d618131c09f7a5f0c776fb02c426b78b1fdd

32 files changed, 741 insertions, 1205 deletions

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/ShellInitImpl.java b/libs/WindowManager/Shell/src/com/android/wm/shell/ShellInitImpl.java
index 38f7465f3f45..f6a3e7fb54d9 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/ShellInitImpl.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/ShellInitImpl.java
@@ -27,6 +27,7 @@ import com.android.wm.shell.common.annotations.ExternalThread;
 import com.android.wm.shell.draganddrop.DragAndDropController;
 import com.android.wm.shell.freeform.FreeformTaskListener;
 import com.android.wm.shell.fullscreen.FullscreenTaskListener;
+import com.android.wm.shell.fullscreen.FullscreenUnfoldController;
 import com.android.wm.shell.kidsmode.KidsModeTaskOrganizer;
 import com.android.wm.shell.pip.phone.PipTouchHandler;
 import com.android.wm.shell.recents.RecentTasksController;
@@ -34,7 +35,6 @@ import com.android.wm.shell.splitscreen.SplitScreenController;
 import com.android.wm.shell.startingsurface.StartingWindowController;
 import com.android.wm.shell.transition.DefaultMixedHandler;
 import com.android.wm.shell.transition.Transitions;
-import com.android.wm.shell.unfold.UnfoldAnimationController;
 import com.android.wm.shell.unfold.UnfoldTransitionHandler;
 
 import java.util.Optional;
@@ -55,7 +55,7 @@ public class ShellInitImpl {
     private final Optional<SplitScreenController> mSplitScreenOptional;
     private final Optional<PipTouchHandler> mPipTouchHandlerOptional;
     private final FullscreenTaskListener mFullscreenTaskListener;
-    private final Optional<UnfoldAnimationController> mUnfoldController;
+    private final Optional<FullscreenUnfoldController> mFullscreenUnfoldController;
     private final Optional<UnfoldTransitionHandler> mUnfoldTransitionHandler;
     private final Optional<FreeformTaskListener> mFreeformTaskListenerOptional;
     private final ShellExecutor mMainExecutor;
@@ -76,7 +76,7 @@ public class ShellInitImpl {
             Optional<SplitScreenController> splitScreenOptional,
             Optional<PipTouchHandler> pipTouchHandlerOptional,
             FullscreenTaskListener fullscreenTaskListener,
-            Optional<UnfoldAnimationController> unfoldAnimationController,
+            Optional<FullscreenUnfoldController> fullscreenUnfoldTransitionController,
             Optional<UnfoldTransitionHandler> unfoldTransitionHandler,
             Optional<FreeformTaskListener> freeformTaskListenerOptional,
             Optional<RecentTasksController> recentTasks,
@@ -93,7 +93,7 @@ public class ShellInitImpl {
         mSplitScreenOptional = splitScreenOptional;
         mFullscreenTaskListener = fullscreenTaskListener;
         mPipTouchHandlerOptional = pipTouchHandlerOptional;
-        mUnfoldController = unfoldAnimationController;
+        mFullscreenUnfoldController = fullscreenUnfoldTransitionController;
         mUnfoldTransitionHandler = unfoldTransitionHandler;
         mFreeformTaskListenerOptional = freeformTaskListenerOptional;
         mRecentTasks = recentTasks;
@@ -146,7 +146,7 @@ public class ShellInitImpl {
                 mShellTaskOrganizer.addListenerForType(
                         f, ShellTaskOrganizer.TASK_LISTENER_TYPE_FREEFORM));
 
-        mUnfoldController.ifPresent(UnfoldAnimationController::init);
+        mFullscreenUnfoldController.ifPresent(FullscreenUnfoldController::init);
         mRecentTasks.ifPresent(RecentTasksController::init);
 
         // Initialize kids mode task organizer
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
index e9d24fbf4d0a..31f0ef0192ae 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
@@ -55,7 +55,6 @@ import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.compatui.CompatUIController;
 import com.android.wm.shell.recents.RecentTasksController;
 import com.android.wm.shell.startingsurface.StartingWindowController;
-import com.android.wm.shell.unfold.UnfoldAnimationController;
 
 import java.io.PrintWriter;
 import java.util.ArrayList;
@@ -180,41 +179,33 @@ public class ShellTaskOrganizer extends TaskOrganizer implements
     private final Optional<RecentTasksController> mRecentTasks;
 
     @Nullable
-    private final UnfoldAnimationController mUnfoldAnimationController;
-
-    @Nullable
     private RunningTaskInfo mLastFocusedTaskInfo;
 
     public ShellTaskOrganizer(ShellExecutor mainExecutor, Context context) {
         this(null /* taskOrganizerController */, mainExecutor, context, null /* compatUI */,
-                Optional.empty() /* unfoldAnimationController */,
                 Optional.empty() /* recentTasksController */);
     }
 
     public ShellTaskOrganizer(ShellExecutor mainExecutor, Context context, @Nullable
             CompatUIController compatUI) {
         this(null /* taskOrganizerController */, mainExecutor, context, compatUI,
-                Optional.empty() /* unfoldAnimationController */,
                 Optional.empty() /* recentTasksController */);
     }
 
     public ShellTaskOrganizer(ShellExecutor mainExecutor, Context context, @Nullable
             CompatUIController compatUI,
-            Optional<UnfoldAnimationController> unfoldAnimationController,
             Optional<RecentTasksController> recentTasks) {
         this(null /* taskOrganizerController */, mainExecutor, context, compatUI,
-                unfoldAnimationController, recentTasks);
+                recentTasks);
     }
 
     @VisibleForTesting
     protected ShellTaskOrganizer(ITaskOrganizerController taskOrganizerController,
             ShellExecutor mainExecutor, Context context, @Nullable CompatUIController compatUI,
-            Optional<UnfoldAnimationController> unfoldAnimationController,
             Optional<RecentTasksController> recentTasks) {
         super(taskOrganizerController, mainExecutor);
         mCompatUI = compatUI;
         mRecentTasks = recentTasks;
-        mUnfoldAnimationController = unfoldAnimationController.orElse(null);
         if (compatUI != null) {
             compatUI.setCompatUICallback(this);
         }
@@ -446,9 +437,6 @@ public class ShellTaskOrganizer extends TaskOrganizer implements
         if (listener != null) {
             listener.onTaskAppeared(info.getTaskInfo(), info.getLeash());
         }
-        if (mUnfoldAnimationController != null) {
-            mUnfoldAnimationController.onTaskAppeared(info.getTaskInfo(), info.getLeash());
-        }
         notifyLocusVisibilityIfNeeded(info.getTaskInfo());
         notifyCompatUI(info.getTaskInfo(), listener);
     }
@@ -470,11 +458,6 @@ public class ShellTaskOrganizer extends TaskOrganizer implements
     public void onTaskInfoChanged(RunningTaskInfo taskInfo) {
         synchronized (mLock) {
             ProtoLog.v(WM_SHELL_TASK_ORG, "Task info changed taskId=%d", taskInfo.taskId);
-
-            if (mUnfoldAnimationController != null) {
-                mUnfoldAnimationController.onTaskInfoChanged(taskInfo);
-            }
-
             final TaskAppearedInfo data = mTasks.get(taskInfo.taskId);
             final TaskListener oldListener = getTaskListener(data.getTaskInfo());
             final TaskListener newListener = getTaskListener(taskInfo);
@@ -524,10 +507,6 @@ public class ShellTaskOrganizer extends TaskOrganizer implements
     public void onTaskVanished(RunningTaskInfo taskInfo) {
         synchronized (mLock) {
             ProtoLog.v(WM_SHELL_TASK_ORG, "Task vanished taskId=%d", taskInfo.taskId);
-            if (mUnfoldAnimationController != null) {
-                mUnfoldAnimationController.onTaskVanished(taskInfo);
-            }
-
             final int taskId = taskInfo.taskId;
             final TaskListener listener = getTaskListener(mTasks.get(taskId).getTaskInfo());
             mTasks.remove(taskId);
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 ce3ae6e7bfc6..c94455d9151a 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
@@ -178,11 +178,6 @@ public final class SplitLayout implements DisplayInsetsController.OnInsetsChange
         return outBounds;
     }
 
-    /** Gets root bounds of the whole split layout */
-    public Rect getRootBounds() {
-        return new Rect(mRootBounds);
-    }
-
     /** Gets bounds of divider window with screen based coordinate. */
     public Rect getDividerBounds() {
         return new Rect(mDividerBounds);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
index a57b0e0ab241..1d10bbe37438 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
@@ -64,6 +64,7 @@ import com.android.wm.shell.draganddrop.DragAndDrop;
 import com.android.wm.shell.draganddrop.DragAndDropController;
 import com.android.wm.shell.freeform.FreeformTaskListener;
 import com.android.wm.shell.fullscreen.FullscreenTaskListener;
+import com.android.wm.shell.fullscreen.FullscreenUnfoldController;
 import com.android.wm.shell.hidedisplaycutout.HideDisplayCutout;
 import com.android.wm.shell.hidedisplaycutout.HideDisplayCutoutController;
 import com.android.wm.shell.kidsmode.KidsModeTaskOrganizer;
@@ -87,7 +88,6 @@ import com.android.wm.shell.tasksurfacehelper.TaskSurfaceHelperController;
 import com.android.wm.shell.transition.ShellTransitions;
 import com.android.wm.shell.transition.Transitions;
 import com.android.wm.shell.unfold.ShellUnfoldProgressProvider;
-import com.android.wm.shell.unfold.UnfoldAnimationController;
 import com.android.wm.shell.unfold.UnfoldTransitionHandler;
 
 import java.util.Optional;
@@ -176,11 +176,9 @@ public abstract class WMShellBaseModule {
     static ShellTaskOrganizer provideShellTaskOrganizer(@ShellMainThread ShellExecutor mainExecutor,
             Context context,
             CompatUIController compatUI,
-            Optional<UnfoldAnimationController> unfoldAnimationController,
             Optional<RecentTasksController> recentTasksOptional
     ) {
-        return new ShellTaskOrganizer(mainExecutor, context, compatUI, unfoldAnimationController,
-                recentTasksOptional);
+        return new ShellTaskOrganizer(mainExecutor, context, compatUI, recentTasksOptional);
     }
 
     @WMSingleton
@@ -192,12 +190,10 @@ public abstract class WMShellBaseModule {
             SyncTransactionQueue syncTransactionQueue,
             DisplayController displayController,
             DisplayInsetsController displayInsetsController,
-            Optional<UnfoldAnimationController> unfoldAnimationController,
             Optional<RecentTasksController> recentTasksOptional
     ) {
         return new KidsModeTaskOrganizer(mainExecutor, mainHandler, context, syncTransactionQueue,
-                displayController, displayInsetsController, unfoldAnimationController,
-                recentTasksOptional);
+                displayController, displayInsetsController, recentTasksOptional);
     }
 
     @WMSingleton
@@ -294,11 +290,13 @@ public abstract class WMShellBaseModule {
     static FullscreenTaskListener provideFullscreenTaskListener(
             @DynamicOverride Optional<FullscreenTaskListener> fullscreenTaskListener,
             SyncTransactionQueue syncQueue,
+            Optional<FullscreenUnfoldController> optionalFullscreenUnfoldController,
             Optional<RecentTasksController> recentTasksOptional) {
         if (fullscreenTaskListener.isPresent()) {
             return fullscreenTaskListener.get();
         } else {
-            return new FullscreenTaskListener(syncQueue, recentTasksOptional);
+            return new FullscreenTaskListener(syncQueue, optionalFullscreenUnfoldController,
+                    recentTasksOptional);
         }
     }
 
@@ -312,13 +310,12 @@ public abstract class WMShellBaseModule {
     // Workaround for dynamic overriding with a default implementation, see {@link DynamicOverride}
     @BindsOptionalOf
     @DynamicOverride
-    abstract UnfoldAnimationController optionalUnfoldController();
+    abstract FullscreenUnfoldController optionalFullscreenUnfoldController();
 
     @WMSingleton
     @Provides
-    static Optional<UnfoldAnimationController> provideUnfoldController(
-            @DynamicOverride Lazy<Optional<UnfoldAnimationController>>
-                    fullscreenUnfoldController,
+    static Optional<FullscreenUnfoldController> provideFullscreenUnfoldController(
+            @DynamicOverride Lazy<Optional<FullscreenUnfoldController>> fullscreenUnfoldController,
             Optional<ShellUnfoldProgressProvider> progressProvider) {
         if (progressProvider.isPresent()
                 && progressProvider.get() != ShellUnfoldProgressProvider.NO_PROVIDER) {
@@ -643,7 +640,7 @@ public abstract class WMShellBaseModule {
             Optional<SplitScreenController> splitScreenOptional,
             Optional<PipTouchHandler> pipTouchHandlerOptional,
             FullscreenTaskListener fullscreenTaskListener,
-            Optional<UnfoldAnimationController> unfoldAnimationController,
+            Optional<FullscreenUnfoldController> appUnfoldTransitionController,
             Optional<UnfoldTransitionHandler> unfoldTransitionHandler,
             Optional<FreeformTaskListener> freeformTaskListener,
             Optional<RecentTasksController> recentTasksOptional,
@@ -660,7 +657,7 @@ public abstract class WMShellBaseModule {
                 splitScreenOptional,
                 pipTouchHandlerOptional,
                 fullscreenTaskListener,
-                unfoldAnimationController,
+                appUnfoldTransitionController,
                 unfoldTransitionHandler,
                 freeformTaskListener,
                 recentTasksOptional,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
index 8d2b5f766f4e..c5453d4394b8 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
@@ -45,6 +45,7 @@ import com.android.wm.shell.common.annotations.ShellBackgroundThread;
 import com.android.wm.shell.common.annotations.ShellMainThread;
 import com.android.wm.shell.draganddrop.DragAndDropController;
 import com.android.wm.shell.freeform.FreeformTaskListener;
+import com.android.wm.shell.fullscreen.FullscreenUnfoldController;
 import com.android.wm.shell.onehanded.OneHandedController;
 import com.android.wm.shell.pip.Pip;
 import com.android.wm.shell.pip.PipAnimationController;
@@ -66,22 +67,17 @@ import com.android.wm.shell.pip.phone.PipMotionHelper;
 import com.android.wm.shell.pip.phone.PipTouchHandler;
 import com.android.wm.shell.recents.RecentTasksController;
 import com.android.wm.shell.splitscreen.SplitScreenController;
+import com.android.wm.shell.splitscreen.StageTaskUnfoldController;
 import com.android.wm.shell.transition.Transitions;
-import com.android.wm.shell.unfold.UnfoldAnimationController;
 import com.android.wm.shell.unfold.ShellUnfoldProgressProvider;
 import com.android.wm.shell.unfold.UnfoldBackgroundController;
 import com.android.wm.shell.unfold.UnfoldTransitionHandler;
 import com.android.wm.shell.unfold.animation.FullscreenUnfoldTaskAnimator;
-import com.android.wm.shell.unfold.animation.SplitTaskUnfoldAnimator;
-import com.android.wm.shell.unfold.animation.UnfoldTaskAnimator;
-import com.android.wm.shell.unfold.qualifier.UnfoldTransition;
-import com.android.wm.shell.unfold.qualifier.UnfoldShellTransition;
 
-import java.util.ArrayList;
-import java.util.List;
 import java.util.Optional;
 
-import dagger.Binds;
+import javax.inject.Provider;
+
 import dagger.Lazy;
 import dagger.Module;
 import dagger.Provides;
@@ -95,7 +91,7 @@ import dagger.Provides;
  * dependencies should go into {@link WMShellBaseModule}.
  */
 @Module(includes = WMShellBaseModule.class)
-public abstract class WMShellModule {
+public class WMShellModule {
 
     //
     // Bubbles
@@ -176,11 +172,12 @@ public abstract class WMShellModule {
             DisplayImeController displayImeController,
             DisplayInsetsController displayInsetsController, Transitions transitions,
             TransactionPool transactionPool, IconProvider iconProvider,
-            Optional<RecentTasksController> recentTasks) {
+            Optional<RecentTasksController> recentTasks,
+            Provider<Optional<StageTaskUnfoldController>> stageTaskUnfoldControllerProvider) {
         return new SplitScreenController(shellTaskOrganizer, syncQueue, context,
                 rootTaskDisplayAreaOrganizer, mainExecutor, displayController, displayImeController,
                 displayInsetsController, transitions, transactionPool, iconProvider,
-                recentTasks);
+                recentTasks, stageTaskUnfoldControllerProvider);
     }
 
     //
@@ -327,77 +324,62 @@ public abstract class WMShellModule {
     //
     // Unfold transition
     //
+
     @WMSingleton
     @Provides
     @DynamicOverride
-    static UnfoldAnimationController provideUnfoldAnimationController(
+    static FullscreenUnfoldController provideFullscreenUnfoldController(
             Optional<ShellUnfoldProgressProvider> progressProvider,
-            TransactionPool transactionPool,
-            @UnfoldTransition SplitTaskUnfoldAnimator splitAnimator,
-            FullscreenUnfoldTaskAnimator fullscreenAnimator,
-            Lazy<Optional<UnfoldTransitionHandler>> unfoldTransitionHandler,
+            Optional<UnfoldTransitionHandler> unfoldTransitionHandler,
+            FullscreenUnfoldTaskAnimator fullscreenUnfoldTaskAnimator,
+            UnfoldBackgroundController unfoldBackgroundController,
             @ShellMainThread ShellExecutor mainExecutor
     ) {
-        final List<UnfoldTaskAnimator> animators = new ArrayList<>();
-        animators.add(splitAnimator);
-        animators.add(fullscreenAnimator);
-
-        return new UnfoldAnimationController(
-                        transactionPool,
-                        progressProvider.get(),
-                        animators,
-                        unfoldTransitionHandler,
-                        mainExecutor
-                );
+        return new FullscreenUnfoldController(mainExecutor,
+                unfoldBackgroundController, progressProvider.get(),
+                unfoldTransitionHandler.get(), fullscreenUnfoldTaskAnimator);
     }
 
-
     @Provides
     static FullscreenUnfoldTaskAnimator provideFullscreenUnfoldTaskAnimator(
             Context context,
-            UnfoldBackgroundController unfoldBackgroundController,
-            DisplayInsetsController displayInsetsController
-    ) {
-        return new FullscreenUnfoldTaskAnimator(context, unfoldBackgroundController,
-                displayInsetsController);
-    }
-
-    @Provides
-    static SplitTaskUnfoldAnimator provideSplitTaskUnfoldAnimatorBase(
-            Context context,
-            UnfoldBackgroundController backgroundController,
-            @ShellMainThread ShellExecutor executor,
-            Lazy<Optional<SplitScreenController>> splitScreenOptional,
             DisplayInsetsController displayInsetsController
     ) {
-        return new SplitTaskUnfoldAnimator(context, executor, splitScreenOptional,
-                backgroundController, displayInsetsController);
+        return new FullscreenUnfoldTaskAnimator(context, displayInsetsController);
     }
 
     @WMSingleton
-    @UnfoldShellTransition
-    @Binds
-    abstract SplitTaskUnfoldAnimator provideShellSplitTaskUnfoldAnimator(
-            SplitTaskUnfoldAnimator splitTaskUnfoldAnimator);
-
-    @WMSingleton
-    @UnfoldTransition
-    @Binds
-    abstract SplitTaskUnfoldAnimator provideSplitTaskUnfoldAnimator(
-            SplitTaskUnfoldAnimator splitTaskUnfoldAnimator);
-
-    @WMSingleton
     @Provides
     @DynamicOverride
     static UnfoldTransitionHandler provideUnfoldTransitionHandler(
             Optional<ShellUnfoldProgressProvider> progressProvider,
             FullscreenUnfoldTaskAnimator animator,
-            @UnfoldShellTransition SplitTaskUnfoldAnimator unfoldAnimator,
+            UnfoldBackgroundController backgroundController,
             TransactionPool transactionPool,
             Transitions transitions,
             @ShellMainThread ShellExecutor executor) {
         return new UnfoldTransitionHandler(progressProvider.get(), animator,
-                unfoldAnimator, transactionPool, executor, transitions);
+                transactionPool, backgroundController, executor, transitions);
+    }
+
+    @Provides
+    static Optional<StageTaskUnfoldController> provideStageTaskUnfoldController(
+            Optional<ShellUnfoldProgressProvider> progressProvider,
+            Context context,
+            TransactionPool transactionPool,
+            Lazy<UnfoldBackgroundController> unfoldBackgroundController,
+            DisplayInsetsController displayInsetsController,
+            @ShellMainThread ShellExecutor mainExecutor
+    ) {
+        return progressProvider.map(shellUnfoldTransitionProgressProvider ->
+                new StageTaskUnfoldController(
+                        context,
+                        transactionPool,
+                        shellUnfoldTransitionProgressProvider,
+                        displayInsetsController,
+                        unfoldBackgroundController.get(),
+                        mainExecutor
+                ));
     }
 
     @WMSingleton
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/fullscreen/FullscreenTaskListener.java b/libs/WindowManager/Shell/src/com/android/wm/shell/fullscreen/FullscreenTaskListener.java
index 79e363bcdb41..1fc1215b6cea 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/fullscreen/FullscreenTaskListener.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/fullscreen/FullscreenTaskListener.java
@@ -16,13 +16,17 @@
 
 package com.android.wm.shell.fullscreen;
 
+import static android.app.WindowConfiguration.ACTIVITY_TYPE_HOME;
+
 import static com.android.wm.shell.ShellTaskOrganizer.TASK_LISTENER_TYPE_FULLSCREEN;
 import static com.android.wm.shell.ShellTaskOrganizer.taskListenerTypeToString;
 
 import android.app.ActivityManager.RunningTaskInfo;
+import android.app.TaskInfo;
 import android.graphics.Point;
 import android.util.Slog;
 import android.util.SparseArray;
+import android.util.SparseBooleanArray;
 import android.view.SurfaceControl;
 
 import androidx.annotation.NonNull;
@@ -44,17 +48,22 @@ public class FullscreenTaskListener implements ShellTaskOrganizer.TaskListener {
     private static final String TAG = "FullscreenTaskListener";
 
     private final SyncTransactionQueue mSyncQueue;
+    private final FullscreenUnfoldController mFullscreenUnfoldController;
     private final Optional<RecentTasksController> mRecentTasksOptional;
 
     private final SparseArray<TaskData> mDataByTaskId = new SparseArray<>();
+    private final AnimatableTasksListener mAnimatableTasksListener = new AnimatableTasksListener();
 
-    public FullscreenTaskListener(SyncTransactionQueue syncQueue) {
-        this(syncQueue, Optional.empty());
+    public FullscreenTaskListener(SyncTransactionQueue syncQueue,
+            Optional<FullscreenUnfoldController> unfoldController) {
+        this(syncQueue, unfoldController, Optional.empty());
     }
 
     public FullscreenTaskListener(SyncTransactionQueue syncQueue,
+            Optional<FullscreenUnfoldController> unfoldController,
             Optional<RecentTasksController> recentTasks) {
         mSyncQueue = syncQueue;
+        mFullscreenUnfoldController = unfoldController.orElse(null);
         mRecentTasksOptional = recentTasks;
     }
 
@@ -67,6 +76,7 @@ public class FullscreenTaskListener implements ShellTaskOrganizer.TaskListener {
                 taskInfo.taskId);
         final Point positionInParent = taskInfo.positionInParent;
         mDataByTaskId.put(taskInfo.taskId, new TaskData(leash, positionInParent));
+        mAnimatableTasksListener.onTaskAppeared(taskInfo);
 
         if (Transitions.ENABLE_SHELL_TRANSITIONS) return;
         mSyncQueue.runInSync(t -> {
@@ -84,6 +94,8 @@ public class FullscreenTaskListener implements ShellTaskOrganizer.TaskListener {
 
     @Override
     public void onTaskInfoChanged(RunningTaskInfo taskInfo) {
+        mAnimatableTasksListener.onTaskInfoChanged(taskInfo);
+
         if (Transitions.ENABLE_SHELL_TRANSITIONS) return;
 
         updateRecentsForVisibleFullscreenTask(taskInfo);
@@ -105,6 +117,7 @@ public class FullscreenTaskListener implements ShellTaskOrganizer.TaskListener {
             return;
         }
 
+        mAnimatableTasksListener.onTaskVanished(taskInfo);
         mDataByTaskId.remove(taskInfo.taskId);
 
         ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TASK_ORG, "Fullscreen Task Vanished: #%d",
@@ -162,4 +175,65 @@ public class FullscreenTaskListener implements ShellTaskOrganizer.TaskListener {
             this.positionInParent = positionInParent;
         }
     }
+
+    class AnimatableTasksListener {
+        private final SparseBooleanArray mTaskIds = new SparseBooleanArray();
+
+        public void onTaskAppeared(RunningTaskInfo taskInfo) {
+            final boolean isApplicable = isAnimatable(taskInfo);
+            if (isApplicable) {
+                mTaskIds.put(taskInfo.taskId, true);
+
+                if (mFullscreenUnfoldController != null) {
+                    SurfaceControl leash = mDataByTaskId.get(taskInfo.taskId).surface;
+                    mFullscreenUnfoldController.onTaskAppeared(taskInfo, leash);
+                }
+            }
+        }
+
+        public void onTaskInfoChanged(RunningTaskInfo taskInfo) {
+            final boolean isCurrentlyApplicable = mTaskIds.get(taskInfo.taskId);
+            final boolean isApplicable = isAnimatable(taskInfo);
+
+            if (isCurrentlyApplicable) {
+                if (isApplicable) {
+                    // Still applicable, send update
+                    if (mFullscreenUnfoldController != null) {
+                        mFullscreenUnfoldController.onTaskInfoChanged(taskInfo);
+                    }
+                } else {
+                    // Became inapplicable
+                    if (mFullscreenUnfoldController != null) {
+                        mFullscreenUnfoldController.onTaskVanished(taskInfo);
+                    }
+                    mTaskIds.put(taskInfo.taskId, false);
+                }
+            } else {
+                if (isApplicable) {
+                    // Became applicable
+                    mTaskIds.put(taskInfo.taskId, true);
+
+                    if (mFullscreenUnfoldController != null) {
+                        SurfaceControl leash = mDataByTaskId.get(taskInfo.taskId).surface;
+                        mFullscreenUnfoldController.onTaskAppeared(taskInfo, leash);
+                    }
+                }
+            }
+        }
+
+        public void onTaskVanished(RunningTaskInfo taskInfo) {
+            final boolean isCurrentlyApplicable = mTaskIds.get(taskInfo.taskId);
+            if (isCurrentlyApplicable && mFullscreenUnfoldController != null) {
+                mFullscreenUnfoldController.onTaskVanished(taskInfo);
+            }
+            mTaskIds.put(taskInfo.taskId, false);
+        }
+
+        private boolean isAnimatable(TaskInfo taskInfo) {
+            // Filter all visible tasks that are not launcher tasks
+            // We do not animate launcher as it handles the animation by itself
+            return taskInfo != null && taskInfo.isVisible && taskInfo.getConfiguration()
+                    .windowConfiguration.getActivityType() != ACTIVITY_TYPE_HOME;
+        }
+    }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/fullscreen/FullscreenUnfoldController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/fullscreen/FullscreenUnfoldController.java
new file mode 100644
index 000000000000..99f15f65f74e
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/fullscreen/FullscreenUnfoldController.java
@@ -0,0 +1,135 @@
+/*
+ * 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.fullscreen;
+
+import static android.app.WindowConfiguration.WINDOWING_MODE_PINNED;
+
+import android.annotation.NonNull;
+import android.app.ActivityManager;
+import android.view.SurfaceControl;
+
+import com.android.wm.shell.unfold.ShellUnfoldProgressProvider;
+import com.android.wm.shell.unfold.ShellUnfoldProgressProvider.UnfoldListener;
+import com.android.wm.shell.unfold.UnfoldBackgroundController;
+import com.android.wm.shell.unfold.UnfoldTransitionHandler;
+import com.android.wm.shell.unfold.animation.FullscreenUnfoldTaskAnimator;
+
+import java.util.concurrent.Executor;
+
+/**
+ * Controls full screen app unfold transition: animating cropping window and scaling when
+ * folding or unfolding a foldable device.
+ *
+ * - When Shell transitions are disabled (legacy mode) this controller animates task surfaces
+ *   when doing both fold and unfold.
+ *
+ * - When Shell transitions are enabled this controller animates the surfaces only when
+ *   folding a foldable device. It's not done as a shell transition because we are not committed
+ *   to the display size WM changes yet.
+ *   In this case unfolding is handled by
+ *   {@link com.android.wm.shell.unfold.UnfoldTransitionHandler}.
+ */
+public final class FullscreenUnfoldController implements UnfoldListener {
+
+    private final Executor mExecutor;
+    private final ShellUnfoldProgressProvider mProgressProvider;
+    private final UnfoldBackgroundController mBackgroundController;
+    private final SurfaceControl.Transaction mTransaction = new SurfaceControl.Transaction();
+    private final FullscreenUnfoldTaskAnimator mAnimator;
+    private final UnfoldTransitionHandler mUnfoldTransitionHandler;
+
+    private boolean mShouldHandleAnimation = false;
+
+    public FullscreenUnfoldController(
+            @NonNull Executor executor,
+            @NonNull UnfoldBackgroundController backgroundController,
+            @NonNull ShellUnfoldProgressProvider progressProvider,
+            @NonNull UnfoldTransitionHandler unfoldTransitionHandler,
+            @NonNull FullscreenUnfoldTaskAnimator animator
+    ) {
+        mExecutor = executor;
+        mProgressProvider = progressProvider;
+        mBackgroundController = backgroundController;
+        mUnfoldTransitionHandler = unfoldTransitionHandler;
+        mAnimator = animator;
+    }
+
+    /**
+     * Initializes the controller
+     */
+    public void init() {
+        mAnimator.init();
+        mProgressProvider.addListener(mExecutor, this);
+    }
+
+    @Override
+    public void onStateChangeStarted() {
+        mShouldHandleAnimation = !mUnfoldTransitionHandler.willHandleTransition();
+    }
+
+    @Override
+    public void onStateChangeProgress(float progress) {
+        if (!mAnimator.hasActiveTasks() || !mShouldHandleAnimation) return;
+
+        mBackgroundController.ensureBackground(mTransaction);
+        mAnimator.applyAnimationProgress(progress, mTransaction);
+        mTransaction.apply();
+    }
+
+    @Override
+    public void onStateChangeFinished() {
+        if (!mShouldHandleAnimation) {
+            return;
+        }
+
+        mShouldHandleAnimation = false;
+        mAnimator.resetAllSurfaces(mTransaction);
+        mBackgroundController.removeBackground(mTransaction);
+        mTransaction.apply();
+    }
+
+    /**
+     * Called when a new matching task appeared
+     */
+    public void onTaskAppeared(ActivityManager.RunningTaskInfo taskInfo, SurfaceControl leash) {
+        mAnimator.addTask(taskInfo, leash);
+    }
+
+    /**
+     * Called when matching task changed
+     */
+    public void onTaskInfoChanged(ActivityManager.RunningTaskInfo taskInfo) {
+        mAnimator.onTaskInfoChanged(taskInfo);
+    }
+
+    /**
+     * Called when matching task vanished
+     */
+    public void onTaskVanished(ActivityManager.RunningTaskInfo taskInfo) {
+        // PiP task has its own cleanup path, ignore surface reset to avoid conflict.
+        if (taskInfo.getWindowingMode() != WINDOWING_MODE_PINNED) {
+            mAnimator.resetSurface(taskInfo, mTransaction);
+        }
+        mAnimator.removeTask(taskInfo);
+
+        if (!mAnimator.hasActiveTasks()) {
+            mBackgroundController.removeBackground(mTransaction);
+        }
+
+        mTransaction.apply();
+    }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/kidsmode/KidsModeTaskOrganizer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/kidsmode/KidsModeTaskOrganizer.java
index 2c8ba0970ccc..b4c87b6cbf95 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/kidsmode/KidsModeTaskOrganizer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/kidsmode/KidsModeTaskOrganizer.java
@@ -51,7 +51,6 @@ import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.common.SyncTransactionQueue;
 import com.android.wm.shell.recents.RecentTasksController;
 import com.android.wm.shell.startingsurface.StartingWindowController;
-import com.android.wm.shell.unfold.UnfoldAnimationController;
 
 import java.io.PrintWriter;
 import java.util.List;
@@ -147,11 +146,9 @@ public class KidsModeTaskOrganizer extends ShellTaskOrganizer {
             SyncTransactionQueue syncTransactionQueue,
             DisplayController displayController,
             DisplayInsetsController displayInsetsController,
-            Optional<UnfoldAnimationController> unfoldAnimationController,
             Optional<RecentTasksController> recentTasks,
             KidsModeSettingsObserver kidsModeSettingsObserver) {
-        super(taskOrganizerController, mainExecutor, context, /* compatUI= */ null,
-                unfoldAnimationController, recentTasks);
+        super(taskOrganizerController, mainExecutor, context, /* compatUI= */ null, recentTasks);
         mContext = context;
         mMainHandler = mainHandler;
         mSyncQueue = syncTransactionQueue;
@@ -167,9 +164,8 @@ public class KidsModeTaskOrganizer extends ShellTaskOrganizer {
             SyncTransactionQueue syncTransactionQueue,
             DisplayController displayController,
             DisplayInsetsController displayInsetsController,
-            Optional<UnfoldAnimationController> unfoldAnimationController,
             Optional<RecentTasksController> recentTasks) {
-        super(mainExecutor, context, /* compatUI= */ null, unfoldAnimationController, recentTasks);
+        super(mainExecutor, context, /* compatUI= */ null, recentTasks);
         mContext = context;
         mMainHandler = mainHandler;
         mSyncQueue = syncTransactionQueue;
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 2bfa5db502ce..ae5e075c4d3f 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
@@ -16,6 +16,7 @@
 
 package com.android.wm.shell.splitscreen;
 
+import android.annotation.Nullable;
 import android.content.Context;
 import android.view.SurfaceSession;
 import android.window.WindowContainerToken;
@@ -37,9 +38,10 @@ class MainStage extends StageTaskListener {
 
     MainStage(Context context, ShellTaskOrganizer taskOrganizer, int displayId,
             StageListenerCallbacks callbacks, SyncTransactionQueue syncQueue,
-            SurfaceSession surfaceSession, IconProvider iconProvider) {
-        super(context, taskOrganizer, displayId, callbacks, syncQueue, surfaceSession,
-                iconProvider);
+            SurfaceSession surfaceSession, IconProvider iconProvider,
+            @Nullable StageTaskUnfoldController stageTaskUnfoldController) {
+        super(context, taskOrganizer, displayId, callbacks, syncQueue, surfaceSession, iconProvider,
+                stageTaskUnfoldController);
     }
 
     boolean isActive() {
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 f92a0d3901b9..d55619f5e5ed 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
@@ -16,6 +16,7 @@
 
 package com.android.wm.shell.splitscreen;
 
+import android.annotation.Nullable;
 import android.app.ActivityManager;
 import android.content.Context;
 import android.view.SurfaceSession;
@@ -37,9 +38,10 @@ class SideStage extends StageTaskListener {
 
     SideStage(Context context, ShellTaskOrganizer taskOrganizer, int displayId,
             StageListenerCallbacks callbacks, SyncTransactionQueue syncQueue,
-            SurfaceSession surfaceSession, IconProvider iconProvider) {
-        super(context, taskOrganizer, displayId, callbacks, syncQueue, surfaceSession,
-                iconProvider);
+            SurfaceSession surfaceSession, IconProvider iconProvider,
+            @Nullable StageTaskUnfoldController stageTaskUnfoldController) {
+        super(context, taskOrganizer, displayId, callbacks, syncQueue, surfaceSession, iconProvider,
+                stageTaskUnfoldController);
     }
 
     boolean removeAllTasks(WindowContainerTransaction wct, boolean toTop) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreen.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreen.java
index 29b6311e5041..448773ae9ea2 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreen.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreen.java
@@ -18,7 +18,6 @@ package com.android.wm.shell.splitscreen;
 
 import android.annotation.IntDef;
 import android.annotation.NonNull;
-import android.graphics.Rect;
 
 import com.android.wm.shell.common.annotations.ExternalThread;
 import com.android.wm.shell.common.split.SplitScreenConstants.SplitPosition;
@@ -59,7 +58,6 @@ public interface SplitScreen {
     interface SplitScreenListener {
         default void onStagePositionChanged(@StageType int stage, @SplitPosition int position) {}
         default void onTaskStageChanged(int taskId, @StageType int stage, boolean visible) {}
-        default void onSplitBoundsChanged(Rect rootBounds, Rect mainBounds, Rect sideBounds) {}
         default void onSplitVisibilityChanged(boolean visible) {}
     }
 
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 e7958ee4d700..0d976d4a81eb 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
@@ -86,6 +86,8 @@ import java.util.Arrays;
 import java.util.Optional;
 import java.util.concurrent.Executor;
 
+import javax.inject.Provider;
+
 /**
  * Class manages split-screen multitasking mode and implements the main interface
  * {@link SplitScreen}.
@@ -136,6 +138,7 @@ public class SplitScreenController implements DragAndDropPolicy.Starter,
     private final SplitscreenEventLogger mLogger;
     private final IconProvider mIconProvider;
     private final Optional<RecentTasksController> mRecentTasksOptional;
+    private final Provider<Optional<StageTaskUnfoldController>> mUnfoldControllerProvider;
 
     private StageCoordinator mStageCoordinator;
     // Only used for the legacy recents animation from splitscreen to allow the tasks to be animated
@@ -149,7 +152,8 @@ public class SplitScreenController implements DragAndDropPolicy.Starter,
             DisplayImeController displayImeController,
             DisplayInsetsController displayInsetsController,
             Transitions transitions, TransactionPool transactionPool, IconProvider iconProvider,
-            Optional<RecentTasksController> recentTasks) {
+            Optional<RecentTasksController> recentTasks,
+            Provider<Optional<StageTaskUnfoldController>> unfoldControllerProvider) {
         mTaskOrganizer = shellTaskOrganizer;
         mSyncQueue = syncQueue;
         mContext = context;
@@ -160,6 +164,7 @@ public class SplitScreenController implements DragAndDropPolicy.Starter,
         mDisplayInsetsController = displayInsetsController;
         mTransitions = transitions;
         mTransactionPool = transactionPool;
+        mUnfoldControllerProvider = unfoldControllerProvider;
         mLogger = new SplitscreenEventLogger();
         mIconProvider = iconProvider;
         mRecentTasksOptional = recentTasks;
@@ -185,7 +190,7 @@ public class SplitScreenController implements DragAndDropPolicy.Starter,
             mStageCoordinator = new StageCoordinator(mContext, DEFAULT_DISPLAY, mSyncQueue,
                     mTaskOrganizer, mDisplayController, mDisplayImeController,
                     mDisplayInsetsController, mTransitions, mTransactionPool, mLogger,
-                    mIconProvider, mMainExecutor, mRecentTasksOptional);
+                    mIconProvider, mMainExecutor, mRecentTasksOptional, mUnfoldControllerProvider);
         }
     }
 
@@ -221,14 +226,6 @@ public class SplitScreenController implements DragAndDropPolicy.Starter,
                 new WindowContainerTransaction());
     }
 
-    /**
-     * Update surfaces of the split screen layout based on the current state
-     * @param transaction to write the updates to
-     */
-    public void updateSplitScreenSurfaces(SurfaceControl.Transaction transaction) {
-        mStageCoordinator.updateSurfaces(transaction);
-    }
-
     private boolean moveToStage(int taskId, @StageType int stageType,
             @SplitPosition int stagePosition, WindowContainerTransaction wct) {
         final ActivityManager.RunningTaskInfo task = mTaskOrganizer.getRunningTaskInfo(taskId);
@@ -542,17 +539,6 @@ public class SplitScreenController implements DragAndDropPolicy.Starter,
             }
 
             @Override
-            public void onSplitBoundsChanged(Rect rootBounds, Rect mainBounds, Rect sideBounds) {
-                for (int i = 0; i < mExecutors.size(); i++) {
-                    final int index = i;
-                    mExecutors.valueAt(index).execute(() -> {
-                        mExecutors.keyAt(index).onSplitBoundsChanged(rootBounds, mainBounds,
-                                sideBounds);
-                    });
-                }
-            }
-
-            @Override
             public void onSplitVisibilityChanged(boolean visible) {
                 for (int i = 0; i < mExecutors.size(); i++) {
                     final int index = i;
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 357a0e63386f..3cf8a45310ef 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
@@ -117,6 +117,8 @@ import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
 
+import javax.inject.Provider;
+
 /**
  * Coordinates the staging (visibility, sizing, ...) of the split-screen {@link MainStage} and
  * {@link SideStage} stages.
@@ -143,8 +145,10 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
 
     private final MainStage mMainStage;
     private final StageListenerImpl mMainStageListener = new StageListenerImpl();
+    private final StageTaskUnfoldController mMainUnfoldController;
     private final SideStage mSideStage;
     private final StageListenerImpl mSideStageListener = new StageListenerImpl();
+    private final StageTaskUnfoldController mSideUnfoldController;
     private final DisplayLayout mDisplayLayout;
     @SplitPosition
     private int mSideStagePosition = SPLIT_POSITION_BOTTOM_OR_RIGHT;
@@ -205,7 +209,8 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
             DisplayInsetsController displayInsetsController, Transitions transitions,
             TransactionPool transactionPool, SplitscreenEventLogger logger,
             IconProvider iconProvider, ShellExecutor mainExecutor,
-            Optional<RecentTasksController> recentTasks) {
+            Optional<RecentTasksController> recentTasks,
+            Provider<Optional<StageTaskUnfoldController>> unfoldControllerProvider) {
         mContext = context;
         mDisplayId = displayId;
         mSyncQueue = syncQueue;
@@ -213,7 +218,8 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         mLogger = logger;
         mMainExecutor = mainExecutor;
         mRecentTasks = recentTasks;
-
+        mMainUnfoldController = unfoldControllerProvider.get().orElse(null);
+        mSideUnfoldController = unfoldControllerProvider.get().orElse(null);
         taskOrganizer.createRootTask(displayId, WINDOWING_MODE_FULLSCREEN, this /* listener */);
 
         mMainStage = new MainStage(
@@ -223,7 +229,8 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
                 mMainStageListener,
                 mSyncQueue,
                 mSurfaceSession,
-                iconProvider);
+                iconProvider,
+                mMainUnfoldController);
         mSideStage = new SideStage(
                 mContext,
                 mTaskOrganizer,
@@ -231,7 +238,8 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
                 mSideStageListener,
                 mSyncQueue,
                 mSurfaceSession,
-                iconProvider);
+                iconProvider,
+                mSideUnfoldController);
         mDisplayController = displayController;
         mDisplayImeController = displayImeController;
         mDisplayInsetsController = displayInsetsController;
@@ -254,7 +262,8 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
             DisplayInsetsController displayInsetsController, SplitLayout splitLayout,
             Transitions transitions, TransactionPool transactionPool,
             SplitscreenEventLogger logger, ShellExecutor mainExecutor,
-            Optional<RecentTasksController> recentTasks) {
+            Optional<RecentTasksController> recentTasks,
+            Provider<Optional<StageTaskUnfoldController>> unfoldControllerProvider) {
         mContext = context;
         mDisplayId = displayId;
         mSyncQueue = syncQueue;
@@ -268,6 +277,8 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         mSplitLayout = splitLayout;
         mSplitTransitions = new SplitScreenTransitions(transactionPool, transitions,
                 this::onTransitionAnimationComplete, this);
+        mMainUnfoldController = unfoldControllerProvider.get().orElse(null);
+        mSideUnfoldController = unfoldControllerProvider.get().orElse(null);
         mLogger = logger;
         mMainExecutor = mainExecutor;
         mRecentTasks = recentTasks;
@@ -619,7 +630,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
                 onLayoutSizeChanged(mSplitLayout);
             } else {
                 updateWindowBounds(mSplitLayout, wct);
-                sendOnBoundsChanged();
+                updateUnfoldBounds();
             }
         }
     }
@@ -849,10 +860,6 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         listener.onStagePositionChanged(STAGE_TYPE_MAIN, getMainStagePosition());
         listener.onStagePositionChanged(STAGE_TYPE_SIDE, getSideStagePosition());
         listener.onSplitVisibilityChanged(isSplitScreenVisible());
-        if (mSplitLayout != null) {
-            listener.onSplitBoundsChanged(mSplitLayout.getRootBounds(), getMainStageBounds(),
-                    getSideStageBounds());
-        }
         mSideStage.onSplitScreenListenerRegistered(listener, STAGE_TYPE_SIDE);
         mMainStage.onSplitScreenListenerRegistered(listener, STAGE_TYPE_MAIN);
     }
@@ -865,14 +872,6 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         }
     }
 
-    private void sendOnBoundsChanged() {
-        if (mSplitLayout == null) return;
-        for (int i = mListeners.size() - 1; i >= 0; --i) {
-            mListeners.get(i).onSplitBoundsChanged(mSplitLayout.getRootBounds(),
-                    getMainStageBounds(), getSideStageBounds());
-        }
-    }
-
     private void onStageChildTaskStatusChanged(StageListenerImpl stageListener, int taskId,
             boolean present, boolean visible) {
         int stage;
@@ -936,7 +935,12 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
             final SplitScreen.SplitScreenListener l = mListeners.get(i);
             l.onSplitVisibilityChanged(mDividerVisible);
         }
-        sendOnBoundsChanged();
+
+        if (mMainUnfoldController != null && mSideUnfoldController != null) {
+            mMainUnfoldController.onSplitVisibilityChanged(mDividerVisible);
+            mSideUnfoldController.onSplitVisibilityChanged(mDividerVisible);
+            updateUnfoldBounds();
+        }
     }
 
     @Override
@@ -957,6 +961,11 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
             mDisplayInsetsController.addInsetsChangedListener(mDisplayId, mSplitLayout);
         }
 
+        if (mMainUnfoldController != null && mSideUnfoldController != null) {
+            mMainUnfoldController.init();
+            mSideUnfoldController.init();
+        }
+
         onRootTaskAppeared();
     }
 
@@ -970,8 +979,13 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         mRootTaskInfo = taskInfo;
         if (mSplitLayout != null
                 && mSplitLayout.updateConfiguration(mRootTaskInfo.configuration)
-                && mMainStage.isActive()
-                && !ENABLE_SHELL_TRANSITIONS) {
+                && 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;
+            }
             // Clear the divider remote animating flag as the divider will be re-rendered to apply
             // the new rotation config.
             mIsDividerRemoteAnimating = false;
@@ -1221,7 +1235,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
     public void onLayoutSizeChanged(SplitLayout layout) {
         final WindowContainerTransaction wct = new WindowContainerTransaction();
         updateWindowBounds(layout, wct);
-        sendOnBoundsChanged();
+        updateUnfoldBounds();
         mSyncQueue.queue(wct);
         mSyncQueue.runInSync(t -> {
             setResizingSplits(false /* resizing */);
@@ -1232,6 +1246,15 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         mLogger.logResize(mSplitLayout.getDividerPositionAsFraction());
     }
 
+    private void updateUnfoldBounds() {
+        if (mMainUnfoldController != null && mSideUnfoldController != null) {
+            mMainUnfoldController.onLayoutChanged(getMainStageBounds(), getMainStagePosition(),
+                    isLandscape());
+            mSideUnfoldController.onLayoutChanged(getSideStageBounds(), getSideStagePosition(),
+                    isLandscape());
+        }
+    }
+
     private boolean isLandscape() {
         return mSplitLayout.isLandscape();
     }
@@ -1311,11 +1334,6 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         mDisplayLayout.set(mDisplayController.getDisplayLayout(displayId));
     }
 
-    void updateSurfaces(SurfaceControl.Transaction transaction) {
-        updateSurfaceBounds(mSplitLayout, transaction, /* applyResizingOffset */ false);
-        mSplitLayout.update(transaction);
-    }
-
     private void onDisplayChange(int displayId, int fromRotation, int toRotation,
             @Nullable DisplayAreaInfo newDisplayAreaInfo, WindowContainerTransaction wct) {
         if (!mMainStage.isActive()) return;
@@ -1328,7 +1346,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
             mSplitLayout.updateConfiguration(newDisplayAreaInfo.configuration);
         }
         updateWindowBounds(mSplitLayout, wct);
-        sendOnBoundsChanged();
+        updateUnfoldBounds();
     }
 
     private void onFoldedStateChanged(boolean folded) {
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 86ffacc4eb50..eb513384d822 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
@@ -26,6 +26,7 @@ import static android.app.WindowConfiguration.WINDOWING_MODE_UNDEFINED;
 import static com.android.wm.shell.transition.Transitions.ENABLE_SHELL_TRANSITIONS;
 
 import android.annotation.CallSuper;
+import android.annotation.Nullable;
 import android.app.ActivityManager;
 import android.content.Context;
 import android.graphics.Point;
@@ -94,14 +95,18 @@ class StageTaskListener implements ShellTaskOrganizer.TaskListener {
     // TODO(b/204308910): Extracts SplitDecorManager related code to common package.
     private SplitDecorManager mSplitDecorManager;
 
+    private final StageTaskUnfoldController mStageTaskUnfoldController;
+
     StageTaskListener(Context context, ShellTaskOrganizer taskOrganizer, int displayId,
             StageListenerCallbacks callbacks, SyncTransactionQueue syncQueue,
-            SurfaceSession surfaceSession, IconProvider iconProvider) {
+            SurfaceSession surfaceSession, IconProvider iconProvider,
+            @Nullable StageTaskUnfoldController stageTaskUnfoldController) {
         mContext = context;
         mCallbacks = callbacks;
         mSyncQueue = syncQueue;
         mSurfaceSession = surfaceSession;
         mIconProvider = iconProvider;
+        mStageTaskUnfoldController = stageTaskUnfoldController;
         taskOrganizer.createRootTask(displayId, WINDOWING_MODE_MULTI_WINDOW, this);
     }
 
@@ -202,6 +207,10 @@ class StageTaskListener implements ShellTaskOrganizer.TaskListener {
             throw new IllegalArgumentException(this + "\n Unknown task: " + taskInfo
                     + "\n mRootTaskInfo: " + mRootTaskInfo);
         }
+
+        if (mStageTaskUnfoldController != null) {
+            mStageTaskUnfoldController.onTaskAppeared(taskInfo, leash);
+        }
     }
 
     @Override
@@ -269,6 +278,10 @@ class StageTaskListener implements ShellTaskOrganizer.TaskListener {
             throw new IllegalArgumentException(this + "\n Unknown task: " + taskInfo
                     + "\n mRootTaskInfo: " + mRootTaskInfo);
         }
+
+        if (mStageTaskUnfoldController != null) {
+            mStageTaskUnfoldController.onTaskVanished(taskInfo);
+        }
     }
 
     @Override
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/animation/SplitTaskUnfoldAnimator.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskUnfoldController.java
index 6e10ebe94c5d..59eecb5db136 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/animation/SplitTaskUnfoldAnimator.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskUnfoldController.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2022 The Android Open Source Project
+ * 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.
@@ -14,19 +14,17 @@
  * limitations under the License.
  */
 
-package com.android.wm.shell.unfold.animation;
+package com.android.wm.shell.splitscreen;
 
-import static android.app.WindowConfiguration.WINDOWING_MODE_MULTI_WINDOW;
 import static android.view.Display.DEFAULT_DISPLAY;
 
 import static com.android.wm.shell.common.split.SplitScreenConstants.SPLIT_POSITION_TOP_OR_LEFT;
 import static com.android.wm.shell.common.split.SplitScreenConstants.SPLIT_POSITION_UNDEFINED;
-import static com.android.wm.shell.splitscreen.SplitScreen.STAGE_TYPE_MAIN;
-import static com.android.wm.shell.splitscreen.SplitScreen.STAGE_TYPE_UNDEFINED;
 
 import android.animation.RectEvaluator;
 import android.animation.TypeEvaluator;
-import android.app.TaskInfo;
+import android.annotation.NonNull;
+import android.app.ActivityManager;
 import android.content.Context;
 import android.graphics.Insets;
 import android.graphics.Rect;
@@ -34,130 +32,67 @@ import android.util.SparseArray;
 import android.view.InsetsSource;
 import android.view.InsetsState;
 import android.view.SurfaceControl;
-import android.view.SurfaceControl.Transaction;
 
 import com.android.internal.policy.ScreenDecorationsUtils;
 import com.android.wm.shell.common.DisplayInsetsController;
+import com.android.wm.shell.common.DisplayInsetsController.OnInsetsChangedListener;
+import com.android.wm.shell.common.TransactionPool;
 import com.android.wm.shell.common.split.SplitScreenConstants.SplitPosition;
-import com.android.wm.shell.splitscreen.SplitScreen;
-import com.android.wm.shell.splitscreen.SplitScreen.SplitScreenListener;
-import com.android.wm.shell.splitscreen.SplitScreenController;
-import com.android.wm.shell.unfold.UnfoldAnimationController;
+import com.android.wm.shell.unfold.ShellUnfoldProgressProvider;
+import com.android.wm.shell.unfold.ShellUnfoldProgressProvider.UnfoldListener;
 import com.android.wm.shell.unfold.UnfoldBackgroundController;
 
-import java.util.Optional;
 import java.util.concurrent.Executor;
 
-import dagger.Lazy;
-
 /**
- * This helper class contains logic that calculates scaling and cropping parameters
- * for the folding/unfolding animation. As an input it receives TaskInfo objects and
- * surfaces leashes and as an output it could fill surface transactions with required
- * transformations.
- *
- * This class is used by
- * {@link com.android.wm.shell.unfold.UnfoldTransitionHandler} and
- * {@link UnfoldAnimationController}.
- * They use independent instances of SplitTaskUnfoldAnimator.
+ * Controls transformations of the split screen task surfaces in response
+ * to the unfolding/folding action on foldable devices
  */
-public class SplitTaskUnfoldAnimator implements UnfoldTaskAnimator,
-        DisplayInsetsController.OnInsetsChangedListener, SplitScreenListener {
+public class StageTaskUnfoldController implements UnfoldListener, OnInsetsChangedListener {
 
     private static final TypeEvaluator<Rect> RECT_EVALUATOR = new RectEvaluator(new Rect());
     private static final float CROPPING_START_MARGIN_FRACTION = 0.05f;
 
-    private final Executor mExecutor;
-    private final DisplayInsetsController mDisplayInsetsController;
     private final SparseArray<AnimationContext> mAnimationContextByTaskId = new SparseArray<>();
+    private final ShellUnfoldProgressProvider mUnfoldProgressProvider;
+    private final DisplayInsetsController mDisplayInsetsController;
+    private final UnfoldBackgroundController mBackgroundController;
+    private final Executor mExecutor;
     private final int mExpandedTaskBarHeight;
     private final float mWindowCornerRadiusPx;
-    private final Lazy<Optional<SplitScreenController>> mSplitScreenController;
-    private final UnfoldBackgroundController mUnfoldBackgroundController;
-
-    private final Rect mMainStageBounds = new Rect();
-    private final Rect mSideStageBounds = new Rect();
-    private final Rect mRootStageBounds = new Rect();
+    private final Rect mStageBounds = new Rect();
+    private final TransactionPool mTransactionPool;
 
     private InsetsSource mTaskbarInsetsSource;
-
-    @SplitPosition
-    private int mMainStagePosition = SPLIT_POSITION_UNDEFINED;
-    @SplitPosition
-    private int mSideStagePosition = SPLIT_POSITION_UNDEFINED;
-
-    public SplitTaskUnfoldAnimator(Context context, Executor executor,
-            Lazy<Optional<SplitScreenController>> splitScreenController,
-            UnfoldBackgroundController unfoldBackgroundController,
-            DisplayInsetsController displayInsetsController) {
-        mDisplayInsetsController = displayInsetsController;
+    private boolean mBothStagesVisible;
+
+    public StageTaskUnfoldController(@NonNull Context context,
+            @NonNull TransactionPool transactionPool,
+            @NonNull ShellUnfoldProgressProvider unfoldProgressProvider,
+            @NonNull DisplayInsetsController displayInsetsController,
+            @NonNull UnfoldBackgroundController backgroundController,
+            @NonNull Executor executor) {
+        mUnfoldProgressProvider = unfoldProgressProvider;
+        mTransactionPool = transactionPool;
         mExecutor = executor;
-        mUnfoldBackgroundController = unfoldBackgroundController;
-        mSplitScreenController = splitScreenController;
+        mBackgroundController = backgroundController;
+        mDisplayInsetsController = displayInsetsController;
+        mWindowCornerRadiusPx = ScreenDecorationsUtils.getWindowCornerRadius(context);
         mExpandedTaskBarHeight = context.getResources().getDimensionPixelSize(
                 com.android.internal.R.dimen.taskbar_frame_height);
-        mWindowCornerRadiusPx = ScreenDecorationsUtils.getWindowCornerRadius(context);
-    }
-
-    /** Initializes the animator, this should be called only once */
-    @Override
-    public void init() {
-        mDisplayInsetsController.addInsetsChangedListener(DEFAULT_DISPLAY, this);
-    }
-
-    /**
-     * Starts listening for split-screen changes and gets initial split-screen
-     * layout information through the listener
-     */
-    @Override
-    public void start() {
-        mSplitScreenController.get().get().asSplitScreen()
-                .registerSplitScreenListener(this, mExecutor);
     }
 
     /**
-     * Stops listening for the split-screen layout changes
+     * Initializes the controller, starts listening for the external events
      */
-    @Override
-    public void stop() {
-        mSplitScreenController.get().get().asSplitScreen()
-                .unregisterSplitScreenListener(this);
+    public void init() {
+        mUnfoldProgressProvider.addListener(mExecutor, this);
+        mDisplayInsetsController.addInsetsChangedListener(DEFAULT_DISPLAY, this);
     }
 
     @Override
     public void insetsChanged(InsetsState insetsState) {
         mTaskbarInsetsSource = insetsState.getSource(InsetsState.ITYPE_EXTRA_NAVIGATION_BAR);
-        updateContexts();
-    }
-
-    @Override
-    public void onTaskStageChanged(int taskId, int stage, boolean visible) {
-        final AnimationContext context = mAnimationContextByTaskId.get(taskId);
-        if (context != null) {
-            context.mStageType = stage;
-            context.update();
-        }
-    }
-
-    @Override
-    public void onStagePositionChanged(int stage, int position) {
-        if (stage == STAGE_TYPE_MAIN) {
-            mMainStagePosition = position;
-        } else {
-            mSideStagePosition = position;
-        }
-        updateContexts();
-    }
-
-    @Override
-    public void onSplitBoundsChanged(Rect rootBounds, Rect mainBounds, Rect sideBounds) {
-        mRootStageBounds.set(rootBounds);
-        mMainStageBounds.set(mainBounds);
-        mSideStageBounds.set(sideBounds);
-        updateContexts();
-    }
-
-    private void updateContexts() {
         for (int i = mAnimationContextByTaskId.size() - 1; i >= 0; i--) {
             AnimationContext context = mAnimationContextByTaskId.valueAt(i);
             context.update();
@@ -165,73 +100,44 @@ public class SplitTaskUnfoldAnimator implements UnfoldTaskAnimator,
     }
 
     /**
-     * Register a split task in the animator
-     * @param taskInfo info of the task
-     * @param leash the surface of the task
+     * Called when split screen task appeared
+     * @param taskInfo info for the appeared task
+     * @param leash surface leash for the appeared task
      */
-    @Override
-    public void onTaskAppeared(TaskInfo taskInfo, SurfaceControl leash) {
+    public void onTaskAppeared(ActivityManager.RunningTaskInfo taskInfo, SurfaceControl leash) {
+        // Only handle child task surface here.
+        if (!taskInfo.hasParentTask()) return;
+
         AnimationContext context = new AnimationContext(leash);
         mAnimationContextByTaskId.put(taskInfo.taskId, context);
     }
 
     /**
-     * Unregister the task from the unfold animation
-     * @param taskInfo info of the task
-     */
-    @Override
-    public void onTaskVanished(TaskInfo taskInfo) {
-        mAnimationContextByTaskId.remove(taskInfo.taskId);
-    }
-
-    @Override
-    public boolean isApplicableTask(TaskInfo taskInfo) {
-        return taskInfo.hasParentTask()
-                && taskInfo.isVisible
-                && taskInfo.realActivity != null // to filter out parents created by organizer
-                && taskInfo.getWindowingMode() == WINDOWING_MODE_MULTI_WINDOW;
-    }
-
-    /**
-     * Clear all registered tasks
+     * Called when a split screen task vanished
+     * @param taskInfo info for the vanished task
      */
-    @Override
-    public void clearTasks() {
-        mAnimationContextByTaskId.clear();
-    }
+    public void onTaskVanished(ActivityManager.RunningTaskInfo taskInfo) {
+        if (!taskInfo.hasParentTask()) return;
 
-    /**
-     * Reset transformations of the task that could have been applied by the animator
-     * @param taskInfo task to reset
-     * @param transaction a transaction to write the changes to
-     */
-    @Override
-    public void resetSurface(TaskInfo taskInfo, Transaction transaction) {
         AnimationContext context = mAnimationContextByTaskId.get(taskInfo.taskId);
         if (context != null) {
+            final SurfaceControl.Transaction transaction = mTransactionPool.acquire();
             resetSurface(transaction, context);
+            transaction.apply();
+            mTransactionPool.release(transaction);
         }
+        mAnimationContextByTaskId.remove(taskInfo.taskId);
     }
 
-    /**
-     * Reset all surface transformation that could have been introduced by the animator
-     * @param transaction to write changes to
-     */
     @Override
-    public void resetAllSurfaces(Transaction transaction) {
-        for (int i = mAnimationContextByTaskId.size() - 1; i >= 0; i--) {
-            final AnimationContext context = mAnimationContextByTaskId.valueAt(i);
-            resetSurface(transaction, context);
-        }
-    }
+    public void onStateChangeProgress(float progress) {
+        if (mAnimationContextByTaskId.size() == 0 || !mBothStagesVisible) return;
+
+        final SurfaceControl.Transaction transaction = mTransactionPool.acquire();
+        mBackgroundController.ensureBackground(transaction);
 
-    @Override
-    public void applyAnimationProgress(float progress, Transaction transaction) {
         for (int i = mAnimationContextByTaskId.size() - 1; i >= 0; i--) {
             AnimationContext context = mAnimationContextByTaskId.valueAt(i);
-            if (context.mStageType == STAGE_TYPE_UNDEFINED) {
-                continue;
-            }
 
             context.mCurrentCropRect.set(RECT_EVALUATOR
                     .evaluate(progress, context.mStartCropRect, context.mEndCropRect));
@@ -239,25 +145,53 @@ public class SplitTaskUnfoldAnimator implements UnfoldTaskAnimator,
             transaction.setWindowCrop(context.mLeash, context.mCurrentCropRect)
                     .setCornerRadius(context.mLeash, mWindowCornerRadiusPx);
         }
+
+        transaction.apply();
+
+        mTransactionPool.release(transaction);
     }
 
     @Override
-    public void prepareStartTransaction(Transaction transaction) {
-        mUnfoldBackgroundController.ensureBackground(transaction);
-        mSplitScreenController.get().get().updateSplitScreenSurfaces(transaction);
+    public void onStateChangeFinished() {
+        resetTransformations();
     }
 
-    @Override
-    public void prepareFinishTransaction(Transaction transaction) {
-        mUnfoldBackgroundController.removeBackground(transaction);
+    /**
+     * Called when split screen visibility changes
+     * @param bothStagesVisible true if both stages of the split screen are visible
+     */
+    public void onSplitVisibilityChanged(boolean bothStagesVisible) {
+        mBothStagesVisible = bothStagesVisible;
+        if (!bothStagesVisible) {
+            resetTransformations();
+        }
     }
 
     /**
-     * @return true if there are tasks to animate
+     * Called when split screen stage bounds changed
+     * @param bounds new bounds for this stage
      */
-    @Override
-    public boolean hasActiveTasks() {
-        return mAnimationContextByTaskId.size() > 0;
+    public void onLayoutChanged(Rect bounds, @SplitPosition int splitPosition,
+            boolean isLandscape) {
+        mStageBounds.set(bounds);
+
+        for (int i = mAnimationContextByTaskId.size() - 1; i >= 0; i--) {
+            final AnimationContext context = mAnimationContextByTaskId.valueAt(i);
+            context.update(splitPosition, isLandscape);
+        }
+    }
+
+    private void resetTransformations() {
+        final SurfaceControl.Transaction transaction = mTransactionPool.acquire();
+
+        for (int i = mAnimationContextByTaskId.size() - 1; i >= 0; i--) {
+            final AnimationContext context = mAnimationContextByTaskId.valueAt(i);
+            resetSurface(transaction, context);
+        }
+        mBackgroundController.removeBackground(transaction);
+        transaction.apply();
+
+        mTransactionPool.release(transaction);
     }
 
     private void resetSurface(SurfaceControl.Transaction transaction, AnimationContext context) {
@@ -268,24 +202,26 @@ public class SplitTaskUnfoldAnimator implements UnfoldTaskAnimator,
 
     private class AnimationContext {
         final SurfaceControl mLeash;
-
         final Rect mStartCropRect = new Rect();
         final Rect mEndCropRect = new Rect();
         final Rect mCurrentCropRect = new Rect();
 
-        @SplitScreen.StageType
-        int mStageType = STAGE_TYPE_UNDEFINED;
+        private @SplitPosition int mSplitPosition = SPLIT_POSITION_UNDEFINED;
+        private boolean mIsLandscape = false;
 
         private AnimationContext(SurfaceControl leash) {
-            mLeash = leash;
+            this.mLeash = leash;
             update();
         }
 
-        private void update() {
-            final Rect stageBounds = mStageType == STAGE_TYPE_MAIN
-                    ? mMainStageBounds : mSideStageBounds;
+        private void update(@SplitPosition int splitPosition, boolean isLandscape) {
+            this.mSplitPosition = splitPosition;
+            this.mIsLandscape = isLandscape;
+            update();
+        }
 
-            mStartCropRect.set(stageBounds);
+        private void update() {
+            mStartCropRect.set(mStageBounds);
 
             boolean taskbarExpanded = isTaskbarExpanded();
             if (taskbarExpanded) {
@@ -303,8 +239,7 @@ public class SplitTaskUnfoldAnimator implements UnfoldTaskAnimator,
 
             // Sides adjacent to split bar or task bar are not be animated.
             Insets margins;
-            final boolean isLandscape = mRootStageBounds.width() > mRootStageBounds.height();
-            if (isLandscape) { // Left and right splits.
+            if (mIsLandscape) { // Left and right splits.
                 margins = getLandscapeMargins(margin, taskbarExpanded);
             } else { // Top and bottom splits.
                 margins = getPortraitMargins(margin, taskbarExpanded);
@@ -316,9 +251,7 @@ public class SplitTaskUnfoldAnimator implements UnfoldTaskAnimator,
             int left = margin;
             int right = margin;
             int bottom = taskbarExpanded ? 0 : margin; // Taskbar margin.
-            final int splitPosition = mStageType == STAGE_TYPE_MAIN
-                    ? mMainStagePosition : mSideStagePosition;
-            if (splitPosition == SPLIT_POSITION_TOP_OR_LEFT) {
+            if (mSplitPosition == SPLIT_POSITION_TOP_OR_LEFT) {
                 right = 0; // Divider margin.
             } else {
                 left = 0; // Divider margin.
@@ -329,9 +262,7 @@ public class SplitTaskUnfoldAnimator implements UnfoldTaskAnimator,
         private Insets getPortraitMargins(int margin, boolean taskbarExpanded) {
             int bottom = margin;
             int top = margin;
-            final int splitPosition = mStageType == STAGE_TYPE_MAIN
-                    ? mMainStagePosition : mSideStagePosition;
-            if (splitPosition == SPLIT_POSITION_TOP_OR_LEFT) {
+            if (mSplitPosition == SPLIT_POSITION_TOP_OR_LEFT) {
                 bottom = 0; // Divider margin.
             } else { // Bottom split.
                 top = 0; // Divider margin.
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/UnfoldAnimationController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/UnfoldAnimationController.java
deleted file mode 100644
index 530d47416665..000000000000
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/UnfoldAnimationController.java
+++ /dev/null
@@ -1,219 +0,0 @@
-/*
- * 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.wm.shell.unfold;
-
-import android.annotation.NonNull;
-import android.app.ActivityManager.RunningTaskInfo;
-import android.app.TaskInfo;
-import android.util.SparseArray;
-import android.view.SurfaceControl;
-
-import com.android.wm.shell.common.TransactionPool;
-import com.android.wm.shell.unfold.ShellUnfoldProgressProvider.UnfoldListener;
-import com.android.wm.shell.unfold.animation.UnfoldTaskAnimator;
-
-import java.util.List;
-import java.util.Optional;
-import java.util.concurrent.Executor;
-
-import dagger.Lazy;
-
-/**
- * Manages fold/unfold animations of tasks on foldable devices.
- * When folding or unfolding a foldable device we play animations that
- * transform task cropping/scaling/rounded corners.
- *
- * This controller manages:
- *  1) Folding/unfolding when Shell transitions disabled
- *  2) Folding when Shell transitions enabled, unfolding is managed by
- *    {@link com.android.wm.shell.unfold.UnfoldTransitionHandler}
- */
-public class UnfoldAnimationController implements UnfoldListener {
-
-    private final ShellUnfoldProgressProvider mUnfoldProgressProvider;
-    private final Executor mExecutor;
-    private final TransactionPool mTransactionPool;
-    private final List<UnfoldTaskAnimator> mAnimators;
-    private final Lazy<Optional<UnfoldTransitionHandler>> mUnfoldTransitionHandler;
-
-    private final SparseArray<SurfaceControl> mTaskSurfaces = new SparseArray<>();
-    private final SparseArray<UnfoldTaskAnimator> mAnimatorsByTaskId = new SparseArray<>();
-
-    public UnfoldAnimationController(@NonNull TransactionPool transactionPool,
-            @NonNull ShellUnfoldProgressProvider unfoldProgressProvider,
-            @NonNull List<UnfoldTaskAnimator> animators,
-            @NonNull Lazy<Optional<UnfoldTransitionHandler>> unfoldTransitionHandler,
-            @NonNull Executor executor) {
-        mUnfoldProgressProvider = unfoldProgressProvider;
-        mUnfoldTransitionHandler = unfoldTransitionHandler;
-        mTransactionPool = transactionPool;
-        mExecutor = executor;
-        mAnimators = animators;
-    }
-
-    /**
-     * Initializes the controller, starts listening for the external events
-     */
-    public void init() {
-        mUnfoldProgressProvider.addListener(mExecutor, this);
-
-        for (int i = 0; i < mAnimators.size(); i++) {
-            final UnfoldTaskAnimator animator = mAnimators.get(i);
-            animator.init();
-            animator.start();
-        }
-    }
-
-    /**
-     * Called when a task appeared
-     * @param taskInfo info for the appeared task
-     * @param leash surface leash for the appeared task
-     */
-    public void onTaskAppeared(RunningTaskInfo taskInfo, SurfaceControl leash) {
-        mTaskSurfaces.put(taskInfo.taskId, leash);
-
-        // Find the first matching animator
-        for (int i = 0; i < mAnimators.size(); i++) {
-            final UnfoldTaskAnimator animator = mAnimators.get(i);
-            if (animator.isApplicableTask(taskInfo)) {
-                mAnimatorsByTaskId.put(taskInfo.taskId, animator);
-                animator.onTaskAppeared(taskInfo, leash);
-                break;
-            }
-        }
-    }
-
-    /**
-     * Called when task info changed
-     * @param taskInfo info for the changed task
-     */
-    public void onTaskInfoChanged(RunningTaskInfo taskInfo) {
-        final UnfoldTaskAnimator animator = mAnimatorsByTaskId.get(taskInfo.taskId);
-        final boolean isCurrentlyApplicable = animator != null;
-
-        if (isCurrentlyApplicable) {
-            final boolean isApplicable = animator.isApplicableTask(taskInfo);
-            if (isApplicable) {
-                // Still applicable, send update
-                animator.onTaskChanged(taskInfo);
-            } else {
-                // Became inapplicable
-                resetTask(animator, taskInfo);
-                animator.onTaskVanished(taskInfo);
-                mAnimatorsByTaskId.remove(taskInfo.taskId);
-            }
-        } else {
-            // Find the first matching animator
-            for (int i = 0; i < mAnimators.size(); i++) {
-                final UnfoldTaskAnimator currentAnimator = mAnimators.get(i);
-                if (currentAnimator.isApplicableTask(taskInfo)) {
-                    // Became applicable
-                    mAnimatorsByTaskId.put(taskInfo.taskId, currentAnimator);
-
-                    SurfaceControl leash = mTaskSurfaces.get(taskInfo.taskId);
-                    currentAnimator.onTaskAppeared(taskInfo, leash);
-                    break;
-                }
-            }
-        }
-    }
-
-    /**
-     * Called when a task vanished
-     * @param taskInfo info for the vanished task
-     */
-    public void onTaskVanished(RunningTaskInfo taskInfo) {
-        mTaskSurfaces.remove(taskInfo.taskId);
-
-        final UnfoldTaskAnimator animator = mAnimatorsByTaskId.get(taskInfo.taskId);
-        final boolean isCurrentlyApplicable = animator != null;
-
-        if (isCurrentlyApplicable) {
-            resetTask(animator, taskInfo);
-            animator.onTaskVanished(taskInfo);
-            mAnimatorsByTaskId.remove(taskInfo.taskId);
-        }
-    }
-
-    @Override
-    public void onStateChangeStarted() {
-        if (mUnfoldTransitionHandler.get().get().willHandleTransition()) {
-            return;
-        }
-
-        SurfaceControl.Transaction transaction = null;
-        for (int i = 0; i < mAnimators.size(); i++) {
-            final UnfoldTaskAnimator animator = mAnimators.get(i);
-            if (animator.hasActiveTasks()) {
-                if (transaction == null) transaction = mTransactionPool.acquire();
-                animator.prepareStartTransaction(transaction);
-            }
-        }
-
-        if (transaction != null) {
-            transaction.apply();
-            mTransactionPool.release(transaction);
-        }
-    }
-
-    @Override
-    public void onStateChangeProgress(float progress) {
-        if (mUnfoldTransitionHandler.get().get().willHandleTransition()) {
-            return;
-        }
-
-        SurfaceControl.Transaction transaction = null;
-        for (int i = 0; i < mAnimators.size(); i++) {
-            final UnfoldTaskAnimator animator = mAnimators.get(i);
-            if (animator.hasActiveTasks()) {
-                if (transaction == null) transaction = mTransactionPool.acquire();
-                animator.applyAnimationProgress(progress, transaction);
-            }
-        }
-
-        if (transaction != null) {
-            transaction.apply();
-            mTransactionPool.release(transaction);
-        }
-    }
-
-    @Override
-    public void onStateChangeFinished() {
-        if (mUnfoldTransitionHandler.get().get().willHandleTransition()) {
-            return;
-        }
-
-        final SurfaceControl.Transaction transaction = mTransactionPool.acquire();
-
-        for (int i = 0; i < mAnimators.size(); i++) {
-            final UnfoldTaskAnimator animator = mAnimators.get(i);
-            animator.resetAllSurfaces(transaction);
-            animator.prepareFinishTransaction(transaction);
-        }
-
-        transaction.apply();
-
-        mTransactionPool.release(transaction);
-    }
-
-    private void resetTask(UnfoldTaskAnimator animator, TaskInfo taskInfo) {
-        final SurfaceControl.Transaction transaction = mTransactionPool.acquire();
-        animator.resetSurface(taskInfo, transaction);
-        transaction.apply();
-        mTransactionPool.release(transaction);
-    }
-}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/UnfoldTransitionHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/UnfoldTransitionHandler.java
index 9bf32faa12bd..8e45e7d36b86 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/UnfoldTransitionHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/UnfoldTransitionHandler.java
@@ -16,6 +16,8 @@
 
 package com.android.wm.shell.unfold;
 
+import static android.app.WindowConfiguration.ACTIVITY_TYPE_HOME;
+import static android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN;
 import static android.view.WindowManager.TRANSIT_CHANGE;
 
 import android.os.IBinder;
@@ -33,22 +35,19 @@ import com.android.wm.shell.transition.Transitions.TransitionFinishCallback;
 import com.android.wm.shell.transition.Transitions.TransitionHandler;
 import com.android.wm.shell.unfold.ShellUnfoldProgressProvider.UnfoldListener;
 import com.android.wm.shell.unfold.animation.FullscreenUnfoldTaskAnimator;
-import com.android.wm.shell.unfold.animation.SplitTaskUnfoldAnimator;
-import com.android.wm.shell.unfold.animation.UnfoldTaskAnimator;
 
-import java.util.ArrayList;
-import java.util.List;
 import java.util.concurrent.Executor;
 
 /**
  * Transition handler that is responsible for animating app surfaces when unfolding of foldable
  * devices. It does not handle the folding animation, which is done in
- * {@link UnfoldAnimationController}.
+ * {@link com.android.wm.shell.fullscreen.FullscreenUnfoldController}.
  */
 public class UnfoldTransitionHandler implements TransitionHandler, UnfoldListener {
 
     private final ShellUnfoldProgressProvider mUnfoldProgressProvider;
     private final Transitions mTransitions;
+    private final UnfoldBackgroundController mUnfoldBackgroundController;
     private final Executor mExecutor;
     private final TransactionPool mTransactionPool;
 
@@ -57,26 +56,22 @@ public class UnfoldTransitionHandler implements TransitionHandler, UnfoldListene
     @Nullable
     private IBinder mTransition;
 
-    private final List<UnfoldTaskAnimator> mAnimators = new ArrayList<>();
+    private final FullscreenUnfoldTaskAnimator mFullscreenAnimator;
 
     public UnfoldTransitionHandler(ShellUnfoldProgressProvider unfoldProgressProvider,
-            FullscreenUnfoldTaskAnimator fullscreenUnfoldAnimator,
-            SplitTaskUnfoldAnimator splitUnfoldTaskAnimator,
-            TransactionPool transactionPool,
+            FullscreenUnfoldTaskAnimator animator, TransactionPool transactionPool,
+            UnfoldBackgroundController unfoldBackgroundController,
             Executor executor, Transitions transitions) {
         mUnfoldProgressProvider = unfoldProgressProvider;
+        mFullscreenAnimator = animator;
         mTransactionPool = transactionPool;
+        mUnfoldBackgroundController = unfoldBackgroundController;
         mExecutor = executor;
         mTransitions = transitions;
-
-        mAnimators.add(splitUnfoldTaskAnimator);
-        mAnimators.add(fullscreenUnfoldAnimator);
     }
 
     public void init() {
-        for (int i = 0; i < mAnimators.size(); i++) {
-            mAnimators.get(i).init();
-        }
+        mFullscreenAnimator.init();
         mTransitions.addHandler(this);
         mUnfoldProgressProvider.addListener(mExecutor, this);
     }
@@ -88,67 +83,44 @@ public class UnfoldTransitionHandler implements TransitionHandler, UnfoldListene
             @NonNull TransitionFinishCallback finishCallback) {
         if (transition != mTransition) return false;
 
-        for (int i = 0; i < mAnimators.size(); i++) {
-            final UnfoldTaskAnimator animator = mAnimators.get(i);
-            animator.clearTasks();
-
-            info.getChanges().forEach(change -> {
-                if (change.getTaskInfo() != null
-                        && change.getMode() == TRANSIT_CHANGE
-                        && animator.isApplicableTask(change.getTaskInfo())) {
-                    animator.onTaskAppeared(change.getTaskInfo(), change.getLeash());
-                }
-            });
-
-            if (animator.hasActiveTasks()) {
-                animator.prepareStartTransaction(startTransaction);
-                animator.prepareFinishTransaction(finishTransaction);
-                animator.start();
+        mUnfoldBackgroundController.ensureBackground(startTransaction);
+        startTransaction.apply();
+
+        mFullscreenAnimator.clearTasks();
+        info.getChanges().forEach(change -> {
+            final boolean allowedToAnimate = change.getTaskInfo() != null
+                    && change.getTaskInfo().isVisible()
+                    && change.getTaskInfo().getWindowingMode() == WINDOWING_MODE_FULLSCREEN
+                    && change.getTaskInfo().getActivityType() != ACTIVITY_TYPE_HOME
+                    && change.getMode() == TRANSIT_CHANGE;
+
+            if (allowedToAnimate) {
+                mFullscreenAnimator.addTask(change.getTaskInfo(), change.getLeash());
             }
-        }
+        });
 
-        startTransaction.apply();
+        mFullscreenAnimator.resetAllSurfaces(finishTransaction);
+        mUnfoldBackgroundController.removeBackground(finishTransaction);
         mFinishCallback = finishCallback;
         return true;
     }
 
     @Override
     public void onStateChangeProgress(float progress) {
-        if (mTransition == null) return;
-
-        SurfaceControl.Transaction transaction = null;
-
-        for (int i = 0; i < mAnimators.size(); i++) {
-            final UnfoldTaskAnimator animator = mAnimators.get(i);
-
-            if (animator.hasActiveTasks()) {
-                if (transaction == null) {
-                    transaction = mTransactionPool.acquire();
-                }
-
-                animator.applyAnimationProgress(progress, transaction);
-            }
-        }
-
-        if (transaction != null) {
-            transaction.apply();
-            mTransactionPool.release(transaction);
-        }
+        final SurfaceControl.Transaction transaction = mTransactionPool.acquire();
+        mFullscreenAnimator.applyAnimationProgress(progress, transaction);
+        transaction.apply();
+        mTransactionPool.release(transaction);
     }
 
     @Override
     public void onStateChangeFinished() {
-        if (mFinishCallback == null) return;
-
-        for (int i = 0; i < mAnimators.size(); i++) {
-            final UnfoldTaskAnimator animator = mAnimators.get(i);
-            animator.clearTasks();
-            animator.stop();
+        if (mFinishCallback != null) {
+            mFinishCallback.onTransitionFinished(null, null);
+            mFinishCallback = null;
+            mTransition = null;
+            mFullscreenAnimator.clearTasks();
         }
-
-        mFinishCallback.onTransitionFinished(null, null);
-        mFinishCallback = null;
-        mTransition = null;
     }
 
     @Nullable
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/animation/FullscreenUnfoldTaskAnimator.java b/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/animation/FullscreenUnfoldTaskAnimator.java
index eab82f00e962..6ec5512c22ec 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/animation/FullscreenUnfoldTaskAnimator.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/animation/FullscreenUnfoldTaskAnimator.java
@@ -16,14 +16,11 @@
 
 package com.android.wm.shell.unfold.animation;
 
-import static android.app.WindowConfiguration.ACTIVITY_TYPE_HOME;
-import static android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN;
 import static android.util.MathUtils.lerp;
 import static android.view.Display.DEFAULT_DISPLAY;
 
 import android.animation.RectEvaluator;
 import android.animation.TypeEvaluator;
-import android.annotation.NonNull;
 import android.app.TaskInfo;
 import android.content.Context;
 import android.graphics.Matrix;
@@ -36,8 +33,6 @@ import android.view.SurfaceControl.Transaction;
 
 import com.android.internal.policy.ScreenDecorationsUtils;
 import com.android.wm.shell.common.DisplayInsetsController;
-import com.android.wm.shell.unfold.UnfoldAnimationController;
-import com.android.wm.shell.unfold.UnfoldBackgroundController;
 
 /**
  * This helper class contains logic that calculates scaling and cropping parameters
@@ -47,10 +42,10 @@ import com.android.wm.shell.unfold.UnfoldBackgroundController;
  *
  * This class is used by
  * {@link com.android.wm.shell.unfold.UnfoldTransitionHandler} and
- * {@link UnfoldAnimationController}. They use independent
+ * {@link com.android.wm.shell.fullscreen.FullscreenUnfoldController}. They use independent
  * instances of FullscreenUnfoldTaskAnimator.
  */
-public class FullscreenUnfoldTaskAnimator implements UnfoldTaskAnimator,
+public class FullscreenUnfoldTaskAnimator implements
         DisplayInsetsController.OnInsetsChangedListener {
 
     private static final float[] FLOAT_9 = new float[9];
@@ -65,15 +60,12 @@ public class FullscreenUnfoldTaskAnimator implements UnfoldTaskAnimator,
     private final int mExpandedTaskBarHeight;
     private final float mWindowCornerRadiusPx;
     private final DisplayInsetsController mDisplayInsetsController;
-    private final UnfoldBackgroundController mBackgroundController;
 
     private InsetsSource mTaskbarInsetsSource;
 
     public FullscreenUnfoldTaskAnimator(Context context,
-            @NonNull UnfoldBackgroundController backgroundController,
             DisplayInsetsController displayInsetsController) {
         mDisplayInsetsController = displayInsetsController;
-        mBackgroundController = backgroundController;
         mExpandedTaskBarHeight = context.getResources().getDimensionPixelSize(
                 com.android.internal.R.dimen.taskbar_frame_height);
         mWindowCornerRadiusPx = ScreenDecorationsUtils.getWindowCornerRadius(context);
@@ -96,32 +88,27 @@ public class FullscreenUnfoldTaskAnimator implements UnfoldTaskAnimator,
         return mAnimationContextByTaskId.size() > 0;
     }
 
-    @Override
-    public void onTaskAppeared(TaskInfo taskInfo, SurfaceControl leash) {
+    public void addTask(TaskInfo taskInfo, SurfaceControl leash) {
         AnimationContext animationContext = new AnimationContext(leash, mTaskbarInsetsSource,
                 taskInfo);
         mAnimationContextByTaskId.put(taskInfo.taskId, animationContext);
     }
 
-    @Override
-    public void onTaskChanged(TaskInfo taskInfo) {
+    public void onTaskInfoChanged(TaskInfo taskInfo) {
         AnimationContext animationContext = mAnimationContextByTaskId.get(taskInfo.taskId);
         if (animationContext != null) {
             animationContext.update(mTaskbarInsetsSource, taskInfo);
         }
     }
 
-    @Override
-    public void onTaskVanished(TaskInfo taskInfo) {
+    public void removeTask(TaskInfo taskInfo) {
         mAnimationContextByTaskId.remove(taskInfo.taskId);
     }
 
-    @Override
     public void clearTasks() {
         mAnimationContextByTaskId.clear();
     }
 
-    @Override
     public void resetSurface(TaskInfo taskInfo, Transaction transaction) {
         final AnimationContext context = mAnimationContextByTaskId.get(taskInfo.taskId);
         if (context != null) {
@@ -129,7 +116,6 @@ public class FullscreenUnfoldTaskAnimator implements UnfoldTaskAnimator,
         }
     }
 
-    @Override
     public void applyAnimationProgress(float progress, Transaction transaction) {
         if (mAnimationContextByTaskId.size() == 0) return;
 
@@ -146,29 +132,11 @@ public class FullscreenUnfoldTaskAnimator implements UnfoldTaskAnimator,
             transaction.setWindowCrop(context.mLeash, context.mCurrentCropRect)
                     .setMatrix(context.mLeash, context.mMatrix, FLOAT_9)
                     .setCornerRadius(context.mLeash, mWindowCornerRadiusPx)
-                    .show(context.mLeash);
+                    .show(context.mLeash)
+            ;
         }
     }
 
-    @Override
-    public void prepareStartTransaction(Transaction transaction) {
-        mBackgroundController.ensureBackground(transaction);
-    }
-
-    @Override
-    public void prepareFinishTransaction(Transaction transaction) {
-        mBackgroundController.removeBackground(transaction);
-    }
-
-    @Override
-    public boolean isApplicableTask(TaskInfo taskInfo) {
-        return taskInfo != null && taskInfo.isVisible()
-                && taskInfo.realActivity != null // to filter out parents created by organizer
-                && taskInfo.getWindowingMode() == WINDOWING_MODE_FULLSCREEN
-                && taskInfo.getActivityType() != ACTIVITY_TYPE_HOME;
-    }
-
-    @Override
     public void resetAllSurfaces(Transaction transaction) {
         for (int i = mAnimationContextByTaskId.size() - 1; i >= 0; i--) {
             final AnimationContext context = mAnimationContextByTaskId.valueAt(i);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/animation/UnfoldTaskAnimator.java b/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/animation/UnfoldTaskAnimator.java
deleted file mode 100644
index e1e366301b46..000000000000
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/animation/UnfoldTaskAnimator.java
+++ /dev/null
@@ -1,117 +0,0 @@
-/*
- * 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.wm.shell.unfold.animation;
-
-import android.app.TaskInfo;
-import android.view.SurfaceControl;
-import android.view.SurfaceControl.Transaction;
-
-/**
- * Interface for classes that handle animations of tasks when folding or unfolding
- * foldable devices.
- */
-public interface UnfoldTaskAnimator {
-    /**
-     * Initializes the animator, this should be called once in the lifetime of the animator
-     */
-    default void init() {}
-
-    /**
-     * Starts the animator, it might start listening for some events from the system.
-     * Applying animation should be done only when animator is started.
-     * Animator could be started/stopped several times.
-     */
-    default void start() {}
-
-    /**
-     * Stops the animator, it could unsubscribe from system events.
-     */
-    default void stop() {}
-
-    /**
-     * If this method returns true then task updates will be propagated to
-     * the animator using the onTaskAppeared/Changed/Vanished callbacks.
-     * @return true if this task should be animated by this animator
-     */
-    default boolean isApplicableTask(TaskInfo taskInfo) {
-        return false;
-    }
-
-    /**
-     * Called whenever a task applicable to this animator appeared
-     * (isApplicableTask returns true for this task)
-     *
-     * @param taskInfo info of the appeared task
-     * @param leash surface of the task
-     */
-    default void onTaskAppeared(TaskInfo taskInfo, SurfaceControl leash) {}
-
-    /**
-     * Called whenever a task applicable to this animator changed
-     * @param taskInfo info of the changed task
-     */
-    default void onTaskChanged(TaskInfo taskInfo) {}
-
-    /**
-     * Called whenever a task applicable to this animator vanished
-     * @param taskInfo info of the vanished task
-     */
-    default void onTaskVanished(TaskInfo taskInfo) {}
-
-    /**
-     * @return true if there tasks that could be potentially animated
-     */
-    default boolean hasActiveTasks() {
-        return false;
-    }
-
-    /**
-     * Clears all registered tasks in the animator
-     */
-    default void clearTasks() {}
-
-    /**
-     * Apply task surfaces transformations based on the current unfold progress
-     * @param progress unfold transition progress
-     * @param transaction to write changes to
-     */
-    default void applyAnimationProgress(float progress, Transaction transaction) {}
-
-    /**
-     * Apply task surfaces transformations that should be set before starting the animation
-     * @param transaction to write changes to
-     */
-    default void prepareStartTransaction(Transaction transaction) {}
-
-    /**
-     * Apply task surfaces transformations that should be set after finishing the animation
-     * @param transaction to write changes to
-     */
-    default void prepareFinishTransaction(Transaction transaction) {}
-
-    /**
-     * Resets task surface to its initial transformation
-     * @param transaction to write changes to
-     */
-    default void resetSurface(TaskInfo taskInfo, Transaction transaction) {}
-
-    /**
-     * Resets all task surfaces to their initial transformations
-     * @param transaction to write changes to
-     */
-    default void resetAllSurfaces(Transaction transaction) {}
-}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/qualifier/UnfoldShellTransition.java b/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/qualifier/UnfoldShellTransition.java
deleted file mode 100644
index 4c868305dcdb..000000000000
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/qualifier/UnfoldShellTransition.java
+++ /dev/null
@@ -1,29 +0,0 @@
-/*
- * 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.wm.shell.unfold.qualifier;
-
-import java.lang.annotation.Retention;
-import java.lang.annotation.RetentionPolicy;
-
-import javax.inject.Qualifier;
-
-/**
- * Indicates that this class is used for the shell unfold transition
- */
-@Qualifier
-@Retention(RetentionPolicy.RUNTIME)
-public @interface UnfoldShellTransition {}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/qualifier/UnfoldTransition.java b/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/qualifier/UnfoldTransition.java
deleted file mode 100644
index 4d2b3e6f899b..000000000000
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/qualifier/UnfoldTransition.java
+++ /dev/null
@@ -1,30 +0,0 @@
-/*
- * 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.wm.shell.unfold.qualifier;
-
-import java.lang.annotation.Retention;
-import java.lang.annotation.RetentionPolicy;
-
-import javax.inject.Qualifier;
-
-/**
- * Indicates that this class is used for unfold transition implemented
- * without using Shell transitions
- */
-@Qualifier
-@Retention(RetentionPolicy.RUNTIME)
-public @interface UnfoldTransition {}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellTaskOrganizerTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellTaskOrganizerTests.java
index 0b53c4069c3f..a6caefe6d3e7 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellTaskOrganizerTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellTaskOrganizerTests.java
@@ -133,7 +133,7 @@ public class ShellTaskOrganizerTests {
                     .when(mTaskOrganizerController).registerTaskOrganizer(any());
         } catch (RemoteException e) {}
         mOrganizer = spy(new ShellTaskOrganizer(mTaskOrganizerController, mTestExecutor, mContext,
-                mCompatUI, Optional.empty(), Optional.empty()));
+                mCompatUI, Optional.empty()));
     }
 
     @Test
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/fullscreen/FullscreenTaskListenerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/fullscreen/FullscreenTaskListenerTest.java
new file mode 100644
index 000000000000..4523e2c9cba5
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/fullscreen/FullscreenTaskListenerTest.java
@@ -0,0 +1,166 @@
+/*
+ * 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.fullscreen;
+
+import static android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD;
+
+import static org.junit.Assume.assumeFalse;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.inOrder;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+import static org.mockito.Mockito.when;
+
+import android.app.ActivityManager.RunningTaskInfo;
+import android.app.WindowConfiguration;
+import android.content.res.Configuration;
+import android.graphics.Point;
+import android.os.SystemProperties;
+import android.view.SurfaceControl;
+
+import androidx.test.filters.SmallTest;
+
+import com.android.wm.shell.common.SyncTransactionQueue;
+import com.android.wm.shell.recents.RecentTasksController;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.InOrder;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+import java.util.Optional;
+
+@SmallTest
+public class FullscreenTaskListenerTest {
+    private static final boolean ENABLE_SHELL_TRANSITIONS =
+            SystemProperties.getBoolean("persist.wm.debug.shell_transit", false);
+
+    @Mock
+    private SyncTransactionQueue mSyncQueue;
+    @Mock
+    private FullscreenUnfoldController mUnfoldController;
+    @Mock
+    private RecentTasksController mRecentTasksController;
+    @Mock
+    private SurfaceControl mSurfaceControl;
+
+    private Optional<FullscreenUnfoldController> mFullscreenUnfoldController;
+
+    private FullscreenTaskListener mListener;
+
+    @Before
+    public void setup() {
+        MockitoAnnotations.initMocks(this);
+        mFullscreenUnfoldController = Optional.of(mUnfoldController);
+        mListener = new FullscreenTaskListener(mSyncQueue, mFullscreenUnfoldController,
+                Optional.empty());
+    }
+
+    @Test
+    public void testAnimatableTaskAppeared_notifiesUnfoldController() {
+        assumeFalse(ENABLE_SHELL_TRANSITIONS);
+        RunningTaskInfo info = createTaskInfo(/* visible */ true, /* taskId */ 0);
+
+        mListener.onTaskAppeared(info, mSurfaceControl);
+
+        verify(mUnfoldController).onTaskAppeared(eq(info), any());
+    }
+
+    @Test
+    public void testMultipleAnimatableTasksAppeared_notifiesUnfoldController() {
+        assumeFalse(ENABLE_SHELL_TRANSITIONS);
+        RunningTaskInfo animatable1 = createTaskInfo(/* visible */ true, /* taskId */ 0);
+        RunningTaskInfo animatable2 = createTaskInfo(/* visible */ true, /* taskId */ 1);
+
+        mListener.onTaskAppeared(animatable1, mSurfaceControl);
+        mListener.onTaskAppeared(animatable2, mSurfaceControl);
+
+        InOrder order = inOrder(mUnfoldController);
+        order.verify(mUnfoldController).onTaskAppeared(eq(animatable1), any());
+        order.verify(mUnfoldController).onTaskAppeared(eq(animatable2), any());
+    }
+
+    @Test
+    public void testNonAnimatableTaskAppeared_doesNotNotifyUnfoldController() {
+        assumeFalse(ENABLE_SHELL_TRANSITIONS);
+        RunningTaskInfo info = createTaskInfo(/* visible */ false, /* taskId */ 0);
+
+        mListener.onTaskAppeared(info, mSurfaceControl);
+
+        verifyNoMoreInteractions(mUnfoldController);
+    }
+
+    @Test
+    public void testNonAnimatableTaskChanged_doesNotNotifyUnfoldController() {
+        assumeFalse(ENABLE_SHELL_TRANSITIONS);
+        RunningTaskInfo info = createTaskInfo(/* visible */ false, /* taskId */ 0);
+        mListener.onTaskAppeared(info, mSurfaceControl);
+
+        mListener.onTaskInfoChanged(info);
+
+        verifyNoMoreInteractions(mUnfoldController);
+    }
+
+    @Test
+    public void testNonAnimatableTaskVanished_doesNotNotifyUnfoldController() {
+        assumeFalse(ENABLE_SHELL_TRANSITIONS);
+        RunningTaskInfo info = createTaskInfo(/* visible */ false, /* taskId */ 0);
+        mListener.onTaskAppeared(info, mSurfaceControl);
+
+        mListener.onTaskVanished(info);
+
+        verifyNoMoreInteractions(mUnfoldController);
+    }
+
+    @Test
+    public void testAnimatableTaskBecameInactive_notifiesUnfoldController() {
+        assumeFalse(ENABLE_SHELL_TRANSITIONS);
+        RunningTaskInfo animatableTask = createTaskInfo(/* visible */ true, /* taskId */ 0);
+        mListener.onTaskAppeared(animatableTask, mSurfaceControl);
+        RunningTaskInfo notAnimatableTask = createTaskInfo(/* visible */ false, /* taskId */ 0);
+
+        mListener.onTaskInfoChanged(notAnimatableTask);
+
+        verify(mUnfoldController).onTaskVanished(eq(notAnimatableTask));
+    }
+
+    @Test
+    public void testAnimatableTaskVanished_notifiesUnfoldController() {
+        assumeFalse(ENABLE_SHELL_TRANSITIONS);
+        RunningTaskInfo taskInfo = createTaskInfo(/* visible */ true, /* taskId */ 0);
+        mListener.onTaskAppeared(taskInfo, mSurfaceControl);
+
+        mListener.onTaskVanished(taskInfo);
+
+        verify(mUnfoldController).onTaskVanished(eq(taskInfo));
+    }
+
+    private RunningTaskInfo createTaskInfo(boolean visible, int taskId) {
+        final RunningTaskInfo info = spy(new RunningTaskInfo());
+        info.isVisible = visible;
+        info.positionInParent = new Point();
+        when(info.getWindowingMode()).thenReturn(WindowConfiguration.WINDOWING_MODE_FULLSCREEN);
+        final Configuration configuration = new Configuration();
+        configuration.windowConfiguration.setActivityType(ACTIVITY_TYPE_STANDARD);
+        when(info.getConfiguration()).thenReturn(configuration);
+        info.taskId = taskId;
+        return info;
+    }
+}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/kidsmode/KidsModeTaskOrganizerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/kidsmode/KidsModeTaskOrganizerTest.java
index 1eadeed7cd3b..440a6f8fb59a 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/kidsmode/KidsModeTaskOrganizerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/kidsmode/KidsModeTaskOrganizerTest.java
@@ -88,7 +88,7 @@ public class KidsModeTaskOrganizerTest {
         // NOTE: KidsModeTaskOrganizer should have a null CompatUIController.
         mOrganizer = spy(new KidsModeTaskOrganizer(mTaskOrganizerController, mTestExecutor,
                 mHandler, mContext, mSyncTransactionQueue, mDisplayController,
-                mDisplayInsetsController, Optional.empty(), Optional.empty(), mObserver));
+                mDisplayInsetsController, Optional.empty(), mObserver));
         mOrganizer.initialize(mStartingWindowController);
         doReturn(mTransaction).when(mOrganizer).getWindowContainerTransaction();
         doReturn(new InsetsState()).when(mDisplayController).getInsetsState(DEFAULT_DISPLAY);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/recents/RecentTasksControllerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/recents/RecentTasksControllerTest.java
index 9191b1564de2..50f6bd7b4927 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/recents/RecentTasksControllerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/recents/RecentTasksControllerTest.java
@@ -80,7 +80,7 @@ public class RecentTasksControllerTest extends ShellTestCase {
         mRecentTasksController = spy(new RecentTasksController(mContext, mTaskStackListener,
                 mMainExecutor));
         mShellTaskOrganizer = new ShellTaskOrganizer(mMainExecutor, mContext,
-                null /* sizeCompatUI */, Optional.empty(), Optional.of(mRecentTasksController));
+                null /* sizeCompatUI */, Optional.of(mRecentTasksController));
     }
 
     @Test
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/MainStageTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/MainStageTests.java
index 68cb57c14d8c..0639ad5d0a62 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/MainStageTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/MainStageTests.java
@@ -61,7 +61,7 @@ public class MainStageTests extends ShellTestCase {
         MockitoAnnotations.initMocks(this);
         mRootTaskInfo = new TestRunningTaskInfoBuilder().build();
         mMainStage = new MainStage(mContext, mTaskOrganizer, DEFAULT_DISPLAY, mCallbacks,
-                mSyncQueue, mSurfaceSession, mIconProvider);
+                mSyncQueue, mSurfaceSession, mIconProvider, null);
         mMainStage.onTaskAppeared(mRootTaskInfo, mRootLeash);
     }
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SideStageTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SideStageTests.java
index 3b42a48b5a40..a31aa58bdc26 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SideStageTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SideStageTests.java
@@ -66,7 +66,7 @@ public class SideStageTests extends ShellTestCase {
         MockitoAnnotations.initMocks(this);
         mRootTask = new TestRunningTaskInfoBuilder().build();
         mSideStage = new SideStage(mContext, mTaskOrganizer, DEFAULT_DISPLAY, mCallbacks,
-                mSyncQueue, mSurfaceSession, mIconProvider);
+                mSyncQueue, mSurfaceSession, mIconProvider, null);
         mSideStage.onTaskAppeared(mRootTask, mRootLeash);
     }
 
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 a67853cfe745..eb9d3a11d285 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
@@ -40,6 +40,8 @@ import com.android.wm.shell.transition.Transitions;
 
 import java.util.Optional;
 
+import javax.inject.Provider;
+
 public class SplitTestUtils {
 
     static SplitLayout createMockSplitLayout() {
@@ -72,10 +74,12 @@ public class SplitTestUtils {
                 DisplayInsetsController insetsController, SplitLayout splitLayout,
                 Transitions transitions, TransactionPool transactionPool,
                 SplitscreenEventLogger logger, ShellExecutor mainExecutor,
-                Optional<RecentTasksController> recentTasks) {
+                Optional<RecentTasksController> recentTasks,
+                Provider<Optional<StageTaskUnfoldController>> unfoldController) {
             super(context, displayId, syncQueue, taskOrganizer, mainStage,
                     sideStage, displayController, imeController, insetsController, splitLayout,
-                    transitions, transactionPool, logger, mainExecutor, recentTasks);
+                    transitions, transactionPool, logger, mainExecutor, recentTasks,
+                    unfoldController);
 
             // Prepare root task for testing.
             mRootTask = new TestRunningTaskInfoBuilder().build();
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 304ca66dd3bb..b52690487944 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
@@ -118,16 +118,16 @@ public class SplitTransitionTests extends ShellTestCase {
         mSplitLayout = SplitTestUtils.createMockSplitLayout();
         mMainStage = new MainStage(mContext, mTaskOrganizer, DEFAULT_DISPLAY, mock(
                 StageTaskListener.StageListenerCallbacks.class), mSyncQueue, mSurfaceSession,
-                mIconProvider);
+                mIconProvider, null);
         mMainStage.onTaskAppeared(new TestRunningTaskInfoBuilder().build(), createMockSurface());
         mSideStage = new SideStage(mContext, mTaskOrganizer, DEFAULT_DISPLAY, mock(
                 StageTaskListener.StageListenerCallbacks.class), mSyncQueue, mSurfaceSession,
-                mIconProvider);
+                mIconProvider, null);
         mSideStage.onTaskAppeared(new TestRunningTaskInfoBuilder().build(), createMockSurface());
         mStageCoordinator = new SplitTestUtils.TestStageCoordinator(mContext, DEFAULT_DISPLAY,
                 mSyncQueue, mTaskOrganizer, mMainStage, mSideStage, mDisplayController,
                 mDisplayImeController, mDisplayInsetsController, mSplitLayout, mTransitions,
-                mTransactionPool, mLogger, mMainExecutor, Optional.empty());
+                mTransactionPool, mLogger, mMainExecutor, Optional.empty(), Optional::empty);
         mSplitScreenTransitions = mStageCoordinator.getSplitTransitions();
         doAnswer((Answer<IBinder>) invocation -> mock(IBinder.class))
                 .when(mTransitions).startTransition(anyInt(), any(), any());
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 af2c495c85c5..42d998f6b0ee 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,7 +34,6 @@ import static org.mockito.ArgumentMatchers.eq;
 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;
@@ -59,7 +58,6 @@ import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.common.SyncTransactionQueue;
 import com.android.wm.shell.common.TransactionPool;
 import com.android.wm.shell.common.split.SplitLayout;
-import com.android.wm.shell.splitscreen.SplitScreen.SplitScreenListener;
 import com.android.wm.shell.transition.Transitions;
 
 import org.junit.Before;
@@ -70,6 +68,8 @@ import org.mockito.MockitoAnnotations;
 
 import java.util.Optional;
 
+import javax.inject.Provider;
+
 /**
  * Tests for {@link StageCoordinator}
  */
@@ -85,6 +85,10 @@ public class StageCoordinatorTests extends ShellTestCase {
     @Mock
     private SideStage mSideStage;
     @Mock
+    private StageTaskUnfoldController mMainUnfoldController;
+    @Mock
+    private StageTaskUnfoldController mSideUnfoldController;
+    @Mock
     private SplitLayout mSplitLayout;
     @Mock
     private DisplayController mDisplayController;
@@ -103,7 +107,6 @@ public class StageCoordinatorTests extends ShellTestCase {
 
     private final Rect mBounds1 = new Rect(10, 20, 30, 40);
     private final Rect mBounds2 = new Rect(5, 10, 15, 20);
-    private final Rect mRootBounds = new Rect(0, 0, 45, 60);
 
     private SurfaceSession mSurfaceSession = new SurfaceSession();
     private SurfaceControl mRootLeash;
@@ -116,12 +119,11 @@ public class StageCoordinatorTests extends ShellTestCase {
         mStageCoordinator = spy(new StageCoordinator(mContext, DEFAULT_DISPLAY, mSyncQueue,
                 mTaskOrganizer, mMainStage, mSideStage, mDisplayController, mDisplayImeController,
                 mDisplayInsetsController, mSplitLayout, mTransitions, mTransactionPool, mLogger,
-                mMainExecutor, Optional.empty()));
+                mMainExecutor, Optional.empty(), new UnfoldControllerProvider()));
         doNothing().when(mStageCoordinator).updateActivityOptions(any(), anyInt());
 
         when(mSplitLayout.getBounds1()).thenReturn(mBounds1);
         when(mSplitLayout.getBounds2()).thenReturn(mBounds2);
-        when(mSplitLayout.getRootBounds()).thenReturn(mRootBounds);
         when(mSplitLayout.isLandscape()).thenReturn(false);
 
         mRootTask = new TestRunningTaskInfoBuilder().build();
@@ -166,6 +168,13 @@ public class StageCoordinatorTests extends ShellTestCase {
     }
 
     @Test
+    public void testRootTaskAppeared_initializesUnfoldControllers() {
+        verify(mMainUnfoldController).init();
+        verify(mSideUnfoldController).init();
+        verify(mStageCoordinator).onRootTaskAppeared();
+    }
+
+    @Test
     public void testRootTaskInfoChanged_updatesSplitLayout() {
         mStageCoordinator.onTaskInfoChanged(mRootTask);
 
@@ -175,25 +184,26 @@ public class StageCoordinatorTests extends ShellTestCase {
     @Test
     public void testLayoutChanged_topLeftSplitPosition_updatesUnfoldStageBounds() {
         mStageCoordinator.setSideStagePosition(SPLIT_POSITION_TOP_OR_LEFT, null);
-        final SplitScreenListener listener = mock(SplitScreenListener.class);
-        mStageCoordinator.registerSplitScreenListener(listener);
-        clearInvocations(listener);
+        clearInvocations(mMainUnfoldController, mSideUnfoldController);
 
         mStageCoordinator.onLayoutSizeChanged(mSplitLayout);
 
-        verify(listener).onSplitBoundsChanged(mRootBounds, mBounds2, mBounds1);
+        verify(mMainUnfoldController).onLayoutChanged(mBounds2, SPLIT_POSITION_BOTTOM_OR_RIGHT,
+                false);
+        verify(mSideUnfoldController).onLayoutChanged(mBounds1, SPLIT_POSITION_TOP_OR_LEFT, false);
     }
 
     @Test
     public void testLayoutChanged_bottomRightSplitPosition_updatesUnfoldStageBounds() {
         mStageCoordinator.setSideStagePosition(SPLIT_POSITION_BOTTOM_OR_RIGHT, null);
-        final SplitScreenListener listener = mock(SplitScreenListener.class);
-        mStageCoordinator.registerSplitScreenListener(listener);
-        clearInvocations(listener);
+        clearInvocations(mMainUnfoldController, mSideUnfoldController);
 
         mStageCoordinator.onLayoutSizeChanged(mSplitLayout);
 
-        verify(listener).onSplitBoundsChanged(mRootBounds, mBounds1, mBounds2);
+        verify(mMainUnfoldController).onLayoutChanged(mBounds1, SPLIT_POSITION_TOP_OR_LEFT,
+                false);
+        verify(mSideUnfoldController).onLayoutChanged(mBounds2, SPLIT_POSITION_BOTTOM_OR_RIGHT,
+                false);
     }
 
     @Test
@@ -304,4 +314,20 @@ public class StageCoordinatorTests extends ShellTestCase {
 
         verify(mSplitLayout).applySurfaceChanges(any(), any(), any(), any(), any(), eq(false));
     }
+
+    private class UnfoldControllerProvider implements
+            Provider<Optional<StageTaskUnfoldController>> {
+
+        private boolean isMain = true;
+
+        @Override
+        public Optional<StageTaskUnfoldController> get() {
+            if (isMain) {
+                isMain = false;
+                return Optional.of(mMainUnfoldController);
+            } else {
+                return Optional.of(mSideUnfoldController);
+            }
+        }
+    }
 }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageTaskListenerTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageTaskListenerTests.java
index 5ee8bf3006a3..157c30bcb6c7 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageTaskListenerTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageTaskListenerTests.java
@@ -25,6 +25,7 @@ import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assume.assumeFalse;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
@@ -71,6 +72,8 @@ public final class StageTaskListenerTests extends ShellTestCase {
     private SyncTransactionQueue mSyncQueue;
     @Mock
     private IconProvider mIconProvider;
+    @Mock
+    private StageTaskUnfoldController mStageTaskUnfoldController;
     @Captor
     private ArgumentCaptor<SyncTransactionQueue.TransactionRunnable> mRunnableCaptor;
     private SurfaceSession mSurfaceSession = new SurfaceSession();
@@ -89,7 +92,8 @@ public final class StageTaskListenerTests extends ShellTestCase {
                 mCallbacks,
                 mSyncQueue,
                 mSurfaceSession,
-                mIconProvider);
+                mIconProvider,
+                mStageTaskUnfoldController);
         mRootTask = new TestRunningTaskInfoBuilder().build();
         mRootTask.parentTaskId = INVALID_TASK_ID;
         mSurfaceControl = new SurfaceControl.Builder(mSurfaceSession).setName("test").build();
@@ -126,6 +130,30 @@ public final class StageTaskListenerTests extends ShellTestCase {
         verify(mCallbacks).onStatusChanged(eq(mRootTask.isVisible), eq(true));
     }
 
+    @Test
+    public void testTaskAppeared_notifiesUnfoldListener() {
+        assumeFalse(ENABLE_SHELL_TRANSITIONS);
+        final ActivityManager.RunningTaskInfo task =
+                new TestRunningTaskInfoBuilder().setParentTaskId(mRootTask.taskId).build();
+
+        mStageTaskListener.onTaskAppeared(task, mSurfaceControl);
+
+        verify(mStageTaskUnfoldController).onTaskAppeared(eq(task), eq(mSurfaceControl));
+    }
+
+    @Test
+    public void testTaskVanished_notifiesUnfoldListener() {
+        assumeFalse(ENABLE_SHELL_TRANSITIONS);
+        final ActivityManager.RunningTaskInfo task =
+                new TestRunningTaskInfoBuilder().setParentTaskId(mRootTask.taskId).build();
+        mStageTaskListener.onTaskAppeared(task, mSurfaceControl);
+        clearInvocations(mStageTaskUnfoldController);
+
+        mStageTaskListener.onTaskVanished(task);
+
+        verify(mStageTaskUnfoldController).onTaskVanished(eq(task));
+    }
+
     @Test(expected = IllegalArgumentException.class)
     public void testUnknownTaskVanished() {
         final ActivityManager.RunningTaskInfo task = new TestRunningTaskInfoBuilder().build();
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/unfold/UnfoldAnimationControllerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/unfold/UnfoldAnimationControllerTest.java
deleted file mode 100644
index 798208956180..000000000000
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/unfold/UnfoldAnimationControllerTest.java
+++ /dev/null
@@ -1,341 +0,0 @@
-/*
- * 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.wm.shell.unfold;
-
-import static com.android.wm.shell.unfold.UnfoldAnimationControllerTest.TestUnfoldTaskAnimator.UNSET;
-
-import static com.google.common.truth.Truth.assertThat;
-
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
-
-import android.app.ActivityManager.RunningTaskInfo;
-import android.app.TaskInfo;
-import android.testing.AndroidTestingRunner;
-import android.view.SurfaceControl;
-import android.view.SurfaceControl.Transaction;
-
-import com.android.wm.shell.ShellTestCase;
-import com.android.wm.shell.TestRunningTaskInfoBuilder;
-import com.android.wm.shell.TestShellExecutor;
-import com.android.wm.shell.common.TransactionPool;
-import com.android.wm.shell.unfold.animation.UnfoldTaskAnimator;
-
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.mockito.Mock;
-import org.mockito.MockitoAnnotations;
-
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Map;
-import java.util.Optional;
-import java.util.Set;
-import java.util.concurrent.Executor;
-import java.util.function.Predicate;
-
-/**
- * Tests for {@link UnfoldAnimationController}.
- *
- * Build/Install/Run:
- * atest WMShellUnitTests:UnfoldAnimationControllerTest
- */
-@RunWith(AndroidTestingRunner.class)
-public class UnfoldAnimationControllerTest extends ShellTestCase {
-
-    @Mock
-    private TransactionPool mTransactionPool;
-    @Mock
-    private UnfoldTransitionHandler mUnfoldTransitionHandler;
-    @Mock
-    private SurfaceControl mLeash;
-
-    private UnfoldAnimationController mUnfoldAnimationController;
-
-    private final TestShellUnfoldProgressProvider mProgressProvider =
-            new TestShellUnfoldProgressProvider();
-    private final TestShellExecutor mShellExecutor = new TestShellExecutor();
-
-    private final TestUnfoldTaskAnimator mTaskAnimator1 = new TestUnfoldTaskAnimator();
-    private final TestUnfoldTaskAnimator mTaskAnimator2 = new TestUnfoldTaskAnimator();
-
-    @Before
-    public void setUp() {
-        MockitoAnnotations.initMocks(this);
-        when(mTransactionPool.acquire()).thenReturn(mock(SurfaceControl.Transaction.class));
-
-        final List<UnfoldTaskAnimator> animators = new ArrayList<>();
-        animators.add(mTaskAnimator1);
-        animators.add(mTaskAnimator2);
-        mUnfoldAnimationController = new UnfoldAnimationController(
-                mTransactionPool,
-                mProgressProvider,
-                animators,
-                () -> Optional.of(mUnfoldTransitionHandler),
-                mShellExecutor
-        );
-    }
-
-    @Test
-    public void testAppearedMatchingTask_appliesUnfoldProgress() {
-        mTaskAnimator1.setTaskMatcher((info) -> info.getWindowingMode() == 2);
-        RunningTaskInfo taskInfo = new TestRunningTaskInfoBuilder()
-                .setWindowingMode(2).build();
-
-        mUnfoldAnimationController.onTaskAppeared(taskInfo, mLeash);
-
-        mUnfoldAnimationController.onStateChangeProgress(0.5f);
-        assertThat(mTaskAnimator1.mLastAppliedProgress).isEqualTo(0.5f);
-    }
-
-    @Test
-    public void testAppearedMatchingTaskTwoDifferentAnimators_appliesUnfoldProgressToBoth() {
-        mTaskAnimator1.setTaskMatcher((info) -> info.getWindowingMode() == 1);
-        mTaskAnimator2.setTaskMatcher((info) -> info.getWindowingMode() == 2);
-        RunningTaskInfo taskInfo1 = new TestRunningTaskInfoBuilder()
-                .setWindowingMode(1).build();
-        RunningTaskInfo taskInfo2 = new TestRunningTaskInfoBuilder()
-                .setWindowingMode(2).build();
-
-        mUnfoldAnimationController.onTaskAppeared(taskInfo1, mLeash);
-        mUnfoldAnimationController.onTaskAppeared(taskInfo2, mLeash);
-
-        mUnfoldAnimationController.onStateChangeProgress(0.5f);
-        assertThat(mTaskAnimator1.mLastAppliedProgress).isEqualTo(0.5f);
-        assertThat(mTaskAnimator2.mLastAppliedProgress).isEqualTo(0.5f);
-    }
-
-    @Test
-    public void testAppearedNonMatchingTask_doesNotApplyUnfoldProgress() {
-        mTaskAnimator1.setTaskMatcher((info) -> info.getWindowingMode() == 2);
-        RunningTaskInfo taskInfo = new TestRunningTaskInfoBuilder()
-                .setWindowingMode(0).build();
-
-        mUnfoldAnimationController.onTaskAppeared(taskInfo, mLeash);
-
-        mUnfoldAnimationController.onStateChangeProgress(0.5f);
-        assertThat(mTaskAnimator1.mLastAppliedProgress).isEqualTo(UNSET);
-    }
-
-    @Test
-    public void testAppearedAndChangedToNonMatchingTask_doesNotApplyUnfoldProgress() {
-        mTaskAnimator1.setTaskMatcher((info) -> info.getWindowingMode() == 2);
-        RunningTaskInfo taskInfo = new TestRunningTaskInfoBuilder()
-                .setWindowingMode(2).build();
-
-        mUnfoldAnimationController.onTaskAppeared(taskInfo, mLeash);
-        taskInfo.configuration.windowConfiguration.setWindowingMode(0);
-        mUnfoldAnimationController.onTaskInfoChanged(taskInfo);
-
-        mUnfoldAnimationController.onStateChangeProgress(0.5f);
-        assertThat(mTaskAnimator1.mLastAppliedProgress).isEqualTo(UNSET);
-    }
-
-    @Test
-    public void testAppearedAndChangedToNonMatchingTaskAndBack_appliesUnfoldProgress() {
-        mTaskAnimator1.setTaskMatcher((info) -> info.getWindowingMode() == 2);
-        RunningTaskInfo taskInfo = new TestRunningTaskInfoBuilder()
-                .setWindowingMode(2).build();
-
-        mUnfoldAnimationController.onTaskAppeared(taskInfo, mLeash);
-        taskInfo.configuration.windowConfiguration.setWindowingMode(0);
-        mUnfoldAnimationController.onTaskInfoChanged(taskInfo);
-        taskInfo.configuration.windowConfiguration.setWindowingMode(2);
-        mUnfoldAnimationController.onTaskInfoChanged(taskInfo);
-
-        mUnfoldAnimationController.onStateChangeProgress(0.5f);
-        assertThat(mTaskAnimator1.mLastAppliedProgress).isEqualTo(0.5f);
-    }
-
-    @Test
-    public void testAppearedNonMatchingTaskAndChangedToMatching_appliesUnfoldProgress() {
-        mTaskAnimator1.setTaskMatcher((info) -> info.getWindowingMode() == 2);
-        RunningTaskInfo taskInfo = new TestRunningTaskInfoBuilder()
-                .setWindowingMode(0).build();
-
-        mUnfoldAnimationController.onTaskAppeared(taskInfo, mLeash);
-        taskInfo.configuration.windowConfiguration.setWindowingMode(2);
-        mUnfoldAnimationController.onTaskInfoChanged(taskInfo);
-
-        mUnfoldAnimationController.onStateChangeProgress(0.5f);
-        assertThat(mTaskAnimator1.mLastAppliedProgress).isEqualTo(0.5f);
-    }
-
-    @Test
-    public void testAppearedMatchingTaskAndChanged_appliesUnfoldProgress() {
-        mTaskAnimator1.setTaskMatcher((info) -> info.getWindowingMode() == 2);
-        RunningTaskInfo taskInfo = new TestRunningTaskInfoBuilder()
-                .setWindowingMode(2).build();
-
-        mUnfoldAnimationController.onTaskAppeared(taskInfo, mLeash);
-        mUnfoldAnimationController.onTaskInfoChanged(taskInfo);
-
-        mUnfoldAnimationController.onStateChangeProgress(0.5f);
-        assertThat(mTaskAnimator1.mLastAppliedProgress).isEqualTo(0.5f);
-    }
-
-    @Test
-    public void testShellTransitionRunning_doesNotApplyUnfoldProgress() {
-        when(mUnfoldTransitionHandler.willHandleTransition()).thenReturn(true);
-        mTaskAnimator1.setTaskMatcher((info) -> info.getWindowingMode() == 2);
-        RunningTaskInfo taskInfo = new TestRunningTaskInfoBuilder()
-                .setWindowingMode(2).build();
-
-        mUnfoldAnimationController.onTaskAppeared(taskInfo, mLeash);
-
-        mUnfoldAnimationController.onStateChangeProgress(0.5f);
-        assertThat(mTaskAnimator1.mLastAppliedProgress).isEqualTo(UNSET);
-    }
-
-    @Test
-    public void testApplicableTaskDisappeared_resetsSurface() {
-        mTaskAnimator1.setTaskMatcher((info) -> info.getWindowingMode() == 2);
-        RunningTaskInfo taskInfo = new TestRunningTaskInfoBuilder()
-                .setWindowingMode(2).build();
-        mUnfoldAnimationController.onTaskAppeared(taskInfo, mLeash);
-        assertThat(mTaskAnimator1.mResetTasks).doesNotContain(taskInfo.taskId);
-
-        mUnfoldAnimationController.onTaskVanished(taskInfo);
-
-        assertThat(mTaskAnimator1.mResetTasks).contains(taskInfo.taskId);
-    }
-
-    @Test
-    public void testNonApplicableTaskAppearedDisappeared_doesNotResetSurface() {
-        mTaskAnimator1.setTaskMatcher((info) -> info.getWindowingMode() == 2);
-        RunningTaskInfo taskInfo = new TestRunningTaskInfoBuilder()
-                .setWindowingMode(0).build();
-
-        mUnfoldAnimationController.onTaskAppeared(taskInfo, mLeash);
-        mUnfoldAnimationController.onTaskVanished(taskInfo);
-
-        assertThat(mTaskAnimator1.mResetTasks).doesNotContain(taskInfo.taskId);
-    }
-
-    @Test
-    public void testInit_initsAndStartsAnimators() {
-        mUnfoldAnimationController.init();
-
-        assertThat(mTaskAnimator1.mInitialized).isTrue();
-        assertThat(mTaskAnimator1.mStarted).isTrue();
-    }
-
-    private static class TestShellUnfoldProgressProvider implements ShellUnfoldProgressProvider,
-            ShellUnfoldProgressProvider.UnfoldListener {
-
-        private final List<UnfoldListener> mListeners = new ArrayList<>();
-
-        @Override
-        public void addListener(Executor executor, UnfoldListener listener) {
-            mListeners.add(listener);
-        }
-
-        @Override
-        public void onStateChangeStarted() {
-            mListeners.forEach(UnfoldListener::onStateChangeStarted);
-        }
-
-        @Override
-        public void onStateChangeProgress(float progress) {
-            mListeners.forEach(unfoldListener -> unfoldListener.onStateChangeProgress(progress));
-        }
-
-        @Override
-        public void onStateChangeFinished() {
-            mListeners.forEach(UnfoldListener::onStateChangeFinished);
-        }
-    }
-
-    public static class TestUnfoldTaskAnimator implements UnfoldTaskAnimator {
-
-        public static final float UNSET = -1f;
-        private Predicate<TaskInfo> mTaskMatcher = (info) -> false;
-
-        Map<Integer, TaskInfo> mTasksMap = new HashMap<>();
-        Set<Integer> mResetTasks = new HashSet<>();
-
-        boolean mInitialized = false;
-        boolean mStarted = false;
-        float mLastAppliedProgress = UNSET;
-
-        @Override
-        public void init() {
-            mInitialized = true;
-        }
-
-        @Override
-        public void start() {
-            mStarted = true;
-        }
-
-        @Override
-        public void stop() {
-            mStarted = false;
-        }
-
-        @Override
-        public boolean isApplicableTask(TaskInfo taskInfo) {
-            return mTaskMatcher.test(taskInfo);
-        }
-
-        @Override
-        public void applyAnimationProgress(float progress, Transaction transaction) {
-            mLastAppliedProgress = progress;
-        }
-
-        public void setTaskMatcher(Predicate<TaskInfo> taskMatcher) {
-            mTaskMatcher = taskMatcher;
-        }
-
-        @Override
-        public void onTaskAppeared(TaskInfo taskInfo, SurfaceControl leash) {
-            mTasksMap.put(taskInfo.taskId, taskInfo);
-        }
-
-        @Override
-        public void onTaskVanished(TaskInfo taskInfo) {
-            mTasksMap.remove(taskInfo.taskId);
-        }
-
-        @Override
-        public void onTaskChanged(TaskInfo taskInfo) {
-            mTasksMap.put(taskInfo.taskId, taskInfo);
-        }
-
-        @Override
-        public void resetSurface(TaskInfo taskInfo, Transaction transaction) {
-            mResetTasks.add(taskInfo.taskId);
-        }
-
-        @Override
-        public void resetAllSurfaces(Transaction transaction) {
-            mTasksMap.values().forEach((t) -> mResetTasks.add(t.taskId));
-        }
-
-        @Override
-        public boolean hasActiveTasks() {
-            return mTasksMap.size() > 0;
-        }
-
-        public List<TaskInfo> getCurrentTasks() {
-            return new ArrayList<>(mTasksMap.values());
-        }
-    }
-}