Fold component getters into a single method generic over references

Addresses the common footgun of calling get::<&T> for symmetry with queries
Ralith · Jul 11, 2022 · 799d409 · 799d409
1 parent 3c01a4e
commit 799d409
Show file tree

Hide file tree

Showing 12 changed files with 153 additions and 119 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,6 +1,9 @@
 # Unreleased
 
 ### Changed
+- Many generic methods that previously took a `Component` now instead take either a
+  `ComponentRef<'a>` or a `Query` to improve consistency with queries and address a common footgun:
+  - `World::get` and `EntityRef::get` now take shared or unique references to component types
 - `SerializeContext` traits now take their serializer arguments by value, and must call `end()`
   themselves.
 

diff --git a/README.md b/README.md
@@ -21,8 +21,8 @@ for (id, (number, &flag)) in world.query_mut::<(&mut i32, &bool)>() {
   if flag { *number *= 2; }
 }
 // Random access is simple and safe
-assert_eq!(*world.get::<i32>(a).unwrap(), 246);
-assert_eq!(*world.get::<i32>(b).unwrap(), 42);
+assert_eq!(*world.get::<&i32>(a).unwrap(), 246);
+assert_eq!(*world.get::<&i32>(b).unwrap(), 42);
 ```
 
 ### Why ECS?

diff --git a/examples/ffa_simulation.rs b/examples/ffa_simulation.rs
@@ -102,7 +102,7 @@ fn system_fire_at_closest(world: &mut World) {
         */
 
         // Or like this:
-        let mut hp1 = world.get_mut::<Health>(closest).unwrap();
+        let mut hp1 = world.get::<&mut Health>(closest).unwrap();
 
         // Is target unit still alive?
         if hp1.0 > 0 {

diff --git a/examples/format.rs b/examples/format.rs
@@ -3,7 +3,7 @@
 
 fn format_entity(entity: hecs::EntityRef<'_>) -> String {
     fn fmt<T: hecs::Component + std::fmt::Display>(entity: hecs::EntityRef<'_>) -> Option<String> {
-        Some(entity.get::<T>()?.to_string())
+        Some(entity.get::<&T>()?.to_string())
     }
 
     const FUNCTIONS: &[&dyn Fn(hecs::EntityRef<'_>) -> Option<String>] =

diff --git a/src/command_buffer.rs b/src/command_buffer.rs
@@ -26,7 +26,7 @@ use crate::{Bundle, Entity};
 /// let mut cmd = CommandBuffer::new();
 /// cmd.insert(entity, (true, 42));
 /// cmd.run_on(&mut world); // cmd can now be reused
-/// assert_eq!(*world.get::<i32>(entity).unwrap(), 42);
+/// assert_eq!(*world.get::<&i32>(entity).unwrap(), 42);
 /// ```
 pub struct CommandBuffer {
     entities: Vec<EntityIndex>,

diff --git a/src/entity_builder.rs b/src/entity_builder.rs
@@ -26,8 +26,8 @@ use crate::{align, Component, DynamicBundle};
 /// let mut builder = EntityBuilder::new();
 /// builder.add(123).add("abc");
 /// let e = world.spawn(builder.build()); // builder can now be reused
-/// assert_eq!(*world.get::<i32>(e).unwrap(), 123);
-/// assert_eq!(*world.get::<&str>(e).unwrap(), "abc");
+/// assert_eq!(*world.get::<&i32>(e).unwrap(), 123);
+/// assert_eq!(*world.get::<&&str>(e).unwrap(), "abc");
 /// ```
 #[derive(Default)]
 pub struct EntityBuilder {
@@ -142,10 +142,10 @@ impl Drop for BuiltEntity<'_> {
 /// let bundle = builder.build();
 /// let e = world.spawn(&bundle);
 /// let f = world.spawn(&bundle); // `&bundle` can be used many times
-/// assert_eq!(*world.get::<i32>(e).unwrap(), 123);
-/// assert_eq!(*world.get::<&str>(e).unwrap(), "abc");
-/// assert_eq!(*world.get::<i32>(f).unwrap(), 123);
-/// assert_eq!(*world.get::<&str>(f).unwrap(), "abc");
+/// assert_eq!(*world.get::<&i32>(e).unwrap(), 123);
+/// assert_eq!(*world.get::<&&str>(e).unwrap(), "abc");
+/// assert_eq!(*world.get::<&i32>(f).unwrap(), 123);
+/// assert_eq!(*world.get::<&&str>(f).unwrap(), "abc");
 /// ```
 #[derive(Clone, Default)]
 pub struct EntityBuilderClone {

diff --git a/src/entity_ref.rs b/src/entity_ref.rs
@@ -33,19 +33,24 @@ impl<'a> EntityRef<'a> {
         self.archetype.has::<T>()
     }
 
-    /// Borrow the component of type `T`, if it exists
+    /// Borrow a single component, if it exists
     ///
-    /// Panics if the component is already uniquely borrowed from another entity with the same
-    /// components.
-    pub fn get<T: Component>(&self) -> Option<Ref<'a, T>> {
-        Some(unsafe { Ref::new(self.archetype, self.index).ok()? })
-    }
-
-    /// Uniquely borrow the component of type `T`, if it exists
+    /// `T` must be a shared or unique reference to a component type.
+    ///
+    /// # Example
+    /// ```
+    /// # use hecs::*;
+    /// let mut world = World::new();
+    /// let a = world.spawn((42, "abc"));
+    /// let e = world.entity(a).unwrap();
+    /// *e.get::<&mut i32>().unwrap() = 17;
+    /// assert_eq!(*e.get::<&i32>().unwrap(), 17);
+    /// ```
     ///
-    /// Panics if the component is already borrowed from another entity with the same components.
-    pub fn get_mut<T: Component>(&self) -> Option<RefMut<'a, T>> {
-        Some(unsafe { RefMut::new(self.archetype, self.index).ok()? })
+    /// Panics if `T` is a unique reference and the component is already borrowed, or if the
+    /// component is already uniquely borrowed.
+    pub fn get<T: ComponentRef<'a>>(&self) -> Option<T::Ref> {
+        T::get_component(*self)
     }
 
     /// Run a query against this entity
@@ -184,3 +189,40 @@ impl<'a, T: Component> DerefMut for RefMut<'a, T> {
         unsafe { self.target.as_mut() }
     }
 }
+
+/// `&T` or `&mut T` where `T` is some component type
+///
+/// The interface of this trait is a private implementation detail.
+pub trait ComponentRef<'a> {
+    /// Smart pointer to a component of the referenced type
+    #[doc(hidden)]
+    type Ref;
+
+    /// Component type referenced by `Ref`
+    #[doc(hidden)]
+    type Component: Component;
+
+    /// Fetch the component from `entity`
+    #[doc(hidden)]
+    fn get_component(entity: EntityRef<'a>) -> Option<Self::Ref>;
+}
+
+impl<'a, T: Component> ComponentRef<'a> for &'a T {
+    type Ref = Ref<'a, T>;
+
+    type Component = T;
+
+    fn get_component(entity: EntityRef<'a>) -> Option<Self::Ref> {
+        Some(unsafe { Ref::new(entity.archetype, entity.index).ok()? })
+    }
+}
+
+impl<'a, T: Component> ComponentRef<'a> for &'a mut T {
+    type Ref = RefMut<'a, T>;
+
+    type Component = T;
+
+    fn get_component(entity: EntityRef<'a>) -> Option<Self::Ref> {
+        Some(unsafe { RefMut::new(entity.archetype, entity.index).ok()? })
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
@@ -28,8 +28,8 @@
 //!   if flag { *number *= 2; }
 //! }
 //! // Random access is simple and safe
-//! assert_eq!(*world.get::<i32>(a).unwrap(), 246);
-//! assert_eq!(*world.get::<i32>(b).unwrap(), 42);
+//! assert_eq!(*world.get::<&i32>(a).unwrap(), 246);
+//! assert_eq!(*world.get::<&i32>(b).unwrap(), 42);
 //! ```
 
 #![warn(missing_docs)]
@@ -79,7 +79,7 @@ pub use column::{Column, ColumnMut};
 pub use command_buffer::CommandBuffer;
 pub use entities::{Entity, NoSuchEntity};
 pub use entity_builder::{BuiltEntity, BuiltEntityClone, EntityBuilder, EntityBuilderClone};
-pub use entity_ref::{EntityRef, Ref, RefMut};
+pub use entity_ref::{ComponentRef, EntityRef, Ref, RefMut};
 pub use query::{
     Access, Batch, BatchedIter, Or, PreparedQuery, PreparedQueryBorrow, PreparedQueryIter,
     PreparedView, Query, QueryBorrow, QueryItem, QueryIter, QueryMut, QueryShared, Satisfies, View,

diff --git a/src/serialize/column.rs b/src/serialize/column.rs
@@ -757,7 +757,7 @@ mod tests {
     impl PartialEq for SerWorld {
         fn eq(&self, other: &Self) -> bool {
             fn same_components<T: Component + PartialEq>(x: &EntityRef, y: &EntityRef) -> bool {
-                x.get::<T>().as_ref().map(|x| &**x) == y.get::<T>().as_ref().map(|x| &**x)
+                x.get::<&T>().as_ref().map(|x| &**x) == y.get::<&T>().as_ref().map(|x| &**x)
             }
 
             for (x, y) in self.0.iter().zip(other.0.iter()) {
@@ -779,8 +779,8 @@ mod tests {
                     (
                         e.entity(),
                         (
-                            e.get::<Position>().map(|x| *x),
-                            e.get::<Velocity>().map(|x| *x),
+                            e.get::<&Position>().map(|x| *x),
+                            e.get::<&Velocity>().map(|x| *x),
                         ),
                     )
                 }))

diff --git a/src/serialize/row.rs b/src/serialize/row.rs
@@ -82,7 +82,7 @@ pub fn try_serialize<T: Component + Serialize, K: Serialize + ?Sized, S: Seriali
     key: &K,
     map: &mut S,
 ) -> Result<(), S::Error> {
-    if let Some(x) = entity.get::<T>() {
+    if let Some(x) = entity.get::<&T>() {
         map.serialize_key(key)?;
         map.serialize_value(&*x)?;
     }
@@ -270,7 +270,7 @@ mod tests {
     impl PartialEq for SerWorld {
         fn eq(&self, other: &Self) -> bool {
             fn same_components<T: Component + PartialEq>(x: &EntityRef, y: &EntityRef) -> bool {
-                x.get::<T>().as_ref().map(|x| &**x) == y.get::<T>().as_ref().map(|x| &**x)
+                x.get::<&T>().as_ref().map(|x| &**x) == y.get::<&T>().as_ref().map(|x| &**x)
             }
 
             for (x, y) in self.0.iter().zip(other.0.iter()) {
@@ -292,8 +292,8 @@ mod tests {
                     (
                         e.entity(),
                         (
-                            e.get::<Position>().map(|x| *x),
-                            e.get::<Velocity>().map(|x| *x),
+                            e.get::<&Position>().map(|x| *x),
+                            e.get::<&Velocity>().map(|x| *x),
                         ),
                     )
                 }))

diff --git a/src/world.rs b/src/world.rs
@@ -24,9 +24,8 @@ use crate::alloc::boxed::Box;
 use crate::archetype::{Archetype, TypeIdMap, TypeInfo};
 use crate::entities::{Entities, EntityMeta, Location, ReserveEntitiesIterator};
 use crate::{
-    Bundle, Column, ColumnBatch, ColumnMut, DynamicBundle, Entity, EntityRef, Fetch,
-    MissingComponent, NoSuchEntity, Query, QueryBorrow, QueryItem, QueryMut, QueryOne, Ref, RefMut,
-    TakenEntity,
+    Bundle, Column, ColumnBatch, ColumnMut, ComponentRef, DynamicBundle, Entity, EntityRef, Fetch,
+    MissingComponent, NoSuchEntity, Query, QueryBorrow, QueryItem, QueryMut, QueryOne, TakenEntity,
 };
 
 /// An unordered collection of entities, each having any number of distinctly typed components
@@ -187,7 +186,7 @@ impl World {
     /// let mut world = World::new();
     /// let entities = world.spawn_batch((0..1_000).map(|i| (i, "abc"))).collect::<Vec<_>>();
     /// for i in 0..1_000 {
-    ///     assert_eq!(*world.get::<i32>(entities[i]).unwrap(), i as i32);
+    ///     assert_eq!(*world.get::<&i32>(entities[i]).unwrap(), i as i32);
     /// }
     /// ```
     pub fn spawn_batch<I>(&mut self, iter: I) -> SpawnBatchIter<'_, I::IntoIter>
@@ -469,25 +468,15 @@ impl World {
         unsafe { Ok(fetch.get(loc.index as usize)) }
     }
 
-    /// Borrow the `T` component of `entity`
-    ///
-    /// Panics if the component is already uniquely borrowed from another entity with the same
-    /// components.
-    pub fn get<T: Component>(&self, entity: Entity) -> Result<Ref<'_, T>, ComponentError> {
-        Ok(self
-            .entity(entity)?
-            .get()
-            .ok_or_else(MissingComponent::new::<T>)?)
-    }
-
-    /// Uniquely borrow the `T` component of `entity`
-    ///
-    /// Panics if the component is already borrowed from another entity with the same components.
-    pub fn get_mut<T: Component>(&self, entity: Entity) -> Result<RefMut<'_, T>, ComponentError> {
+    /// Short-hand for [`entity`](Self::entity) followed by [`EntityRef::get`]
+    pub fn get<'a, T: ComponentRef<'a>>(
+        &'a self,
+        entity: Entity,
+    ) -> Result<T::Ref, ComponentError> {
         Ok(self
             .entity(entity)?
-            .get_mut()
-            .ok_or_else(MissingComponent::new::<T>)?)
+            .get::<T>()
+            .ok_or_else(MissingComponent::new::<T::Component>)?)
     }
 
     /// Access an entity regardless of its component types
@@ -547,8 +536,8 @@ impl World {
     /// let mut world = World::new();
     /// let e = world.spawn((123, "abc"));
     /// world.insert(e, (456, true));
-    /// assert_eq!(*world.get::<i32>(e).unwrap(), 456);
-    /// assert_eq!(*world.get::<bool>(e).unwrap(), true);
+    /// assert_eq!(*world.get::<&i32>(e).unwrap(), 456);
+    /// assert_eq!(*world.get::<&bool>(e).unwrap(), true);
     /// ```
     pub fn insert(
         &mut self,
@@ -666,9 +655,9 @@ impl World {
     /// let mut world = World::new();
     /// let e = world.spawn((123, "abc", true));
     /// assert_eq!(world.remove::<(i32, &str)>(e), Ok((123, "abc")));
-    /// assert!(world.get::<i32>(e).is_err());
-    /// assert!(world.get::<&str>(e).is_err());
-    /// assert_eq!(*world.get::<bool>(e).unwrap(), true);
+    /// assert!(world.get::<&i32>(e).is_err());
+    /// assert!(world.get::<&&str>(e).is_err());
+    /// assert_eq!(*world.get::<&bool>(e).unwrap(), true);
     /// ```
     pub fn remove<T: Bundle + 'static>(&mut self, entity: Entity) -> Result<T, ComponentError> {
         self.flush();
@@ -1389,13 +1378,13 @@ mod tests {
     fn reuse_populated() {
         let mut world = World::new();
         let a = world.spawn((42,));
-        assert_eq!(*world.get::<i32>(a).unwrap(), 42);
+        assert_eq!(*world.get::<&i32>(a).unwrap(), 42);
         world.despawn(a).unwrap();
         let b = world.spawn((true,));
         assert_eq!(a.id, b.id);
         assert_ne!(a.generation, b.generation);
-        assert!(world.get::<i32>(b).is_err());
-        assert!(*world.get::<bool>(b).unwrap());
+        assert!(world.get::<&i32>(b).is_err());
+        assert!(*world.get::<&bool>(b).unwrap());
     }
 
     #[test]