Consolidate isInTransition API and enable wildcard for it

The isInTransition[fromState/toState] API has a large API surface. This
CL consolidates it into just isInTransition and also adds functionality
such that scene edges with a wildcard can also be queried.

Test: None
Bug: b/330311871
Flag: com.android.systemui.scene_container
Change-Id: I0caddb1eb39a3cf7619df185e7d8658c6ae7e7a1
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorTest.kt
index 99cccb2..5756bca 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorTest.kt
@@ -40,7 +40,10 @@
 import com.android.systemui.keyguard.shared.model.TransitionState.STARTED
 import com.android.systemui.keyguard.shared.model.TransitionStep
 import com.android.systemui.kosmos.testScope
+import com.android.systemui.scene.data.repository.Idle
+import com.android.systemui.scene.data.repository.Transition
 import com.android.systemui.scene.data.repository.sceneContainerRepository
+import com.android.systemui.scene.data.repository.setSceneTransition
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
@@ -527,7 +530,13 @@
     @DisableSceneContainer
     fun isInTransitionToState() =
         testScope.runTest {
-            val results by collectValues(underTest.isInTransitionToState(GONE))
+            val results by
+                collectValues(
+                    underTest.isInTransition(
+                        edge = Edge.create(OFF, OFF),
+                        edgeWithoutSceneContainer = Edge.create(to = LOCKSCREEN)
+                    )
+                )
 
             sendSteps(
                 TransitionStep(AOD, DOZING, 0f, STARTED),
@@ -543,7 +552,7 @@
                 )
 
             sendSteps(
-                TransitionStep(DOZING, GONE, 0f, STARTED),
+                TransitionStep(DOZING, LOCKSCREEN, 0f, STARTED),
             )
 
             assertThat(results)
@@ -555,7 +564,7 @@
                 )
 
             sendSteps(
-                TransitionStep(DOZING, GONE, 0f, RUNNING),
+                TransitionStep(DOZING, LOCKSCREEN, 0f, RUNNING),
             )
 
             assertThat(results)
@@ -567,7 +576,7 @@
                 )
 
             sendSteps(
-                TransitionStep(DOZING, GONE, 0f, FINISHED),
+                TransitionStep(DOZING, LOCKSCREEN, 0f, FINISHED),
             )
 
             assertThat(results)
@@ -580,9 +589,9 @@
                 )
 
             sendSteps(
-                TransitionStep(GONE, DOZING, 0f, STARTED),
-                TransitionStep(GONE, DOZING, 0f, RUNNING),
-                TransitionStep(GONE, DOZING, 1f, FINISHED),
+                TransitionStep(LOCKSCREEN, DOZING, 0f, STARTED),
+                TransitionStep(LOCKSCREEN, DOZING, 0f, RUNNING),
+                TransitionStep(LOCKSCREEN, DOZING, 1f, FINISHED),
             )
 
             assertThat(results)
@@ -595,8 +604,8 @@
                 )
 
             sendSteps(
-                TransitionStep(DOZING, GONE, 0f, STARTED),
-                TransitionStep(DOZING, GONE, 0f, RUNNING),
+                TransitionStep(DOZING, LOCKSCREEN, 0f, STARTED),
+                TransitionStep(DOZING, LOCKSCREEN, 0f, RUNNING),
             )
 
             assertThat(results)
@@ -611,9 +620,96 @@
         }
 
     @Test
+    @EnableSceneContainer
+    fun isInTransitionFromScene() =
+        testScope.runTest {
+            val results by
+                collectValues(underTest.isInTransition(edge = Edge.create(Scenes.Lockscreen, null)))
+
+            kosmos.setSceneTransition(Transition(from = Scenes.Gone, to = Scenes.Lockscreen))
+            kosmos.setSceneTransition(Idle(Scenes.Lockscreen))
+
+            assertThat(results)
+                .isEqualTo(
+                    listOf(
+                        false,
+                    )
+                )
+
+            kosmos.setSceneTransition(Transition(from = Scenes.Lockscreen, to = Scenes.Shade))
+
+            assertThat(results)
+                .isEqualTo(
+                    listOf(
+                        false,
+                        true,
+                    )
+                )
+
+            kosmos.setSceneTransition(Idle(Scenes.Shade))
+
+            assertThat(results)
+                .isEqualTo(
+                    listOf(
+                        false,
+                        true,
+                        false,
+                    )
+                )
+        }
+
+    @Test
+    @EnableSceneContainer
+    fun isInTransitionToScene() =
+        testScope.runTest {
+            val results by
+                collectValues(underTest.isInTransition(edge = Edge.create(null, Scenes.Lockscreen)))
+
+            kosmos.setSceneTransition(Transition(from = Scenes.Gone, to = Scenes.Lockscreen))
+            kosmos.setSceneTransition(Idle(Scenes.Lockscreen))
+            kosmos.setSceneTransition(Transition(from = Scenes.Lockscreen, to = Scenes.Gone))
+
+            assertThat(results)
+                .isEqualTo(
+                    listOf(
+                        false,
+                        true,
+                        false,
+                    )
+                )
+        }
+
+    @Test
+    @EnableSceneContainer
+    fun isInTransitionStateToScene() =
+        testScope.runTest {
+            val results by
+                collectValues(underTest.isInTransition(edge = Edge.create(AOD, Scenes.Gone)))
+
+            kosmos.setSceneTransition(Transition(from = Scenes.Lockscreen, to = Scenes.Gone))
+
+            sendSteps(
+                TransitionStep(AOD, UNDEFINED, 0f, STARTED),
+                TransitionStep(AOD, UNDEFINED, 0.5f, RUNNING),
+                TransitionStep(AOD, UNDEFINED, 1f, FINISHED),
+            )
+
+            kosmos.setSceneTransition(Idle(Scenes.Gone))
+
+            assertThat(results)
+                .isEqualTo(
+                    listOf(
+                        false,
+                        true,
+                        false,
+                    )
+                )
+        }
+
+    @Test
     fun isInTransitionFromState() =
         testScope.runTest {
-            val results by collectValues(underTest.isInTransitionFromState(DOZING))
+            val results by collectValues(underTest.isInTransition(Edge.create(from = DOZING)))
 
             sendSteps(
                 TransitionStep(AOD, DOZING, 0f, STARTED),
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/InWindowLauncherUnlockAnimationInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/InWindowLauncherUnlockAnimationInteractor.kt
index 8ec831c..28d4ba96 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/InWindowLauncherUnlockAnimationInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/InWindowLauncherUnlockAnimationInteractor.kt
@@ -20,7 +20,9 @@
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.keyguard.data.repository.InWindowLauncherUnlockAnimationRepository
 import com.android.systemui.keyguard.data.repository.KeyguardSurfaceBehindRepository
-import com.android.systemui.keyguard.shared.model.KeyguardState
+import com.android.systemui.keyguard.shared.model.Edge
+import com.android.systemui.keyguard.shared.model.KeyguardState.GONE
+import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.shared.system.ActivityManagerWrapper
 import com.android.systemui.shared.system.smartspace.SmartspaceState
 import javax.inject.Inject
@@ -50,7 +52,10 @@
      */
     val transitioningToGoneWithInWindowAnimation: StateFlow<Boolean> =
         transitionInteractor
-            .isInTransitionToState(KeyguardState.GONE)
+            .isInTransition(
+                edge = Edge.create(to = Scenes.Gone),
+                edgeWithoutSceneContainer = Edge.create(to = GONE)
+            )
             .map { transitioningToGone -> transitioningToGone && isLauncherUnderneath() }
             .stateIn(scope, SharingStarted.Eagerly, false)
 
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 c65dc30..2766b71 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
@@ -180,7 +180,7 @@
                 val fromScene =
                     when (edge) {
                         is Edge.StateToState -> edge.from?.mapToSceneContainerScene()
-                        is Edge.StateToScene -> edge.from.mapToSceneContainerScene()
+                        is Edge.StateToScene -> edge.from?.mapToSceneContainerScene()
                         is Edge.SceneToState -> edge.from
                     }
 
@@ -188,7 +188,7 @@
                     when (edge) {
                         is Edge.StateToState -> edge.to?.mapToSceneContainerScene()
                         is Edge.StateToScene -> edge.to
-                        is Edge.SceneToState -> edge.to.mapToSceneContainerScene()
+                        is Edge.SceneToState -> edge.to?.mapToSceneContainerScene()
                     }
 
                 fun SceneKey?.isLockscreenOrNull() = this == Scenes.Lockscreen || this == null
@@ -450,16 +450,6 @@
         }
     }
 
-    /** Whether we're in a transition to the given [KeyguardState], but haven't yet completed it. */
-    fun isInTransitionToState(
-        state: KeyguardState,
-    ): Flow<Boolean> {
-        return transition(Edge.create(from = null, to = state))
-            .mapLatest { it.transitionState.isTransitioning() }
-            .onStart { emit(false) }
-            .distinctUntilChanged()
-    }
-
     /**
      * Whether we're in a transition to and from the given [KeyguardState]s, but haven't yet
      * completed it.
@@ -469,23 +459,23 @@
      */
     fun isInTransition(edge: Edge, edgeWithoutSceneContainer: Edge? = null): Flow<Boolean> {
         return if (SceneContainerFlag.isEnabled) {
-                transition(edge)
+                if (edge.isSceneWildcardEdge()) {
+                    sceneInteractor.get().transitionState.map {
+                        when (edge) {
+                            is Edge.StateToState ->
+                                throw IllegalStateException("Should not be reachable.")
+                            is Edge.SceneToState -> it.isTransitioning(from = edge.from)
+                            is Edge.StateToScene -> it.isTransitioning(to = edge.to)
+                        }
+                    }
+                } else {
+                    transition(edge).mapLatest { it.transitionState.isTransitioning() }
+                }
             } else {
-                transition(edgeWithoutSceneContainer ?: edge)
+                transition(edgeWithoutSceneContainer ?: edge).mapLatest {
+                    it.transitionState.isTransitioning()
+                }
             }
-            .mapLatest { it.transitionState.isTransitioning() }
-            .onStart { emit(false) }
-            .distinctUntilChanged()
-    }
-
-    /**
-     * Whether we're in a transition out of the given [KeyguardState], but haven't yet completed it.
-     */
-    fun isInTransitionFromState(
-        state: KeyguardState,
-    ): Flow<Boolean> {
-        return transition(Edge.create(from = state, to = null))
-            .mapLatest { it.transitionState.isTransitioning() }
             .onStart { emit(false) }
             .distinctUntilChanged()
     }
@@ -494,7 +484,7 @@
      * Whether we're in a transition to a [KeyguardState] that matches the given predicate, but
      * haven't yet completed it.
      *
-     * If you only care about a single state, instead use the optimized [isInTransitionToState].
+     * If you only care about a single state, instead use the optimized [isInTransition].
      */
     fun isInTransitionToStateWhere(
         stateMatcher: (KeyguardState) -> Boolean,
@@ -506,7 +496,7 @@
      * Whether we're in a transition out of a [KeyguardState] that matches the given predicate, but
      * haven't yet completed it.
      *
-     * If you only care about a single state, instead use the optimized [isInTransitionFromState].
+     * If you only care about a single state, instead use the optimized [isInTransition].
      */
     fun isInTransitionFromStateWhere(
         stateMatcher: (KeyguardState) -> Boolean,
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/WindowManagerLockscreenVisibilityInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/WindowManagerLockscreenVisibilityInteractor.kt
index 350527a..8ba09bd 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/WindowManagerLockscreenVisibilityInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/WindowManagerLockscreenVisibilityInteractor.kt
@@ -21,6 +21,7 @@
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.deviceentry.domain.interactor.DeviceEntryInteractor
 import com.android.systemui.keyguard.shared.model.BiometricUnlockMode
+import com.android.systemui.keyguard.shared.model.Edge
 import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.keyguard.shared.model.TransitionState
 import com.android.systemui.scene.domain.interactor.SceneInteractor
@@ -132,7 +133,10 @@
                 .distinctUntilChanged()
         } else {
             combine(
-                    transitionInteractor.isInTransitionToState(KeyguardState.GONE),
+                    transitionInteractor.isInTransition(
+                        edge = Edge.create(to = Scenes.Gone),
+                        edgeWithoutSceneContainer = Edge.create(to = KeyguardState.GONE)
+                    ),
                     transitionInteractor.finishedKeyguardState,
                     surfaceBehindInteractor.isAnimatingSurface,
                     notificationLaunchAnimationInteractor.isLaunchAnimationRunning,
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/shared/model/Edge.kt b/packages/SystemUI/src/com/android/systemui/keyguard/shared/model/Edge.kt
index 4f516f5..c1e8d22 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/shared/model/Edge.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/shared/model/Edge.kt
@@ -23,11 +23,6 @@
 /**
  * Represents an edge either between two Keyguard Transition Framework states (KTF) or a KTF state
  * and a scene container scene. Passing [null] in either [from] or [to] indicates a wildcard.
- *
- * Wildcards are not allowed for transitions involving a scene. Use [sceneInteractor] directly
- * instead. Reason: [TransitionStep]s are not emitted for every edge leading into/out of a scene.
- * For example: Lockscreen -> Gone would be emitted as LOCKSCREEN -> UNDEFINED but Bouncer -> Gone
- * would not emit anything.
  */
 sealed class Edge {
 
@@ -59,17 +54,6 @@
                     Please remove or port this edge to scene container."""
                         .trimIndent(),
                 )
-            } else if ((fromChanged && to == null) || (toChanged && from == null)) {
-                Log.e(
-                    TAG,
-                    """
-                    The edge ${from?.name} => ${to?.name} was automatically converted to
-                    ${mappedFrom?.name} => ${mappedTo?.name}. Wildcards are not allowed together
-                    with UNDEFINED because it will only be tracking edges leading in and out of
-                    the Lockscreen scene but miss others. Please remove or port this edge."""
-                        .trimIndent(),
-                    Exception()
-                )
             } else if (fromChanged || toChanged) {
                 Log.w(
                     TAG,
@@ -90,23 +74,32 @@
         }
     }
 
+    fun isSceneWildcardEdge(): Boolean {
+        return when (this) {
+            is StateToState -> false
+            is SceneToState -> to == null
+            is StateToScene -> from == null
+        }
+    }
+
     data class StateToState(val from: KeyguardState?, val to: KeyguardState?) : Edge() {
+
         init {
             check(!(from == null && to == null)) { "to and from can't both be null" }
         }
     }
 
-    data class StateToScene(val from: KeyguardState, val to: SceneKey) : Edge()
+    data class StateToScene(val from: KeyguardState? = null, val to: SceneKey) : Edge()
 
-    data class SceneToState(val from: SceneKey, val to: KeyguardState) : Edge()
+    data class SceneToState(val from: SceneKey, val to: KeyguardState? = null) : Edge()
 
     companion object {
         private const val TAG = "Edge"
 
         fun create(from: KeyguardState? = null, to: KeyguardState? = null) = StateToState(from, to)
 
-        fun create(from: KeyguardState, to: SceneKey) = StateToScene(from, to)
+        fun create(from: KeyguardState? = null, to: SceneKey) = StateToScene(from, to)
 
-        fun create(from: SceneKey, to: KeyguardState) = SceneToState(from, to)
+        fun create(from: SceneKey, to: KeyguardState? = null) = SceneToState(from, to)
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModel.kt
index f8a9310..aaec69f 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModel.kt
@@ -142,8 +142,8 @@
         combine(
                 keyguardTransitionInteractor.isFinishedInState(LOCKSCREEN).onStart { emit(false) },
                 anyOf(
-                    keyguardTransitionInteractor.isInTransitionToState(LOCKSCREEN),
-                    keyguardTransitionInteractor.isInTransitionFromState(LOCKSCREEN),
+                    keyguardTransitionInteractor.isInTransition(Edge.create(to = LOCKSCREEN)),
+                    keyguardTransitionInteractor.isInTransition(Edge.create(from = LOCKSCREEN)),
                 ),
             ) { onLockscreen, transitioningToOrFromLockscreen ->
                 onLockscreen || transitioningToOrFromLockscreen
diff --git a/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/ShadeInteractorImpl.kt b/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/ShadeInteractorImpl.kt
index 06a8d18..4014512 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/ShadeInteractorImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/ShadeInteractorImpl.kt
@@ -21,6 +21,7 @@
 import com.android.systemui.keyguard.data.repository.KeyguardRepository
 import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor
 import com.android.systemui.keyguard.shared.model.DozeStateModel
+import com.android.systemui.keyguard.shared.model.Edge
 import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.power.domain.interactor.PowerInteractor
 import com.android.systemui.statusbar.disableflags.data.repository.DisableFlagsRepository
@@ -85,7 +86,7 @@
     override val isShadeTouchable: Flow<Boolean> =
         combine(
             powerInteractor.isAsleep,
-            keyguardTransitionInteractor.isInTransitionToState(KeyguardState.AOD),
+            keyguardTransitionInteractor.isInTransition(Edge.create(to = KeyguardState.AOD)),
             keyguardRepository.dozeTransitionModel.map { it.to == DozeStateModel.DOZE_PULSING },
         ) { isAsleep, goingToSleep, isPulsing ->
             when {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModel.kt
index 3393321..6dfaec9 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModel.kt
@@ -249,8 +249,14 @@
                     state == GLANCEABLE_HUB
                 },
                 anyOf(
-                    keyguardTransitionInteractor.isInTransitionToState(GLANCEABLE_HUB),
-                    keyguardTransitionInteractor.isInTransitionFromState(GLANCEABLE_HUB),
+                    keyguardTransitionInteractor.isInTransition(
+                        edge = Edge.create(to = Scenes.Communal),
+                        edgeWithoutSceneContainer = Edge.create(to = GLANCEABLE_HUB)
+                    ),
+                    keyguardTransitionInteractor.isInTransition(
+                        edge = Edge.create(from = Scenes.Communal),
+                        edgeWithoutSceneContainer = Edge.create(from = GLANCEABLE_HUB)
+                    ),
                 ),
             ) { isOnGlanceableHub, transitioningToOrFromHub ->
                 isOnGlanceableHub || transitioningToOrFromHub
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerBaseTest.java b/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerBaseTest.java
index 043dba1..041adea 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerBaseTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerBaseTest.java
@@ -421,6 +421,8 @@
         mPowerInteractor = keyguardInteractorDeps.getPowerInteractor();
         when(mKeyguardTransitionInteractor.isInTransitionToStateWhere(any())).thenReturn(
                 MutableStateFlow(false));
+        when(mKeyguardTransitionInteractor.isInTransition(any(), any()))
+                .thenReturn(emptyFlow());
         when(mKeyguardTransitionInteractor.getCurrentKeyguardState()).thenReturn(
                 MutableSharedFlow(0, 0, BufferOverflow.SUSPEND));
         when(mDeviceEntryFaceAuthInteractor.isBypassEnabled()).thenReturn(MutableStateFlow(false));
@@ -428,8 +430,6 @@
                 mock(DeviceEntryUdfpsInteractor.class);
         when(deviceEntryUdfpsInteractor.isUdfpsSupported()).thenReturn(MutableStateFlow(false));
 
-        when(mKeyguardTransitionInteractor.isInTransitionToState(any())).thenReturn(emptyFlow());
-
         mShadeInteractor = new ShadeInteractorImpl(
                 mTestScope.getBackgroundScope(),
                 mKosmos.getDeviceProvisioningInteractor(),
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerWithCoroutinesTest.kt b/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerWithCoroutinesTest.kt
index e1ee358..d47bd47 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerWithCoroutinesTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerWithCoroutinesTest.kt
@@ -79,10 +79,9 @@
 
         verify(mView, atLeastOnce()).addView(viewCaptor.capture(), anyInt())
         val userSwitcherStub =
-            CollectionUtils.find(
-                viewCaptor.getAllValues(),
-                { view -> view.getId() == R.id.keyguard_user_switcher_stub }
-            )
+            CollectionUtils.find(viewCaptor.allValues) { view ->
+                view.id == R.id.keyguard_user_switcher_stub
+            }
         assertThat(userSwitcherStub).isNotNull()
         assertThat(userSwitcherStub).isInstanceOf(ViewStub::class.java)
     }