1use bevy_color::Color;
2use bevy_ecs::{
3 component::Mutable,
4 prelude::*,
5 system::{Query, SystemParam},
6};
7
8use crate::{TextColor, TextFont, TextSpan};
9
10pub trait TextSpanAccess: Component<Mutability = Mutable> {
12 fn read_span(&self) -> &str;
14 fn write_span(&mut self) -> &mut String;
16}
17
18pub trait TextRoot: TextSpanAccess + From<String> {}
20
21pub trait TextSpanComponent: TextSpanAccess + From<String> {}
23
24#[derive(Resource, Default)]
26pub struct TextIterScratch {
27 stack: Vec<(&'static Children, usize)>,
28}
29
30impl TextIterScratch {
31 fn take<'a>(&mut self) -> Vec<(&'a Children, usize)> {
32 core::mem::take(&mut self.stack)
33 .into_iter()
34 .map(|_| -> (&Children, usize) { unreachable!() })
35 .collect()
36 }
37
38 fn recover(&mut self, mut stack: Vec<(&Children, usize)>) {
39 stack.clear();
40 self.stack = stack
41 .into_iter()
42 .map(|_| -> (&'static Children, usize) { unreachable!() })
43 .collect();
44 }
45}
46
47#[derive(SystemParam)]
51pub struct TextReader<'w, 's, R: TextRoot> {
52 scratch: Local<'s, TextIterScratch>,
54 roots: Query<
55 'w,
56 's,
57 (
58 &'static R,
59 &'static TextFont,
60 &'static TextColor,
61 Option<&'static Children>,
62 ),
63 >,
64 spans: Query<
65 'w,
66 's,
67 (
68 &'static TextSpan,
69 &'static TextFont,
70 &'static TextColor,
71 Option<&'static Children>,
72 ),
73 >,
74}
75
76impl<'w, 's, R: TextRoot> TextReader<'w, 's, R> {
77 pub fn iter(&mut self, root_entity: Entity) -> TextSpanIter<'_, R> {
79 let stack = self.scratch.take();
80
81 TextSpanIter {
82 scratch: &mut self.scratch,
83 root_entity: Some(root_entity),
84 stack,
85 roots: &self.roots,
86 spans: &self.spans,
87 }
88 }
89
90 pub fn get(
92 &mut self,
93 root_entity: Entity,
94 index: usize,
95 ) -> Option<(Entity, usize, &str, &TextFont, Color)> {
96 self.iter(root_entity).nth(index)
97 }
98
99 pub fn get_text(&mut self, root_entity: Entity, index: usize) -> Option<&str> {
101 self.get(root_entity, index).map(|(_, _, text, _, _)| text)
102 }
103
104 pub fn get_font(&mut self, root_entity: Entity, index: usize) -> Option<&TextFont> {
106 self.get(root_entity, index).map(|(_, _, _, font, _)| font)
107 }
108
109 pub fn get_color(&mut self, root_entity: Entity, index: usize) -> Option<Color> {
111 self.get(root_entity, index)
112 .map(|(_, _, _, _, color)| color)
113 }
114
115 pub fn text(&mut self, root_entity: Entity, index: usize) -> &str {
119 self.get_text(root_entity, index).unwrap()
120 }
121
122 pub fn font(&mut self, root_entity: Entity, index: usize) -> &TextFont {
126 self.get_font(root_entity, index).unwrap()
127 }
128
129 pub fn color(&mut self, root_entity: Entity, index: usize) -> Color {
133 self.get_color(root_entity, index).unwrap()
134 }
135}
136
137pub struct TextSpanIter<'a, R: TextRoot> {
143 scratch: &'a mut TextIterScratch,
144 root_entity: Option<Entity>,
145 stack: Vec<(&'a Children, usize)>,
147 roots: &'a Query<
148 'a,
149 'a,
150 (
151 &'static R,
152 &'static TextFont,
153 &'static TextColor,
154 Option<&'static Children>,
155 ),
156 >,
157 spans: &'a Query<
158 'a,
159 'a,
160 (
161 &'static TextSpan,
162 &'static TextFont,
163 &'static TextColor,
164 Option<&'static Children>,
165 ),
166 >,
167}
168
169impl<'a, R: TextRoot> Iterator for TextSpanIter<'a, R> {
170 type Item = (Entity, usize, &'a str, &'a TextFont, Color);
172 fn next(&mut self) -> Option<Self::Item> {
173 if let Some(root_entity) = self.root_entity.take() {
175 if let Ok((text, text_font, color, maybe_children)) = self.roots.get(root_entity) {
176 if let Some(children) = maybe_children {
177 self.stack.push((children, 0));
178 }
179 return Some((root_entity, 0, text.read_span(), text_font, color.0));
180 }
181 return None;
182 }
183
184 loop {
186 let (children, idx) = self.stack.last_mut()?;
187
188 loop {
189 let Some(child) = children.get(*idx) else {
190 break;
191 };
192
193 *idx += 1;
195
196 let entity = *child;
197 let Ok((span, text_font, color, maybe_children)) = self.spans.get(entity) else {
198 continue;
199 };
200
201 let depth = self.stack.len();
202 if let Some(children) = maybe_children {
203 self.stack.push((children, 0));
204 }
205 return Some((entity, depth, span.read_span(), text_font, color.0));
206 }
207
208 self.stack.pop();
210 }
211 }
212}
213
214impl<'a, R: TextRoot> Drop for TextSpanIter<'a, R> {
215 fn drop(&mut self) {
216 let stack = core::mem::take(&mut self.stack);
218 self.scratch.recover(stack);
219 }
220}
221
222#[derive(SystemParam)]
226pub struct TextWriter<'w, 's, R: TextRoot> {
227 scratch: ResMut<'w, TextIterScratch>,
229 roots: Query<
230 'w,
231 's,
232 (
233 &'static mut R,
234 &'static mut TextFont,
235 &'static mut TextColor,
236 ),
237 Without<TextSpan>,
238 >,
239 spans: Query<
240 'w,
241 's,
242 (
243 &'static mut TextSpan,
244 &'static mut TextFont,
245 &'static mut TextColor,
246 ),
247 Without<R>,
248 >,
249 children: Query<'w, 's, &'static Children>,
250}
251
252impl<'w, 's, R: TextRoot> TextWriter<'w, 's, R> {
253 pub fn get(
255 &mut self,
256 root_entity: Entity,
257 index: usize,
258 ) -> Option<(
259 Entity,
260 usize,
261 Mut<'_, String>,
262 Mut<'_, TextFont>,
263 Mut<'_, TextColor>,
264 )> {
265 if index == 0 {
267 let (text, font, color) = self.roots.get_mut(root_entity).ok()?;
268 return Some((
269 root_entity,
270 0,
271 text.map_unchanged(|t| t.write_span()),
272 font,
273 color,
274 ));
275 }
276
277 let mut stack: Vec<(&Children, usize)> = self.scratch.take();
279 if let Ok(children) = self.children.get(root_entity) {
280 stack.push((children, 0));
281 }
282
283 let mut count = 1;
285 let (depth, entity) = 'l: loop {
286 let Some((children, idx)) = stack.last_mut() else {
287 self.scratch.recover(stack);
288 return None;
289 };
290
291 loop {
292 let Some(child) = children.get(*idx) else {
293 stack.pop();
295 break;
296 };
297
298 *idx += 1;
300
301 if !self.spans.contains(*child) {
302 continue;
303 };
304 count += 1;
305
306 if count - 1 == index {
307 let depth = stack.len();
308 self.scratch.recover(stack);
309 break 'l (depth, *child);
310 }
311
312 if let Ok(children) = self.children.get(*child) {
313 stack.push((children, 0));
314 break;
315 }
316 }
317 };
318
319 let (text, font, color) = self.spans.get_mut(entity).unwrap();
321 Some((
322 entity,
323 depth,
324 text.map_unchanged(|t| t.write_span()),
325 font,
326 color,
327 ))
328 }
329
330 pub fn get_text(&mut self, root_entity: Entity, index: usize) -> Option<Mut<'_, String>> {
332 self.get(root_entity, index).map(|(_, _, text, ..)| text)
333 }
334
335 pub fn get_font(&mut self, root_entity: Entity, index: usize) -> Option<Mut<'_, TextFont>> {
337 self.get(root_entity, index).map(|(_, _, _, font, _)| font)
338 }
339
340 pub fn get_color(&mut self, root_entity: Entity, index: usize) -> Option<Mut<'_, TextColor>> {
342 self.get(root_entity, index)
343 .map(|(_, _, _, _, color)| color)
344 }
345
346 pub fn text(&mut self, root_entity: Entity, index: usize) -> Mut<'_, String> {
350 self.get_text(root_entity, index).unwrap()
351 }
352
353 pub fn font(&mut self, root_entity: Entity, index: usize) -> Mut<'_, TextFont> {
357 self.get_font(root_entity, index).unwrap()
358 }
359
360 pub fn color(&mut self, root_entity: Entity, index: usize) -> Mut<'_, TextColor> {
364 self.get_color(root_entity, index).unwrap()
365 }
366
367 pub fn for_each(
369 &mut self,
370 root_entity: Entity,
371 mut callback: impl FnMut(Entity, usize, Mut<String>, Mut<TextFont>, Mut<TextColor>),
372 ) {
373 self.for_each_until(root_entity, |a, b, c, d, e| {
374 (callback)(a, b, c, d, e);
375 true
376 });
377 }
378
379 pub fn for_each_text(&mut self, root_entity: Entity, mut callback: impl FnMut(Mut<String>)) {
381 self.for_each(root_entity, |_, _, text, _, _| {
382 (callback)(text);
383 });
384 }
385
386 pub fn for_each_font(&mut self, root_entity: Entity, mut callback: impl FnMut(Mut<TextFont>)) {
388 self.for_each(root_entity, |_, _, _, font, _| {
389 (callback)(font);
390 });
391 }
392
393 pub fn for_each_color(
395 &mut self,
396 root_entity: Entity,
397 mut callback: impl FnMut(Mut<TextColor>),
398 ) {
399 self.for_each(root_entity, |_, _, _, _, color| {
400 (callback)(color);
401 });
402 }
403
404 pub fn for_each_until(
409 &mut self,
410 root_entity: Entity,
411 mut callback: impl FnMut(Entity, usize, Mut<String>, Mut<TextFont>, Mut<TextColor>) -> bool,
412 ) {
413 let Ok((text, font, color)) = self.roots.get_mut(root_entity) else {
415 return;
416 };
417 if !(callback)(
418 root_entity,
419 0,
420 text.map_unchanged(|t| t.write_span()),
421 font,
422 color,
423 ) {
424 return;
425 }
426
427 let mut stack: Vec<(&Children, usize)> = self.scratch.take();
429 if let Ok(children) = self.children.get(root_entity) {
430 stack.push((children, 0));
431 }
432
433 loop {
435 let depth = stack.len();
436 let Some((children, idx)) = stack.last_mut() else {
437 self.scratch.recover(stack);
438 return;
439 };
440
441 loop {
442 let Some(child) = children.get(*idx) else {
443 stack.pop();
445 break;
446 };
447
448 *idx += 1;
450
451 let entity = *child;
452 let Ok((text, font, color)) = self.spans.get_mut(entity) else {
453 continue;
454 };
455
456 if !(callback)(
457 entity,
458 depth,
459 text.map_unchanged(|t| t.write_span()),
460 font,
461 color,
462 ) {
463 self.scratch.recover(stack);
464 return;
465 }
466
467 if let Ok(children) = self.children.get(entity) {
468 stack.push((children, 0));
469 break;
470 }
471 }
472 }
473 }
474}