bevy_easy_stats/
commands.rs

1use std::marker::PhantomData;
2
3use bevy::{
4    ecs::system::{EntityCommand, EntityCommands},
5    prelude::{Commands, Component, Entity, World},
6};
7
8use crate::{stat_modification::ModificationType, StatData, StatIdentifier, Stats};
9
10/// Make changes to an entities stats in a deferred patter using commands.
11pub struct ModifyStatEntityCommands<
12    'a,
13    StatCollection: AsMut<Stats> + Send + Sync + 'static + Component,
14> {
15    target_entity: Entity,
16    target_component: PhantomData<StatCollection>,
17    commands: Commands<'a, 'a>,
18}
19
20impl<'a, StatCollection: AsMut<Stats> + Send + Sync + 'static + Component>
21    ModifyStatEntityCommands<'a, StatCollection>
22{
23    /// Return the entity id that is targeted for stat changes
24    pub fn id(&self) -> Entity {
25        self.target_entity
26    }
27
28    /// Change the target entity
29    pub fn entity(&mut self, entity: Entity) {
30        self.target_entity = entity
31    }
32
33    /// Return a reference to the commands object contained
34    pub fn commands(&mut self) -> &mut Commands<'a, 'a> {
35        &mut self.commands
36    }
37}
38
39impl<StatCollection: AsMut<Stats> + Send + Sync + 'static + Component>
40    ModifyStatEntityCommands<'_, StatCollection>
41{
42    /// Get entity commands for the targeted entity
43    pub fn entity_commands(&mut self) -> EntityCommands<'_> {
44        let id = self.id();
45        self.commands().entity(id)
46    }
47
48    /// Queue a command to perform an add with the given [`StatData`] to the targeted [`StatIdentifier`]
49    pub fn add(
50        &mut self,
51        stat_id: impl StatIdentifier + 'static + Send + Sync,
52        stat_data: impl StatData,
53    ) -> &mut Self {
54        self.entity_commands()
55            .queue(modify_entity_stat::<StatCollection>(
56                stat_id,
57                ModificationType::add(stat_data),
58            ));
59        self
60    }
61
62    /// Queue a command to perform a sub with the given [`StatData`] to the targeted [`StatIdentifier`]
63    pub fn sub(
64        &mut self,
65        stat_id: impl StatIdentifier + 'static + Send + Sync,
66        stat_data: impl StatData,
67    ) -> &mut Self {
68        self.entity_commands()
69            .queue(modify_entity_stat::<StatCollection>(
70                stat_id,
71                ModificationType::sub(stat_data),
72            ));
73        self
74    }
75
76    /// Queue a command to perform a set with the given [`StatData`] to the targeted [`StatIdentifier`]
77    pub fn set(
78        &mut self,
79        stat_id: impl StatIdentifier + 'static + Send + Sync,
80        stat_data: impl StatData,
81    ) -> &mut Self {
82        self.entity_commands()
83            .queue(modify_entity_stat::<StatCollection>(
84                stat_id,
85                ModificationType::set(stat_data),
86            ));
87        self
88    }
89
90    /// Queue a command to perform a remove to the targeted [`StatIdentifier`]
91    pub fn remove(&mut self, stat_id: impl StatIdentifier + 'static + Send + Sync) -> &mut Self {
92        self.entity_commands()
93            .queue(modify_entity_stat::<StatCollection>(
94                stat_id,
95                ModificationType::remove(),
96            ));
97        self
98    }
99
100    /// Queue a command to perform a reset to the targeted [`StatIdentifier`]
101    pub fn reset(&mut self, stat_id: impl StatIdentifier + 'static + Send + Sync) -> &mut Self {
102        self.entity_commands()
103            .queue(modify_entity_stat::<StatCollection>(
104                stat_id,
105                ModificationType::reset(),
106            ));
107        self
108    }
109}
110
111pub trait StatCommandsExt {
112    /// Returns a [`ModifyStatEntityCommands`] object for the given entity
113    fn entity_stats<StatCollection: AsMut<Stats> + Send + Sync + 'static + Component>(
114        &mut self,
115        entity: Entity,
116    ) -> ModifyStatEntityCommands<'_, StatCollection>;
117
118    /// Modify a single stat on an entity
119    fn modify_stat<StatCollection: AsMut<Stats> + Send + Sync + 'static + Component>(
120        &mut self,
121        entity: Entity,
122        stat_id: impl StatIdentifier + 'static + Send + Sync,
123        modification_type: ModificationType,
124    );
125}
126
127impl<'a> StatCommandsExt for Commands<'a, 'a> {
128    /// Get access to an object to modify entity stats
129    fn entity_stats<StatCollection: AsMut<Stats> + Send + Sync + 'static + Component>(
130        &mut self,
131        entity: Entity,
132    ) -> ModifyStatEntityCommands<'_, StatCollection> {
133        ModifyStatEntityCommands {
134            target_entity: entity,
135            target_component: PhantomData,
136            commands: self.reborrow(),
137        }
138    }
139
140    /// Modify a single stat on an entity
141    fn modify_stat<StatCollection: AsMut<Stats> + Send + Sync + 'static + Component>(
142        &mut self,
143        entity: Entity,
144        stat_id: impl StatIdentifier + 'static + Send + Sync,
145        modification_type: ModificationType,
146    ) {
147        self.entity(entity)
148            .modify_stat::<StatCollection>(stat_id, modification_type);
149    }
150}
151
152pub trait StatEntityCommandsExt {
153    /// Get a special command pattern object to affect changes to an entity
154    fn entity_stats<StatCollection: AsMut<Stats> + Send + Sync + 'static + Component>(
155        &mut self,
156    ) -> ModifyStatEntityCommands<'_, StatCollection>;
157
158    /// Modify a single stat on an entity
159    fn modify_stat<StatCollection: AsMut<Stats> + Send + Sync + 'static + Component>(
160        &mut self,
161        stat_id: impl StatIdentifier + 'static + Send + Sync,
162        modification_type: ModificationType,
163    );
164}
165
166impl<'a> StatEntityCommandsExt for EntityCommands<'a> {
167    fn modify_stat<StatCollection: AsMut<Stats> + Send + Sync + 'static + Component>(
168        &mut self,
169        stat_id: impl StatIdentifier + 'static + Send + Sync,
170        modification_type: ModificationType,
171    ) {
172        self.queue(modify_entity_stat::<StatCollection>(
173            stat_id,
174            modification_type,
175        ));
176    }
177
178    fn entity_stats<StatCollection: AsMut<Stats> + Send + Sync + 'static + Component>(
179        &mut self,
180    ) -> ModifyStatEntityCommands<'_, StatCollection> {
181        ModifyStatEntityCommands {
182            target_entity: self.id(),
183            target_component: PhantomData,
184            commands: self.commands(),
185        }
186    }
187}
188
189fn modify_entity_stat<StatCollection: AsMut<Stats> + Send + Sync + 'static + Component>(
190    stat_id: impl StatIdentifier + 'static + Send + Sync,
191    modification_type: ModificationType,
192) -> impl EntityCommand {
193    move |entity: Entity, world: &mut World| {
194        if let Ok(mut entity_mut) = world.get_entity_mut(entity) {
195            if let Some(mut stat_collection) = entity_mut.get_mut::<StatCollection>() {
196                let stats = stat_collection.as_mut().as_mut();
197
198                match modification_type {
199                    ModificationType::Add(data) => {
200                        stats.add_to_stat_manual(stat_id.identifier(), data)
201                    }
202                    ModificationType::Sub(data) => {
203                        stats.sub_from_stat_manual(stat_id.identifier(), data)
204                    }
205                    ModificationType::Remove => stats.remove_stat_manual(stat_id.identifier()),
206                    ModificationType::Set(data) => {
207                        stats.set_stat_manual(stat_id.identifier(), data)
208                    }
209                    ModificationType::Reset => stats.reset_stat_manual(stat_id.identifier()),
210                }
211            }
212        }
213    }
214}
215
216#[cfg(test)]
217mod tests {
218    use super::*;
219
220    #[derive(Hash)]
221    pub struct EnemiesKilled;
222
223    impl StatIdentifier for EnemiesKilled {
224        fn identifier(&self) -> &'static str {
225            "Enemies Killed"
226        }
227    }
228
229    #[derive(Component)]
230    pub struct EntityStats {
231        stats: Stats,
232    }
233
234    impl AsMut<Stats> for EntityStats {
235        fn as_mut(&mut self) -> &mut Stats {
236            &mut self.stats
237        }
238    }
239
240    #[test]
241    fn entity_commands() {
242        let mut world = World::new();
243        let entity = world
244            .spawn(EntityStats {
245                stats: Stats::new(),
246            })
247            .id();
248
249        let mut commands = world.commands();
250
251        let mut stats = commands.entity_stats::<EntityStats>(entity);
252        stats.add(EnemiesKilled, 5u64);
253        stats.add(EnemiesKilled, 10u64);
254        world.flush();
255
256        assert_eq!(
257            *world
258                .entity(entity)
259                .get::<EntityStats>()
260                .unwrap()
261                .stats
262                .get_stat_downcast::<u64>(&EnemiesKilled)
263                .unwrap(),
264            15u64
265        );
266
267        let mut commands = world.commands();
268        let mut stats = commands.entity_stats::<EntityStats>(entity);
269        stats.reset(EnemiesKilled);
270        stats.add(EnemiesKilled, 15u64);
271        stats.sub(EnemiesKilled, 5u64);
272        stats.sub(EnemiesKilled, 7u64);
273        world.flush();
274
275        assert_eq!(
276            *world
277                .entity(entity)
278                .get::<EntityStats>()
279                .unwrap()
280                .stats
281                .get_stat_downcast::<u64>(&EnemiesKilled)
282                .unwrap(),
283            3u64
284        );
285
286        let mut commands = world.commands();
287        let mut stats = commands.entity_stats::<EntityStats>(entity);
288        stats.set(EnemiesKilled, 7u64);
289        world.flush();
290
291        assert_eq!(
292            *world
293                .entity(entity)
294                .get::<EntityStats>()
295                .unwrap()
296                .stats
297                .get_stat_downcast::<u64>(&EnemiesKilled)
298                .unwrap(),
299            7u64
300        );
301
302        let mut commands = world.commands();
303        let mut stats = commands.entity_stats::<EntityStats>(entity);
304        stats.remove(EnemiesKilled);
305        world.flush();
306
307        assert_eq!(
308            world
309                .entity(entity)
310                .get::<EntityStats>()
311                .unwrap()
312                .stats
313                .get_stat_downcast::<u64>(&EnemiesKilled),
314            None
315        );
316    }
317}