Merge changes from topic "keyguard-occlude-race-transitions" into udc-dev

* changes:
  All keyguard occludes go in-band with transitions
  Call finishedCallback after transitioning away from DREAM
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
index 18bd78f..2d82c50 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
@@ -139,6 +139,7 @@
 import com.android.systemui.flags.Flags;
 import com.android.systemui.flags.SystemPropertiesHelper;
 import com.android.systemui.keyguard.dagger.KeyguardModule;
+import com.android.systemui.keyguard.shared.model.TransitionStep;
 import com.android.systemui.keyguard.ui.viewmodel.DreamingToLockscreenTransitionViewModel;
 import com.android.systemui.log.SessionTracker;
 import com.android.systemui.navigationbar.NavigationModeController;
@@ -538,6 +539,8 @@
 
     private CentralSurfaces mCentralSurfaces;
 
+    private IRemoteAnimationFinishedCallback mUnoccludeFromDreamFinishedCallback;
+
     private final DeviceConfig.OnPropertiesChangedListener mOnPropertiesChangedListener =
             new DeviceConfig.OnPropertiesChangedListener() {
             @Override
@@ -1164,6 +1167,7 @@
                             getRemoteSurfaceAlphaApplier().accept(0.0f);
                             mDreamingToLockscreenTransitionViewModel.get()
                                     .startTransition();
+                            mUnoccludeFromDreamFinishedCallback = finishedCallback;
                             return;
                         }
 
@@ -1243,6 +1247,19 @@
         };
     }
 
+    private Consumer<TransitionStep> getFinishedCallbackConsumer() {
+        return (TransitionStep step) -> {
+            if (mUnoccludeFromDreamFinishedCallback == null) return;
+            try {
+                mUnoccludeFromDreamFinishedCallback.onAnimationFinished();
+                mUnoccludeFromDreamFinishedCallback = null;
+            } catch (RemoteException e) {
+                Log.e(TAG, "Wasn't able to callback", e);
+            }
+            mInteractionJankMonitor.end(CUJ_LOCKSCREEN_OCCLUSION);
+        };
+    }
+
     private DeviceConfigProxy mDeviceConfig;
     private DozeParameters mDozeParameters;
 
@@ -1502,6 +1519,9 @@
                 collectFlow(viewRootImpl.getView(),
                         mDreamingToLockscreenTransitionViewModel.get().getDreamOverlayAlpha(),
                         getRemoteSurfaceAlphaApplier(), mMainDispatcher);
+                collectFlow(viewRootImpl.getView(),
+                        mDreamingToLockscreenTransitionViewModel.get().getTransitionEnded(),
+                        getFinishedCallbackConsumer(), mMainDispatcher);
             }
         }
         // Most services aren't available until the system reaches the ready state, so we
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractor.kt
index 42f12f8..41a81a8 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractor.kt
@@ -49,11 +49,15 @@
 ) {
     /** (any)->GONE transition information */
     val anyStateToGoneTransition: Flow<TransitionStep> =
-        repository.transitions.filter { step -> step.to == KeyguardState.GONE }
+        repository.transitions.filter { step -> step.to == GONE }
 
     /** (any)->AOD transition information */
     val anyStateToAodTransition: Flow<TransitionStep> =
-        repository.transitions.filter { step -> step.to == KeyguardState.AOD }
+        repository.transitions.filter { step -> step.to == AOD }
+
+    /** DREAMING->(any) transition information. */
+    val fromDreamingTransition: Flow<TransitionStep> =
+        repository.transitions.filter { step -> step.from == DREAMING }
 
     /** AOD->LOCKSCREEN transition information. */
     val aodToLockscreenTransition: Flow<TransitionStep> = repository.transition(AOD, LOCKSCREEN)
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/DreamingToLockscreenTransitionViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/DreamingToLockscreenTransitionViewModel.kt
index 9ca4bd6..e24d326 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/DreamingToLockscreenTransitionViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/DreamingToLockscreenTransitionViewModel.kt
@@ -48,7 +48,7 @@
         )
 
     val transitionEnded =
-        keyguardTransitionInteractor.dreamingToLockscreenTransition.filter { step ->
+        keyguardTransitionInteractor.fromDreamingTransition.filter { step ->
             step.transitionState == TransitionState.FINISHED ||
                 step.transitionState == TransitionState.CANCELED
         }
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 c280538..7a501a8 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardViewMediatorTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardViewMediatorTest.java
@@ -186,6 +186,8 @@
         when(mStatusBarKeyguardViewManager.getViewRootImpl()).thenReturn(testViewRoot);
         when(mDreamingToLockscreenTransitionViewModel.getDreamOverlayAlpha())
                 .thenReturn(mock(Flow.class));
+        when(mDreamingToLockscreenTransitionViewModel.getTransitionEnded())
+                .thenReturn(mock(Flow.class));
         mNotificationShadeWindowController = new NotificationShadeWindowControllerImpl(mContext,
                 mWindowManager, mActivityManager, mDozeParameters, mStatusBarStateController,
                 mConfigurationController, mViewMediator, mKeyguardBypassController,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/DreamingToLockscreenTransitionViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/DreamingToLockscreenTransitionViewModelTest.kt
index a341346..ab994b7 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/DreamingToLockscreenTransitionViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/DreamingToLockscreenTransitionViewModelTest.kt
@@ -22,8 +22,16 @@
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.keyguard.data.repository.FakeKeyguardTransitionRepository
 import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor
-import com.android.systemui.keyguard.shared.model.KeyguardState
+import com.android.systemui.keyguard.shared.model.KeyguardState.AOD
+import com.android.systemui.keyguard.shared.model.KeyguardState.DOZING
+import com.android.systemui.keyguard.shared.model.KeyguardState.DREAMING
+import com.android.systemui.keyguard.shared.model.KeyguardState.GONE
+import com.android.systemui.keyguard.shared.model.KeyguardState.LOCKSCREEN
 import com.android.systemui.keyguard.shared.model.TransitionState
+import com.android.systemui.keyguard.shared.model.TransitionState.CANCELED
+import com.android.systemui.keyguard.shared.model.TransitionState.FINISHED
+import com.android.systemui.keyguard.shared.model.TransitionState.RUNNING
+import com.android.systemui.keyguard.shared.model.TransitionState.STARTED
 import com.android.systemui.keyguard.shared.model.TransitionStep
 import com.android.systemui.util.mockito.mock
 import com.google.common.collect.Range
@@ -60,7 +68,7 @@
             val job =
                 underTest.dreamOverlayTranslationY(pixels).onEach { values.add(it) }.launchIn(this)
 
-            repository.sendTransitionStep(step(0f, TransitionState.STARTED))
+            repository.sendTransitionStep(step(0f, STARTED))
             repository.sendTransitionStep(step(0f))
             repository.sendTransitionStep(step(0.3f))
             repository.sendTransitionStep(step(0.5f))
@@ -82,7 +90,7 @@
             val job = underTest.dreamOverlayAlpha.onEach { values.add(it) }.launchIn(this)
 
             // Should start running here...
-            repository.sendTransitionStep(step(0f, TransitionState.STARTED))
+            repository.sendTransitionStep(step(0f, STARTED))
             repository.sendTransitionStep(step(0f))
             repository.sendTransitionStep(step(0.1f))
             repository.sendTransitionStep(step(0.5f))
@@ -104,7 +112,7 @@
 
             val job = underTest.lockscreenAlpha.onEach { values.add(it) }.launchIn(this)
 
-            repository.sendTransitionStep(step(0f, TransitionState.STARTED))
+            repository.sendTransitionStep(step(0f, STARTED))
             repository.sendTransitionStep(step(0f))
             repository.sendTransitionStep(step(0.1f))
             repository.sendTransitionStep(step(0.2f))
@@ -126,7 +134,7 @@
             val job =
                 underTest.lockscreenTranslationY(pixels).onEach { values.add(it) }.launchIn(this)
 
-            repository.sendTransitionStep(step(0f, TransitionState.STARTED))
+            repository.sendTransitionStep(step(0f, STARTED))
             repository.sendTransitionStep(step(0f))
             repository.sendTransitionStep(step(0.3f))
             repository.sendTransitionStep(step(0.5f))
@@ -138,13 +146,44 @@
             job.cancel()
         }
 
-    private fun step(
-        value: Float,
-        state: TransitionState = TransitionState.RUNNING
-    ): TransitionStep {
+    @Test
+    fun transitionEnded() =
+        runTest(UnconfinedTestDispatcher()) {
+            val values = mutableListOf<TransitionStep>()
+
+            val job = underTest.transitionEnded.onEach { values.add(it) }.launchIn(this)
+
+            repository.sendTransitionStep(TransitionStep(DOZING, DREAMING, 0.0f, STARTED))
+            repository.sendTransitionStep(TransitionStep(DOZING, DREAMING, 1.0f, FINISHED))
+
+            repository.sendTransitionStep(TransitionStep(DREAMING, LOCKSCREEN, 0.0f, STARTED))
+            repository.sendTransitionStep(TransitionStep(DREAMING, LOCKSCREEN, 0.1f, RUNNING))
+            repository.sendTransitionStep(TransitionStep(DREAMING, LOCKSCREEN, 1.0f, FINISHED))
+
+            repository.sendTransitionStep(TransitionStep(LOCKSCREEN, DREAMING, 0.0f, STARTED))
+            repository.sendTransitionStep(TransitionStep(LOCKSCREEN, DREAMING, 0.5f, RUNNING))
+            repository.sendTransitionStep(TransitionStep(LOCKSCREEN, DREAMING, 1.0f, FINISHED))
+
+            repository.sendTransitionStep(TransitionStep(DREAMING, GONE, 0.0f, STARTED))
+            repository.sendTransitionStep(TransitionStep(DREAMING, GONE, 0.5f, RUNNING))
+            repository.sendTransitionStep(TransitionStep(DREAMING, GONE, 1.0f, CANCELED))
+
+            repository.sendTransitionStep(TransitionStep(DREAMING, AOD, 0.0f, STARTED))
+            repository.sendTransitionStep(TransitionStep(DREAMING, AOD, 1.0f, FINISHED))
+
+            assertThat(values.size).isEqualTo(3)
+            values.forEach {
+                assertThat(it.transitionState == FINISHED || it.transitionState == CANCELED)
+                    .isTrue()
+            }
+
+            job.cancel()
+        }
+
+    private fun step(value: Float, state: TransitionState = RUNNING): TransitionStep {
         return TransitionStep(
-            from = KeyguardState.DREAMING,
-            to = KeyguardState.LOCKSCREEN,
+            from = DREAMING,
+            to = LOCKSCREEN,
             value = value,
             transitionState = state,
             ownerName = "DreamingToLockscreenTransitionViewModelTest"
diff --git a/services/core/java/com/android/server/policy/PhoneWindowManager.java b/services/core/java/com/android/server/policy/PhoneWindowManager.java
index 3132c5d..cbded89 100644
--- a/services/core/java/com/android/server/policy/PhoneWindowManager.java
+++ b/services/core/java/com/android/server/policy/PhoneWindowManager.java
@@ -3570,19 +3570,17 @@
     }
 
     @Override
-    public void onKeyguardOccludedChangedLw(boolean occluded, boolean waitAppTransition) {
-        if (mKeyguardDelegate != null && waitAppTransition) {
+    public void onKeyguardOccludedChangedLw(boolean occluded) {
+        if (mKeyguardDelegate != null) {
             mPendingKeyguardOccluded = occluded;
             mKeyguardOccludedChanged = true;
-        } else {
-            setKeyguardOccludedLw(occluded);
         }
     }
 
     @Override
     public int applyKeyguardOcclusionChange() {
         if (DEBUG_KEYGUARD) Slog.d(TAG, "transition/occluded commit occluded="
-                + mPendingKeyguardOccluded);
+                + mPendingKeyguardOccluded + " changed=" + mKeyguardOccludedChanged);
 
         // TODO(b/276433230): Explicitly save before/after for occlude state in each
         // Transition so we don't need to update SysUI every time.
diff --git a/services/core/java/com/android/server/policy/WindowManagerPolicy.java b/services/core/java/com/android/server/policy/WindowManagerPolicy.java
index 887f946..03a7bd3 100644
--- a/services/core/java/com/android/server/policy/WindowManagerPolicy.java
+++ b/services/core/java/com/android/server/policy/WindowManagerPolicy.java
@@ -169,7 +169,7 @@
      *
      * @param occluded Whether Keyguard is currently occluded or not.
      */
-    void onKeyguardOccludedChangedLw(boolean occluded, boolean waitAppTransition);
+    void onKeyguardOccludedChangedLw(boolean occluded);
 
     /**
      * Commit any queued changes to keyguard occlude status that had been deferred during the
diff --git a/services/core/java/com/android/server/wm/KeyguardController.java b/services/core/java/com/android/server/wm/KeyguardController.java
index ad9c3b2..83fd725 100644
--- a/services/core/java/com/android/server/wm/KeyguardController.java
+++ b/services/core/java/com/android/server/wm/KeyguardController.java
@@ -418,13 +418,17 @@
             return;
         }
 
-        final boolean waitAppTransition = isKeyguardLocked(displayId);
-        mWindowManager.mPolicy.onKeyguardOccludedChangedLw(isDisplayOccluded(DEFAULT_DISPLAY),
-                waitAppTransition);
-        if (waitAppTransition) {
-            mService.deferWindowLayout();
-            try {
-                if (isDisplayOccluded(DEFAULT_DISPLAY)) {
+        final TransitionController tc = mRootWindowContainer.mTransitionController;
+
+        final boolean occluded = isDisplayOccluded(displayId);
+        final boolean performTransition = isKeyguardLocked(displayId);
+        final boolean executeTransition = performTransition && !tc.isCollecting();
+
+        mWindowManager.mPolicy.onKeyguardOccludedChangedLw(occluded);
+        mService.deferWindowLayout();
+        try {
+            if (isKeyguardLocked(displayId)) {
+                if (occluded) {
                     mRootWindowContainer.getDefaultDisplay().requestTransitionAndLegacyPrepare(
                             TRANSIT_KEYGUARD_OCCLUDE,
                             TRANSIT_FLAG_KEYGUARD_OCCLUDING,
@@ -434,11 +438,19 @@
                             TRANSIT_KEYGUARD_UNOCCLUDE,
                             TRANSIT_FLAG_KEYGUARD_UNOCCLUDING);
                 }
-                updateKeyguardSleepToken(DEFAULT_DISPLAY);
-                mWindowManager.executeAppTransition();
-            } finally {
-                mService.continueWindowLayout();
+            } else {
+                if (tc.inTransition()) {
+                    tc.mStateValidators.add(mWindowManager.mPolicy::applyKeyguardOcclusionChange);
+                } else {
+                    mWindowManager.mPolicy.applyKeyguardOcclusionChange();
+                }
             }
+            updateKeyguardSleepToken(displayId);
+            if (performTransition && executeTransition) {
+                mWindowManager.executeAppTransition();
+            }
+        } finally {
+            mService.continueWindowLayout();
         }
     }
 
@@ -485,6 +497,9 @@
         }
     }
 
+    /**
+     * @return true if Keyguard is occluded or the device is dreaming.
+     */
     boolean isDisplayOccluded(int displayId) {
         return getDisplayState(displayId).mOccluded;
     }
diff --git a/services/core/java/com/android/server/wm/Transition.java b/services/core/java/com/android/server/wm/Transition.java
index 738869b..2c582c1 100644
--- a/services/core/java/com/android/server/wm/Transition.java
+++ b/services/core/java/com/android/server/wm/Transition.java
@@ -27,6 +27,7 @@
 import static android.view.Display.DEFAULT_DISPLAY;
 import static android.view.Display.INVALID_DISPLAY;
 import static android.view.WindowManager.INPUT_CONSUMER_RECENTS_ANIMATION;
+import static android.view.WindowManager.KEYGUARD_VISIBILITY_TRANSIT_FLAGS;
 import static android.view.WindowManager.LayoutParams.ROTATION_ANIMATION_SEAMLESS;
 import static android.view.WindowManager.LayoutParams.ROTATION_ANIMATION_UNSPECIFIED;
 import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_STARTING;
@@ -2659,7 +2660,7 @@
     }
 
     private void validateKeyguardOcclusion() {
-        if ((mFlags & TRANSIT_FLAG_KEYGUARD_LOCKED) != 0) {
+        if ((mFlags & KEYGUARD_VISIBILITY_TRANSIT_FLAGS) != 0) {
             mController.mStateValidators.add(
                 mController.mAtm.mWindowManager.mPolicy::applyKeyguardOcclusionChange);
         }
diff --git a/services/tests/wmtests/src/com/android/server/wm/TestWindowManagerPolicy.java b/services/tests/wmtests/src/com/android/server/wm/TestWindowManagerPolicy.java
index adf3f39..bd111ad 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TestWindowManagerPolicy.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TestWindowManagerPolicy.java
@@ -233,7 +233,7 @@
     }
 
     @Override
-    public void onKeyguardOccludedChangedLw(boolean occluded, boolean waitAppTransition) {
+    public void onKeyguardOccludedChangedLw(boolean occluded) {
     }
 
     public void setSafeMode(boolean safeMode) {