Fix updating root constraints (#1387)

Since `owner.size` doesn't really observed during composition setting
scene size might not trigger invalidations

- Manually run `updateRootConstraints` instead of marking size as state
- Pick extra conditions from
[aosp/2598886](https://android-review.googlesource.com/c/platform/frameworks/support/+/2598886)

Fixes stretching reported in
JetBrains/compose-multiplatform#4850

## Testing

Run `MultiLayerComposeSceneTest`

## Release Notes

### Fixes - iOS

- Fix missing invalidations during native view resize

compose/ui/ui/src/skikoMain/kotlin/androidx/compose/ui/node/RootNodeOwner.skiko.kt

@@ -127,12 +127,13 @@ internal class RootNodeOwner(
         }
     val owner: Owner = OwnerImpl(layoutDirection, coroutineContext)
     val semanticsOwner = SemanticsOwner(owner.root)
-    var size by mutableStateOf(size)
+    var size: IntSize? = size
+        set(value) {
+            field = value
+            onRootConstrainsChanged(value?.toConstraints())
+        }
     var density by mutableStateOf(density)
 
-    private val constraints
-        get() = size?.toConstraints() ?: Constraints()
-
     private var _layoutDirection by mutableStateOf(layoutDirection)
     var layoutDirection: LayoutDirection
         get() = _layoutDirection
@@ -152,6 +153,7 @@ internal class RootNodeOwner(
         snapshotObserver.startObserving()
         owner.root.attach(owner)
         platformContext.rootForTestListener?.onRootForTestCreated(rootForTest)
+        onRootConstrainsChanged(size?.toConstraints())
     }
 
     fun dispose() {
@@ -209,6 +211,13 @@ internal class RootNodeOwner(
         owner.root.modifier = rootModifier then modifier
     }
 
+    private fun onRootConstrainsChanged(constraints: Constraints?) {
+        measureAndLayoutDelegate.updateRootConstraintsWithInfinityCheck(constraints)
+        if (measureAndLayoutDelegate.hasPendingMeasureOrLayout) {
+            snapshotInvalidationTracker.requestMeasureAndLayout()
+        }
+    }
+
     @OptIn(InternalCoreApi::class)
     fun onPointerInput(event: PointerInputEvent) {
         if (event.button != null) {
@@ -313,22 +322,32 @@ internal class RootNodeOwner(
         }
 
         override fun measureAndLayout(sendPointerUpdate: Boolean) {
-            measureAndLayoutDelegate.updateRootConstraintsWithInfinityCheck(constraints)
-            val rootNodeResized = measureAndLayoutDelegate.measureAndLayout {
-                if (sendPointerUpdate) {
-                    inputHandler.onPointerUpdate()
+            // only run the logic when we have something pending
+            if (measureAndLayoutDelegate.hasPendingMeasureOrLayout ||
+                measureAndLayoutDelegate.hasPendingOnPositionedCallbacks
+            ) {
+                trace("RootNodeOwner:measureAndLayout") {
+                    val resend = if (sendPointerUpdate) inputHandler::onPointerUpdate else null
+                    val rootNodeResized = measureAndLayoutDelegate.measureAndLayout(resend)
+                    if (rootNodeResized) {
+                        snapshotInvalidationTracker.requestDraw()
+                    }
+                    measureAndLayoutDelegate.dispatchOnPositionedCallbacks()
                 }
             }
-            if (rootNodeResized) {
-                snapshotInvalidationTracker.requestDraw()
-            }
-            measureAndLayoutDelegate.dispatchOnPositionedCallbacks()
         }
 
         override fun measureAndLayout(layoutNode: LayoutNode, constraints: Constraints) {
-            measureAndLayoutDelegate.measureAndLayout(layoutNode, constraints)
-            inputHandler.onPointerUpdate()
-            measureAndLayoutDelegate.dispatchOnPositionedCallbacks()
+            trace("RootNodeOwner:measureAndLayout") {
+                measureAndLayoutDelegate.measureAndLayout(layoutNode, constraints)
+                inputHandler.onPointerUpdate()
+                // only dispatch the callbacks if we don't have other nodes to process as otherwise
+                // we will have one more measureAndLayout() pass anyway in the same frame.
+                // it allows us to not traverse the hierarchy twice.
+                if (!measureAndLayoutDelegate.hasPendingMeasureOrLayout) {
+                    measureAndLayoutDelegate.dispatchOnPositionedCallbacks()
+                }
+            }
         }
 
         override fun forceMeasureTheSubtree(layoutNode: LayoutNode, affectsLookahead: Boolean) {
@@ -562,12 +581,23 @@ internal const val LargeDimension = ConstraintsMinNonFocusMask - 1
  * and pass constraint large enough instead
  */
 private fun MeasureAndLayoutDelegate.updateRootConstraintsWithInfinityCheck(
-    constraints: Constraints
+    constraints: Constraints?
 ) {
-    val maxWidth = if (constraints.hasBoundedWidth) constraints.maxWidth else LargeDimension
-    val maxHeight = if (constraints.hasBoundedHeight) constraints.maxHeight else LargeDimension
     updateRootConstraints(
-        Constraints(constraints.minWidth, maxWidth, constraints.minHeight, maxHeight)
+        constraints = Constraints(
+            minWidth = constraints?.minWidth ?: 0,
+            maxWidth = if (constraints != null && constraints.hasBoundedWidth) {
+                constraints.maxWidth
+            } else {
+                LargeDimension
+            },
+            minHeight = constraints?.minHeight ?: 0,
+            maxHeight = if (constraints != null && constraints.hasBoundedHeight) {
+                constraints.maxHeight
+            } else {
+                LargeDimension
+            }
+        )
     )
 }
 

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

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2024 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.scene
+
+import androidx.compose.foundation.layout.Box
+import androidx.compose.foundation.layout.fillMaxSize
+import androidx.compose.ui.Modifier
+import androidx.compose.ui.unit.IntSize
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.runTest
+
+class MultiLayerComposeSceneTest {
+
+    @Test
+    fun sceneSizeChangeTriggersInvalidation() = runTest(StandardTestDispatcher()) {
+        var invalidationCount = 0
+        val scene = MultiLayerComposeScene(
+            size = IntSize(100, 100),
+            invalidate = { invalidationCount++ }
+        )
+        try {
+            scene.setContent { Box(Modifier.fillMaxSize()) }
+
+            assertEquals(1, invalidationCount)
+            scene.size = IntSize(120, 120)
+            assertEquals(2, invalidationCount)
+        } finally {
+            scene.close()
+        }
+    }
+}