From 1702c7f21968f822ea77e0adacfb93a75f067d7c Mon Sep 17 00:00:00 2001
From: Ivan Matkov <ivan.matkov@jetbrains.com>
Date: Wed, 22 May 2024 19:37:37 +0200
Subject: [PATCH] Copy DelegatableNode.visitSubtreeIf

---
 ...Container.kt => InteropContainer.skiko.kt} | 11 ---
 .../node/TraversableNodeInDrawOrder.skiko.kt  | 89 +++++++++++++++++++
 2 files changed, 89 insertions(+), 11 deletions(-)
 rename compose/ui/ui/src/skikoMain/kotlin/androidx/compose/ui/node/{InteropContainer.kt => InteropContainer.skiko.kt} (91%)
 create mode 100644 compose/ui/ui/src/skikoMain/kotlin/androidx/compose/ui/node/TraversableNodeInDrawOrder.skiko.kt

diff --git a/compose/ui/ui/src/skikoMain/kotlin/androidx/compose/ui/node/InteropContainer.kt b/compose/ui/ui/src/skikoMain/kotlin/androidx/compose/ui/node/InteropContainer.skiko.kt
similarity index 91%
rename from compose/ui/ui/src/skikoMain/kotlin/androidx/compose/ui/node/InteropContainer.kt
rename to compose/ui/ui/src/skikoMain/kotlin/androidx/compose/ui/node/InteropContainer.skiko.kt
index 48e3b47c68198..cb20de18a73de 100644
--- a/compose/ui/ui/src/skikoMain/kotlin/androidx/compose/ui/node/InteropContainer.kt
+++ b/compose/ui/ui/src/skikoMain/kotlin/androidx/compose/ui/node/InteropContainer.skiko.kt
@@ -18,7 +18,6 @@ package androidx.compose.ui.node
 
 import androidx.compose.runtime.Composable
 import androidx.compose.ui.Modifier
-import androidx.compose.ui.areObjectsOfSameType
 import androidx.compose.ui.layout.LayoutCoordinates
 import androidx.compose.ui.layout.OverlayLayout
 
@@ -58,16 +57,6 @@ internal fun <T> InteropContainer<T>.countInteropComponentsBefore(nativeView: T)
     return componentsBefore
 }
 
-private fun <T> T.traverseDescendantsInDrawOrder(block: (T) -> Boolean) where T : TraversableNode {
-    visitSubtreeIf(Nodes.Traversable, zOrder = true) {
-        if (this.traverseKey == it.traverseKey && areObjectsOfSameType(this, it)) {
-            @Suppress("UNCHECKED_CAST")
-            if (!block(it as T)) return
-        }
-        true
-    }
-}
-
 /**
  * Wrapper of Compose content that might contain interop views. It adds a helper modifier to root
  * that allows traversing interop views in the tree with the right order.
diff --git a/compose/ui/ui/src/skikoMain/kotlin/androidx/compose/ui/node/TraversableNodeInDrawOrder.skiko.kt b/compose/ui/ui/src/skikoMain/kotlin/androidx/compose/ui/node/TraversableNodeInDrawOrder.skiko.kt
new file mode 100644
index 0000000000000..4c32b1ce93ec1
--- /dev/null
+++ b/compose/ui/ui/src/skikoMain/kotlin/androidx/compose/ui/node/TraversableNodeInDrawOrder.skiko.kt
@@ -0,0 +1,89 @@
+/*
+ * 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.node
+
+import androidx.compose.runtime.collection.MutableVector
+import androidx.compose.runtime.collection.mutableVectorOf
+import androidx.compose.ui.Modifier
+import androidx.compose.ui.areObjectsOfSameType
+
+internal fun <T> T.traverseDescendantsInDrawOrder(block: (T) -> Boolean) where T : TraversableNode {
+    visitSubtreeIf(Nodes.Traversable, zOrder = true) {
+        if (this.traverseKey == it.traverseKey && areObjectsOfSameType(this, it)) {
+            @Suppress("UNCHECKED_CAST")
+            if (!block(it as T)) return
+        }
+        true
+    }
+}
+
+// TODO: Remove a copy once aosp/3045224 merged
+
+private inline fun <reified T> DelegatableNode.visitSubtreeIf(
+    type: NodeKind<T>,
+    zOrder: Boolean = false,
+    block: (T) -> Boolean
+) = visitSubtreeIf(type.mask, zOrder) foo@{ node ->
+    node.dispatchForKind(type) {
+        if (!block(it)) return@foo false
+    }
+    true
+}
+
+private fun LayoutNode.getChildren(zOrder: Boolean) =
+    if (zOrder) {
+        zSortedChildren
+    } else {
+        _children
+    }
+
+private fun MutableVector<Modifier.Node>.addLayoutNodeChildren(
+    node: Modifier.Node,
+    zOrder: Boolean,
+) {
+    node.requireLayoutNode().getChildren(zOrder).forEachReversed {
+        add(it.nodes.head)
+    }
+}
+
+private inline fun DelegatableNode.visitSubtreeIf(
+    mask: Int,
+    zOrder: Boolean,
+    block: (Modifier.Node) -> Boolean
+) {
+    check(node.isAttached) { "visitSubtreeIf called on an unattached node" }
+    val branches = mutableVectorOf<Modifier.Node>()
+    val child = node.child
+    if (child == null)
+        branches.addLayoutNodeChildren(node, zOrder)
+    else
+        branches.add(child)
+    outer@ while (branches.isNotEmpty()) {
+        val branch = branches.removeAt(branches.size - 1)
+        if (branch.aggregateChildKindSet and mask != 0) {
+            var node: Modifier.Node? = branch
+            while (node != null) {
+                if (node.kindSet and mask != 0) {
+                    val diveDeeper = block(node)
+                    if (!diveDeeper) continue@outer
+                }
+                node = node.child
+            }
+        }
+        branches.addLayoutNodeChildren(branch, zOrder)
+    }
+}