Test entity labels, fixed corner cases, changed interface

crates/bevy_core/src/label.rs

@@ -47,6 +47,10 @@ impl Labels {
         self.labels.insert(label.into());
     }
 
+    pub fn remove<T: Into<Cow<'static, str>>>(&mut self, label: T) {
+        self.labels.remove(&label.into());
+    }
+
     pub fn iter(&self) -> impl Iterator<Item = &str> {
         self.labels.iter().map(|label| label.deref())
     }
@@ -60,37 +64,182 @@ 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<EntityLabels>,
-    // TODO: use change tracking when add/remove events are added
-    // mut query: Query<(Entity, Changed<Labels>)>,
+    // the system runs in an early stage and so can't use a Changed<Labels> filter
     query: Query<(Entity, &Labels)>,
 ) {
     let entity_labels = entity_labels.deref_mut();
+
+    for entity in query.removed::<Labels>() {
+        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);
             }
         }
 
         for added_label in labels.labels.difference(&current_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::<EntityLabels>().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::<Labels>(e1).unwrap().insert("shalau");
+        schedule.initialize_and_run(&mut world, &mut resources);
+
+        let entity_labels = resources.get::<EntityLabels>().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::<Labels>(e1).unwrap().remove("holy");
+        schedule.initialize_and_run(&mut world, &mut resources);
+
+        let entity_labels = resources.get::<EntityLabels>().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::<EntityLabels>().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::<Labels>(e1).unwrap();
+        schedule.initialize_and_run(&mut world, &mut resources);
+
+        let entity_labels = resources.get::<EntityLabels>().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::<EntityLabels>().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::<EntityLabels>().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");
+    }
+}

crates/bevy_core/src/lib.rs

@@ -36,6 +36,7 @@ impl Plugin for CorePlugin {
             .init_resource::<EntityLabels>()
             .init_resource::<FixedTimesteps>()
             .register_type::<Option<String>>()
+            .register_type::<Labels>()
             .register_type::<Range<f32>>()
             .register_type::<Timer>()
             .add_system_to_stage(stage::FIRST, time_system.system())