Add commands + get/mut methods on entity_ref/mut

crates/bevy_ecs/src/system/commands.rs

@@ -276,6 +276,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)]
@@ -408,6 +426,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 {
@@ -505,4 +551,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());
+ }
 }

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 {

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>)*) => {
  $(