Correctly return child compose scenes created for Window and DialogWi…

…ndow in ComposeScene.roots (for test)
JetBrains · Jul 20, 2023 · 3d4bd82 · 3d4bd82
1 parent 0fdc2a1
commit 3d4bd82
Show file tree

Hide file tree

Showing 9 changed files with 104 additions and 10 deletions.
diff --git a/compose/ui/ui-test-junit4/src/desktopTest/kotlin/androidx/compose/ui/test/AssertionsTest.kt b/compose/ui/ui-test-junit4/src/desktopTest/kotlin/androidx/compose/ui/test/AssertionsTest.kt
@@ -20,7 +20,10 @@ import androidx.compose.foundation.focusable
 import androidx.compose.foundation.layout.*
 import androidx.compose.material.*
 import androidx.compose.runtime.LaunchedEffect
+import androidx.compose.runtime.getValue
+import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
+import androidx.compose.runtime.setValue
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.focus.FocusRequester
 import androidx.compose.ui.focus.focusRequester
@@ -31,6 +34,8 @@ import androidx.compose.ui.semantics.semantics
 import androidx.compose.ui.semantics.stateDescription
 import androidx.compose.ui.test.junit4.createComposeRule
 import androidx.compose.ui.unit.dp
+import androidx.compose.ui.window.DialogWindow
+import androidx.compose.ui.window.Window
 import kotlin.test.assertFails
 import org.junit.Rule
 import org.junit.Test
@@ -542,4 +547,48 @@ class AssertionsTest {
             rule.onAllNodesWithTag("tag").assertAll(hasText("World", substring = true))
         }
     }
+
+    @Test
+    fun testNodeInDialogWindow() {
+        var show by mutableStateOf(true)
+        rule.setContent {
+            if (show) {
+                DialogWindow(
+                    onCloseRequest = {},
+                ) {
+                    Text(
+                        text = "Text",
+                        modifier = Modifier.testTag("tag")
+                    )
+                }
+            }
+        }
+
+        rule.onNodeWithTag("tag").assertExists()
+
+        show = false
+        rule.onNodeWithTag("tag").assertDoesNotExist()
+    }
+
+    @Test
+    fun testNodeInWindow() {
+        var show by mutableStateOf(true)
+        rule.setContent {
+            if (show) {
+                Window(
+                    onCloseRequest = {},
+                ) {
+                    Text(
+                        text = "Text",
+                        modifier = Modifier.testTag("tag")
+                    )
+                }
+            }
+        }
+
+        rule.onNodeWithTag("tag").assertExists()
+
+        show = false
+        rule.onNodeWithTag("tag").assertDoesNotExist()
+    }
 }
diff --git a/compose/ui/ui/src/desktopMain/kotlin/androidx/compose/ui/awt/ComposeDialog.desktop.kt b/compose/ui/ui/src/desktopMain/kotlin/androidx/compose/ui/awt/ComposeDialog.desktop.kt
@@ -17,6 +17,7 @@ package androidx.compose.ui.awt
 
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.CompositionLocalContext
+import androidx.compose.ui.ComposeScene
 import androidx.compose.ui.ExperimentalComposeUiApi
 import androidx.compose.ui.input.key.KeyEvent
 import androidx.compose.ui.window.DialogWindowScope
@@ -41,6 +42,9 @@ class ComposeDialog : JDialog {
     private val skiaLayerAnalytics: SkiaLayerAnalytics
     private val delegate: ComposeWindowDelegate
 
+    internal val scene: ComposeScene
+        get() = delegate.scene
+
     constructor(
         owner: Window?,
         modalityType: ModalityType = ModalityType.MODELESS,

diff --git a/compose/ui/ui/src/desktopMain/kotlin/androidx/compose/ui/awt/ComposeWindow.desktop.kt b/compose/ui/ui/src/desktopMain/kotlin/androidx/compose/ui/awt/ComposeWindow.desktop.kt
@@ -17,6 +17,7 @@ package androidx.compose.ui.awt
 
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.CompositionLocalContext
+import androidx.compose.ui.ComposeScene
 import androidx.compose.ui.ExperimentalComposeUiApi
 import androidx.compose.ui.input.key.KeyEvent
 import androidx.compose.ui.window.FrameWindowScope
@@ -60,6 +61,9 @@ class ComposeWindow @ExperimentalComposeUiApi constructor(
 
     private val delegate = ComposeWindowDelegate(this, ::isUndecorated, skiaLayerAnalytics)
 
+    internal val scene: ComposeScene
+        get() = delegate.scene
+
     // Don't override the accessible context of JFrame, since accessibility work through HardwareLayer
     internal val windowAccessible: Accessible
         get() = delegate.windowAccessible

diff --git a/...ose/ui/ui/src/desktopMain/kotlin/androidx/compose/ui/awt/ComposeWindowDelegate.desktop.kt b/...ose/ui/ui/src/desktopMain/kotlin/androidx/compose/ui/awt/ComposeWindowDelegate.desktop.kt
@@ -19,6 +19,7 @@ package androidx.compose.ui.awt
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.CompositionLocalContext
 import androidx.compose.runtime.CompositionLocalProvider
+import androidx.compose.ui.ComposeScene
 import androidx.compose.ui.ExperimentalComposeUiApi
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.input.key.KeyEvent
@@ -54,6 +55,8 @@ internal class ComposeWindowDelegate(
         get() = requireNotNull(_bridge) {
             "ComposeBridge is disposed"
         }
+    internal val scene: ComposeScene
+        get() = bridge.scene
     internal val windowAccessible: Accessible
         get() = bridge.sceneAccessible
     val undecoratedWindowResizer = UndecoratedWindowResizer(window)

diff --git a/compose/ui/ui/src/desktopMain/kotlin/androidx/compose/ui/window/DesktopPopup.desktop.kt b/compose/ui/ui/src/desktopMain/kotlin/androidx/compose/ui/window/DesktopPopup.desktop.kt
@@ -37,7 +37,7 @@ import kotlin.math.roundToInt
  */
 @Composable
 private fun rememberCursorPosition(): Offset? {
-    val scene = LocalComposeScene.current
+    val scene = LocalComposeScene.current!!
     return remember { scene.pointerPositions.values.firstOrNull() }
 }
 

diff --git a/compose/ui/ui/src/desktopMain/kotlin/androidx/compose/ui/window/Dialog.desktop.kt b/compose/ui/ui/src/desktopMain/kotlin/androidx/compose/ui/window/Dialog.desktop.kt
@@ -17,12 +17,14 @@
 package androidx.compose.ui.window
 
 import androidx.compose.runtime.Composable
+import androidx.compose.runtime.SideEffect
 import androidx.compose.runtime.Stable
 import androidx.compose.runtime.currentCompositionLocalContext
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.rememberUpdatedState
 import androidx.compose.ui.ExperimentalComposeUiApi
+import androidx.compose.ui.LocalComposeScene
 import androidx.compose.ui.awt.ComposeDialog
 import androidx.compose.ui.graphics.painter.Painter
 import androidx.compose.ui.input.key.KeyEvent
@@ -338,16 +340,21 @@ fun DialogWindow(
     val windowExceptionHandlerFactory by rememberUpdatedState(
         LocalWindowExceptionHandlerFactory.current
     )
+    val parentScene = LocalComposeScene.current
     AwtWindow(
         visible = visible,
         create = {
             create().apply {
+                parentScene?.addChildScene(scene)
                 this.compositionLocalContext = compositionLocalContext
                 this.exceptionHandler = windowExceptionHandlerFactory.exceptionHandler(this)
                 setContent(onPreviewKeyEvent, onKeyEvent, content)
             }
         },
-        dispose = dispose,
+        dispose = {
+            parentScene?.removeChildScene(it.scene)
+            dispose(it)
+        },
         update = {
             it.compositionLocalContext = compositionLocalContext
             it.exceptionHandler = windowExceptionHandlerFactory.exceptionHandler(it)
@@ -364,7 +371,7 @@ fun DialogWindow(
             if (!wasDisplayable && it.isDisplayable) {
                 it.contentPane.paint(it.contentPane.graphics)
             }
-        }
+        },
     )
 }
 

diff --git a/compose/ui/ui/src/desktopMain/kotlin/androidx/compose/ui/window/Window.desktop.kt b/compose/ui/ui/src/desktopMain/kotlin/androidx/compose/ui/window/Window.desktop.kt
@@ -25,6 +25,7 @@ import androidx.compose.runtime.remember
 import androidx.compose.runtime.rememberCompositionContext
 import androidx.compose.runtime.rememberUpdatedState
 import androidx.compose.ui.ExperimentalComposeUiApi
+import androidx.compose.ui.LocalComposeScene
 import androidx.compose.ui.awt.ComposeWindow
 import androidx.compose.ui.graphics.painter.Painter
 import androidx.compose.ui.input.key.KeyEvent
@@ -397,16 +398,21 @@ fun Window(
     val windowExceptionHandlerFactory by rememberUpdatedState(
         LocalWindowExceptionHandlerFactory.current
     )
+    val parentScene = LocalComposeScene.current
     AwtWindow(
         visible = visible,
         create = {
             create().apply {
+                parentScene?.addChildScene(scene)
                 this.compositionLocalContext = compositionLocalContext
                 this.exceptionHandler = windowExceptionHandlerFactory.exceptionHandler(this)
                 setContent(onPreviewKeyEvent, onKeyEvent, content)
             }
         },
-        dispose = dispose,
+        dispose = {
+            parentScene?.removeChildScene(it.scene)
+            dispose(it)
+        },
         update = {
             it.compositionLocalContext = compositionLocalContext
             it.exceptionHandler = windowExceptionHandlerFactory.exceptionHandler(it)
@@ -423,7 +429,7 @@ fun Window(
             if (!wasDisplayable && it.isDisplayable) {
                 it.contentPane.paint(it.contentPane.graphics)
             }
-        }
+        },
     )
 }
 

diff --git a/compose/ui/ui/src/skikoMain/kotlin/androidx/compose/ui/ComposeScene.skiko.kt b/compose/ui/ui/src/skikoMain/kotlin/androidx/compose/ui/ComposeScene.skiko.kt
@@ -37,9 +37,7 @@ import org.jetbrains.skia.Canvas
 import org.jetbrains.skiko.SkiaLayer
 import org.jetbrains.skiko.currentNanoTime
 
-internal val LocalComposeScene = staticCompositionLocalOf<ComposeScene> {
-    error("CompositionLocal LocalComposeScene not provided")
-}
+internal val LocalComposeScene = staticCompositionLocalOf<ComposeScene?> { null }
 
 /**
  * A virtual container that encapsulates Compose UI content. UI content can be constructed via
@@ -108,6 +106,16 @@ class ComposeScene internal constructor(
         invalidate,
     )
 
+    private val childScenes = mutableSetOf<ComposeScene>()
+
+    fun addChildScene(scene: ComposeScene) {
+        childScenes.add(scene)
+    }
+
+    fun removeChildScene(scene: ComposeScene) {
+        childScenes.remove(scene)
+    }
+
     private var isInvalidationDisabled = false
 
     @Volatile
@@ -175,12 +183,25 @@ class ComposeScene internal constructor(
         }
     }
 
+    /**
+     * Adds this scene's [RootForTest]s, including any of its child scenes, into the given set.
+     */
+    private fun addRootsForTestTo(target: MutableSet<RootForTest>) {
+        target.addAll(owners)
+        for (child in childScenes) {
+            child.addRootsForTestTo(target)
+        }
+    }
+
     /**
      * All currently registered [RootForTest]s. After calling [setContent] the first root
      * will be added. If there is an any [Popup] is present in the content, it will be added as
      * another [RootForTest]
      */
-    val roots: Set<RootForTest> get() = owners.toSet()
+    val roots: Set<RootForTest>
+        get() = buildSet(owners.size) {
+            addRootsForTestTo(this)
+        }
 
     private val defaultPointerStateTracker = DefaultPointerStateTracker()
 

diff --git a/compose/ui/ui/src/skikoMain/kotlin/androidx/compose/ui/window/PopupLayout.skiko.kt b/compose/ui/ui/src/skikoMain/kotlin/androidx/compose/ui/window/PopupLayout.skiko.kt
@@ -41,7 +41,7 @@ internal fun PopupLayout(
     onKeyEvent: ((KeyEvent) -> Boolean) = { false },
     content: @Composable () -> Unit
 ) {
-    val scene = LocalComposeScene.current
+    val scene = LocalComposeScene.current!!
     val density = LocalDensity.current
     val layoutDirection = LocalLayoutDirection.current