Synchronize screen turning on and unfold overlay

Changes unfold overlay implementation from a window
to surface control view host. To make sure that the
unfold overlay is drawn by the time when we remove
screen blocker we synchronously apply a transaction
with drawn overlay and apply another empty transaction
with vsyncId+1.

Also these changes disable the unfold transition when
using power button by filtering only first screen turning
on events after unfolding the device.

Bug: 197538198
Test: manual fold/unfolds
Test: killing SysUI process, checking rotation animation, magnification
Test: atest com.android.systemui.unfold.updates.DeviceFoldStateProviderTest
Change-Id: I8e0bc635b041595602145b313548b14fcabd157e

18 files changed, 788 insertions, 84 deletions

diff --git a/core/java/android/view/SurfaceControlViewHost.java b/core/java/android/view/SurfaceControlViewHost.java
index 11b161ad3cb2..a6c5042db275 100644
--- a/core/java/android/view/SurfaceControlViewHost.java
+++ b/core/java/android/view/SurfaceControlViewHost.java
@@ -292,11 +292,18 @@ public class SurfaceControlViewHost {
      */
     @TestApi
     public void relayout(WindowManager.LayoutParams attrs) {
+        relayout(attrs, SurfaceControl.Transaction::apply);
+    }
+
+    /**
+     * Forces relayout and draw and allows to set a custom callback when it is finished
+     * @hide
+     */
+    public void relayout(WindowManager.LayoutParams attrs,
+            WindowlessWindowManager.ResizeCompleteCallback callback) {
         mViewRoot.setLayoutParams(attrs, false);
         mViewRoot.setReportNextDraw();
-        mWm.setCompletionCallback(mViewRoot.mWindow.asBinder(), (SurfaceControl.Transaction t) -> {
-            t.apply();
-        });
+        mWm.setCompletionCallback(mViewRoot.mWindow.asBinder(), callback);
     }
 
     /**
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/RootDisplayAreaOrganizer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/RootDisplayAreaOrganizer.java
new file mode 100644
index 000000000000..14ba9df93f24
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/RootDisplayAreaOrganizer.java
@@ -0,0 +1,114 @@
+/*
+ * Copyright (C) 2020 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;
+
+import android.util.SparseArray;
+import android.view.SurfaceControl;
+import android.window.DisplayAreaAppearedInfo;
+import android.window.DisplayAreaInfo;
+import android.window.DisplayAreaOrganizer;
+
+import androidx.annotation.NonNull;
+
+import java.io.PrintWriter;
+import java.util.List;
+import java.util.concurrent.Executor;
+
+/** Display area organizer for the root display areas */
+public class RootDisplayAreaOrganizer extends DisplayAreaOrganizer {
+
+    private static final String TAG = RootDisplayAreaOrganizer.class.getSimpleName();
+
+    /** {@link DisplayAreaInfo} list, which is mapped by display IDs. */
+    private final SparseArray<DisplayAreaInfo> mDisplayAreasInfo = new SparseArray<>();
+    /** Display area leashes, which is mapped by display IDs. */
+    private final SparseArray<SurfaceControl> mLeashes = new SparseArray<>();
+
+    public RootDisplayAreaOrganizer(Executor executor) {
+        super(executor);
+        List<DisplayAreaAppearedInfo> infos = registerOrganizer(FEATURE_ROOT);
+        for (int i = infos.size() - 1; i >= 0; --i) {
+            onDisplayAreaAppeared(infos.get(i).getDisplayAreaInfo(), infos.get(i).getLeash());
+        }
+    }
+
+    public void attachToDisplayArea(int displayId, SurfaceControl.Builder b) {
+        final SurfaceControl sc = mLeashes.get(displayId);
+        if (sc != null) {
+            b.setParent(sc);
+        }
+    }
+
+    @Override
+    public void onDisplayAreaAppeared(@NonNull DisplayAreaInfo displayAreaInfo,
+            @NonNull SurfaceControl leash) {
+        if (displayAreaInfo.featureId != FEATURE_ROOT) {
+            throw new IllegalArgumentException(
+                    "Unknown feature: " + displayAreaInfo.featureId
+                            + "displayAreaInfo:" + displayAreaInfo);
+        }
+
+        final int displayId = displayAreaInfo.displayId;
+        if (mDisplayAreasInfo.get(displayId) != null) {
+            throw new IllegalArgumentException(
+                    "Duplicate DA for displayId: " + displayId
+                            + " displayAreaInfo:" + displayAreaInfo
+                            + " mDisplayAreasInfo.get():" + mDisplayAreasInfo.get(displayId));
+        }
+
+        mDisplayAreasInfo.put(displayId, displayAreaInfo);
+        mLeashes.put(displayId, leash);
+    }
+
+    @Override
+    public void onDisplayAreaVanished(@NonNull DisplayAreaInfo displayAreaInfo) {
+        final int displayId = displayAreaInfo.displayId;
+        if (mDisplayAreasInfo.get(displayId) == null) {
+            throw new IllegalArgumentException(
+                    "onDisplayAreaVanished() Unknown DA displayId: " + displayId
+                            + " displayAreaInfo:" + displayAreaInfo
+                            + " mDisplayAreasInfo.get():" + mDisplayAreasInfo.get(displayId));
+        }
+
+        mDisplayAreasInfo.remove(displayId);
+    }
+
+    @Override
+    public void onDisplayAreaInfoChanged(@NonNull DisplayAreaInfo displayAreaInfo) {
+        final int displayId = displayAreaInfo.displayId;
+        if (mDisplayAreasInfo.get(displayId) == null) {
+            throw new IllegalArgumentException(
+                    "onDisplayAreaInfoChanged() Unknown DA displayId: " + displayId
+                            + " displayAreaInfo:" + displayAreaInfo
+                            + " mDisplayAreasInfo.get():" + mDisplayAreasInfo.get(displayId));
+        }
+
+        mDisplayAreasInfo.put(displayId, displayAreaInfo);
+    }
+
+    public void dump(@NonNull PrintWriter pw, String prefix) {
+        final String innerPrefix = prefix + "  ";
+        final String childPrefix = innerPrefix + "  ";
+        pw.println(prefix + this);
+    }
+
+    @Override
+    public String toString() {
+        return TAG + "#" + mDisplayAreasInfo.size();
+    }
+
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/displayareahelper/DisplayAreaHelper.java b/libs/WindowManager/Shell/src/com/android/wm/shell/displayareahelper/DisplayAreaHelper.java
new file mode 100644
index 000000000000..defbd5af01d9
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/displayareahelper/DisplayAreaHelper.java
@@ -0,0 +1,39 @@
+/*
+ * 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.displayareahelper;
+
+import android.view.SurfaceControl;
+
+import java.util.function.Consumer;
+
+/**
+ * Interface that allows to perform various display area related actions
+ */
+public interface DisplayAreaHelper {
+
+    /**
+     * Updates SurfaceControl builder to reparent it to the root display area
+     * @param displayId id of the display to which root display area it should be reparented to
+     * @param builder surface control builder that should be updated
+     * @param onUpdated callback that is invoked after updating the builder, called on
+     *                  the shell main thread
+     */
+    default void attachToRootDisplayArea(int displayId, SurfaceControl.Builder builder,
+            Consumer<SurfaceControl.Builder> onUpdated) {
+    }
+
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/displayareahelper/DisplayAreaHelperController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/displayareahelper/DisplayAreaHelperController.java
new file mode 100644
index 000000000000..ef9ad6d10e6b
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/displayareahelper/DisplayAreaHelperController.java
@@ -0,0 +1,45 @@
+/*
+ * 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.displayareahelper;
+
+import android.view.SurfaceControl;
+
+import com.android.wm.shell.RootDisplayAreaOrganizer;
+
+import java.util.concurrent.Executor;
+import java.util.function.Consumer;
+
+public class DisplayAreaHelperController implements DisplayAreaHelper {
+
+    private final Executor mExecutor;
+    private final RootDisplayAreaOrganizer mRootDisplayAreaOrganizer;
+
+    public DisplayAreaHelperController(Executor executor,
+            RootDisplayAreaOrganizer rootDisplayAreaOrganizer) {
+        mExecutor = executor;
+        mRootDisplayAreaOrganizer = rootDisplayAreaOrganizer;
+    }
+
+    @Override
+    public void attachToRootDisplayArea(int displayId, SurfaceControl.Builder builder,
+            Consumer<SurfaceControl.Builder> onUpdated) {
+        mExecutor.execute(() -> {
+            mRootDisplayAreaOrganizer.attachToDisplayArea(displayId, builder);
+            onUpdated.accept(builder);
+        });
+    }
+}
diff --git a/packages/SystemUI/shared/src/com/android/systemui/unfold/progress/PhysicsBasedUnfoldTransitionProgressProvider.kt b/packages/SystemUI/shared/src/com/android/systemui/unfold/progress/PhysicsBasedUnfoldTransitionProgressProvider.kt
index 10e6c2bfea3a..90f5998053b8 100644
--- a/packages/SystemUI/shared/src/com/android/systemui/unfold/progress/PhysicsBasedUnfoldTransitionProgressProvider.kt
+++ b/packages/SystemUI/shared/src/com/android/systemui/unfold/progress/PhysicsBasedUnfoldTransitionProgressProvider.kt
@@ -16,6 +16,7 @@
 package com.android.systemui.unfold.progress
 
 import android.os.Handler
+import android.util.Log
 import android.util.MathUtils.saturate
 import androidx.dynamicanimation.animation.DynamicAnimation
 import androidx.dynamicanimation.animation.FloatPropertyCompat
@@ -92,7 +93,14 @@ internal class PhysicsBasedUnfoldTransitionProgressProvider(
                 }
             }
             FOLD_UPDATE_FINISH_FULL_OPEN -> {
-                cancelTransition(endValue = 1f, animate = true)
+                // Do not cancel if we haven't started the transition yet.
+                // This could happen when we fully unfolded the device before the screen
+                // became available. In this case we start and immediately cancel the animation
+                // in FOLD_UPDATE_UNFOLDED_SCREEN_AVAILABLE event handler, so we don't need to
+                // cancel it here.
+                if (isTransitionRunning) {
+                    cancelTransition(endValue = 1f, animate = true)
+                }
             }
             FOLD_UPDATE_FINISH_CLOSED -> {
                 cancelTransition(endValue = 0f, animate = false)
@@ -101,6 +109,10 @@ internal class PhysicsBasedUnfoldTransitionProgressProvider(
                 startTransition(startValue = 1f)
             }
         }
+
+        if (DEBUG) {
+            Log.d(TAG, "onFoldUpdate = $update")
+        }
     }
 
     private fun cancelTransition(endValue: Float, animate: Boolean) {
@@ -118,6 +130,10 @@ internal class PhysicsBasedUnfoldTransitionProgressProvider(
             listeners.forEach {
                 it.onTransitionFinished()
             }
+
+            if (DEBUG) {
+                Log.d(TAG, "onTransitionFinished")
+            }
         }
     }
 
@@ -137,6 +153,10 @@ internal class PhysicsBasedUnfoldTransitionProgressProvider(
             it.onTransitionStarted()
         }
         isTransitionRunning = true
+
+        if (DEBUG) {
+            Log.d(TAG, "onTransitionStarted")
+        }
     }
 
     private fun startTransition(startValue: Float) {
@@ -189,6 +209,9 @@ internal class PhysicsBasedUnfoldTransitionProgressProvider(
     }
 }
 
+private const val TAG = "PhysicsBasedUnfoldTransitionProgressProvider"
+private const val DEBUG = true
+
 private const val TRANSITION_TIMEOUT_MILLIS = 2000L
 private const val SPRING_STIFFNESS = 200.0f
 private const val MINIMAL_VISIBLE_CHANGE = 0.001f
diff --git a/packages/SystemUI/shared/src/com/android/systemui/unfold/updates/DeviceFoldStateProvider.kt b/packages/SystemUI/shared/src/com/android/systemui/unfold/updates/DeviceFoldStateProvider.kt
index c37ab06abb60..35e2b30d0a39 100644
--- a/packages/SystemUI/shared/src/com/android/systemui/unfold/updates/DeviceFoldStateProvider.kt
+++ b/packages/SystemUI/shared/src/com/android/systemui/unfold/updates/DeviceFoldStateProvider.kt
@@ -25,7 +25,7 @@ import com.android.systemui.unfold.updates.hinge.FULLY_OPEN_DEGREES
 import com.android.systemui.unfold.updates.hinge.HingeAngleProvider
 import java.util.concurrent.Executor
 
-internal class DeviceFoldStateProvider(
+class DeviceFoldStateProvider(
     context: Context,
     private val hingeAngleProvider: HingeAngleProvider,
     private val screenStatusProvider: ScreenStatusProvider,
@@ -43,6 +43,7 @@ internal class DeviceFoldStateProvider(
     private val foldStateListener = FoldStateListener(context)
 
     private var isFolded = false
+    private var isUnfoldHandled = true
 
     override fun start() {
         deviceStateManager.registerCallback(
@@ -104,6 +105,7 @@ internal class DeviceFoldStateProvider(
                 lastFoldUpdate = FOLD_UPDATE_FINISH_CLOSED
                 outputListeners.forEach { it.onFoldUpdate(FOLD_UPDATE_FINISH_CLOSED) }
                 hingeAngleProvider.stop()
+                isUnfoldHandled = false
             } else {
                 lastFoldUpdate = FOLD_UPDATE_START_OPENING
                 outputListeners.forEach { it.onFoldUpdate(FOLD_UPDATE_START_OPENING) }
@@ -115,8 +117,15 @@ internal class DeviceFoldStateProvider(
         ScreenStatusProvider.ScreenListener {
 
         override fun onScreenTurnedOn() {
-            if (!isFolded) {
+            // Trigger this event only if we are unfolded and this is the first screen
+            // turned on event since unfold started. This prevents running the animation when
+            // turning on the internal display using the power button.
+            // Initially isUnfoldHandled is true so it will be reset to false *only* when we
+            // receive 'folded' event. If SystemUI started when device is already folded it will
+            // still receive 'folded' event on startup.
+            if (!isFolded && !isUnfoldHandled) {
                 outputListeners.forEach { it.onFoldUpdate(FOLD_UPDATE_UNFOLDED_SCREEN_AVAILABLE) }
+                isUnfoldHandled = true
             }
         }
     }
diff --git a/packages/SystemUI/shared/src/com/android/systemui/unfold/updates/FoldStateProvider.kt b/packages/SystemUI/shared/src/com/android/systemui/unfold/updates/FoldStateProvider.kt
index 11984b93638e..643ece353522 100644
--- a/packages/SystemUI/shared/src/com/android/systemui/unfold/updates/FoldStateProvider.kt
+++ b/packages/SystemUI/shared/src/com/android/systemui/unfold/updates/FoldStateProvider.kt
@@ -24,7 +24,7 @@ import com.android.systemui.statusbar.policy.CallbackController
  * Allows to subscribe to main events related to fold/unfold process such as hinge angle update,
  * start folding/unfolding, screen availability
  */
-internal interface FoldStateProvider : CallbackController<FoldUpdatesListener> {
+interface FoldStateProvider : CallbackController<FoldUpdatesListener> {
     fun start()
     fun stop()
 
diff --git a/packages/SystemUI/shared/src/com/android/systemui/unfold/updates/hinge/HingeAngleProvider.kt b/packages/SystemUI/shared/src/com/android/systemui/unfold/updates/hinge/HingeAngleProvider.kt
index 2520d356b721..6f524560de99 100644
--- a/packages/SystemUI/shared/src/com/android/systemui/unfold/updates/hinge/HingeAngleProvider.kt
+++ b/packages/SystemUI/shared/src/com/android/systemui/unfold/updates/hinge/HingeAngleProvider.kt
@@ -9,7 +9,7 @@ import com.android.systemui.statusbar.policy.CallbackController
  * For foldable devices usually 0 corresponds to fully closed (folded) state and
  * 180 degrees corresponds to fully open (flat) state
  */
-internal interface HingeAngleProvider : CallbackController<Consumer<Float>> {
+interface HingeAngleProvider : CallbackController<Consumer<Float>> {
     fun start()
     fun stop()
 }
diff --git a/packages/SystemUI/src/com/android/systemui/SystemUIFactory.java b/packages/SystemUI/src/com/android/systemui/SystemUIFactory.java
index a28223de2bf1..c64f416f9672 100644
--- a/packages/SystemUI/src/com/android/systemui/SystemUIFactory.java
+++ b/packages/SystemUI/src/com/android/systemui/SystemUIFactory.java
@@ -118,6 +118,7 @@ public class SystemUIFactory {
                     .setTaskViewFactory(mWMComponent.getTaskViewFactory())
                     .setTransitions(mWMComponent.getTransitions())
                     .setStartingSurface(mWMComponent.getStartingSurface())
+                    .setDisplayAreaHelper(mWMComponent.getDisplayAreaHelper())
                     .setTaskSurfaceHelper(mWMComponent.getTaskSurfaceHelper());
         } else {
             // TODO: Call on prepareSysUIComponentBuilder but not with real components. Other option
@@ -133,6 +134,7 @@ public class SystemUIFactory {
                     .setAppPairs(Optional.ofNullable(null))
                     .setTaskViewFactory(Optional.ofNullable(null))
                     .setTransitions(Transitions.createEmptyForTesting())
+                    .setDisplayAreaHelper(Optional.ofNullable(null))
                     .setStartingSurface(Optional.ofNullable(null))
                     .setTaskSurfaceHelper(Optional.ofNullable(null));
         }
diff --git a/packages/SystemUI/src/com/android/systemui/dagger/SysUIComponent.java b/packages/SystemUI/src/com/android/systemui/dagger/SysUIComponent.java
index 0fdc4d8e86a7..aac03f8551fc 100644
--- a/packages/SystemUI/src/com/android/systemui/dagger/SysUIComponent.java
+++ b/packages/SystemUI/src/com/android/systemui/dagger/SysUIComponent.java
@@ -30,6 +30,7 @@ import com.android.wm.shell.ShellCommandHandler;
 import com.android.wm.shell.TaskViewFactory;
 import com.android.wm.shell.apppairs.AppPairs;
 import com.android.wm.shell.bubbles.Bubbles;
+import com.android.wm.shell.displayareahelper.DisplayAreaHelper;
 import com.android.wm.shell.hidedisplaycutout.HideDisplayCutout;
 import com.android.wm.shell.legacysplitscreen.LegacySplitScreen;
 import com.android.wm.shell.onehanded.OneHanded;
@@ -96,6 +97,9 @@ public interface SysUIComponent {
         Builder setStartingSurface(Optional<StartingSurface> s);
 
         @BindsInstance
+        Builder setDisplayAreaHelper(Optional<DisplayAreaHelper> h);
+
+        @BindsInstance
         Builder setTaskSurfaceHelper(Optional<TaskSurfaceHelper> t);
 
         SysUIComponent build();
diff --git a/packages/SystemUI/src/com/android/systemui/dagger/WMComponent.java b/packages/SystemUI/src/com/android/systemui/dagger/WMComponent.java
index 442d351729d7..618c26b89196 100644
--- a/packages/SystemUI/src/com/android/systemui/dagger/WMComponent.java
+++ b/packages/SystemUI/src/com/android/systemui/dagger/WMComponent.java
@@ -27,6 +27,7 @@ import com.android.wm.shell.ShellInit;
 import com.android.wm.shell.TaskViewFactory;
 import com.android.wm.shell.apppairs.AppPairs;
 import com.android.wm.shell.bubbles.Bubbles;
+import com.android.wm.shell.displayareahelper.DisplayAreaHelper;
 import com.android.wm.shell.hidedisplaycutout.HideDisplayCutout;
 import com.android.wm.shell.legacysplitscreen.LegacySplitScreen;
 import com.android.wm.shell.onehanded.OneHanded;
@@ -105,5 +106,8 @@ public interface WMComponent {
     Optional<StartingSurface> getStartingSurface();
 
     @WMSingleton
+    Optional<DisplayAreaHelper> getDisplayAreaHelper();
+
+    @WMSingleton
     Optional<TaskSurfaceHelper> getTaskSurfaceHelper();
 }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
index c5694398cdc6..be84a82234a1 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
@@ -119,12 +119,15 @@ import com.android.systemui.statusbar.phone.StatusBar;
 import com.android.systemui.statusbar.phone.UnlockedScreenOffAnimationController;
 import com.android.systemui.statusbar.policy.KeyguardStateController;
 import com.android.systemui.statusbar.policy.UserSwitcherController;
+import com.android.systemui.unfold.UnfoldLightRevealOverlayAnimation;
+import com.android.systemui.unfold.config.UnfoldTransitionConfig;
 import com.android.systemui.util.DeviceConfigProxy;
 
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
 import java.util.ArrayList;
 import java.util.concurrent.Executor;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import dagger.Lazy;
 
@@ -815,6 +818,10 @@ public class KeyguardViewMediator extends SystemUI implements Dumpable,
     private DeviceConfigProxy mDeviceConfig;
     private DozeParameters mDozeParameters;
 
+    private final UnfoldTransitionConfig mUnfoldTransitionConfig;
+    private final Lazy<UnfoldLightRevealOverlayAnimation> mUnfoldLightRevealAnimation;
+    private final AtomicInteger mPendingDrawnTasks = new AtomicInteger();
+
     private final KeyguardStateController mKeyguardStateController;
     private final Lazy<KeyguardUnlockAnimationController> mKeyguardUnlockAnimationControllerLazy;
     private boolean mWallpaperSupportsAmbientMode;
@@ -837,6 +844,8 @@ public class KeyguardViewMediator extends SystemUI implements Dumpable,
             NavigationModeController navigationModeController,
             KeyguardDisplayManager keyguardDisplayManager,
             DozeParameters dozeParameters,
+            UnfoldTransitionConfig unfoldTransitionConfig,
+            Lazy<UnfoldLightRevealOverlayAnimation> unfoldLightRevealOverlayAnimation,
             SysuiStatusBarStateController statusBarStateController,
             KeyguardStateController keyguardStateController,
             Lazy<KeyguardUnlockAnimationController> keyguardUnlockAnimationControllerLazy,
@@ -870,6 +879,8 @@ public class KeyguardViewMediator extends SystemUI implements Dumpable,
                     mInGestureNavigationMode = QuickStepContract.isGesturalMode(mode);
                 }));
         mDozeParameters = dozeParameters;
+        mUnfoldTransitionConfig = unfoldTransitionConfig;
+        mUnfoldLightRevealAnimation = unfoldLightRevealOverlayAnimation;
         mStatusBarStateController = statusBarStateController;
         statusBarStateController.addCallback(this);
 
@@ -2552,6 +2563,24 @@ public class KeyguardViewMediator extends SystemUI implements Dumpable,
         Trace.beginSection("KeyguardViewMediator#handleNotifyScreenTurningOn");
         synchronized (KeyguardViewMediator.this) {
             if (DEBUG) Log.d(TAG, "handleNotifyScreenTurningOn");
+
+            if (mUnfoldTransitionConfig.isEnabled()) {
+                mPendingDrawnTasks.set(2); // unfold overlay and keyguard drawn
+
+                mUnfoldLightRevealAnimation.get()
+                        .onScreenTurningOn(() -> {
+                            if (mPendingDrawnTasks.decrementAndGet() == 0) {
+                                try {
+                                    callback.onDrawn();
+                                } catch (RemoteException e) {
+                                    Slog.w(TAG, "Exception calling onDrawn():", e);
+                                }
+                            }
+                        });
+            } else {
+                mPendingDrawnTasks.set(1); // only keyguard drawn
+            }
+
             mKeyguardViewControllerLazy.get().onScreenTurningOn();
             if (callback != null) {
                 if (mWakeAndUnlocking) {
@@ -2582,10 +2611,12 @@ public class KeyguardViewMediator extends SystemUI implements Dumpable,
 
     private void notifyDrawn(final IKeyguardDrawnCallback callback) {
         Trace.beginSection("KeyguardViewMediator#notifyDrawn");
-        try {
-            callback.onDrawn();
-        } catch (RemoteException e) {
-            Slog.w(TAG, "Exception calling onDrawn():", e);
+        if (mPendingDrawnTasks.decrementAndGet() == 0) {
+            try {
+                callback.onDrawn();
+            } catch (RemoteException e) {
+                Slog.w(TAG, "Exception calling onDrawn():", e);
+            }
         }
         Trace.endSection();
     }
@@ -2738,6 +2769,7 @@ public class KeyguardViewMediator extends SystemUI implements Dumpable,
         pw.print("  mHideAnimationRun: "); pw.println(mHideAnimationRun);
         pw.print("  mPendingReset: "); pw.println(mPendingReset);
         pw.print("  mPendingLock: "); pw.println(mPendingLock);
+        pw.print("  mPendingDrawnTasks: "); pw.println(mPendingDrawnTasks.get());
         pw.print("  mWakeAndUnlocking: "); pw.println(mWakeAndUnlocking);
         pw.print("  mDrawnCallback: "); pw.println(mDrawnCallback);
     }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/dagger/KeyguardModule.java b/packages/SystemUI/src/com/android/systemui/keyguard/dagger/KeyguardModule.java
index 8a383b974d77..11d4aac9dc27 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/dagger/KeyguardModule.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/dagger/KeyguardModule.java
@@ -55,6 +55,8 @@ import com.android.systemui.statusbar.phone.StatusBar;
 import com.android.systemui.statusbar.phone.UnlockedScreenOffAnimationController;
 import com.android.systemui.statusbar.policy.KeyguardStateController;
 import com.android.systemui.statusbar.policy.UserSwitcherController;
+import com.android.systemui.unfold.UnfoldLightRevealOverlayAnimation;
+import com.android.systemui.unfold.config.UnfoldTransitionConfig;
 import com.android.systemui.util.DeviceConfigProxy;
 import com.android.systemui.util.sensors.AsyncSensorManager;
 import com.android.systemui.util.settings.GlobalSettings;
@@ -99,6 +101,8 @@ public class KeyguardModule {
             NavigationModeController navigationModeController,
             KeyguardDisplayManager keyguardDisplayManager,
             DozeParameters dozeParameters,
+            UnfoldTransitionConfig unfoldTransitionConfig,
+            Lazy<UnfoldLightRevealOverlayAnimation> unfoldLightRevealOverlayAnimation,
             SysuiStatusBarStateController statusBarStateController,
             KeyguardStateController keyguardStateController,
             Lazy<KeyguardUnlockAnimationController> keyguardUnlockAnimationController,
@@ -121,6 +125,8 @@ public class KeyguardModule {
                 navigationModeController,
                 keyguardDisplayManager,
                 dozeParameters,
+                unfoldTransitionConfig,
+                unfoldLightRevealOverlayAnimation,
                 statusBarStateController,
                 keyguardStateController,
                 keyguardUnlockAnimationController,
diff --git a/packages/SystemUI/src/com/android/systemui/unfold/UnfoldLightRevealOverlayAnimation.kt b/packages/SystemUI/src/com/android/systemui/unfold/UnfoldLightRevealOverlayAnimation.kt
index 3c3cc64a49cd..f0760d4e2187 100644
--- a/packages/SystemUI/src/com/android/systemui/unfold/UnfoldLightRevealOverlayAnimation.kt
+++ b/packages/SystemUI/src/com/android/systemui/unfold/UnfoldLightRevealOverlayAnimation.kt
@@ -19,13 +19,27 @@ import android.content.Context
 import android.graphics.PixelFormat
 import android.hardware.devicestate.DeviceStateManager
 import android.hardware.devicestate.DeviceStateManager.FoldStateListener
+import android.hardware.display.DisplayManager
+import android.os.Handler
+import android.os.Trace
+import android.view.Choreographer
+import android.view.Display
+import android.view.DisplayInfo
 import android.view.Surface
+import android.view.SurfaceControl
+import android.view.SurfaceControlViewHost
+import android.view.SurfaceSession
 import android.view.WindowManager
+import android.view.WindowlessWindowManager
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Main
-import com.android.systemui.unfold.UnfoldTransitionProgressProvider.TransitionProgressListener
+import com.android.systemui.dagger.qualifiers.UiBackground
+import com.android.systemui.statusbar.LightRevealEffect
 import com.android.systemui.statusbar.LightRevealScrim
 import com.android.systemui.statusbar.LinearLightRevealEffect
+import com.android.systemui.unfold.UnfoldTransitionProgressProvider.TransitionProgressListener
+import com.android.wm.shell.displayareahelper.DisplayAreaHelper
+import java.util.Optional
 import java.util.concurrent.Executor
 import java.util.function.Consumer
 import javax.inject.Inject
@@ -34,53 +48,148 @@ import javax.inject.Inject
 class UnfoldLightRevealOverlayAnimation @Inject constructor(
     private val context: Context,
     private val deviceStateManager: DeviceStateManager,
+    private val displayManager: DisplayManager,
     private val unfoldTransitionProgressProvider: UnfoldTransitionProgressProvider,
+    private val displayAreaHelper: Optional<DisplayAreaHelper>,
     @Main private val executor: Executor,
-    private val windowManager: WindowManager
+    @Main private val handler: Handler,
+    @UiBackground private val backgroundExecutor: Executor
 ) {
 
     private val transitionListener = TransitionListener()
+    private val displayListener = DisplayChangeListener()
+
+    private lateinit var wwm: WindowlessWindowManager
+    private lateinit var unfoldedDisplayInfo: DisplayInfo
+    private lateinit var overlayContainer: SurfaceControl
+
+    private var root: SurfaceControlViewHost? = null
     private var scrimView: LightRevealScrim? = null
+    private var isFolded: Boolean = false
+    private var isUnfoldHandled: Boolean = true
+
+    private var currentRotation: Int = context.display!!.rotation
 
     fun init() {
         deviceStateManager.registerCallback(executor, FoldListener())
         unfoldTransitionProgressProvider.addCallback(transitionListener)
-    }
 
-    private inner class TransitionListener : TransitionProgressListener {
+        val containerBuilder = SurfaceControl.Builder(SurfaceSession())
+            .setContainerLayer()
+            .setName("unfold-overlay-container")
 
-        override fun onTransitionProgress(progress: Float) {
-            scrimView?.revealAmount = progress
-        }
+        displayAreaHelper.get().attachToRootDisplayArea(Display.DEFAULT_DISPLAY,
+            containerBuilder) { builder ->
+            executor.execute {
+                overlayContainer = builder.build()
 
-        override fun onTransitionFinished() {
-            removeOverlayView()
+                SurfaceControl.Transaction()
+                    .setLayer(overlayContainer, Integer.MAX_VALUE)
+                    .show(overlayContainer)
+                    .apply()
+
+                wwm = WindowlessWindowManager(context.resources.configuration,
+                    overlayContainer, null)
+            }
         }
 
-        override fun onTransitionStarted() {
-            // When unfolding the view is added earlier, add view for folding case
-            if (scrimView == null) {
-                addOverlayView()
+        displayManager.registerDisplayListener(displayListener, handler,
+            DisplayManager.EVENT_FLAG_DISPLAY_CHANGED)
+
+        // Get unfolded display size immediately as 'current display info' might be
+        // not up-to-date during unfolding
+        unfoldedDisplayInfo = getUnfoldedDisplayInfo()
+    }
+
+    /**
+     * Called when screen starts turning on, the contents of the screen might not be visible yet.
+     * This method reports back that the overlay is ready in [onOverlayReady] callback.
+     *
+     * @param onOverlayReady callback when the overlay is drawn and visible on the screen
+     * @see [com.android.systemui.keyguard.KeyguardViewMediator]
+     */
+    fun onScreenTurningOn(onOverlayReady: Runnable) {
+        Trace.beginSection("UnfoldLightRevealOverlayAnimation#onScreenTurningOn")
+        try {
+            // Add the view only if we are unfolding and this is the first screen on
+            if (!isFolded && !isUnfoldHandled) {
+                addView(onOverlayReady)
+                isUnfoldHandled = true
+            } else {
+                // No unfold transition, immediately report that overlay is ready
+                ensureOverlayRemoved()
+                onOverlayReady.run()
             }
+        } finally {
+            Trace.endSection()
         }
     }
 
-    private inner class FoldListener : FoldStateListener(context, Consumer { isFolded ->
-        if (isFolded) {
-            removeOverlayView()
-        } else {
-            // Add overlay view before starting the transition as soon as we unfolded the device
-            addOverlayView()
+    private fun addView(onOverlayReady: Runnable? = null) {
+        if (!::wwm.isInitialized) {
+            // Surface overlay is not created yet on the first SysUI launch
+            onOverlayReady?.run()
+            return
         }
-    })
 
-    private fun addOverlayView() {
+        ensureOverlayRemoved()
+
+        val newRoot = SurfaceControlViewHost(context, context.display!!, wwm, false)
+        val newView = LightRevealScrim(context, null)
+            .apply {
+                revealEffect = createLightRevealEffect()
+                isScrimOpaqueChangedListener = Consumer {}
+                revealAmount = 0f
+            }
+
+        val params = getLayoutParams()
+        newRoot.setView(newView, params)
+
+        onOverlayReady?.let { callback ->
+            Trace.beginAsyncSection(
+                "UnfoldLightRevealOverlayAnimation#relayout", 0)
+
+            newRoot.relayout(params) { transaction ->
+                val vsyncId = Choreographer.getSfInstance().vsyncId
+
+                backgroundExecutor.execute {
+                    // Apply the transaction that contains the first frame of the overlay
+                    // synchronously and apply another empty transaction with
+                    // 'vsyncId + 1' to make sure that it is actually displayed on
+                    // the screen. The second transaction is necessary to remove the screen blocker
+                    // (turn on the brightness) only when the content is actually visible as it
+                    // might be presented only in the next frame.
+                    // See b/197538198
+                    transaction.setFrameTimelineVsync(vsyncId)
+                        .apply(/* sync */true)
+
+                    transaction
+                        .setFrameTimelineVsync(vsyncId + 1)
+                        .apply(/* sync */ true)
+
+                    Trace.endAsyncSection(
+                        "UnfoldLightRevealOverlayAnimation#relayout", 0)
+                    callback.run()
+                }
+            }
+        }
+
+        scrimView = newView
+        root = newRoot
+    }
+
+    private fun getLayoutParams(): WindowManager.LayoutParams {
         val params: WindowManager.LayoutParams = WindowManager.LayoutParams()
-        params.height = WindowManager.LayoutParams.MATCH_PARENT
-        params.width = WindowManager.LayoutParams.MATCH_PARENT
-        params.format = PixelFormat.TRANSLUCENT
 
-        // TODO(b/193801466): create a separate type for this overlay
+        val rotation = context.display!!.rotation
+        val isNatural = rotation == Surface.ROTATION_0 || rotation == Surface.ROTATION_180
+
+        params.height = if (isNatural)
+            unfoldedDisplayInfo.naturalHeight else unfoldedDisplayInfo.naturalWidth
+        params.width = if (isNatural)
+            unfoldedDisplayInfo.naturalWidth else unfoldedDisplayInfo.naturalHeight
+
+        params.format = PixelFormat.TRANSLUCENT
         params.type = WindowManager.LayoutParams.TYPE_DISPLAY_OVERLAY
         params.title = "Unfold Light Reveal Animation"
         params.layoutInDisplayCutoutMode =
@@ -90,41 +199,72 @@ class UnfoldLightRevealOverlayAnimation @Inject constructor(
             or WindowManager.LayoutParams.FLAG_NOT_TOUCHABLE)
         params.setTrustedOverlay()
 
-        val rotation = windowManager.defaultDisplay.rotation
-        val isVerticalFold = rotation == Surface.ROTATION_0 || rotation == Surface.ROTATION_180
+        val packageName: String = context.opPackageName
+        params.packageName = packageName
 
-        val newScrimView = LightRevealScrim(context, null)
-            .apply {
-                revealEffect = LinearLightRevealEffect(isVerticalFold)
-                isScrimOpaqueChangedListener = Consumer {}
-                revealAmount = 0f
-            }
+        return params
+    }
 
-        val packageName: String = newScrimView.context.opPackageName
-        params.packageName = packageName
-        params.hideTimeoutMilliseconds = OVERLAY_HIDE_TIMEOUT_MILLIS
+    private fun createLightRevealEffect(): LightRevealEffect {
+        val isVerticalFold = currentRotation == Surface.ROTATION_0 ||
+            currentRotation == Surface.ROTATION_180
+        return LinearLightRevealEffect(isVertical = isVerticalFold)
+    }
+
+    private fun ensureOverlayRemoved() {
+        root?.release()
+        root = null
+        scrimView = null
+    }
 
-        if (scrimView?.parent != null) {
-            windowManager.removeView(scrimView)
+    private fun getUnfoldedDisplayInfo(): DisplayInfo =
+        displayManager.displays
+            .asSequence()
+            .map { DisplayInfo().apply { it.getDisplayInfo(this) } }
+            .filter { it.type == Display.TYPE_INTERNAL }
+            .maxByOrNull { it.naturalWidth }!!
+
+    private inner class TransitionListener : TransitionProgressListener {
+
+        override fun onTransitionProgress(progress: Float) {
+            scrimView?.revealAmount = progress
         }
 
-        this.scrimView = newScrimView
+        override fun onTransitionFinished() {
+            ensureOverlayRemoved()
+        }
 
-        try {
-            windowManager.addView(scrimView, params)
-        } catch (e: WindowManager.BadTokenException) {
-            e.printStackTrace()
+        override fun onTransitionStarted() {
+            // Add view for folding case (when unfolding the view is added earlier)
+            if (scrimView == null) {
+                addView()
+            }
         }
     }
 
-    private fun removeOverlayView() {
-        scrimView?.let {
-            if (it.parent != null) {
-                windowManager.removeViewImmediate(it)
+    private inner class DisplayChangeListener : DisplayManager.DisplayListener {
+
+        override fun onDisplayChanged(displayId: Int) {
+            val newRotation: Int = context.display!!.rotation
+            if (currentRotation != newRotation) {
+                currentRotation = newRotation
+                scrimView?.revealEffect = createLightRevealEffect()
+                root?.relayout(getLayoutParams())
             }
-            scrimView = null
+        }
+
+        override fun onDisplayAdded(displayId: Int) {
+        }
+
+        override fun onDisplayRemoved(displayId: Int) {
         }
     }
-}
 
-private const val OVERLAY_HIDE_TIMEOUT_MILLIS = 10_000L
+    private inner class FoldListener : FoldStateListener(context, Consumer { isFolded ->
+        if (isFolded) {
+            ensureOverlayRemoved()
+            isUnfoldHandled = false
+        }
+        this.isFolded = isFolded
+    })
+}
diff --git a/packages/SystemUI/src/com/android/systemui/wmshell/WMShellBaseModule.java b/packages/SystemUI/src/com/android/systemui/wmshell/WMShellBaseModule.java
index 514a9035a9ef..5af2cc545d52 100644
--- a/packages/SystemUI/src/com/android/systemui/wmshell/WMShellBaseModule.java
+++ b/packages/SystemUI/src/com/android/systemui/wmshell/WMShellBaseModule.java
@@ -28,6 +28,7 @@ import com.android.internal.logging.UiEventLogger;
 import com.android.internal.statusbar.IStatusBarService;
 import com.android.systemui.dagger.WMComponent;
 import com.android.systemui.dagger.WMSingleton;
+import com.android.wm.shell.RootDisplayAreaOrganizer;
 import com.android.wm.shell.RootTaskDisplayAreaOrganizer;
 import com.android.wm.shell.ShellCommandHandler;
 import com.android.wm.shell.ShellCommandHandlerImpl;
@@ -54,6 +55,8 @@ import com.android.wm.shell.common.TransactionPool;
 import com.android.wm.shell.common.annotations.ShellAnimationThread;
 import com.android.wm.shell.common.annotations.ShellMainThread;
 import com.android.wm.shell.common.annotations.ShellSplashscreenThread;
+import com.android.wm.shell.displayareahelper.DisplayAreaHelper;
+import com.android.wm.shell.displayareahelper.DisplayAreaHelperController;
 import com.android.wm.shell.draganddrop.DragAndDropController;
 import com.android.wm.shell.freeform.FreeformTaskListener;
 import com.android.wm.shell.fullscreen.FullscreenTaskListener;
@@ -342,13 +345,21 @@ public abstract class WMShellBaseModule {
         return taskSurfaceController.map((controller) -> controller.asTaskSurfaceHelper());
     }
 
-    @WMSingleton
     @Provides
     static Optional<TaskSurfaceHelperController> provideTaskSurfaceHelperController(
             ShellTaskOrganizer taskOrganizer, @ShellMainThread ShellExecutor mainExecutor) {
         return Optional.ofNullable(new TaskSurfaceHelperController(taskOrganizer, mainExecutor));
     }
 
+    @WMSingleton
+    @Provides
+    static Optional<DisplayAreaHelper> provideDisplayAreaHelper(
+            @ShellMainThread ShellExecutor mainExecutor,
+            RootDisplayAreaOrganizer rootDisplayAreaOrganizer) {
+        return Optional.ofNullable(new DisplayAreaHelperController(mainExecutor,
+                rootDisplayAreaOrganizer));
+    }
+
     //
     // Pip (optional feature)
     //
@@ -424,6 +435,13 @@ public abstract class WMShellBaseModule {
 
     @WMSingleton
     @Provides
+    static RootDisplayAreaOrganizer provideRootDisplayAreaOrganizer(
+            @ShellMainThread ShellExecutor mainExecutor) {
+        return new RootDisplayAreaOrganizer(mainExecutor);
+    }
+
+    @WMSingleton
+    @Provides
     static Optional<SplitScreen> provideSplitScreen(
             Optional<SplitScreenController> splitScreenController) {
         return splitScreenController.map((controller) -> controller.asSplitScreen());
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardViewMediatorTest.java b/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardViewMediatorTest.java
index 31d70f5c811f..1bb660e4cced 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardViewMediatorTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardViewMediatorTest.java
@@ -34,12 +34,14 @@ import android.app.admin.DevicePolicyManager;
 import android.app.trust.TrustManager;
 import android.os.PowerManager;
 import android.os.PowerManager.WakeLock;
+import android.os.RemoteException;
 import android.telephony.TelephonyManager;
 import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
 
 import androidx.test.filters.SmallTest;
 
+import com.android.internal.policy.IKeyguardDrawnCallback;
 import com.android.internal.widget.LockPatternUtils;
 import com.android.keyguard.KeyguardDisplayManager;
 import com.android.keyguard.KeyguardUpdateMonitor;
@@ -55,6 +57,8 @@ import com.android.systemui.statusbar.phone.StatusBarKeyguardViewManager;
 import com.android.systemui.statusbar.phone.UnlockedScreenOffAnimationController;
 import com.android.systemui.statusbar.policy.KeyguardStateController;
 import com.android.systemui.statusbar.policy.UserSwitcherController;
+import com.android.systemui.unfold.UnfoldLightRevealOverlayAnimation;
+import com.android.systemui.unfold.config.UnfoldTransitionConfig;
 import com.android.systemui.util.DeviceConfigProxy;
 import com.android.systemui.util.DeviceConfigProxyFake;
 import com.android.systemui.util.concurrency.FakeExecutor;
@@ -63,6 +67,7 @@ import com.android.systemui.util.time.FakeSystemClock;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
@@ -85,11 +90,14 @@ public class KeyguardViewMediatorTest extends SysuiTestCase {
     private @Mock NavigationModeController mNavigationModeController;
     private @Mock KeyguardDisplayManager mKeyguardDisplayManager;
     private @Mock DozeParameters mDozeParameters;
+    private @Mock UnfoldTransitionConfig mUnfoldTransitionConfig;
+    private @Mock UnfoldLightRevealOverlayAnimation mUnfoldAnimation;
     private @Mock SysuiStatusBarStateController mStatusBarStateController;
     private @Mock KeyguardStateController mKeyguardStateController;
     private @Mock NotificationShadeDepthController mNotificationShadeDepthController;
     private @Mock KeyguardUnlockAnimationController mKeyguardUnlockAnimationController;
     private @Mock UnlockedScreenOffAnimationController mUnlockedScreenOffAnimationController;
+    private @Mock IKeyguardDrawnCallback mKeyguardDrawnCallback;
     private DeviceConfigProxy mDeviceConfig = new DeviceConfigProxyFake();
     private FakeExecutor mUiBgExecutor = new FakeExecutor(new FakeSystemClock());
 
@@ -120,6 +128,8 @@ public class KeyguardViewMediatorTest extends SysuiTestCase {
                 mNavigationModeController,
                 mKeyguardDisplayManager,
                 mDozeParameters,
+                mUnfoldTransitionConfig,
+                () -> mUnfoldAnimation,
                 mStatusBarStateController,
                 mKeyguardStateController,
                 () -> mKeyguardUnlockAnimationController,
@@ -148,6 +158,33 @@ public class KeyguardViewMediatorTest extends SysuiTestCase {
     }
 
     @Test
+    @TestableLooper.RunWithLooper(setAsMainLooper = true)
+    public void testUnfoldTransitionEnabledDrawnTasksReady_onScreenTurningOn_callsDrawnCallback()
+            throws RemoteException {
+        when(mUnfoldTransitionConfig.isEnabled()).thenReturn(true);
+
+        mViewMediator.onScreenTurningOn(mKeyguardDrawnCallback);
+        TestableLooper.get(this).processAllMessages();
+        onUnfoldOverlayReady();
+
+        // Should be called when both unfold overlay and keyguard drawn ready
+        verify(mKeyguardDrawnCallback).onDrawn();
+    }
+
+    @Test
+    @TestableLooper.RunWithLooper(setAsMainLooper = true)
+    public void testUnfoldTransitionDisabledDrawnTasksReady_onScreenTurningOn_callsDrawnCallback()
+            throws RemoteException {
+        when(mUnfoldTransitionConfig.isEnabled()).thenReturn(false);
+
+        mViewMediator.onScreenTurningOn(mKeyguardDrawnCallback);
+        TestableLooper.get(this).processAllMessages();
+
+        // Should be called when only keyguard drawn
+        verify(mKeyguardDrawnCallback).onDrawn();
+    }
+
+    @Test
     public void testIsAnimatingScreenOff() {
         when(mDozeParameters.shouldControlUnlockedScreenOff()).thenReturn(true);
 
@@ -187,4 +224,11 @@ public class KeyguardViewMediatorTest extends SysuiTestCase {
         // then make sure it comes back
         verify(mStatusBarKeyguardViewManager, atLeast(1)).show(null);
     }
+
+    private void onUnfoldOverlayReady() {
+        ArgumentCaptor<Runnable> overlayReadyCaptor = ArgumentCaptor.forClass(Runnable.class);
+        verify(mUnfoldAnimation).onScreenTurningOn(overlayReadyCaptor.capture());
+        overlayReadyCaptor.getValue().run();
+        TestableLooper.get(this).processAllMessages();
+    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/unfold/updates/DeviceFoldStateProviderTest.kt b/packages/SystemUI/tests/src/com/android/systemui/unfold/updates/DeviceFoldStateProviderTest.kt
new file mode 100644
index 000000000000..a1d9a7b50d81
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/unfold/updates/DeviceFoldStateProviderTest.kt
@@ -0,0 +1,178 @@
+/*
+ * 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.systemui.unfold.updates
+
+import android.hardware.devicestate.DeviceStateManager
+import android.hardware.devicestate.DeviceStateManager.FoldStateListener
+import android.testing.AndroidTestingRunner
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.unfold.updates.hinge.HingeAngleProvider
+import com.android.systemui.unfold.updates.screen.ScreenStatusProvider
+import com.android.systemui.unfold.updates.screen.ScreenStatusProvider.ScreenListener
+import com.android.systemui.util.mockito.any
+import com.google.common.truth.Truth.assertThat
+import org.junit.Assume.assumeTrue
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.ArgumentCaptor
+import org.mockito.Mock
+import org.mockito.Mockito.verify
+import org.mockito.MockitoAnnotations
+
+@RunWith(AndroidTestingRunner::class)
+@SmallTest
+class DeviceFoldStateProviderTest : SysuiTestCase() {
+
+    @Mock
+    private lateinit var hingeAngleProvider: HingeAngleProvider
+
+    @Mock
+    private lateinit var screenStatusProvider: ScreenStatusProvider
+
+    @Mock
+    private lateinit var deviceStateManager: DeviceStateManager
+
+    private lateinit var foldStateProvider: FoldStateProvider
+
+    private val foldUpdates: MutableList<Int> = arrayListOf()
+    private val hingeAngleUpdates: MutableList<Float> = arrayListOf()
+
+    private val foldStateListenerCaptor = ArgumentCaptor.forClass(FoldStateListener::class.java)
+    private var foldedDeviceState: Int = 0
+    private var unfoldedDeviceState: Int = 0
+
+    private val screenOnListenerCaptor = ArgumentCaptor.forClass(ScreenListener::class.java)
+
+    @Before
+    fun setUp() {
+        MockitoAnnotations.initMocks(this)
+        val foldedDeviceStates: IntArray = context.resources.getIntArray(
+            com.android.internal.R.array.config_foldedDeviceStates)
+        assumeTrue("Test should be launched on a foldable device",
+            foldedDeviceStates.isNotEmpty())
+
+        foldedDeviceState = foldedDeviceStates.maxOrNull()!!
+        unfoldedDeviceState = foldedDeviceState + 1
+
+        foldStateProvider = DeviceFoldStateProvider(
+            context,
+            hingeAngleProvider,
+            screenStatusProvider,
+            deviceStateManager,
+            context.mainExecutor
+        )
+
+        foldStateProvider.addCallback(object : FoldStateProvider.FoldUpdatesListener {
+            override fun onHingeAngleUpdate(angle: Float) {
+                hingeAngleUpdates.add(angle)
+            }
+
+            override fun onFoldUpdate(update: Int) {
+                foldUpdates.add(update)
+            }
+        })
+        foldStateProvider.start()
+
+        verify(deviceStateManager).registerCallback(any(), foldStateListenerCaptor.capture())
+        verify(screenStatusProvider).addCallback(screenOnListenerCaptor.capture())
+    }
+
+    @Test
+    fun testOnFolded_emitsFinishClosedEvent() {
+        setFoldState(folded = true)
+
+        assertThat(foldUpdates).containsExactly(FOLD_UPDATE_FINISH_CLOSED)
+    }
+
+    @Test
+    fun testOnUnfolded_emitsStartOpeningEvent() {
+        setFoldState(folded = false)
+
+        assertThat(foldUpdates).containsExactly(FOLD_UPDATE_START_OPENING)
+    }
+
+    @Test
+    fun testOnFolded_stopsHingeAngleProvider() {
+        setFoldState(folded = true)
+
+        verify(hingeAngleProvider).stop()
+    }
+
+    @Test
+    fun testOnUnfolded_startsHingeAngleProvider() {
+        setFoldState(folded = false)
+
+        verify(hingeAngleProvider).start()
+    }
+
+    @Test
+    fun testFirstScreenOnEventWhenFolded_doesNotEmitEvents() {
+        setFoldState(folded = true)
+        foldUpdates.clear()
+
+        fireScreenOnEvent()
+
+        // Power button turn on
+        assertThat(foldUpdates).isEmpty()
+    }
+
+    @Test
+    fun testFirstScreenOnEventWhenUnfolded_doesNotEmitEvents() {
+        setFoldState(folded = false)
+        foldUpdates.clear()
+
+        fireScreenOnEvent()
+
+        assertThat(foldUpdates).isEmpty()
+    }
+
+    @Test
+    fun testFirstScreenOnEventAfterFoldAndUnfold_emitsUnfoldedScreenAvailableEvent() {
+        setFoldState(folded = false)
+        setFoldState(folded = true)
+        fireScreenOnEvent()
+        setFoldState(folded = false)
+        foldUpdates.clear()
+
+        fireScreenOnEvent()
+
+        assertThat(foldUpdates).containsExactly(FOLD_UPDATE_UNFOLDED_SCREEN_AVAILABLE)
+    }
+
+    @Test
+    fun testSecondScreenOnEventWhenUnfolded_doesNotEmitEvents() {
+        setFoldState(folded = false)
+        fireScreenOnEvent()
+        foldUpdates.clear()
+
+        fireScreenOnEvent()
+
+        // No events as this is power button turn on
+        assertThat(foldUpdates).isEmpty()
+    }
+
+    private fun setFoldState(folded: Boolean) {
+        val state = if (folded) foldedDeviceState else unfoldedDeviceState
+        foldStateListenerCaptor.value.onStateChanged(state)
+    }
+
+    private fun fireScreenOnEvent() {
+        screenOnListenerCaptor.value.onScreenTurnedOn()
+    }
+}
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index ff4dd46b1f18..2a598c85414a 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -297,7 +297,7 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
      * The direct child layer of the display to put all non-overlay windows. This is also used for
      * screen rotation animation so that there is a parent layer to put the animation leash.
      */
-    private final SurfaceControl mWindowingLayer;
+    private SurfaceControl mWindowingLayer;
 
     /**
      * The window token of the layer of the hierarchy to mirror, or null if this DisplayContent
@@ -329,7 +329,7 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
     private final ImeContainer mImeWindowsContainer = new ImeContainer(mWmService);
 
     @VisibleForTesting
-    final DisplayAreaPolicy mDisplayAreaPolicy;
+    DisplayAreaPolicy mDisplayAreaPolicy;
 
     private WindowState mTmpWindow;
     private boolean mUpdateImeTarget;
@@ -1104,52 +1104,91 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
         mDividerControllerLocked = new DockedTaskDividerController(this);
         mPinnedTaskController = new PinnedTaskController(mWmService, this);
 
+        final Transaction pendingTransaction = getPendingTransaction();
+        configureSurfaces(pendingTransaction);
+        pendingTransaction.apply();
+
+        // Sets the display content for the children.
+        onDisplayChanged(this);
+        updateDisplayAreaOrganizers();
+
+        mInputMonitor = new InputMonitor(mWmService, this);
+        mInsetsPolicy = new InsetsPolicy(mInsetsStateController, this);
+
+        if (DEBUG_DISPLAY) Slog.v(TAG_WM, "Creating display=" + display);
+
+        mWmService.mDisplayWindowSettings.applySettingsToDisplayLocked(this);
+    }
+
+    @Override
+    void migrateToNewSurfaceControl(Transaction t) {
+        t.remove(mSurfaceControl);
+
+        mLastSurfacePosition.set(0, 0);
+
+        configureSurfaces(t);
+
+        for (int i = 0; i < mChildren.size(); i++)  {
+            SurfaceControl sc = mChildren.get(i).getSurfaceControl();
+            if (sc != null) {
+                t.reparent(sc, mSurfaceControl);
+            }
+        }
+
+        scheduleAnimation();
+    }
+
+    /**
+     * Configures the surfaces hierarchy for DisplayContent
+     * This method always recreates the main surface control but reparents the children
+     * if they are already created.
+     * @param transaction as part of which to perform the configuration
+     */
+    private void configureSurfaces(Transaction transaction) {
         final SurfaceControl.Builder b = mWmService.makeSurfaceBuilder(mSession)
                 .setOpaque(true)
                 .setContainerLayer()
                 .setCallsite("DisplayContent");
-        mSurfaceControl = b.setName("Root").setContainerLayer().build();
+        mSurfaceControl = b.setName(getName()).setContainerLayer().build();
 
-        // Setup the policy and build the display area hierarchy.
-        mDisplayAreaPolicy = mWmService.getDisplayAreaPolicyProvider().instantiate(
-                mWmService, this /* content */, this /* root */, mImeWindowsContainer);
+        if (mDisplayAreaPolicy == null) {
+            // Setup the policy and build the display area hierarchy.
+            // Build the hierarchy only after creating the surface so it is reparented correctly
+            mDisplayAreaPolicy = mWmService.getDisplayAreaPolicyProvider().instantiate(
+                    mWmService, this /* content */, this /* root */,
+                    mImeWindowsContainer);
+        }
 
         final List<DisplayArea<? extends WindowContainer>> areas =
                 mDisplayAreaPolicy.getDisplayAreas(FEATURE_WINDOWED_MAGNIFICATION);
         final DisplayArea<?> area = areas.size() == 1 ? areas.get(0) : null;
+
         if (area != null && area.getParent() == this) {
             // The windowed magnification area should contain all non-overlay windows, so just use
             // it as the windowing layer.
             mWindowingLayer = area.mSurfaceControl;
+            transaction.reparent(mWindowingLayer, mSurfaceControl);
         } else {
             // Need an additional layer for screen level animation, so move the layer containing
             // the windows to the new root.
             mWindowingLayer = mSurfaceControl;
             mSurfaceControl = b.setName("RootWrapper").build();
-            getPendingTransaction().reparent(mWindowingLayer, mSurfaceControl)
+            transaction.reparent(mWindowingLayer, mSurfaceControl)
                     .show(mWindowingLayer);
         }
 
-        mOverlayLayer = b.setName("Display Overlays").setParent(mSurfaceControl).build();
+        if (mOverlayLayer == null) {
+            mOverlayLayer = b.setName("Display Overlays").setParent(mSurfaceControl).build();
+        } else {
+            transaction.reparent(mOverlayLayer, mSurfaceControl);
+        }
 
-        getPendingTransaction()
+        transaction
                 .setLayer(mSurfaceControl, 0)
                 .setLayerStack(mSurfaceControl, mDisplayId)
                 .show(mSurfaceControl)
                 .setLayer(mOverlayLayer, Integer.MAX_VALUE)
                 .show(mOverlayLayer);
-        getPendingTransaction().apply();
-
-        // Sets the display content for the children.
-        onDisplayChanged(this);
-        updateDisplayAreaOrganizers();
-
-        mInputMonitor = new InputMonitor(mWmService, this);
-        mInsetsPolicy = new InsetsPolicy(mInsetsStateController, this);
-
-        if (DEBUG_DISPLAY) Slog.v(TAG_WM, "Creating display=" + display);
-
-        mWmService.mDisplayWindowSettings.applySettingsToDisplayLocked(this);
     }
 
     boolean isReady() {