From 5cd70c9b1d4d3dabc0738a83f5e501b31ffe020d Mon Sep 17 00:00:00 2001
From: Ivan Matkov <ivan.matkov@jetbrains.com>
Date: Wed, 17 Apr 2024 18:00:55 +0200
Subject: [PATCH] Add an option to visit children and subtree according to
 z-order

Currently there is no way to iterate over nodes in draw order

Test: TBD
Change-Id: Ib1e8bc917978c68bfb88050cb2122879b82f6476
---
 .../compose/ui/node/DelegatableNode.kt        | 56 ++++++++++++++-----
 1 file changed, 41 insertions(+), 15 deletions(-)

diff --git a/compose/ui/ui/src/commonMain/kotlin/androidx/compose/ui/node/DelegatableNode.kt b/compose/ui/ui/src/commonMain/kotlin/androidx/compose/ui/node/DelegatableNode.kt
index d5ad56e956f6a..c1d6b0bebd087 100644
--- a/compose/ui/ui/src/commonMain/kotlin/androidx/compose/ui/node/DelegatableNode.kt
+++ b/compose/ui/ui/src/commonMain/kotlin/androidx/compose/ui/node/DelegatableNode.kt
@@ -94,7 +94,11 @@ internal fun DelegatableNode.nearestAncestor(mask: Int): Modifier.Node? {
     return null
 }
 
-internal inline fun DelegatableNode.visitSubtree(mask: Int, block: (Modifier.Node) -> Unit) {
+internal inline fun DelegatableNode.visitSubtree(
+    mask: Int,
+    zOrder: Boolean,
+    block: (Modifier.Node) -> Unit
+) {
     // TODO(lmr): we might want to add some safety wheels to prevent this from being called
     //  while one of the chains is being diffed / updated.
     check(node.isAttached) { "visitSubtree called on an unattached node" }
@@ -118,29 +122,43 @@ internal inline fun DelegatableNode.visitSubtree(mask: Int, block: (Modifier.Nod
             }
         }
         node = null
-        nodes.push(layout._children)
+        nodes.push(layout.getChildren(zOrder))
         layout = if (nodes.isNotEmpty()) nodes.pop() else null
     }
 }
 
-private fun MutableVector<Modifier.Node>.addLayoutNodeChildren(node: Modifier.Node) {
-    node.requireLayoutNode()._children.forEachReversed {
+private fun LayoutNode.getChildren(zOrder: Boolean) =
+    if (zOrder) {
+        _children
+    } else {
+        zSortedChildren
+    }
+
+private fun MutableVector<Modifier.Node>.addLayoutNodeChildren(
+    node: Modifier.Node,
+    zOrder: Boolean,
+) {
+    node.requireLayoutNode().getChildren(zOrder).forEachReversed {
         add(it.nodes.head)
     }
 }
 
-internal inline fun DelegatableNode.visitChildren(mask: Int, block: (Modifier.Node) -> Unit) {
+internal inline fun DelegatableNode.visitChildren(
+    mask: Int,
+    zOrder: Boolean,
+    block: (Modifier.Node) -> Unit
+) {
     check(node.isAttached) { "visitChildren called on an unattached node" }
     val branches = mutableVectorOf<Modifier.Node>()
     val child = node.child
     if (child == null)
-        branches.addLayoutNodeChildren(node)
+        branches.addLayoutNodeChildren(node, zOrder)
     else
         branches.add(child)
     while (branches.isNotEmpty()) {
         val branch = branches.removeAt(branches.lastIndex)
         if (branch.aggregateChildKindSet and mask == 0) {
-            branches.addLayoutNodeChildren(branch)
+            branches.addLayoutNodeChildren(branch, zOrder)
             // none of these nodes match the mask, so don't bother traversing them
             continue
         }
@@ -159,12 +177,16 @@ internal inline fun DelegatableNode.visitChildren(mask: Int, block: (Modifier.No
  * visit the shallow tree of children of a given mask, but if block returns true, we will continue
  * traversing below it
  */
-internal inline fun DelegatableNode.visitSubtreeIf(mask: Int, block: (Modifier.Node) -> Boolean) {
-    check(node.isAttached) { "visitSubtreeIf called on an unattached node" }
+internal inline fun DelegatableNode.visitSubtreeIf(
+    mask: Int,
+    zOrder: Boolean,
+    block: (Modifier.Node) -> Boolean
+) {
+    checkPrecondition(node.isAttached) { "visitSubtreeIf called on an unattached node" }
     val branches = mutableVectorOf<Modifier.Node>()
     val child = node.child
     if (child == null)
-        branches.addLayoutNodeChildren(node)
+        branches.addLayoutNodeChildren(node, zOrder)
     else
         branches.add(child)
     outer@ while (branches.isNotEmpty()) {
@@ -179,7 +201,7 @@ internal inline fun DelegatableNode.visitSubtreeIf(mask: Int, block: (Modifier.N
                 node = node.child
             }
         }
-        branches.addLayoutNodeChildren(branch)
+        branches.addLayoutNodeChildren(branch, zOrder)
     }
 }
 
@@ -267,26 +289,30 @@ internal inline fun <reified T : Any> DelegatableNode.nearestAncestor(type: Node
 
 internal inline fun <reified T> DelegatableNode.visitSubtree(
     type: NodeKind<T>,
+    zOrder: Boolean = false,
     block: (T) -> Unit
-) = visitSubtree(type.mask) { it.dispatchForKind(type, block) }
+) = visitSubtree(type.mask, zOrder) { it.dispatchForKind(type, block) }
 
 internal inline fun <reified T> DelegatableNode.visitChildren(
     type: NodeKind<T>,
+    zOrder: Boolean = false,
     block: (T) -> Unit
-) = visitChildren(type.mask) { it.dispatchForKind(type, block) }
+) = visitChildren(type.mask, zOrder) { it.dispatchForKind(type, block) }
 
 internal inline fun <reified T> DelegatableNode.visitSelfAndChildren(
     type: NodeKind<T>,
+    zOrder: Boolean = false,
     block: (T) -> Unit
 ) {
     node.dispatchForKind(type, block)
-    visitChildren(type.mask) { it.dispatchForKind(type, block) }
+    visitChildren(type.mask, zOrder) { it.dispatchForKind(type, block) }
 }
 
 internal inline fun <reified T> DelegatableNode.visitSubtreeIf(
     type: NodeKind<T>,
+    zOrder: Boolean = false,
     block: (T) -> Boolean
-) = visitSubtreeIf(type.mask) foo@{ node ->
+) = visitSubtreeIf(type.mask, zOrder) foo@{ node ->
     node.dispatchForKind(type) {
         if (!block(it)) return@foo false
     }