Copy DelegatableNode.visitSubtreeIf

JetBrains · May 22, 2024 · 112dbf4 · 112dbf4
1 parent 9bcb9ed
commit 112dbf4
Show file tree

Hide file tree

Showing 2 changed files with 89 additions and 11 deletions.
diff --git a/...roidx/compose/ui/node/InteropContainer.kt → ...compose/ui/node/InteropContainer.skiko.kt b/...roidx/compose/ui/node/InteropContainer.kt → ...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/...e/ui/ui/src/skikoMain/kotlin/androidx/compose/ui/node/TraversableNodeInDrawOrder.skiko.kt b/...e/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)
+    }
+}