Repository: Ralith/hecs
Pull request #314 (2023-04-06 22:15:52)

Add `Ref::map` and `RefMut::map`.
This commit adds analogues of [std::cell::Ref::map](https://doc.rust-lang.org/stable/std/cell/struct.Ref.html#method.map) and [std::cell::RefMut::map](https://doc.rust-lang.org/stable/std/cell/struct.RefMut.html#method.map), which allow the contents of a `Ref`/`RefMut` to be transformed without reborrowing. This can make it easier to write certain kinds of functions which return data borrowed from components.

An example of this kind of code is below:

```rust
trait MyComponentTrait {
    fn get_state(&self) -> &dyn MyState;
}

fn entity_get_state<T: MyComponentTrait>(ent: EntityRef<'_>) -> Option<Ref<'_, dyn MyState>> {
    let component = ent.get::<&'_ T>();
    component.map(|comp_ref| Ref::map(comp_ref, |comp| comp.get_state()))
}

struct MyComponentInstance {
    // state unspecified...
}

impl MyComponentTrait for MyComponmentInstance {
    // implementation unspecified...
}

let mut comp_ty_to_state_getter: HashMap<TypeId, fn(EntityRef<'_>) -> Option<Ref<'_, dyn MyState>> = Default::default();
comp_ty_to_state_getter.insert(TypeId::of::<MyComponmentInstance>(), entity_get_state::<MyComponmentInstance>);
```

From there, when you have an entity, you could iterate over the `TypeId`s of the components in that entity, and then
call use the `comp_ty_to_state_getter` to call common functionality on all of them. Examples of where this could
be useful include serialisation, or visiting component fields to render them in a GUI inspector.

As part of this change, the sizes of the `Ref` and `RefMut` types have increased to store the `TypeId` of the original
component, and the `Archetype` can now be released polymorphically using that `TypeId`.

Let me know if there is anything else I should change or if there are any tests I should add.


Ralith: Thanks for the PR! I agree that this API is well motivated. Because it should be impossible to construct a `Ref` or `RefMut` with an incorrect `state` field wrt. its `archetype`, I think we should skip the extra `TypeId` and instead add an unsafe internal API to `Archetype` to borrow/release by state index.
burtonageo: Done
burtonageo: I'd also like to rebase this onto https://github.com/Ralith/hecs/pull/315 when that's approved, as the `Debug` impls could be expanded to require `?Sized + Debug`
Ralith: I like the refactoring, but bundling it together with the API extension will make this take longer to review. You might also want to rebase to fix the conflict.
burtonageo: Rebased!

Sorry about the extra refactoring. I decided to add a test, which exposed a bug in the impl of `Ref::clone()`, where A. The `#[derive(Clone)]` was putting an extra requirement that `impl<T: Clone> Clone for Ref<T>` which was unnecessary, and also not incrementing the borrow ref counts in the archetype for immutable borrows of the component.

Therefore, I decided to refactor the borrowing permissions into the `ComponentBorrow` and `ComponentBorrowMut` types, so that these could manage the borrowing lifetime from the `Archetype`, and then the `Ref` and `RefMut` types could focus on being useful wrappers over the `NonNull<T>`. This fixed the case in test where the extra clone of the `Ref<T>` would cause a panic due to imbalanced references, and it also simplified the `{Ref, RefMut}::map` implementations, as the borrow could be passed to the new `Ref`, rather than having to forget the old ref.

Thanks for your patience, please lmk if there is anything else I can do. I think this should be everything I want to change in this PR.
burtonageo: Comments all (hopefully) addressed

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 52476ef..1eb2cb0 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,12 @@
+# Unreleased
+
+### Added
+- `Ref::map` and `RefMut::map` to reborrow a subfield of a component, or cast it to an unsized type
+
+### Fixed
+- `Ref::clone()` is now callable when containing types which do not implement `Clone`
+- `Ref::clone()` could cause a panic as internal state wasn't being updated correctly
+
 # 0.10.2
 
 ### Fixed
diff --git a/src/archetype.rs b/src/archetype.rs
index 9d5f38a..9fbec2e 100644
--- a/src/archetype.rs
+++ b/src/archetype.rs
@@ -127,6 +127,12 @@ impl Archetype {
         }
     }
 
+    pub(crate) unsafe fn borrow_raw(&self, state: usize) {
+        if !self.data[state].state.borrow() {
+            panic!("state index {} already borrowed uniquely", state);
+        }
+    }
+
     pub(crate) fn borrow_mut<T: Component>(&self, state: usize) {
         assert_eq!(self.types[state].id, TypeId::of::<T>());
 
@@ -145,6 +151,14 @@ impl Archetype {
         self.data[state].state.release_mut();
     }
 
+    pub(crate) unsafe fn release_raw(&self, state: usize) {
+        self.data[state].state.release();
+    }
+
+    pub(crate) unsafe fn release_raw_mut(&self, state: usize) {
+        self.data[state].state.release_mut();
+    }
+
     /// Number of entities in this archetype
     #[inline]
     pub fn len(&self) -> u32 {
diff --git a/src/entity_ref.rs b/src/entity_ref.rs
index f3339b5..3680935 100644
--- a/src/entity_ref.rs
+++ b/src/entity_ref.rs
@@ -1,6 +1,7 @@
 use core::any::TypeId;
 use core::fmt::{self, Debug, Formatter};
-use core::ops::{Deref, DerefMut};
+use core::marker::PhantomData;
+use core::ops::{Deref, DerefMut, FnOnce};
 use core::ptr::NonNull;
 
 use crate::archetype::Archetype;
@@ -109,12 +110,10 @@ unsafe impl<'a> Send for EntityRef<'a> {}
 unsafe impl<'a> Sync for EntityRef<'a> {}
 
 /// Shared borrow of an entity's component
-#[derive(Clone)]
-pub struct Ref<'a, T: Component> {
-    archetype: &'a Archetype,
-    /// State index for `T` in `archetype`
-    state: usize,
+pub struct Ref<'a, T: ?Sized> {
+    borrow: ComponentBorrow<'a>,
     target: NonNull<T>,
+    _phantom: PhantomData<&'a T>,
 }
 
 impl<'a, T: Component> Ref<'a, T> {
@@ -122,48 +121,80 @@ impl<'a, T: Component> Ref<'a, T> {
         archetype: &'a Archetype,
         index: u32,
     ) -> Result<Self, MissingComponent> {
-        let state = archetype
-            .get_state::<T>()
-            .ok_or_else(MissingComponent::new::<T>)?;
-        let target =
-            NonNull::new_unchecked(archetype.get_base::<T>(state).as_ptr().add(index as usize));
-        archetype.borrow::<T>(state);
+        let (target, borrow) = ComponentBorrow::for_component::<T>(archetype, index)?;
         Ok(Self {
-            archetype,
-            state,
+            borrow,
             target,
+            _phantom: PhantomData,
         })
     }
 }
 
-unsafe impl<T: Component> Send for Ref<'_, T> {}
-unsafe impl<T: Component> Sync for Ref<'_, T> {}
+unsafe impl<T: ?Sized + Sync> Send for Ref<'_, T> {}
+unsafe impl<T: ?Sized + Sync> Sync for Ref<'_, T> {}
 
-impl<'a, T: Component> Drop for Ref<'a, T> {
-    fn drop(&mut self) {
-        self.archetype.release::<T>(self.state);
+impl<'a, T: ?Sized> Ref<'a, T> {
+    /// Transform the `Ref<'_, T>` to point to a part of the borrowed data, e.g.
+    /// a struct field.
+    ///
+    /// The `Ref<'_, T>` is already borrowed, so this cannot fail.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use hecs::{EntityRef, Ref};
+    /// struct Component {
+    ///     member: i32,
+    /// }
+    ///
+    /// # fn example(entity_ref: EntityRef<'_>) {
+    /// let component_ref = entity_ref.get::<&Component>()
+    ///     .expect("Entity does not contain an instance of \"Component\"");
+    /// let member_ref = Ref::map(component_ref, |component| &component.member);
+    /// println!("member = {:?}", *member_ref);
+    /// # }
+    /// ```
+    pub fn map<U: ?Sized, F>(orig: Ref<'a, T>, f: F) -> Ref<'a, U>
+    where
+        F: FnOnce(&T) -> &U,
+    {
+        let target = NonNull::from(f(&*orig));
+        Ref {
+            borrow: orig.borrow,
+            target,
+            _phantom: PhantomData,
+        }
     }
 }
 
-impl<'a, T: Component> Deref for Ref<'a, T> {
+impl<'a, T: ?Sized> Deref for Ref<'a, T> {
     type Target = T;
     fn deref(&self) -> &T {
         unsafe { self.target.as_ref() }
     }
 }
 
-impl<'a, T: Component + Debug> Debug for Ref<'a, T> {
+impl<'a, T: ?Sized + Debug> Debug for Ref<'a, T> {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         Debug::fmt(self.deref(), f)
     }
 }
 
+impl<'a, T: ?Sized> Clone for Ref<'a, T> {
+    fn clone(&self) -> Self {
+        Self {
+            borrow: self.borrow.clone(),
+            target: self.target,
+            _phantom: self._phantom,
+        }
+    }
+}
+
 /// Unique borrow of an entity's component
-pub struct RefMut<'a, T: Component> {
-    archetype: &'a Archetype,
-    /// State index for `T` in `archetype`
-    state: usize,
+pub struct RefMut<'a, T: ?Sized> {
+    borrow: ComponentBorrowMut<'a>,
     target: NonNull<T>,
+    _phantom: PhantomData<&'a mut T>,
 }
 
 impl<'a, T: Component> RefMut<'a, T> {
@@ -171,43 +202,67 @@ impl<'a, T: Component> RefMut<'a, T> {
         archetype: &'a Archetype,
         index: u32,
     ) -> Result<Self, MissingComponent> {
-        let state = archetype
-            .get_state::<T>()
-            .ok_or_else(MissingComponent::new::<T>)?;
-        let target =
-            NonNull::new_unchecked(archetype.get_base::<T>(state).as_ptr().add(index as usize));
-        archetype.borrow_mut::<T>(state);
+        let (target, borrow) = ComponentBorrowMut::for_component::<T>(archetype, index)?;
         Ok(Self {
-            archetype,
-            state,
+            borrow,
             target,
+            _phantom: PhantomData,
         })
     }
 }
 
-unsafe impl<T: Component> Send for RefMut<'_, T> {}
-unsafe impl<T: Component> Sync for RefMut<'_, T> {}
+unsafe impl<T: ?Sized + Send> Send for RefMut<'_, T> {}
+unsafe impl<T: ?Sized + Sync> Sync for RefMut<'_, T> {}
 
-impl<'a, T: Component> Drop for RefMut<'a, T> {
-    fn drop(&mut self) {
-        self.archetype.release_mut::<T>(self.state);
+impl<'a, T: ?Sized> RefMut<'a, T> {
+    /// Transform the `RefMut<'_, T>` to point to a part of the borrowed data, e.g.
+    /// a struct field.
+    ///
+    /// The `RefMut<'_, T>` is already mutably borrowed, so this cannot fail.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use hecs::{EntityRef, RefMut};
+    /// struct Component {
+    ///     member: i32,
+    /// }
+    ///
+    /// # fn example(entity_ref: EntityRef<'_>) {
+    /// let component_ref = entity_ref.get::<&mut Component>()
+    ///     .expect("Entity does not contain an instance of \"Component\"");
+    /// let mut member_ref = RefMut::map(component_ref, |component| &mut component.member);
+    /// *member_ref = 21;
+    /// println!("member = {:?}", *member_ref);
+    /// # }
+    /// ```
+    pub fn map<U: ?Sized, F>(mut orig: RefMut<'a, T>, f: F) -> RefMut<'a, U>
+    where
+        F: FnOnce(&mut T) -> &mut U,
+    {
+        let target = NonNull::from(f(&mut *orig));
+        RefMut {
+            borrow: orig.borrow,
+            target,
+            _phantom: PhantomData,
+        }
     }
 }
 
-impl<'a, T: Component> Deref for RefMut<'a, T> {
+impl<'a, T: ?Sized> Deref for RefMut<'a, T> {
     type Target = T;
     fn deref(&self) -> &T {
         unsafe { self.target.as_ref() }
     }
 }
 
-impl<'a, T: Component> DerefMut for RefMut<'a, T> {
+impl<'a, T: ?Sized> DerefMut for RefMut<'a, T> {
     fn deref_mut(&mut self) -> &mut T {
         unsafe { self.target.as_mut() }
     }
 }
 
-impl<'a, T: Component + Debug> Debug for RefMut<'a, T> {
+impl<'a, T: ?Sized + Debug> Debug for RefMut<'a, T> {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         Debug::fmt(self.deref(), f)
     }
@@ -294,3 +349,85 @@ impl<'a, T: Component> ComponentRef<'a> for &'a mut T {
 pub trait ComponentRefShared<'a>: ComponentRef<'a> {}
 
 impl<'a, T: Component> ComponentRefShared<'a> for &'a T {}
+
+struct ComponentBorrow<'a> {
+    archetype: &'a Archetype,
+    /// State index for the borrowed component in the `archetype`.
+    state: usize,
+}
+
+impl<'a> ComponentBorrow<'a> {
+    // This method is unsafe as if the `index` is out of bounds,
+    // then this will cause undefined behavior as the returned
+    // `target` will point to undefined memory.
+    unsafe fn for_component<T: Component>(
+        archetype: &'a Archetype,
+        index: u32,
+    ) -> Result<(NonNull<T>, Self), MissingComponent> {
+        let state = archetype
+            .get_state::<T>()
+            .ok_or_else(MissingComponent::new::<T>)?;
+
+        let target =
+            NonNull::new_unchecked(archetype.get_base::<T>(state).as_ptr().add(index as usize));
+
+        archetype.borrow::<T>(state);
+
+        Ok((target, Self { archetype, state }))
+    }
+}
+
+impl<'a> Clone for ComponentBorrow<'a> {
+    fn clone(&self) -> Self {
+        unsafe {
+            self.archetype.borrow_raw(self.state);
+        }
+        Self {
+            archetype: self.archetype,
+            state: self.state,
+        }
+    }
+}
+
+impl<'a> Drop for ComponentBorrow<'a> {
+    fn drop(&mut self) {
+        unsafe {
+            self.archetype.release_raw(self.state);
+        }
+    }
+}
+
+struct ComponentBorrowMut<'a> {
+    archetype: &'a Archetype,
+    /// State index for the borrowed component in the `archetype`.
+    state: usize,
+}
+
+impl<'a> ComponentBorrowMut<'a> {
+    // This method is unsafe as if the `index` is out of bounds,
+    // then this will cause undefined behavior as the returned
+    // `target` will point to undefined memory.
+    unsafe fn for_component<T: Component>(
+        archetype: &'a Archetype,
+        index: u32,
+    ) -> Result<(NonNull<T>, Self), MissingComponent> {
+        let state = archetype
+            .get_state::<T>()
+            .ok_or_else(MissingComponent::new::<T>)?;
+
+        let target =
+            NonNull::new_unchecked(archetype.get_base::<T>(state).as_ptr().add(index as usize));
+
+        archetype.borrow_mut::<T>(state);
+
+        Ok((target, Self { archetype, state }))
+    }
+}
+
+impl<'a> Drop for ComponentBorrowMut<'a> {
+    fn drop(&mut self) {
+        unsafe {
+            self.archetype.release_raw_mut(self.state);
+        }
+    }
+}


Test patch:
diff --git a/tests/tests.rs b/tests/tests.rs
index 122da18..0543b35 100644
--- a/tests/tests.rs
+++ b/tests/tests.rs
@@ -935,3 +935,34 @@ fn empty_archetype_conflict() {
         for _ in world.query::<(&mut i32, &bool)>().iter() {}
     }
 }
+
+#[test]
+fn component_ref_map() {
+    struct TestComponent {
+        id: i32,
+    }
+
+    let mut world = World::new();
+    let e = world.spawn((TestComponent { id: 21 },));
+
+    let e_ref = world.entity(e).unwrap();
+    {
+        let comp = e_ref.get::<&'_ TestComponent>().unwrap();
+        // Test that no unbalanced releases occur when cloning refs.
+        let _comp2 = comp.clone();
+        let id = Ref::map(comp, |c| &c.id);
+        assert_eq!(*id, 21);
+    }
+
+    {
+        let comp = e_ref.get::<&'_ mut TestComponent>().unwrap();
+        let mut id = RefMut::map(comp, |c| &mut c.id);
+        *id = 31;
+    }
+
+    {
+        let comp = e_ref.get::<&'_ TestComponent>().unwrap();
+        let id = Ref::map(comp, |c| &c.id);
+        assert_eq!(*id, 31);
+    }
+}