Fix invalidation during removing the layer (#3146)

[CMP-10360](https://youtrack.jetbrains.com/issue/CMP-10360) Dialog hide
animation hangs on the last frame

Regression after #3096, not released yet

## Release Notes
N/A

Author: MatkovIvan

compose/ui/ui/src/desktopTest/kotlin/androidx/compose/ui/ComposeSceneTest.kt

@@ -468,7 +468,7 @@ class ComposeSceneTest {
         screenshotRule.snap(surface, "frame4_change_height")
 
         // see https://youtrack.jetbrains.com/issue/CMP-2171, we have extra rendered frames here
-        skipRenders()
+        skipRendersUntilIdle()
 
         assertFalse(hasRenders())
     }

compose/ui/ui/src/desktopTest/kotlin/androidx/compose/ui/platform/RenderingTestScope.kt

@@ -26,8 +26,10 @@ import androidx.compose.ui.scene.SingleComposeSceneRenderingScope
 import androidx.compose.ui.unit.Density
 import androidx.compose.ui.unit.IntSize
 import kotlin.coroutines.CoroutineContext
+import kotlin.time.Duration.Companion.milliseconds
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.withTimeout
 import kotlinx.coroutines.yield
 import org.jetbrains.skia.Surface
 import org.jetbrains.skiko.FrameDispatcher
@@ -37,13 +39,16 @@ internal fun renderingTest(
     width: Int,
     height: Int,
     context: CoroutineContext = MainUIDispatcher,
+    timeoutMillis: Long = 10000,
     block: suspend RenderingTestScope.() -> Unit
 ) = runBlocking(MainUIDispatcher) {
-    val scope = RenderingTestScope(width, height, context)
-    try {
-        scope.block()
-    } finally {
-        scope.dispose()
+    withTimeout(timeoutMillis.milliseconds) {
+        val scope = RenderingTestScope(width, height, context)
+        try {
+            scope.block()
+        } finally {
+            scope.dispose()
+        }
     }
 }
 
@@ -101,9 +106,17 @@ internal class RenderingTestScope(
         onRender.await()
     }
 
-    suspend fun skipRenders() {
-        repeat(1000) {
-            yield()
+    suspend fun skipRendersUntilIdle(maxFrames: Int = 1000) {
+        var frames = 0
+        while (frames < maxFrames) {
+            currentTimeMillis += 16
+            if (!hasRenders()) {
+                yield()
+                if (!hasRenders()) {
+                    return
+                }
+            }
+            frames++
         }
     }
 

compose/ui/ui/src/desktopTest/kotlin/androidx/compose/ui/window/DesktopDialogTest.kt

@@ -0,0 +1,62 @@
+/*
+ * Copyright 2026 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package androidx.compose.ui.window
+
+import androidx.compose.foundation.layout.Box
+import androidx.compose.foundation.layout.size
+import androidx.compose.runtime.getValue
+import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.setValue
+import androidx.compose.ui.Modifier
+import androidx.compose.ui.graphics.Color
+import androidx.compose.ui.graphics.toComposeImageBitmap
+import androidx.compose.ui.graphics.toPixelMap
+import androidx.compose.ui.platform.RenderingTestScope
+import androidx.compose.ui.platform.renderingTest
+import androidx.compose.ui.unit.dp
+import kotlin.test.assertEquals
+import org.junit.Test
+
+class DesktopDialogTest {
+
+    @Test
+    fun scrimDisappearsAfterDialogHideAnimation() = renderingTest(width = 200, height = 200) {
+        var showDialog by mutableStateOf(true)
+
+        setContent {
+            if (showDialog) {
+                Dialog(onDismissRequest = {}) {
+                    Box(Modifier.size(50.dp))
+                }
+            }
+        }
+
+        // Settle the shown state (the appearance animation also runs through the frame loop).
+        awaitNextRender()
+        skipRendersUntilIdle()
+        assertEquals(Color.Black.copy(alpha = 0.6f), colorOfCornerPixel())
+
+        // Dismiss the dialog and let the on-demand loop run the hide animation to completion.
+        showDialog = false
+        skipRendersUntilIdle()
+
+        assertEquals(Color.Transparent, colorOfCornerPixel())
+    }
+
+    private fun RenderingTestScope.colorOfCornerPixel(): Color =
+        surface.makeImageSnapshot().toComposeImageBitmap().toPixelMap()[0, 0]
+}

compose/ui/ui/src/skikoMain/kotlin/androidx/compose/ui/scene/BaseComposeScene.skiko.kt

@@ -38,6 +38,7 @@ import androidx.compose.ui.input.rotary.RotaryScrollEvent
 import androidx.compose.ui.platform.FrameRecomposer
 import androidx.compose.ui.platform.ProvidePlatformCompositionLocals
 import androidx.compose.ui.util.trace
+import kotlin.concurrent.Volatile
 
 /**
  * BaseComposeScene is an internal abstract class that implements the ComposeScene interface.
@@ -82,12 +83,25 @@ internal abstract class BaseComposeScene(
             }
         }
 
-    protected fun invokeInvalidationCallbacks() {
+    @Volatile
+    protected var hasForcedLayout: Boolean = false
+        private set
+
+    @Volatile
+    protected var hasForcedDraw: Boolean = false
+        private set
+
+    protected fun invokeInvalidationCallbacks(
+        forceLayout: Boolean = false,
+        forceDraw: Boolean = false,
+    ) {
+        hasForcedLayout = hasForcedLayout || forceLayout
+        hasForcedDraw = hasForcedDraw || forceDraw
         if (isInvalidationDisabled || isClosed || composition == null) return
-        if (hasPendingMeasureOrLayout) {
+        if (hasForcedLayout || hasPendingMeasureOrLayout) {
             invalidateLayout()
         }
-        if (hasPendingDraw) {
+        if (hasForcedDraw || hasPendingDraw) {
             invalidateDraw()
         }
     }
@@ -139,6 +153,7 @@ internal abstract class BaseComposeScene(
 
     override fun measureAndLayout() {
         if (isClosed) return
+        hasForcedLayout = false
 
         postponeInvalidation("BaseComposeScene:measureAndLayout") {
             doMeasureAndLayout()
@@ -154,6 +169,7 @@ internal abstract class BaseComposeScene(
 
     override fun draw(canvas: Canvas) {
         if (isClosed) return
+        hasForcedDraw = false
 
         postponeInvalidation("BaseComposeScene:draw") {
             // FIXME: Remove applying the global snapshot here.

compose/ui/ui/src/skikoMain/kotlin/androidx/compose/ui/scene/CanvasLayersComposeScene.skiko.kt

@@ -227,11 +227,11 @@ private class CanvasLayersComposeSceneImpl(
     }
 
     override val hasPendingMeasureOrLayout: Boolean
-        get() = mainOwner.hasPendingMeasureOrLayout
+        get() = hasForcedLayout || mainOwner.hasPendingMeasureOrLayout
             || layers.fastAny { it.owner.hasPendingMeasureOrLayout }
 
     override val hasPendingDraw: Boolean
-        get() = mainOwner.hasPendingDraw
+        get() = hasForcedDraw || mainOwner.hasPendingDraw
             || layers.fastAny { it.owner.hasPendingDraw }
 
     override fun createComposition(
@@ -517,7 +517,7 @@ private class CanvasLayersComposeSceneImpl(
         onOwnerAppended(layer.owner)
 
         inputHandler.onPointerUpdate()
-        invokeInvalidationCallbacks()
+        invokeInvalidationCallbacks(forceLayout = true, forceDraw = true)
     }
 
     private fun detachLayer(layer: AttachedComposeSceneLayer) {
@@ -528,7 +528,9 @@ private class CanvasLayersComposeSceneImpl(
         onOwnerRemoved(layer.owner)
 
         inputHandler.onPointerUpdate()
-        invokeInvalidationCallbacks()
+        // A detached layer was composited onto this scene's canvas, so its removal changes
+        // the scene's output even though no remaining owner is dirty.
+        invokeInvalidationCallbacks(forceLayout = true, forceDraw = true)
     }
 
     private fun requestFocus(layer: AttachedComposeSceneLayer) {

compose/ui/ui/src/skikoMain/kotlin/androidx/compose/ui/scene/PlatformLayersComposeScene.skiko.kt

@@ -162,10 +162,10 @@ private class PlatformLayersComposeSceneImpl(
     }
 
     override val hasPendingMeasureOrLayout: Boolean
-        get() = mainOwner.hasPendingMeasureOrLayout
+        get() = hasForcedLayout || mainOwner.hasPendingMeasureOrLayout
 
     override val hasPendingDraw: Boolean
-        get() = mainOwner.hasPendingDraw
+        get() = hasForcedDraw || mainOwner.hasPendingDraw
 
     override fun createComposition(
         parentCompositionContext: CompositionContext,

compose/ui/ui/src/skikoTest/kotlin/androidx/compose/ui/scene/CanvasLayersComposeSceneTest.kt

@@ -20,6 +20,7 @@ import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.fillMaxSize
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.geometry.Offset
+import androidx.compose.ui.graphics.asComposeCanvas
 import androidx.compose.ui.input.pointer.PointerEventType
 import androidx.compose.ui.platform.FrameRecomposer
 import androidx.compose.ui.unit.IntSize
@@ -32,6 +33,7 @@ import kotlin.test.assertFalse
 import kotlin.test.assertTrue
 import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.runTest
+import org.jetbrains.skia.Surface
 
 class CanvasLayersComposeSceneTest {
 
@@ -79,4 +81,40 @@ class CanvasLayersComposeSceneTest {
         }
         frameRecomposer.close()
     }
+
+    @Test
+    fun detachingLayerRequestsDrawPass() = runTest(StandardTestDispatcher()) {
+        var drawInvalidations = 0
+        var layer: ComposeSceneLayer? = null
+        val frameRecomposer = FrameRecomposer(coroutineContext)
+        val surface = Surface.makeRasterN32Premul(100, 100)
+        CanvasLayersComposeScene(
+            frameRecomposer = frameRecomposer,
+            size = IntSize(100, 100),
+            invalidateDraw = { drawInvalidations++ },
+        ).use { scene ->
+            scene.setContent {
+                Box(Modifier.fillMaxSize())
+                layer = rememberComposeSceneLayer(focusable = true)
+            }
+
+            // Settle measure/layout/draw so every owner's pending-draw flag is cleared; otherwise
+            // the close below would invalidate simply because an owner was still dirty.
+            scene.measureAndLayout()
+            scene.draw(surface.canvas.asComposeCanvas())
+            assertFalse(scene.hasPendingMeasureOrLayout)
+            assertFalse(scene.hasPendingDraw)
+
+            val drawInvalidationsBeforeClose = drawInvalidations
+            layer!!.close()
+
+            assertTrue(scene.hasPendingMeasureOrLayout)
+            assertTrue(scene.hasPendingDraw)
+            assertTrue(
+                drawInvalidations > drawInvalidationsBeforeClose,
+                "Detaching a layer must request a draw pass to repaint the scene without it",
+            )
+        }
+        frameRecomposer.close()
+    }
 }