bevy_text/
text_access.rs

1use bevy_color::Color;
2use bevy_ecs::{
3    component::Mutable,
4    prelude::*,
5    system::{Query, SystemParam},
6};
7
8use crate::{TextColor, TextFont, TextSpan};
9
10/// Helper trait for using the [`TextReader`] and [`TextWriter`] system params.
11pub trait TextSpanAccess: Component<Mutability = Mutable> {
12    /// Gets the text span's string.
13    fn read_span(&self) -> &str;
14    /// Gets mutable reference to the text span's string.
15    fn write_span(&mut self) -> &mut String;
16}
17
18/// Helper trait for the root text component in a text block.
19pub trait TextRoot: TextSpanAccess + From<String> {}
20
21/// Helper trait for the text span components in a text block.
22pub trait TextSpanComponent: TextSpanAccess + From<String> {}
23
24/// Scratch buffer used to store intermediate state when iterating over text spans.
25#[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/// System parameter for reading text spans in a text block.
48///
49/// `R` is the root text component.
50#[derive(SystemParam)]
51pub struct TextReader<'w, 's, R: TextRoot> {
52    // This is a local to avoid system ambiguities when TextReaders run in parallel.
53    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    /// Returns an iterator over text spans in a text block, starting with the root entity.
78    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    /// Gets a text span within a text block at a specific index in the flattened span list.
91    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    /// Gets the text value of a text span within a text block at a specific index in the flattened span list.
100    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    /// Gets the [`TextFont`] of a text span within a text block at a specific index in the flattened span list.
105    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    /// Gets the [`TextColor`] of a text span within a text block at a specific index in the flattened span list.
110    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    /// Gets the text value of a text span within a text block at a specific index in the flattened span list.
116    ///
117    /// Panics if there is no span at the requested index.
118    pub fn text(&mut self, root_entity: Entity, index: usize) -> &str {
119        self.get_text(root_entity, index).unwrap()
120    }
121
122    /// Gets the [`TextFont`] of a text span within a text block at a specific index in the flattened span list.
123    ///
124    /// Panics if there is no span at the requested index.
125    pub fn font(&mut self, root_entity: Entity, index: usize) -> &TextFont {
126        self.get_font(root_entity, index).unwrap()
127    }
128
129    /// Gets the [`TextColor`] of a text span within a text block at a specific index in the flattened span list.
130    ///
131    /// Panics if there is no span at the requested index.
132    pub fn color(&mut self, root_entity: Entity, index: usize) -> Color {
133        self.get_color(root_entity, index).unwrap()
134    }
135}
136
137/// Iterator returned by [`TextReader::iter`].
138///
139/// Iterates all spans in a text block according to hierarchy traversal order.
140/// Does *not* flatten interspersed ghost nodes. Only contiguous spans are traversed.
141// TODO: Use this iterator design in UiChildrenIter to reduce allocations.
142pub struct TextSpanIter<'a, R: TextRoot> {
143    scratch: &'a mut TextIterScratch,
144    root_entity: Option<Entity>,
145    /// Stack of (children, next index into children).
146    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    /// Item = (entity in text block, hierarchy depth in the block, span text, span style).
171    type Item = (Entity, usize, &'a str, &'a TextFont, Color);
172    fn next(&mut self) -> Option<Self::Item> {
173        // Root
174        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        // Span
185        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                // Increment to prep the next entity in this stack level.
194                *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            // All children at this stack entry have been iterated.
209            self.stack.pop();
210        }
211    }
212}
213
214impl<'a, R: TextRoot> Drop for TextSpanIter<'a, R> {
215    fn drop(&mut self) {
216        // Return the internal stack.
217        let stack = core::mem::take(&mut self.stack);
218        self.scratch.recover(stack);
219    }
220}
221
222/// System parameter for reading and writing text spans in a text block.
223///
224/// `R` is the root text component, and `S` is the text span component on children.
225#[derive(SystemParam)]
226pub struct TextWriter<'w, 's, R: TextRoot> {
227    // This is a resource because two TextWriters can't run in parallel.
228    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    /// Gets a mutable reference to a text span within a text block at a specific index in the flattened span list.
254    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        // Root
266        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        // Prep stack.
278        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        // Span
284        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                    // All children at this stack entry have been iterated.
294                    stack.pop();
295                    break;
296                };
297
298                // Increment to prep the next entity in this stack level.
299                *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        // Note: We do this outside the loop due to borrow checker limitations.
320        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    /// Gets the text value of a text span within a text block at a specific index in the flattened span list.
331    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    /// Gets the [`TextFont`] of a text span within a text block at a specific index in the flattened span list.
336    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    /// Gets the [`TextColor`] of a text span within a text block at a specific index in the flattened span list.
341    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    /// Gets the text value of a text span within a text block at a specific index in the flattened span list.
347    ///
348    /// Panics if there is no span at the requested index.
349    pub fn text(&mut self, root_entity: Entity, index: usize) -> Mut<'_, String> {
350        self.get_text(root_entity, index).unwrap()
351    }
352
353    /// Gets the [`TextFont`] of a text span within a text block at a specific index in the flattened span list.
354    ///
355    /// Panics if there is no span at the requested index.
356    pub fn font(&mut self, root_entity: Entity, index: usize) -> Mut<'_, TextFont> {
357        self.get_font(root_entity, index).unwrap()
358    }
359
360    /// Gets the [`TextColor`] of a text span within a text block at a specific index in the flattened span list.
361    ///
362    /// Panics if there is no span at the requested index.
363    pub fn color(&mut self, root_entity: Entity, index: usize) -> Mut<'_, TextColor> {
364        self.get_color(root_entity, index).unwrap()
365    }
366
367    /// Invokes a callback on each span in a text block, starting with the root entity.
368    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    /// Invokes a callback on each span's string value in a text block, starting with the root entity.
380    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    /// Invokes a callback on each span's [`TextFont`] in a text block, starting with the root entity.
387    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    /// Invokes a callback on each span's [`TextColor`] in a text block, starting with the root entity.
394    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    /// Invokes a callback on each span in a text block, starting with the root entity.
405    ///
406    /// Traversal will stop when the callback returns `false`.
407    // TODO: find a way to consolidate get and for_each_until, or provide a real iterator. Lifetime issues are challenging here.
408    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        // Root
414        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        // Prep stack.
428        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        // Span
434        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                    // All children at this stack entry have been iterated.
444                    stack.pop();
445                    break;
446                };
447
448                // Increment to prep the next entity in this stack level.
449                *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}