Fix new Expandable animation lag

This CL ensures that the new implementation of the Compose Expandable is
drawing the expandable inside the dialog overlay during dialog launches,
rather than in the Expandable overlay. See b/397904604#comment6 for
details and before/after videos.

Bug: 397904604
Test: Manual
Flag: com.android.systemui.qs_ui_refactor_compose_fragment
Change-Id: I691d17a9c3c8deb74484b744bb40a4a93812bf7d

2 files changed, 8 insertions, 5 deletions

diff --git a/packages/SystemUI/compose/core/src/com/android/compose/animation/Expandable.kt b/packages/SystemUI/compose/core/src/com/android/compose/animation/Expandable.kt
index 873991923e51..1e3c4c9e7c09 100644
--- a/packages/SystemUI/compose/core/src/com/android/compose/animation/Expandable.kt
+++ b/packages/SystemUI/compose/core/src/com/android/compose/animation/Expandable.kt
@@ -314,7 +314,7 @@ private fun Modifier.expandable(
         if (isAnimating) {
             val graphicsLayer = rememberGraphicsLayer()
 
-            FullScreenComposeViewInOverlay { view ->
+            FullScreenComposeViewInOverlay(controller.overlay) { view ->
                 Modifier.then(DrawExpandableInOverlayElement(view, controller, graphicsLayer))
             }
 
@@ -397,7 +397,7 @@ private class DrawExpandableInOverlayNode(
     }
 
     override fun ContentDrawScope.draw() {
-        val state = controller.animatorState ?: return
+        val state = controller.animatorState?.takeIf { it.visible } ?: return
         val topOffset = state.top.toFloat() - composeViewLocationOnScreen[1]
         val leftOffset = state.left.toFloat() - composeViewLocationOnScreen[0]
 
diff --git a/packages/SystemUI/compose/core/src/com/android/compose/ui/graphics/DrawInOverlay.kt b/packages/SystemUI/compose/core/src/com/android/compose/ui/graphics/DrawInOverlay.kt
index 089da4b932b2..d5436d3aa89c 100644
--- a/packages/SystemUI/compose/core/src/com/android/compose/ui/graphics/DrawInOverlay.kt
+++ b/packages/SystemUI/compose/core/src/com/android/compose/ui/graphics/DrawInOverlay.kt
@@ -47,15 +47,18 @@ fun Modifier.drawInOverlay(): Modifier {
 }
 
 @Composable
-internal fun FullScreenComposeViewInOverlay(modifier: (ComposeView) -> Modifier = { Modifier }) {
+internal fun FullScreenComposeViewInOverlay(
+    overlay: ViewGroupOverlay? = null,
+    modifier: (ComposeView) -> Modifier = { Modifier },
+) {
     val context = LocalContext.current
     val localView = LocalView.current
     val compositionContext = rememberCompositionContext()
     val displayMetrics = context.resources.displayMetrics
     val displaySize = IntSize(displayMetrics.widthPixels, displayMetrics.heightPixels)
+    val overlay = overlay ?: localView.rootView.overlay as ViewGroupOverlay
 
-    DisposableEffect(context, localView, compositionContext, displaySize) {
-        val overlay = localView.rootView.overlay as ViewGroupOverlay
+    DisposableEffect(context, localView, overlay, compositionContext, displaySize) {
         val view =
             ComposeView(context).apply {
                 setParentCompositionContext(compositionContext)