Return Ref and RefMut from the borrow queries to enforce borrow checks even more

lyra-ecs/README.md

@@ -12,8 +12,8 @@ I couldn't find anything that fulfilled my needs, specifically an ECS that can s
     - [x] Find some way to fill in the gaps in component columns after entities are deleted
       * This was done by moving the last entity in the column to the gap that was just removed.
   - [x] Grow archetype as it fills up
-  - [ ] Borrow safety of components inside Archetypes
-    * Return `Ref` and `RefMut` from borrow queries
+  - [x] Borrow safety of components inside Archetypes
+    - [x] Return `Ref` and `RefMut` from borrow queries
 - [ ] Querying components from archetypes
   - [x] Views
   - [ ] Mutable views
@@ -25,6 +25,7 @@ I couldn't find anything that fulfilled my needs, specifically an ECS that can s
     * Needed for scripting engines that can create their own components that Rust does not know the types of.
 - [ ] Systems
   - [ ] Dispatchers/Executors
+    - [ ] Execution graph that follows dependencies
 
 <br>
 

lyra-ecs/src/query/borrow.rs

@@ -1,12 +1,12 @@
 use std::{marker::PhantomData, any::TypeId, ptr::NonNull, cell::{Ref, RefCell, RefMut}};
 
-use crate::world::World;
+use crate::{world::World, ComponentColumn};
 
 use super::{Fetch, Query, AsQuery, DefaultQuery};
 
 /// Fetcher for borrowing components from archetypes.
 pub struct FetchBorrow<'a, T> {
-    ptr: Option<Ref<'a, NonNull<u8>>>,
+    col: &'a ComponentColumn,
     size: usize,
     _phantom: PhantomData<&'a T>
 }
@@ -15,22 +15,20 @@ impl<'a, T> Fetch<'a> for FetchBorrow<'a, T>
 where
     T: 'a,
 {
-    type Item = &'a T;
+    type Item = Ref<'a, T>;
 
     fn dangling() -> Self {
-        FetchBorrow {
-            ptr: None,
-            size: 0,
-            _phantom: PhantomData::default()
-        }
+        unreachable!()
     }
 
     unsafe fn get_item(&mut self, entity: crate::world::ArchetypeEntityId) -> Self::Item {
-        let ptr = self.ptr.as_ref().unwrap();
-        let ptr = NonNull::new_unchecked(ptr.as_ptr()
-            .add(entity.0 as usize * self.size))
-            .cast();
-        &*ptr.as_ptr()
+        let ptr = self.col.borrow_ptr();
+        Ref::map(ptr, |ptr| {
+            let ptr = NonNull::new_unchecked(ptr.as_ptr()
+                .add(entity.0 as usize * self.size))
+                .cast();
+            &*ptr.as_ptr()
+        })
     }
 }
 
@@ -60,7 +58,7 @@ impl<T> Query for QueryBorrow<T>
 where
     T: 'static
 {
-    type Item<'a> = &'a T;
+    type Item<'a> = Ref<'a, T>;
 
     type Fetch<'a> = FetchBorrow<'a, T>;
 
@@ -71,10 +69,9 @@ where
     unsafe fn fetch<'a>(&self, _world: &'a World, _arch_id: crate::archetype::ArchetypeId, archetype: &'a crate::archetype::Archetype) -> Self::Fetch<'a> {
         let col = archetype.columns.iter().find(|c| c.info.type_id == self.type_id)
             .expect("You ignored 'can_visit_archetype'!");
-        let col_data = col.borrow_ptr();
 
         FetchBorrow {
-            ptr: Some(col_data),
+            col,
             size: col.info.layout.size(),
             _phantom: PhantomData,
         }
@@ -103,7 +100,7 @@ impl<T: 'static> DefaultQuery for &T {
 
 /// A fetcher for mutably borrowing components from archetypes.
 pub struct FetchBorrowMut<'a, T> {
-    ptr: Option<RefMut<'a, NonNull<u8>>>,
+    col: &'a ComponentColumn,
     size: usize,
     _phantom: PhantomData<&'a T>
 }
@@ -112,22 +109,20 @@ impl<'a, T> Fetch<'a> for FetchBorrowMut<'a, T>
 where
     T: 'a,
 {
-    type Item = &'a mut T;
+    type Item = RefMut<'a, T>;
 
     fn dangling() -> Self {
-        FetchBorrowMut {
-            ptr: None,
-            size: 0,
-            _phantom: PhantomData::default()
-        }
+        unreachable!()
     }
 
     unsafe fn get_item(&mut self, entity: crate::world::ArchetypeEntityId) -> Self::Item {
-        let ptr = self.ptr.as_ref().unwrap();
-        let ptr = NonNull::new_unchecked(ptr.as_ptr()
-            .add(entity.0 as usize * self.size))
-            .cast();
-        &mut *ptr.as_ptr()
+        let ptr = self.col.borrow_mut_ptr();
+        RefMut::map(ptr, |ptr| {
+            let ptr = NonNull::new_unchecked(ptr.as_ptr()
+                .add(entity.0 as usize * self.size))
+                .cast();
+            &mut *ptr.as_ptr()
+        })
     }
 }
 
@@ -157,7 +152,7 @@ impl<T> Query for QueryBorrowMut<T>
 where
     T: 'static
 {
-    type Item<'a> = &'a mut T;
+    type Item<'a> = RefMut<'a, T>;
 
     type Fetch<'a> = FetchBorrowMut<'a, T>;
 
@@ -168,10 +163,9 @@ where
     unsafe fn fetch<'a>(&self, _world: &'a World, _arch_id: crate::archetype::ArchetypeId, archetype: &'a crate::archetype::Archetype) -> Self::Fetch<'a> {
         let col = archetype.columns.iter().find(|c| c.info.type_id == self.type_id)
             .expect("You ignored 'can_visit_archetype'!");
-        let col_data = col.borrow_mut_ptr();
 
         FetchBorrowMut {
-            ptr: Some(col_data),
+            col,
             size: col.info.layout.size(),
             _phantom: PhantomData,
         }
@@ -202,7 +196,7 @@ impl<T: 'static> DefaultQuery for &mut T {
 mod tests {
     use std::{any::TypeId, mem::size_of, marker::PhantomData};
 
-    use crate::{world::{World, Entity, EntityId}, archetype::{Archetype, ArchetypeId}, query::View, tests::Vec2, bundle::Bundle};
+    use crate::{world::{World, Entity, EntityId}, archetype::{Archetype, ArchetypeId}, query::View, tests::Vec2, bundle::Bundle, Fetch};
 
     use super::{QueryBorrow, QueryBorrowMut, FetchBorrowMut};
 
@@ -245,7 +239,7 @@ mod tests {
 
         let mut orig = vec![];
 
-        for v in v.into_iter() {
+        for mut v in v.into_iter() {
             orig.push(v.clone());
             v.x += 10.0;
             v.y += 10.0;
@@ -281,14 +275,16 @@ mod tests {
 
         let col = a.columns.iter().find(|c| c.info.type_id == TypeId::of::<Vec2>()).unwrap();
 
-        let bmut = FetchBorrowMut::<Vec2> {
-            ptr: Some(col.borrow_mut_ptr()),
+        let mut bmut = FetchBorrowMut::<Vec2> {
+            col,
             size: size_of::<Vec2>(),
             _phantom: PhantomData,
         };
+        let item = unsafe { bmut.get_item(crate::ArchetypeEntityId(0)) };
         
+        // ensure only one mutable borrow can be done at a time
         assert!(col.try_borrow_mut_ptr().is_err());
-        drop(bmut);
+        drop(item);
         assert!(col.try_borrow_mut_ptr().is_ok());
     }
 }