From a2bb660ef331e4dcb1a629b44f9c45b7dfd10fd1 Mon Sep 17 00:00:00 2001
From: Jordan Demeulenaere <jdemeulenaere@google.com>
Date: Wed, 4 Sep 2024 16:10:38 +0200
Subject: Add ObservableTransitionState.currentOverlays()

Bug: 353679003
Test: Manual
Flag: com.android.systemui.scene_container
Change-Id: I7c9ecfa2ff61cafe2b13a58be1ffa6653102c666
---
 .../animation/scene/ObservableTransitionState.kt   | 29 +++++++++++++++++-----
 1 file changed, 23 insertions(+), 6 deletions(-)

diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/ObservableTransitionState.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/ObservableTransitionState.kt
index bd21a69e0699..8ae3a1145e24 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/ObservableTransitionState.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/ObservableTransitionState.kt
@@ -46,8 +46,21 @@ sealed interface ObservableTransitionState {
         }
     }
 
+    /** The current overlays. */
+    fun currentOverlays(): Flow<Set<OverlayKey>> {
+        return when (this) {
+            is Idle -> flowOf(currentOverlays)
+            is Transition -> currentOverlays
+        }
+    }
+
     /** No transition/animation is currently running. */
-    data class Idle(val currentScene: SceneKey) : ObservableTransitionState
+    data class Idle
+    @JvmOverloads
+    constructor(
+        val currentScene: SceneKey,
+        val currentOverlays: Set<OverlayKey> = emptySet(),
+    ) : ObservableTransitionState
 
     /** There is a transition animating between two scenes. */
     sealed class Transition(
@@ -94,7 +107,7 @@ sealed interface ObservableTransitionState {
             val fromScene: SceneKey,
             val toScene: SceneKey,
             val currentScene: Flow<SceneKey>,
-            currentOverlays: Flow<Set<OverlayKey>>,
+            currentOverlays: Set<OverlayKey>,
             progress: Flow<Float>,
             isInitiatedByUserInput: Boolean,
             isUserInputOngoing: Flow<Boolean>,
@@ -104,7 +117,7 @@ sealed interface ObservableTransitionState {
             Transition(
                 fromScene,
                 toScene,
-                currentOverlays,
+                flowOf(currentOverlays),
                 progress,
                 isInitiatedByUserInput,
                 isUserInputOngoing,
@@ -169,7 +182,7 @@ sealed interface ObservableTransitionState {
                 isUserInputOngoing: Flow<Boolean>,
                 previewProgress: Flow<Float> = flowOf(0f),
                 isInPreviewStage: Flow<Boolean> = flowOf(false),
-                currentOverlays: Flow<Set<OverlayKey>> = flowOf(emptySet()),
+                currentOverlays: Set<OverlayKey> = emptySet(),
             ): ChangeScene {
                 return ChangeScene(
                     fromScene,
@@ -205,13 +218,17 @@ sealed interface ObservableTransitionState {
 fun SceneTransitionLayoutState.observableTransitionState(): Flow<ObservableTransitionState> {
     return snapshotFlow {
             when (val state = transitionState) {
-                is TransitionState.Idle -> ObservableTransitionState.Idle(state.currentScene)
+                is TransitionState.Idle ->
+                    ObservableTransitionState.Idle(
+                        state.currentScene,
+                        state.currentOverlays,
+                    )
                 is TransitionState.Transition.ChangeScene -> {
                     ObservableTransitionState.Transition.ChangeScene(
                         fromScene = state.fromScene,
                         toScene = state.toScene,
                         currentScene = snapshotFlow { state.currentScene },
-                        currentOverlays = flowOf(state.currentOverlays),
+                        currentOverlays = state.currentOverlays,
                         progress = snapshotFlow { state.progress },
                         isInitiatedByUserInput = state.isInitiatedByUserInput,
                         isUserInputOngoing = snapshotFlow { state.isUserInputOngoing },
-- 
cgit v1.2.3-59-g8ed1b