diff --git a/crates/bevy_core/src/label.rs b/crates/bevy_core/src/label.rs index 26309477e7843f..93f7f8ec8e5339 100644 --- a/crates/bevy_core/src/label.rs +++ b/crates/bevy_core/src/label.rs @@ -47,6 +47,10 @@ impl Labels { self.labels.insert(label.into()); } + pub fn remove>>(&mut self, label: T) { + self.labels.remove(&label.into()); + } + pub fn iter(&self) -> impl Iterator { self.labels.iter().map(|label| label.deref()) } @@ -60,25 +64,36 @@ pub struct EntityLabels { } impl EntityLabels { - pub fn get(&self, label: &str) -> Option<&[Entity]> { + pub fn get(&self, label: &str) -> &[Entity] { self.label_entities .get(label) .map(|entities| entities.as_slice()) + .unwrap_or(&[]) } } pub(crate) fn entity_labels_system( mut entity_labels: ResMut, - // TODO: use change tracking when add/remove events are added - // mut query: Query<(Entity, Changed)>, - query: Query<(Entity, &Labels)>, + query: Query<(Entity, &Labels), Changed>, ) { let entity_labels = entity_labels.deref_mut(); + + for entity in query.removed::() { + if let Some(labels) = entity_labels.entity_labels.get(entity) { + for label in labels.iter() { + if let Some(entities) = entity_labels.label_entities.get_mut(label) { + entities.retain(|e| e != entity); + } + } + } + } + for (entity, labels) in query.iter() { let current_labels = entity_labels .entity_labels .entry(entity) .or_insert_with(HashSet::default); + for removed_label in current_labels.difference(&labels.labels) { if let Some(entities) = entity_labels.label_entities.get_mut(removed_label) { entities.retain(|e| *e != entity); @@ -86,11 +101,144 @@ pub(crate) fn entity_labels_system( } for added_label in labels.labels.difference(¤t_labels) { - if let Some(entities) = entity_labels.label_entities.get_mut(added_label) { - entities.push(entity); - } + entity_labels + .label_entities + .entry(added_label.clone()) + .or_insert_with(Vec::new) + .push(entity); } *current_labels = labels.labels.clone(); } } + +#[cfg(test)] +mod tests { + use super::*; + + fn setup() -> (World, Resources, bevy_ecs::Schedule) { + let world = World::new(); + let mut resources = Resources::default(); + resources.insert(EntityLabels::default()); + let mut schedule = bevy_ecs::Schedule::default(); + schedule.add_stage("test", SystemStage::serial()); + schedule.add_system_to_stage("test", entity_labels_system.system()); + (world, resources, schedule) + } + + fn holy_cow() -> Labels { + Labels::from(["holy", "cow"].iter().cloned()) + } + + fn holy_shamoni() -> Labels { + Labels::from(["holy", "shamoni"].iter().cloned()) + } + + #[test] + fn adds_spawned_entity() { + let (mut world, mut resources, mut schedule) = setup(); + + let e1 = world.spawn((holy_cow(),)); + schedule.initialize_and_run(&mut world, &mut resources); + + let entity_labels = resources.get::().unwrap(); + assert_eq!(entity_labels.get("holy"), &[e1], "holy"); + assert_eq!(entity_labels.get("cow"), &[e1], "cow"); + assert_eq!(entity_labels.get("shalau"), &[], "shalau"); + } + + #[test] + fn add_labels() { + let (mut world, mut resources, mut schedule) = setup(); + let e1 = world.spawn((holy_cow(),)); + schedule.initialize_and_run(&mut world, &mut resources); + + world.get_mut::(e1).unwrap().insert("shalau"); + schedule.initialize_and_run(&mut world, &mut resources); + + let entity_labels = resources.get::().unwrap(); + assert_eq!(entity_labels.get("holy"), &[e1], "holy"); + assert_eq!(entity_labels.get("cow"), &[e1], "cow"); + assert_eq!(entity_labels.get("shalau"), &[e1], "shalau"); + } + + #[test] + fn remove_labels() { + let (mut world, mut resources, mut schedule) = setup(); + let e1 = world.spawn((holy_cow(),)); + schedule.initialize_and_run(&mut world, &mut resources); + + world.get_mut::(e1).unwrap().remove("holy"); + schedule.initialize_and_run(&mut world, &mut resources); + + let entity_labels = resources.get::().unwrap(); + assert_eq!(entity_labels.get("holy"), &[], "holy"); + assert_eq!(entity_labels.get("cow"), &[e1], "cow"); + assert_eq!(entity_labels.get("shalau"), &[], "shalau"); + } + + #[test] + fn removes_despawned_entity() { + let (mut world, mut resources, mut schedule) = setup(); + let e1 = world.spawn((holy_cow(),)); + schedule.initialize_and_run(&mut world, &mut resources); + + world.despawn(e1).unwrap(); + schedule.initialize_and_run(&mut world, &mut resources); + + let entity_labels = resources.get::().unwrap(); + assert_eq!(entity_labels.get("holy"), &[], "holy"); + assert_eq!(entity_labels.get("cow"), &[], "cow"); + assert_eq!(entity_labels.get("shalau"), &[], "shalau"); + } + + #[test] + fn removes_labels_when_component_removed() { + let (mut world, mut resources, mut schedule) = setup(); + let e1 = world.spawn((holy_cow(),)); + schedule.initialize_and_run(&mut world, &mut resources); + + world.remove_one::(e1).unwrap(); + schedule.initialize_and_run(&mut world, &mut resources); + + let entity_labels = resources.get::().unwrap(); + assert_eq!(entity_labels.get("holy"), &[], "holy"); + assert_eq!(entity_labels.get("cow"), &[], "cow"); + assert_eq!(entity_labels.get("shalau"), &[], "shalau"); + } + + #[test] + fn adds_another_spawned_entity() { + let (mut world, mut resources, mut schedule) = setup(); + let e1 = world.spawn((holy_cow(),)); + schedule.initialize_and_run(&mut world, &mut resources); + + let e2 = world.spawn((holy_shamoni(),)); + schedule.initialize_and_run(&mut world, &mut resources); + + let entity_labels = resources.get::().unwrap(); + assert_eq!(entity_labels.get("holy"), &[e1, e2], "holy"); + assert_eq!(entity_labels.get("cow"), &[e1], "cow"); + assert_eq!(entity_labels.get("shamoni"), &[e2], "shamoni"); + assert_eq!(entity_labels.get("shalau"), &[], "shalau"); + } + + #[test] + fn removes_despawned_entity_but_leaves_other() { + let (mut world, mut resources, mut schedule) = setup(); + let e1 = world.spawn((holy_cow(),)); + schedule.initialize_and_run(&mut world, &mut resources); + + let e2 = world.spawn((holy_shamoni(),)); + schedule.initialize_and_run(&mut world, &mut resources); + + world.despawn(e1).unwrap(); + schedule.initialize_and_run(&mut world, &mut resources); + + let entity_labels = resources.get::().unwrap(); + assert_eq!(entity_labels.get("holy"), &[e2], "holy"); + assert_eq!(entity_labels.get("cow"), &[], "cow"); + assert_eq!(entity_labels.get("shamoni"), &[e2], "shamoni"); + assert_eq!(entity_labels.get("shalau"), &[], "shalau"); + } +}