1#[cfg(feature = "ghost_nodes")]
4use crate::ui_node::ComputedUiTargetCamera;
5use crate::Node;
6#[cfg(feature = "ghost_nodes")]
7use bevy_camera::visibility::Visibility;
8use bevy_ecs::{prelude::*, system::SystemParam};
9#[cfg(feature = "ghost_nodes")]
10use bevy_reflect::prelude::*;
11#[cfg(feature = "ghost_nodes")]
12use bevy_transform::prelude::Transform;
13#[cfg(feature = "ghost_nodes")]
14use smallvec::SmallVec;
15#[cfg(feature = "ghost_nodes")]
21#[derive(Component, Debug, Copy, Clone, Reflect)]
22#[cfg_attr(feature = "ghost_nodes", derive(Default))]
23#[reflect(Component, Debug, Clone)]
24#[require(Visibility, Transform, ComputedUiTargetCamera)]
25pub struct GhostNode;
26
27#[cfg(feature = "ghost_nodes")]
28#[derive(SystemParam)]
32pub struct UiRootNodes<'w, 's> {
33 root_node_query: Query<'w, 's, Entity, (With<Node>, Without<ChildOf>)>,
34 root_ghost_node_query: Query<'w, 's, Entity, (With<GhostNode>, Without<ChildOf>)>,
35 all_nodes_query: Query<'w, 's, Entity, With<Node>>,
36 ui_children: UiChildren<'w, 's>,
37}
38
39#[cfg(not(feature = "ghost_nodes"))]
40pub type UiRootNodes<'w, 's> = Query<'w, 's, Entity, (With<Node>, Without<ChildOf>)>;
41
42#[cfg(feature = "ghost_nodes")]
43impl<'w, 's> UiRootNodes<'w, 's> {
44 pub fn iter(&'s self) -> impl Iterator<Item = Entity> + 's {
45 self.root_node_query
46 .iter()
47 .chain(self.root_ghost_node_query.iter().flat_map(|root_ghost| {
48 self.all_nodes_query
49 .iter_many(self.ui_children.iter_ui_children(root_ghost))
50 }))
51 }
52}
53
54#[cfg(feature = "ghost_nodes")]
55#[derive(SystemParam)]
57pub struct UiChildren<'w, 's> {
58 ui_children_query: Query<
59 'w,
60 's,
61 (Option<&'static Children>, Has<GhostNode>),
62 Or<(With<Node>, With<GhostNode>)>,
63 >,
64 changed_children_query: Query<'w, 's, Entity, Changed<Children>>,
65 children_query: Query<'w, 's, &'static Children>,
66 ghost_nodes_query: Query<'w, 's, Entity, With<GhostNode>>,
67 parents_query: Query<'w, 's, &'static ChildOf>,
68}
69
70#[cfg(not(feature = "ghost_nodes"))]
71#[derive(SystemParam)]
73pub struct UiChildren<'w, 's> {
74 ui_children_query: Query<'w, 's, Option<&'static Children>, With<Node>>,
75 changed_children_query: Query<'w, 's, Entity, Changed<Children>>,
76 parents_query: Query<'w, 's, &'static ChildOf>,
77}
78
79#[cfg(feature = "ghost_nodes")]
80impl<'w, 's> UiChildren<'w, 's> {
81 pub fn iter_ui_children(&'s self, entity: Entity) -> UiChildrenIter<'w, 's> {
89 UiChildrenIter {
90 stack: self
91 .ui_children_query
92 .get(entity)
93 .map_or(SmallVec::new(), |(children, _)| {
94 children.into_iter().flatten().rev().copied().collect()
95 }),
96 query: &self.ui_children_query,
97 }
98 }
99
100 pub fn get_parent(&'s self, entity: Entity) -> Option<Entity> {
102 self.parents_query
103 .iter_ancestors(entity)
104 .find(|entity| !self.ghost_nodes_query.contains(*entity))
105 }
106
107 pub fn iter_ghost_nodes(&'s self, entity: Entity) -> Box<dyn Iterator<Item = Entity> + 's> {
109 Box::new(
110 self.children_query
111 .get(entity)
112 .into_iter()
113 .flat_map(|children| {
114 self.ghost_nodes_query
115 .iter_many(children)
116 .flat_map(|entity| {
117 core::iter::once(entity).chain(self.iter_ghost_nodes(entity))
118 })
119 }),
120 )
121 }
122
123 pub fn is_changed(&'s self, entity: Entity) -> bool {
125 self.changed_children_query.contains(entity)
126 || self
127 .iter_ghost_nodes(entity)
128 .any(|entity| self.changed_children_query.contains(entity))
129 }
130
131 pub fn is_ui_node(&'s self, entity: Entity) -> bool {
133 self.ui_children_query.contains(entity)
134 }
135}
136
137#[cfg(not(feature = "ghost_nodes"))]
138impl<'w, 's> UiChildren<'w, 's> {
139 pub fn iter_ui_children(&'s self, entity: Entity) -> impl Iterator<Item = Entity> + 's {
141 self.ui_children_query
142 .get(entity)
143 .ok()
144 .flatten()
145 .map(|children| children.as_ref())
146 .unwrap_or(&[])
147 .iter()
148 .copied()
149 }
150
151 pub fn get_parent(&'s self, entity: Entity) -> Option<Entity> {
153 self.parents_query.get(entity).ok().map(ChildOf::parent)
154 }
155
156 pub fn is_changed(&'s self, entity: Entity) -> bool {
158 self.changed_children_query.contains(entity)
159 }
160
161 pub fn is_ui_node(&'s self, entity: Entity) -> bool {
163 self.ui_children_query.contains(entity)
164 }
165}
166
167#[cfg(feature = "ghost_nodes")]
168pub struct UiChildrenIter<'w, 's> {
169 stack: SmallVec<[Entity; 8]>,
170 query: &'s Query<
171 'w,
172 's,
173 (Option<&'static Children>, Has<GhostNode>),
174 Or<(With<Node>, With<GhostNode>)>,
175 >,
176}
177
178#[cfg(feature = "ghost_nodes")]
179impl<'w, 's> Iterator for UiChildrenIter<'w, 's> {
180 type Item = Entity;
181 fn next(&mut self) -> Option<Self::Item> {
182 loop {
183 let entity = self.stack.pop()?;
184 if let Ok((children, has_ghost_node)) = self.query.get(entity) {
185 if !has_ghost_node {
186 return Some(entity);
187 }
188 if let Some(children) = children {
189 self.stack.extend(children.iter().rev());
190 }
191 }
192 }
193 }
194}
195
196#[cfg(all(test, feature = "ghost_nodes"))]
197mod tests {
198 use bevy_ecs::{
199 prelude::Component,
200 system::{Query, SystemState},
201 world::World,
202 };
203
204 use super::{GhostNode, Node, UiChildren, UiRootNodes};
205
206 #[derive(Component, PartialEq, Debug)]
207 struct A(usize);
208
209 #[test]
210 fn iterate_ui_root_nodes() {
211 let world = &mut World::new();
212
213 world
215 .spawn((A(1), Node::default()))
216 .with_children(|parent| {
217 parent.spawn((A(2), Node::default()));
218 parent
219 .spawn((A(3), GhostNode))
220 .with_child((A(4), Node::default()));
221 });
222
223 world.spawn((A(5), GhostNode)).with_children(|parent| {
225 parent.spawn((A(6), Node::default()));
226 parent
227 .spawn((A(7), GhostNode))
228 .with_child((A(8), Node::default()))
229 .with_child(A(9));
230 });
231
232 let mut system_state = SystemState::<(UiRootNodes, Query<&A>)>::new(world);
233 let (ui_root_nodes, a_query) = system_state.get(world);
234
235 let result: Vec<_> = a_query.iter_many(ui_root_nodes.iter()).collect();
236
237 assert_eq!([&A(1), &A(6), &A(8)], result.as_slice());
238 }
239
240 #[test]
241 fn iterate_ui_children() {
242 let world = &mut World::new();
243
244 let n1 = world.spawn((A(1), Node::default())).id();
245 let n2 = world.spawn((A(2), GhostNode)).id();
246 let n3 = world.spawn((A(3), GhostNode)).id();
247 let n4 = world.spawn((A(4), Node::default())).id();
248 let n5 = world.spawn((A(5), Node::default())).id();
249
250 let n6 = world.spawn((A(6), GhostNode)).id();
251 let n7 = world.spawn((A(7), GhostNode)).id();
252 let n8 = world.spawn((A(8), Node::default())).id();
253 let n9 = world.spawn((A(9), GhostNode)).id();
254 let n10 = world.spawn((A(10), Node::default())).id();
255
256 let no_ui = world.spawn_empty().id();
257
258 world.entity_mut(n1).add_children(&[n2, n3, n4, n6]);
259 world.entity_mut(n2).add_children(&[n5]);
260
261 world.entity_mut(n6).add_children(&[n7, no_ui, n9]);
262 world.entity_mut(n7).add_children(&[n8]);
263 world.entity_mut(n9).add_children(&[n10]);
264
265 let mut system_state = SystemState::<(UiChildren, Query<&A>)>::new(world);
266 let (ui_children, a_query) = system_state.get(world);
267
268 let result: Vec<_> = a_query
269 .iter_many(ui_children.iter_ui_children(n1))
270 .collect();
271
272 assert_eq!([&A(5), &A(4), &A(8), &A(10)], result.as_slice());
273 }
274}