Merge "Fix new Expandable animation lag" into main

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)