Smooth burn-in transitions

GONE->AOD transitions are easily interruptible by a power button push
to keep the phone on, or to prevent unlocking. This adds a smoother
transition from AOD->LOCKSCREEN which is next in the sequence. Rather
than rely on a preset starting point, it begins from the current
y-translation value and interpolates to the final resting point.

Also, adds the ability for animation flows to return both a value and
the transition state. The transition state is frequently used in view
models to know which transition is truly active.

Fixes: 322197793
Test: manual - interrupt GONE->AOD
Test: atest com.android.systemui.keyguard
Flag: ACONFIG com.android.systemui.keyguard_shade_migration_nssl
DEVELOPMENT

Change-Id: If9501500796bee9d421a291bf555077e8eefcc14
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/AodBurnInViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/AodBurnInViewModelTest.kt
index 0543bc2..d52696a 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/AodBurnInViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/AodBurnInViewModelTest.kt
@@ -20,6 +20,7 @@
 
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.systemui.Flags as AConfigFlags
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
@@ -68,6 +69,8 @@
 
     @Before
     fun setUp() {
+        mSetFlagsRule.disableFlags(AConfigFlags.FLAG_MIGRATE_CLOCKS_TO_BLUEPRINT)
+
         MockitoAnnotations.initMocks(this)
         whenever(burnInInteractor.keyguardBurnIn).thenReturn(burnInFlow)
         kosmos.burnInInteractor = burnInInteractor
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/BurnInInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/BurnInInteractor.kt
index cc1cf91..7ae70a9 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/BurnInInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/BurnInInteractor.kt
@@ -74,6 +74,7 @@
                 BurnInModel(translationX, translationY, burnInHelperWrapper.burnInScale())
             }
             .distinctUntilChanged()
+            .stateIn(scope, SharingStarted.Lazily, BurnInModel())
 
     /**
      * Use for max burn-in offsets that are NOT specified in pixels. This flow will recalculate the
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromAodTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromAodTransitionInteractor.kt
index 8fa33ee7..32de8e5 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromAodTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromAodTransitionInteractor.kt
@@ -97,6 +97,8 @@
                         val modeOnCanceled =
                             if (lastStartedStep.from == KeyguardState.LOCKSCREEN) {
                                 TransitionModeOnCanceled.REVERSE
+                            } else if (lastStartedStep.from == KeyguardState.GONE) {
+                                TransitionModeOnCanceled.RESET
                             } else {
                                 TransitionModeOnCanceled.LAST_VALUE
                             }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/KeyguardTransitionAnimationFlow.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/KeyguardTransitionAnimationFlow.kt
index 00b7989..8b278cd 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/KeyguardTransitionAnimationFlow.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/KeyguardTransitionAnimationFlow.kt
@@ -112,6 +112,38 @@
             interpolator: Interpolator = LINEAR,
             name: String? = null
         ): Flow<Float> {
+            return sharedFlowWithState(
+                    duration = duration,
+                    onStep = onStep,
+                    startTime = startTime,
+                    onStart = onStart,
+                    onCancel = onCancel,
+                    onFinish = onFinish,
+                    interpolator = interpolator,
+                    name = name,
+                )
+                .mapNotNull { stateToValue -> stateToValue.value }
+        }
+
+        /**
+         * Transitions will occur over a [transitionDuration] with [TransitionStep]s being emitted
+         * in the range of [0, 1]. View animations should begin and end within a subset of this
+         * range. This function maps the [startTime] and [duration] into [0, 1], when this subset is
+         * valid.
+         *
+         * Will return a [StateToValue], which encompasses the calculated value as well as the
+         * transitionState that is associated with it.
+         */
+        fun sharedFlowWithState(
+            duration: Duration,
+            onStep: (Float) -> Float,
+            startTime: Duration = 0.milliseconds,
+            onStart: (() -> Unit)? = null,
+            onCancel: (() -> Float)? = null,
+            onFinish: (() -> Float)? = null,
+            interpolator: Interpolator = LINEAR,
+            name: String? = null
+        ): Flow<StateToValue> {
             if (!duration.isPositive()) {
                 throw IllegalArgumentException("duration must be a positive number: $duration")
             }
@@ -164,7 +196,6 @@
                         .also { logger.logTransitionStep(name, step, it.value) }
                 }
                 .distinctUntilChanged()
-                .mapNotNull { stateToValue -> stateToValue.value }
         }
 
         /**
@@ -174,9 +205,9 @@
             return sharedFlow(duration = 1.milliseconds, onStep = { value }, onFinish = { value })
         }
     }
-
-    data class StateToValue(
-        val transitionState: TransitionState,
-        val value: Float?,
-    )
 }
+
+data class StateToValue(
+    val transitionState: TransitionState = TransitionState.FINISHED,
+    val value: Float? = 0f,
+)
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardRootViewBinder.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardRootViewBinder.kt
index eb11c3f..9e7c70d 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardRootViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardRootViewBinder.kt
@@ -311,6 +311,12 @@
             }
         }
 
+        if (KeyguardShadeMigrationNssl.isEnabled) {
+            burnInParams.update { current ->
+                current.copy(translationY = { childViews[burnInLayerId]?.translationY })
+            }
+        }
+
         onLayoutChangeListener = OnLayoutChange(viewModel, burnInParams)
         view.addOnLayoutChangeListener(onLayoutChangeListener)
 
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/AodBurnInViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/AodBurnInViewModel.kt
index 828e033..8110de2 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/AodBurnInViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/AodBurnInViewModel.kt
@@ -31,6 +31,10 @@
 import com.android.systemui.keyguard.shared.model.KeyguardState.ALTERNATE_BOUNCER
 import com.android.systemui.keyguard.shared.model.KeyguardState.AOD
 import com.android.systemui.keyguard.shared.model.KeyguardState.GONE
+import com.android.systemui.keyguard.shared.model.TransitionState
+import com.android.systemui.keyguard.shared.model.TransitionState.RUNNING
+import com.android.systemui.keyguard.shared.model.TransitionState.STARTED
+import com.android.systemui.keyguard.ui.StateToValue
 import com.android.systemui.plugins.clocks.ClockController
 import com.android.systemui.res.R
 import javax.inject.Inject
@@ -58,6 +62,7 @@
     private val keyguardInteractor: KeyguardInteractor,
     private val keyguardTransitionInteractor: KeyguardTransitionInteractor,
     private val goneToAodTransitionViewModel: GoneToAodTransitionViewModel,
+    private val aodToLockscreenTransitionViewModel: AodToLockscreenTransitionViewModel,
     private val occludedToLockscreenTransitionViewModel: OccludedToLockscreenTransitionViewModel,
     private val keyguardClockViewModel: KeyguardClockViewModel,
 ) {
@@ -83,21 +88,22 @@
                     burnIn(params).map { it.translationY.toFloat() }.onStart { emit(0f) },
                     goneToAodTransitionViewModel
                         .enterFromTopTranslationY(enterFromTopAmount)
-                        .onStart { emit(0f) },
+                        .onStart { emit(StateToValue()) },
                     occludedToLockscreenTransitionViewModel.lockscreenTranslationY.onStart {
                         emit(0f)
                     },
-                ) {
-                    keyguardTransitionY,
-                    burnInTranslationY,
-                    goneToAodTransitionTranslationY,
-                    occludedToLockscreenTransitionTranslationY ->
-
-                    // All values need to be combined for a smooth translation
-                    keyguardTransitionY +
-                        burnInTranslationY +
-                        goneToAodTransitionTranslationY +
-                        occludedToLockscreenTransitionTranslationY
+                    aodToLockscreenTransitionViewModel.translationY(params.translationY).onStart {
+                        emit(StateToValue())
+                    },
+                ) { keyguardTranslationY, burnInY, goneToAod, occludedToLockscreen, aodToLockscreen
+                    ->
+                    if (isInTransition(aodToLockscreen.transitionState)) {
+                        aodToLockscreen.value ?: 0f
+                    } else if (isInTransition(goneToAod.transitionState)) {
+                        (goneToAod.value ?: 0f) + burnInY
+                    } else {
+                        burnInY + occludedToLockscreen + keyguardTranslationY
+                    }
                 }
             }
             .distinctUntilChanged()
@@ -115,6 +121,10 @@
         }
     }
 
+    private fun isInTransition(state: TransitionState): Boolean {
+        return state == STARTED || state == RUNNING
+    }
+
     private fun burnIn(
         params: BurnInParameters,
     ): Flow<BurnInModel> {
@@ -185,6 +195,8 @@
     val topInset: Int = 0,
     /** Status view top, without translation added in */
     val statusViewTop: Int = 0,
+    /** The current y translation of the view */
+    val translationY: () -> Float? = { null }
 )
 
 /**
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/AodToLockscreenTransitionViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/AodToLockscreenTransitionViewModel.kt
index 266fd02..6d1d3cb 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/AodToLockscreenTransitionViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/AodToLockscreenTransitionViewModel.kt
@@ -16,11 +16,14 @@
 
 package com.android.systemui.keyguard.ui.viewmodel
 
+import android.util.MathUtils
+import com.android.app.animation.Interpolators.FAST_OUT_SLOW_IN
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.deviceentry.domain.interactor.DeviceEntryUdfpsInteractor
 import com.android.systemui.keyguard.domain.interactor.FromAodTransitionInteractor.Companion.TO_LOCKSCREEN_DURATION
 import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.keyguard.ui.KeyguardTransitionAnimationFlow
+import com.android.systemui.keyguard.ui.StateToValue
 import com.android.systemui.keyguard.ui.transitions.DeviceEntryIconTransition
 import javax.inject.Inject
 import kotlin.time.Duration.Companion.milliseconds
@@ -48,6 +51,22 @@
             to = KeyguardState.LOCKSCREEN,
         )
 
+    /**
+     * Begin the transition from wherever the y-translation value is currently. This helps ensure a
+     * smooth transition if a transition in canceled.
+     */
+    fun translationY(currentTranslationY: () -> Float?): Flow<StateToValue> {
+        var startValue = 0f
+        return transitionAnimation.sharedFlowWithState(
+            duration = 500.milliseconds,
+            onStart = {
+                startValue = currentTranslationY() ?: 0f
+                startValue
+            },
+            onStep = { MathUtils.lerp(startValue, 0f, FAST_OUT_SLOW_IN.getInterpolation(it)) },
+        )
+    }
+
     /** Ensure alpha is set to be visible */
     val lockscreenAlpha: Flow<Float> =
         transitionAnimation.sharedFlow(
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/GoneToAodTransitionViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/GoneToAodTransitionViewModel.kt
index ba04fd3..5801793 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/GoneToAodTransitionViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/GoneToAodTransitionViewModel.kt
@@ -22,6 +22,7 @@
 import com.android.systemui.keyguard.domain.interactor.FromGoneTransitionInteractor.Companion.TO_AOD_DURATION
 import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.keyguard.ui.KeyguardTransitionAnimationFlow
+import com.android.systemui.keyguard.ui.StateToValue
 import com.android.systemui.keyguard.ui.transitions.DeviceEntryIconTransition
 import javax.inject.Inject
 import kotlin.time.Duration.Companion.milliseconds
@@ -48,8 +49,8 @@
         )
 
     /** y-translation from the top of the screen for AOD */
-    fun enterFromTopTranslationY(translatePx: Int): Flow<Float> {
-        return transitionAnimation.sharedFlow(
+    fun enterFromTopTranslationY(translatePx: Int): Flow<StateToValue> {
+        return transitionAnimation.sharedFlowWithState(
             startTime = 600.milliseconds,
             duration = 500.milliseconds,
             onStart = { translatePx },
@@ -63,8 +64,8 @@
     /** alpha animation upon entering AOD */
     val enterFromTopAnimationAlpha: Flow<Float> =
         transitionAnimation.sharedFlow(
-            startTime = 600.milliseconds,
-            duration = 500.milliseconds,
+            startTime = 700.milliseconds,
+            duration = 400.milliseconds,
             onStart = { 0f },
             onStep = { it },
         )
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/KeyguardTransitionAnimationFlowTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/KeyguardTransitionAnimationFlowTest.kt
index 2d9d5ed..0e9197e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/KeyguardTransitionAnimationFlowTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/KeyguardTransitionAnimationFlowTest.kt
@@ -181,6 +181,31 @@
         }
 
     @Test
+    fun usesOnStepToDoubleValueWithState() =
+        testScope.runTest {
+            val flow =
+                underTest.sharedFlowWithState(
+                    duration = 1000.milliseconds,
+                    onStep = { it * 2 },
+                )
+            val animationValues by collectLastValue(flow)
+            runCurrent()
+
+            repository.sendTransitionStep(step(0f, TransitionState.STARTED))
+            assertThat(animationValues).isEqualTo(StateToValue(TransitionState.STARTED, 0f))
+            repository.sendTransitionStep(step(0.3f, TransitionState.RUNNING))
+            assertThat(animationValues).isEqualTo(StateToValue(TransitionState.RUNNING, 0.6f))
+            repository.sendTransitionStep(step(0.6f, TransitionState.RUNNING))
+            assertThat(animationValues).isEqualTo(StateToValue(TransitionState.RUNNING, 1.2f))
+            repository.sendTransitionStep(step(0.8f, TransitionState.RUNNING))
+            assertThat(animationValues).isEqualTo(StateToValue(TransitionState.RUNNING, 1.6f))
+            repository.sendTransitionStep(step(1f, TransitionState.RUNNING))
+            assertThat(animationValues).isEqualTo(StateToValue(TransitionState.RUNNING, 2f))
+            repository.sendTransitionStep(step(1f, TransitionState.FINISHED))
+            assertThat(animationValues).isEqualTo(StateToValue(TransitionState.FINISHED, null))
+        }
+
+    @Test
     fun sameFloatValueWithTheSameTransitionStateDoesNotEmitTwice() =
         testScope.runTest {
             val flow =
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/GoneToAodTransitionViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/GoneToAodTransitionViewModelTest.kt
index 1c9c942..bfa8433 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/GoneToAodTransitionViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/GoneToAodTransitionViewModelTest.kt
@@ -26,6 +26,7 @@
 import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.keyguard.shared.model.TransitionState
 import com.android.systemui.keyguard.shared.model.TransitionStep
+import com.android.systemui.keyguard.ui.StateToValue
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.testKosmos
 import com.google.common.collect.Range
@@ -57,17 +58,19 @@
 
             // The animation should only start > .4f way through
             repository.sendTransitionStep(step(0f, TransitionState.STARTED))
-            assertThat(enterFromTopTranslationY).isEqualTo(pixels)
+            assertThat(enterFromTopTranslationY)
+                .isEqualTo(StateToValue(TransitionState.STARTED, pixels))
 
-            repository.sendTransitionStep(step(0.4f))
-            assertThat(enterFromTopTranslationY).isEqualTo(pixels)
+            repository.sendTransitionStep(step(.55f))
+            assertThat(enterFromTopTranslationY!!.value ?: -1f).isIn(Range.closed(pixels, 0f))
 
             repository.sendTransitionStep(step(.85f))
-            assertThat(enterFromTopTranslationY).isIn(Range.closed(pixels, 0f))
+            assertThat(enterFromTopTranslationY!!.value ?: -1f).isIn(Range.closed(pixels, 0f))
 
             // At the end, the translation should be complete and set to zero
             repository.sendTransitionStep(step(1f))
-            assertThat(enterFromTopTranslationY).isEqualTo(0f)
+            assertThat(enterFromTopTranslationY)
+                .isEqualTo(StateToValue(TransitionState.RUNNING, 0f))
         }
 
     @Test
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/AodBurnInViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/AodBurnInViewModelKosmos.kt
index a8f45b0..6f168d4 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/AodBurnInViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/AodBurnInViewModelKosmos.kt
@@ -33,6 +33,7 @@
         keyguardInteractor = keyguardInteractor,
         keyguardTransitionInteractor = keyguardTransitionInteractor,
         goneToAodTransitionViewModel = goneToAodTransitionViewModel,
+        aodToLockscreenTransitionViewModel = aodToLockscreenTransitionViewModel,
         occludedToLockscreenTransitionViewModel = occludedToLockscreenTransitionViewModel,
         keyguardClockViewModel = keyguardClockViewModel,
     )