bevy_yoleck/
populating.rs

1use std::ops::RangeFrom;
2
3use bevy::ecs::query::{QueryData, QueryFilter};
4use bevy::ecs::system::{EntityCommands, SystemParam};
5use bevy::platform::collections::HashMap;
6use bevy::prelude::*;
7
8use crate::entity_management::EntitiesToPopulate;
9
10/// Wrapper for writing queries in populate systems.
11#[derive(SystemParam)]
12pub struct YoleckPopulate<'w, 's, Q: 'static + QueryData, F: 'static + QueryFilter = ()> {
13    entities_to_populate: Res<'w, EntitiesToPopulate>,
14    query: Query<'w, 's, Q, F>,
15    commands: Commands<'w, 's>,
16}
17
18impl<Q: 'static + QueryData, F: 'static + QueryFilter> YoleckPopulate<'_, '_, Q, F> {
19    /// Iterate over the entities that need populating in order to add/update components using
20    /// a Bevy command.
21    pub fn populate(
22        &mut self,
23        mut dlg: impl FnMut(YoleckPopulateContext, EntityCommands, <Q as QueryData>::Item<'_>),
24    ) {
25        for (entity, populate_reason) in self.entities_to_populate.0.iter() {
26            if let Ok(data) = self.query.get_mut(*entity) {
27                let cmd = self.commands.entity(*entity);
28                let context = YoleckPopulateContext {
29                    reason: *populate_reason,
30                };
31                dlg(context, cmd, data);
32            }
33        }
34    }
35}
36
37#[derive(Clone, Copy)]
38pub(crate) enum PopulateReason {
39    EditorInit,
40    EditorUpdate,
41    RealGame,
42}
43
44/// A context for [`YoleckPopulate::populate`].
45pub struct YoleckPopulateContext {
46    pub(crate) reason: PopulateReason,
47}
48
49impl YoleckPopulateContext {
50    /// `true` if the entity is created in editor mode, `false` if created in playtest or actual game.
51    pub fn is_in_editor(&self) -> bool {
52        match self.reason {
53            PopulateReason::EditorInit => true,
54            PopulateReason::EditorUpdate => true,
55            PopulateReason::RealGame => false,
56        }
57    }
58
59    /// `true` if this is this is the first time the entity is populated, `false` if the entity was
60    /// popultated before.
61    pub fn is_first_time(&self) -> bool {
62        match self.reason {
63            PopulateReason::EditorInit => true,
64            PopulateReason::EditorUpdate => false,
65            PopulateReason::RealGame => true,
66        }
67    }
68}
69
70/// See [`YoleckMarking`].
71#[derive(Debug, Component, PartialEq, Eq, Clone, Copy)]
72pub struct YoleckSystemMarker(usize);
73
74#[derive(Resource)]
75struct MarkerGenerator(RangeFrom<usize>);
76
77impl FromWorld for YoleckSystemMarker {
78    fn from_world(world: &mut World) -> Self {
79        let mut marker = world.get_resource_or_insert_with(|| MarkerGenerator(1..));
80        YoleckSystemMarker(marker.0.next().unwrap())
81    }
82}
83
84/// Use to mark child entities created from a specific system.
85///
86/// Using `despawn_descendants` is dangerous because it can delete entities created by other
87/// populate systems in the same frame. Instead, a populate system that wants to despawn its own
88/// child entities should do something like this:
89///
90/// ```no_run
91/// # use bevy::prelude::*;
92/// # use bevy_yoleck::prelude::*;
93/// # use serde::{Deserialize, Serialize};
94/// # #[derive(Default, Clone, PartialEq, Serialize, Deserialize, Component, YoleckComponent)]
95/// # struct MyComponent;
96/// fn populate_system(mut populate: YoleckPopulate<&MyComponent>, marking: YoleckMarking) {
97///     populate.populate(|_ctx, mut cmd, my_component| {
98///         marking.despawn_marked(&mut cmd);
99///         cmd.with_children(|commands| {
100///             let mut child = commands.spawn(marking.marker());
101///             child.insert((
102///                 // relevant Bevy components
103///             ));
104///         });
105///     });
106/// }
107/// ```
108#[derive(SystemParam)]
109pub struct YoleckMarking<'w, 's> {
110    designated_marker: Local<'s, YoleckSystemMarker>,
111    children_query: Query<'w, 's, &'static Children>,
112    marked_query: Query<'w, 's, (&'static ChildOf, &'static YoleckSystemMarker)>,
113}
114
115impl YoleckMarking<'_, '_> {
116    /// Create a marker unique to this system.
117    pub fn marker(&self) -> YoleckSystemMarker {
118        *self.designated_marker
119    }
120
121    /// Despawn all the entities marked by the current system, that are descendants of the entity
122    /// edited by the supplied `cmd`.
123    pub fn despawn_marked(&self, cmd: &mut EntityCommands) {
124        let mut marked_children_map: HashMap<Entity, Vec<Entity>> = Default::default();
125        for child in self.children_query.iter_descendants(cmd.id()) {
126            let Ok((child_of, marker)) = self.marked_query.get(child) else {
127                continue;
128            };
129            if *marker == *self.designated_marker {
130                marked_children_map
131                    .entry(child_of.parent())
132                    .or_default()
133                    .push(child);
134            }
135        }
136        for (parent, children) in marked_children_map {
137            cmd.commands().entity(parent).remove_children(&children);
138            for child in children {
139                cmd.commands().entity(child).despawn();
140            }
141        }
142    }
143}