[flexiglass] Fix HUN shared element transitions

Fixed for Shade -> Lockscreen and QS -> Lockscreen

Bug: 369843768
Test: manually verified shared element transitions look correct
Flag: com.android.systemui.scene_container
Change-Id: I64ebfcf215b99f68e47b460cc9e4532f9bd2af36

2 files changed, 11 insertions, 0 deletions

diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/Notifications.kt b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/Notifications.kt
index 2066c9314bc7..2745f6ea6bb4 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/Notifications.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/Notifications.kt
@@ -259,6 +259,7 @@ fun SceneScope.ConstrainedNotificationStack(
         HeadsUpNotificationSpace(
             stackScrollView = stackScrollView,
             viewModel = viewModel,
+            useHunBounds = { shouldUseLockscreenHunBounds(layoutState.transitionState) },
             modifier = Modifier.align(Alignment.TopCenter),
         )
         NotificationStackCutoffGuideline(
@@ -541,6 +542,7 @@ fun SceneScope.NotificationScrollingStack(
             HeadsUpNotificationSpace(
                 stackScrollView = stackScrollView,
                 viewModel = viewModel,
+                useHunBounds = { !shouldUseLockscreenHunBounds(layoutState.transitionState) },
                 modifier = Modifier.padding(top = topPadding),
             )
         }
@@ -675,6 +677,14 @@ private fun shouldUseLockscreenStackBounds(state: TransitionState): Boolean {
     return state is TransitionState.Idle && state.currentScene == Scenes.Lockscreen
 }
 
+private fun shouldUseLockscreenHunBounds(state: TransitionState): Boolean {
+    return when (state) {
+        is TransitionState.Idle -> state.currentScene == Scenes.Lockscreen
+        is TransitionState.Transition ->
+            state.isTransitioning(from = Scenes.QuickSettings, to = Scenes.Lockscreen)
+    }
+}
+
 private fun shouldAnimateScrimCornerRadius(
     state: SceneTransitionLayoutState,
     shouldPunchHoleBehindScrim: Boolean,
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainerTransitions.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainerTransitions.kt
index 303a6f0d942a..ac58ab5296f6 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainerTransitions.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainerTransitions.kt
@@ -85,6 +85,7 @@ val SceneContainerTransitions = transitions {
     from(Scenes.Shade, to = Scenes.Lockscreen) {
         reversed { lockscreenToShadeTransition() }
         sharedElement(Notifications.Elements.NotificationStackPlaceholder, enabled = false)
+        sharedElement(Notifications.Elements.HeadsUpNotificationPlaceholder, enabled = false)
     }
 
     // Overlay transitions