From 179c2e108ea7358e0fa1f8ff252e5b6ef3180ab3 Mon Sep 17 00:00:00 2001
From: Viktor Gustavsson <villor94@gmail.com>
Date: Sat, 21 Sep 2024 16:28:26 +0200
Subject: [PATCH] Use UiChildren/UiRootNodes in bevy_ui::update

---
 crates/bevy_ui/src/update.rs | 51 +++++++++++++++++-------------------
 1 file changed, 24 insertions(+), 27 deletions(-)

diff --git a/crates/bevy_ui/src/update.rs b/crates/bevy_ui/src/update.rs
index afb24274bbc1ca..2102df55fac1bc 100644
--- a/crates/bevy_ui/src/update.rs
+++ b/crates/bevy_ui/src/update.rs
@@ -1,14 +1,13 @@
 //! This module contains systems that update the UI when something changes
 
-use crate::{CalculatedClip, Display, OverflowAxis, Style, TargetCamera, UiRootNodes};
+use crate::{CalculatedClip, Display, OverflowAxis, Style, TargetCamera, UiChildren, UiRootNodes};
 
 use super::Node;
 use bevy_ecs::{
     entity::Entity,
-    query::{Changed, With, Without},
+    query::{Changed, With},
     system::{Commands, Query},
 };
-use bevy_hierarchy::{Children, Parent};
 use bevy_math::Rect;
 use bevy_transform::components::GlobalTransform;
 use bevy_utils::HashSet;
@@ -18,12 +17,12 @@ pub fn update_clipping_system(
     mut commands: Commands,
     root_nodes: UiRootNodes,
     mut node_query: Query<(&Node, &GlobalTransform, &Style, Option<&mut CalculatedClip>)>,
-    children_query: Query<&Children>,
+    ui_children: UiChildren,
 ) {
     for root_node in root_nodes.iter() {
         update_clipping(
             &mut commands,
-            &children_query,
+            &ui_children,
             &mut node_query,
             root_node,
             None,
@@ -33,7 +32,7 @@ pub fn update_clipping_system(
 
 fn update_clipping(
     commands: &mut Commands,
-    children_query: &Query<&Children>,
+    ui_children: &UiChildren,
     node_query: &mut Query<(&Node, &GlobalTransform, &Style, Option<&mut CalculatedClip>)>,
     entity: Entity,
     mut maybe_inherited_clip: Option<Rect>,
@@ -92,10 +91,8 @@ fn update_clipping(
         Some(maybe_inherited_clip.map_or(node_rect, |c| c.intersect(node_rect)))
     };
 
-    if let Ok(children) = children_query.get(entity) {
-        for &child in children {
-            update_clipping(commands, children_query, node_query, child, children_clip);
-        }
+    for child in ui_children.iter_ui_children(entity) {
+        update_clipping(commands, ui_children, node_query, child, children_clip);
     }
 }
 
@@ -103,11 +100,11 @@ pub fn update_target_camera_system(
     mut commands: Commands,
     changed_root_nodes_query: Query<
         (Entity, Option<&TargetCamera>),
-        (With<Node>, Without<Parent>, Changed<TargetCamera>),
+        (With<Node>, Changed<TargetCamera>),
     >,
-    changed_children_query: Query<(Entity, Option<&TargetCamera>), (With<Node>, Changed<Children>)>,
-    children_query: Query<&Children, With<Node>>,
-    node_query: Query<Option<&TargetCamera>, With<Node>>,
+    node_query: Query<(Entity, Option<&TargetCamera>), With<Node>>,
+    ui_root_nodes: UiRootNodes,
+    ui_children: UiChildren,
 ) {
     // Track updated entities to prevent redundant updates, as `Commands` changes are deferred,
     // and updates done for changed_children_query can overlap with itself or with root_node_query
@@ -115,12 +112,12 @@ pub fn update_target_camera_system(
 
     // Assuming that TargetCamera is manually set on the root node only,
     // update root nodes first, since it implies the biggest change
-    for (root_node, target_camera) in &changed_root_nodes_query {
+    for (root_node, target_camera) in changed_root_nodes_query.iter_many(ui_root_nodes.iter()) {
         update_children_target_camera(
             root_node,
             target_camera,
             &node_query,
-            &children_query,
+            &ui_children,
             &mut commands,
             &mut updated_entities,
         );
@@ -129,12 +126,16 @@ pub fn update_target_camera_system(
     // If the root node TargetCamera was changed, then every child is updated
     // by this point, and iteration will be skipped.
     // Otherwise, update changed children
-    for (parent, target_camera) in &changed_children_query {
+    for (parent, target_camera) in &node_query {
+        if !ui_children.is_changed(parent) {
+            continue;
+        }
+
         update_children_target_camera(
             parent,
             target_camera,
             &node_query,
-            &children_query,
+            &ui_children,
             &mut commands,
             &mut updated_entities,
         );
@@ -144,19 +145,15 @@ pub fn update_target_camera_system(
 fn update_children_target_camera(
     entity: Entity,
     camera_to_set: Option<&TargetCamera>,
-    node_query: &Query<Option<&TargetCamera>, With<Node>>,
-    children_query: &Query<&Children, With<Node>>,
+    node_query: &Query<(Entity, Option<&TargetCamera>), With<Node>>,
+    ui_children: &UiChildren,
     commands: &mut Commands,
     updated_entities: &mut HashSet<Entity>,
 ) {
-    let Ok(children) = children_query.get(entity) else {
-        return;
-    };
-
-    for &child in children {
+    for child in ui_children.iter_ui_children(entity) {
         // Skip if the child has already been updated or update is not needed
         if updated_entities.contains(&child)
-            || camera_to_set == node_query.get(child).ok().flatten()
+            || camera_to_set == node_query.get(child).ok().and_then(|(_, camera)| camera)
         {
             continue;
         }
@@ -175,7 +172,7 @@ fn update_children_target_camera(
             child,
             camera_to_set,
             node_query,
-            children_query,
+            ui_children,
             commands,
             updated_entities,
         );