Add commands + get/mut methods on entity_ref/mut

BoxyUwU · BoxyUwU · commit d35d9b9185cc · 2021-04-17T00:20:52.000+01:00
diff --git a/crates/bevy_ecs/src/system/commands.rs b/crates/bevy_ecs/src/system/commands.rs
@@ -281,6 +281,24 @@ impl<'a, 'b> EntityCommands<'a, 'b> {
     pub fn commands(&mut self) -> &mut Commands<'a> {
         self.commands
     }
+
+    pub fn insert_relation<T: Component>(&mut self, relation: T, target: Entity) -> &mut Self {
+        self.commands.add(InsertRelation {
+            this: self.entity,
+            relation,
+            target,
+        });
+        self
+    }
+
+    pub fn remove_relation<T: Component>(&mut self, target: Entity) -> &mut Self {
+        self.commands.add(RemoveRelation::<T> {
+            this: self.entity,
+            target,
+            p: PhantomData,
+        });
+        self
+    }
 }
 
 #[derive(Debug)]
@@ -413,6 +431,34 @@ impl<T: Component> Command for RemoveResource<T> {
     }
 }
 
+pub struct InsertRelation<T: Component> {
+    this: Entity,
+    relation: T,
+    target: Entity,
+}
+
+impl<T: Component> Command for InsertRelation<T> {
+    fn write(self: Box<Self>, world: &mut World) {
+        world
+            .entity_mut(self.this)
+            .insert_relation(self.relation, self.target);
+    }
+}
+
+pub struct RemoveRelation<T: Component> {
+    this: Entity,
+    target: Entity,
+    p: PhantomData<T>,
+}
+
+impl<T: Component> Command for RemoveRelation<T> {
+    fn write(self: Box<Self>, world: &mut World) {
+        world
+            .entity_mut(self.this)
+            .remove_relation::<T>(self.target);
+    }
+}
+
 #[cfg(test)]
 #[allow(clippy::float_cmp, clippy::approx_constant)]
 mod tests {
@@ -510,4 +556,63 @@ mod tests {
         assert!(!world.contains_resource::<i32>());
         assert!(world.contains_resource::<f64>());
     }
+
+    #[test]
+    fn relations() {
+        struct MyRelation(bool);
+        struct MyRelationTwo(u32);
+
+        let mut world = World::default();
+        let mut queue = CommandQueue::default();
+
+        let t1 = world.spawn().id();
+        let t2 = world.spawn().id();
+        let t3 = world.spawn().id();
+
+        let mut commands = Commands::new(&mut queue, &world);
+        let e1 = commands
+            .spawn()
+            .insert_relation(MyRelation(true), t1)
+            .insert_relation(MyRelation(false), t2)
+            .insert_relation(MyRelationTwo(18), t3)
+            .id();
+        let e2 = commands
+            .spawn()
+            .insert_relation(MyRelation(true), t3)
+            .insert_relation(MyRelationTwo(10), t1)
+            .insert_relation(MyRelation(false), t1)
+            .insert_relation(MyRelation(false), t3)
+            .id();
+        drop(commands);
+        queue.apply(&mut world);
+
+        let e1 = world.entity(e1);
+        assert!(e1.get_relation::<MyRelation>(Some(t1)).unwrap().0 == true);
+        assert!(e1.get_relation::<MyRelation>(Some(t2)).unwrap().0 == false);
+        assert!(e1.get_relation::<MyRelationTwo>(Some(t3)).unwrap().0 == 18);
+        let e1 = e1.id();
+
+        let e2 = world.entity(e2);
+        assert!(e2.get_relation::<MyRelation>(Some(t3)).unwrap().0 == false);
+        assert!(e2.get_relation::<MyRelationTwo>(Some(t1)).unwrap().0 == 10);
+        assert!(e2.get_relation::<MyRelation>(Some(t1)).unwrap().0 == false);
+        let e2 = e2.id();
+
+        let mut commands = Commands::new(&mut queue, &world);
+
+        commands.entity(e1).remove_relation::<MyRelation>(t2);
+        commands.entity(e2).remove_relation::<MyRelation>(t3);
+
+        drop(commands);
+        queue.apply(&mut world);
+
+        assert!(world
+            .entity(e1)
+            .get_relation::<MyRelation>(Some(t2))
+            .is_none());
+        assert!(world
+            .entity(e2)
+            .get_relation::<MyRelation>(Some(t3))
+            .is_none());
+    }
 }
diff --git a/crates/bevy_ecs/src/world/entity_ref.rs b/crates/bevy_ecs/src/world/entity_ref.rs
@@ -75,10 +75,21 @@ impl<'w> EntityRef<'w> {
 
     #[inline]
     pub fn get<T: Component>(&self) -> Option<&'w T> {
+        self.get_relation::<T>(None)
+    }
+
+    #[inline]
+    pub fn get_relation<T: Component>(&self, target: Option<Entity>) -> Option<&'w T> {
         // SAFE: entity location is valid and returned component is of type T
         unsafe {
-            get_component_with_type(self.world, TypeId::of::<T>(), self.entity, self.location)
-                .map(|value| &*value.cast::<T>())
+            get_component_with_type(
+                self.world,
+                TypeId::of::<T>(),
+                target,
+                self.entity,
+                self.location,
+            )
+            .map(|value| &*value.cast::<T>())
         }
     }
 
@@ -91,13 +102,33 @@ impl<'w> EntityRef<'w> {
         last_change_tick: u32,
         change_tick: u32,
     ) -> Option<Mut<'w, T>> {
-        get_component_and_ticks_with_type(self.world, TypeId::of::<T>(), self.entity, self.location)
-            .map(|(value, ticks)| Mut {
-                value: &mut *value.cast::<T>(),
-                component_ticks: &mut *ticks,
-                last_change_tick,
-                change_tick,
-            })
+        // SAFETY: Caller
+        self.get_relation_unchecked_mut(None, last_change_tick, change_tick)
+    }
+
+    /// # Safety
+    /// This allows aliased mutability. You must make sure this call does not result in multiple
+    /// mutable references to the same component
+    #[inline]
+    pub unsafe fn get_relation_unchecked_mut<T: Component>(
+        &self,
+        target: Option<Entity>,
+        last_change_tick: u32,
+        change_tick: u32,
+    ) -> Option<Mut<'w, T>> {
+        get_component_and_ticks_with_type(
+            self.world,
+            TypeId::of::<T>(),
+            target,
+            self.entity,
+            self.location,
+        )
+        .map(|(value, ticks)| Mut {
+            value: &mut *value.cast::<T>(),
+            component_ticks: &mut *ticks,
+            last_change_tick,
+            change_tick,
+        })
     }
 }
 
@@ -165,21 +196,38 @@ impl<'w> EntityMut<'w> {
 
     #[inline]
     pub fn get<T: Component>(&self) -> Option<&'w T> {
+        self.get_relation::<T>(None)
+    }
+
+    #[inline]
+    pub fn get_relation<T: Component>(&self, target: Option<Entity>) -> Option<&'w T> {
         // SAFE: entity location is valid and returned component is of type T
         unsafe {
-            get_component_with_type(self.world, TypeId::of::<T>(), self.entity, self.location)
-                .map(|value| &*value.cast::<T>())
+            get_component_with_type(
+                self.world,
+                TypeId::of::<T>(),
+                target,
+                self.entity,
+                self.location,
+            )
+            .map(|value| &*value.cast::<T>())
         }
     }
 
     #[inline]
     pub fn get_mut<T: Component>(&mut self) -> Option<Mut<'w, T>> {
+        self.get_relation_mut::<T>(None)
+    }
+
+    #[inline]
+    pub fn get_relation_mut<T: Component>(&mut self, target: Option<Entity>) -> Option<Mut<'w, T>> {
         // SAFE: world access is unique, entity location is valid, and returned component is of type
         // T
         unsafe {
             get_component_and_ticks_with_type(
                 self.world,
                 TypeId::of::<T>(),
+                target,
                 self.entity,
                 self.location,
             )
@@ -197,13 +245,30 @@ impl<'w> EntityMut<'w> {
     /// mutable references to the same component
     #[inline]
     pub unsafe fn get_unchecked_mut<T: Component>(&self) -> Option<Mut<'w, T>> {
-        get_component_and_ticks_with_type(self.world, TypeId::of::<T>(), self.entity, self.location)
-            .map(|(value, ticks)| Mut {
-                value: &mut *value.cast::<T>(),
-                component_ticks: &mut *ticks,
-                last_change_tick: self.world.last_change_tick(),
-                change_tick: self.world.read_change_tick(),
-            })
+        self.get_relation_unchecked_mut::<T>(None)
+    }
+
+    /// # Safety
+    /// This allows aliased mutability. You must make sure this call does not result in multiple
+    /// mutable references to the same component
+    #[inline]
+    pub unsafe fn get_relation_unchecked_mut<T: Component>(
+        &self,
+        target: Option<Entity>,
+    ) -> Option<Mut<'w, T>> {
+        get_component_and_ticks_with_type(
+            self.world,
+            TypeId::of::<T>(),
+            target,
+            self.entity,
+            self.location,
+        )
+        .map(|(value, ticks)| Mut {
+            value: &mut *value.cast::<T>(),
+            component_ticks: &mut *ticks,
+            last_change_tick: self.world.last_change_tick(),
+            change_tick: self.world.read_change_tick(),
+        })
     }
 
     // TODO: factor out non-generic part to cut down on monomorphization (just check perf)
@@ -796,28 +861,32 @@ unsafe fn remove_component(
     }
 }
 
+/// Set `target` to None if you just want normal components
 /// # Safety
 /// `entity_location` must be within bounds of an archetype that exists.
 unsafe fn get_component_with_type(
     world: &World,
     type_id: TypeId,
+    target: Option<Entity>,
     entity: Entity,
     location: EntityLocation,
 ) -> Option<*mut u8> {
     let kind = world.components.get_component_kind(type_id)?;
-    get_component(world, kind.id(), None, entity, location)
+    get_component(world, kind.id(), target, entity, location)
 }
 
+/// Set `target` to None if you just want normal components
 /// # Safety
 /// `entity_location` must be within bounds of an archetype that exists.
 pub(crate) unsafe fn get_component_and_ticks_with_type(
     world: &World,
     type_id: TypeId,
+    target: Option<Entity>,
     entity: Entity,
     location: EntityLocation,
 ) -> Option<(*mut u8, *mut ComponentTicks)> {
     let kind_info = world.components.get_component_kind(type_id)?;
-    get_component_and_ticks(world, kind_info.id(), None, entity, location)
+    get_component_and_ticks(world, kind_info.id(), target, entity, location)
 }
 
 fn contains_component_with_type(world: &World, type_id: TypeId, location: EntityLocation) -> bool {
diff --git a/crates/bevy_ecs/src/world/tests.rs b/crates/bevy_ecs/src/world/tests.rs
@@ -351,6 +351,34 @@ fn relation_query_mut_raw(storage_type: StorageType) {
     assert!(was_targeter1 && was_targeter2 && was_targeter3);
 }
 
+#[test]
+fn some_example_code() {
+    #[derive(PartialEq, Eq, Debug)]
+    struct MyRelation;
+
+    let mut world = World::new();
+
+    let target1 = world.spawn().id();
+    let target2 = world.spawn().id();
+    let my_entity = world
+        .spawn()
+        .insert_relation(MyRelation, target1)
+        .insert_relation(MyRelation, target2)
+        .id();
+
+    let mut iterated_entities = Vec::new();
+    let mut query = world.query::<(Entity, &Relation<MyRelation>)>();
+    for (entity, relations) in query.iter_mut(&mut world) {
+        iterated_entities.push(entity);
+        assert_eq!(
+            &relations.collect::<Vec<_>>(),
+            &[(Some(target1), &MyRelation), (Some(target2), &MyRelation)],
+        );
+    }
+
+    assert_eq!(&iterated_entities, &[my_entity]);
+}
+
 macro_rules! query_conflict_tests {
     ($($name:ident => <$param:ty>)*) => {
         $(
