Enforce single mutable borrows of component columns

lyra-ecs/Cargo.toml

@@ -7,6 +7,5 @@ edition = "2021"
 
 [dependencies]
 
-
 [dev-dependencies]
 rand = "0.8.5" # used for tests

lyra-ecs/README.md

@@ -10,13 +10,13 @@ I couldn't find anything that fulfilled my needs, specifically an ECS that can s
   - [x] Despawning entities
     - [x] Delete entities from archetypes
     - [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.
+      * 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
+  - [x] Borrow safety of components inside Archetypes
-      * Its wonk right now; a component can be borrowed mutably from a non-mutable reference.
 - [ ] Querying components from archetypes
   - [x] Views
   - [ ] Mutable views
+    - [ ] Make it possible so that ONLY `ViewMut` can borrow mutably
 - [ ] Resources
 - [ ] Relationships (maybe this can be done through queries, idk)
 - [ ] Dynamic queries

lyra-ecs/src/archetype.rs

@@ -1,27 +1,40 @@
-use std::{any::TypeId, ptr::{NonNull, self}, alloc::{self, Layout, alloc, dealloc}, mem, collections::HashMap};
+use std::{any::TypeId, ptr::{NonNull, self}, alloc::{self, Layout, alloc, dealloc}, mem, collections::HashMap, cell::{RefCell, Ref, RefMut, BorrowError, BorrowMutError}, ops::{DerefMut, Deref}};
 
-use crate::{world::{Entity, ArchetypeEntityId, Record}, bundle::Bundle, component_info::ComponentInfo};
+use crate::{world::{Entity, ArchetypeEntityId}, bundle::Bundle, component_info::ComponentInfo};
 
 pub struct ComponentColumn {
-    pub data: NonNull<u8>,
+    data: RefCell<NonNull<u8>>,
-    pub capacity: usize,
+    len: usize,
+    capacity: usize,
     pub info: ComponentInfo,
-    pub entry_size: usize,
+}
-    pub len: usize,
+
+impl Drop for ComponentColumn {
+    fn drop(&mut self) {
+        let data = self.data.borrow_mut();
+        let data = data.as_ptr();
+
+        unsafe {
+            // layout of current alloc
+            let layout = Layout::from_size_align_unchecked(self.info.layout.size() * self.capacity,
+                self.info.layout.align());
+            dealloc(data, layout);
+        }
+    }
 }
 
 #[allow(dead_code)]
 impl ComponentColumn {
     /// Creates an invalid component column. Do not attempt to use it.
-    pub unsafe fn dangling() -> Self {
+    /* pub unsafe fn dangling() -> Self {
         ComponentColumn {
-            data: NonNull::dangling(),
+            data: RefCell::
             capacity: 0,
             info: ComponentInfo::new::<()>(),
             entry_size: 0,
             len: 0,
         }
-    }
+    } */
     
     pub unsafe fn alloc(component_layout: Layout, capacity: usize) -> NonNull<u8> {
         let new_layout = Layout::from_size_align(
@@ -39,12 +52,10 @@ impl ComponentColumn {
     pub unsafe fn new(info: ComponentInfo, capacity: usize) -> Self {
         let data = ComponentColumn::alloc(info.layout, capacity);
         
-        let size = info.layout.size();
         Self {
-            data,
+            data: RefCell::new(data),
             capacity,
             info,
-            entry_size: size,
             len: 0,
         }
     }
@@ -56,8 +67,12 @@ impl ComponentColumn {
     /// This column must have space to fit the component, if it does not have room it will panic.
     pub unsafe fn set_at(&mut self, entity_index: usize, comp_src: NonNull<u8>) {
         assert!(entity_index < self.capacity);
-        let dest = NonNull::new_unchecked(self.data.as_ptr().add(entity_index * self.entry_size));
+
-        ptr::copy_nonoverlapping(comp_src.as_ptr(), dest.as_ptr(), self.entry_size);
+        let mut data = self.data.borrow_mut();
+        let data = data.deref_mut();
+        
+        let dest = NonNull::new_unchecked(data.as_ptr().add(entity_index * self.info.layout.size()));
+        ptr::copy_nonoverlapping(comp_src.as_ptr(), dest.as_ptr(), self.info.layout.size());
         self.len += 1;
     }
 
@@ -67,8 +82,11 @@ impl ComponentColumn {
     /// 
     /// This column MUST have the entity. If it does not, it WILL NOT panic and will cause UB.
     pub unsafe fn get<T>(&self, entity_index: usize) -> &T {
-        let ptr = NonNull::new_unchecked(self.data.as_ptr()
+        let data = self.data.borrow();
-            .add(entity_index * self.entry_size))
+        let data = data.deref();
+
+        let ptr = NonNull::new_unchecked(data.as_ptr()
+            .add(entity_index * self.info.layout.size()))
             .cast();
         &*ptr.as_ptr()
     }
@@ -79,9 +97,12 @@ impl ComponentColumn {
     /// 
     /// This column must have the entity.
     pub unsafe fn get_mut<T>(&mut self, entity_index: usize) -> &mut T {
-        let p = self.data.as_ptr()
+        let mut data = self.data.borrow_mut();
+        let data = data.deref_mut();
+
+        let p = data.as_ptr()
             .cast::<T>()
-            .add(entity_index * self.entry_size);
+            .add(entity_index * self.info.layout.size());
         &mut *p
     }
 
@@ -95,31 +116,46 @@ impl ComponentColumn {
     /// Will panic if `new_capacity` is less than the current capacity of the column.
     pub unsafe fn grow(&mut self, new_capacity: usize) {
         assert!(new_capacity > self.capacity);
+
+        let mut data = self.data.borrow_mut();
+        //let data = data.deref_mut();
+
         let mut new_ptr = Self::alloc(self.info.layout, new_capacity);
-        ptr::copy_nonoverlapping(self.data.as_ptr(), new_ptr.as_ptr(), self.capacity * self.entry_size);
 
-        let old_layout = Layout::from_size_align_unchecked(
+        if self.len > 0 {
-            self.info.layout.size().checked_mul(self.capacity).unwrap(),
+            ptr::copy_nonoverlapping(data.as_ptr(), new_ptr.as_ptr(), self.len * self.info.layout.size());
-            self.info.layout.align()
+        }
-        );
 
-        mem::swap(&mut self.data, &mut new_ptr); // 'new_ptr' is now the old pointer
+        // dont attempt to free if we weren't able to store anything anyway
-        dealloc(new_ptr.as_ptr(), old_layout);
+        if self.capacity != 0 {
+            let old_layout = Layout::from_size_align_unchecked(
+                self.info.layout.size().checked_mul(self.capacity).unwrap(),
+                self.info.layout.align()
+            );
+
+            mem::swap(data.deref_mut(), &mut new_ptr);
+            dealloc(new_ptr.as_ptr(), old_layout);
+        } else {
+            *data = new_ptr;
+        }
         
         self.capacity = new_capacity;
     }
 
     /// Removes a component from the column, freeing it, and returning the old index of the entity that took its place in the column.
     pub unsafe fn remove_component(&mut self, entity_index: usize) -> Option<usize> {
-        let mut old_comp_ptr = NonNull::new_unchecked(self.data.as_ptr()
+        let mut data = self.data.borrow_mut();
-            .add(entity_index * self.entry_size));
+        let data = data.deref_mut();
+
+        let mut old_comp_ptr = NonNull::new_unchecked(data.as_ptr()
+            .add(entity_index * self.info.layout.size()));
 
         let moved_index = if entity_index != self.len - 1 {
             let moved_index = self.len - 1;
-            let mut new_comp_ptr = NonNull::new_unchecked(self.data.as_ptr()
+            let mut new_comp_ptr = NonNull::new_unchecked(data.as_ptr()
-                .add(moved_index * self.entry_size));
+                .add(moved_index * self.info.layout.size()));
 
-            ptr::copy_nonoverlapping(new_comp_ptr.as_ptr(), old_comp_ptr.as_ptr(), self.entry_size);
+            ptr::copy_nonoverlapping(new_comp_ptr.as_ptr(), old_comp_ptr.as_ptr(), self.info.layout.size());
 
             mem::swap(&mut old_comp_ptr, &mut new_comp_ptr); // new_comp_ptr is now the old ptr
             Some(moved_index)
@@ -129,6 +165,22 @@ impl ComponentColumn {
 
         moved_index
     }
+
+    pub fn borrow_ptr<'a>(&'a self) -> Ref<'a, NonNull<u8>> {
+        self.data.borrow()
+    }
+
+    pub fn borrow_mut_ptr<'a>(&'a self) -> RefMut<'a, NonNull<u8>> {
+        self.data.borrow_mut()
+    }
+
+    pub fn try_borrow_ptr<'a>(&'a self) -> Result<Ref<'a, NonNull<u8>>, BorrowError> {
+        self.data.try_borrow()
+    }
+
+    pub fn try_borrow_mut_ptr<'a>(&'a self) -> Result<RefMut<'a, NonNull<u8>>, BorrowMutError> {
+        self.data.try_borrow_mut()
+    }
 }
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
@@ -252,6 +304,8 @@ impl Archetype {
         for c in self.columns.iter_mut() {
             unsafe { c.grow(new_capacity); }
         }
+
+        self.capacity = new_capacity;
     }
 }
 
@@ -353,19 +407,30 @@ mod tests {
         assert_eq!(vec3.clone(), b2.1);
     }
 
+    /// Tests manual growth of archetypes
     #[test]
-    fn column_growth() {
+    fn archetype_growth() {
+        let info = (Vec2::new(0.0, 0.0),).info();
+        let mut a = Archetype::from_bundle_info(super::ArchetypeId(0), info);
+        a.grow_columns(64);
+        a.grow_columns(128);
+        a.grow_columns(256);
+    }
+
+    /// Tests the auto growth of archetypes when adding entities
+    #[test]
+    fn auto_archetype_growth() {
         let mut rng = rand::thread_rng();
         let bundle_count = rng.gen_range(50..150);
         
-        let mut bundles = vec![];
+        let mut bundles: Vec<(Vec2,)> = vec![];
         bundles.reserve(bundle_count);
 
         let info = (Vec2::new(0.0, 0.0),).info();
         let mut a = Archetype::from_bundle_info(super::ArchetypeId(0), info);
 
         for i in 0..bundle_count {
-            let c = (Vec2::new(rng.gen_range(10.0..3000.0), rng.gen_range(10.0..3000.0)),);
+            let c = (Vec2::rand(),);
             bundles.push(c);
             
             a.add_entity(
@@ -376,6 +441,7 @@ mod tests {
                 c
             );
         }
+        println!("Inserted {} entities", bundle_count);
 
         let col = a.columns.get(0).unwrap();
         for i in 0..bundle_count {
@@ -386,14 +452,8 @@ mod tests {
 
     #[test]
     fn remove_entity() {
-        let mut rng = rand::thread_rng();
+        let bundles = vec![Vec2::rand(), Vec2::rand(), Vec2::rand()];
-        let range = 30.0..1853.0;
+        println!("Bundles: {:?}", bundles);
-        
-        let bundles = vec![
-            ( Vec2::new(rng.gen_range(range.clone()), rng.gen_range(range.clone())), ),
-            ( Vec2::new(rng.gen_range(range.clone()), rng.gen_range(range.clone())), ),
-            ( Vec2::new(rng.gen_range(range.clone()), rng.gen_range(range.clone())), )
-        ];
 
         let info = (Vec2::new(0.0, 0.0),).info();
         let mut a = Archetype::from_bundle_info(super::ArchetypeId(0), info);
@@ -405,7 +465,7 @@ mod tests {
                     id: EntityId(i as u64),
                     generation: 0
                 },
-                bundles[i],
+                (bundles[i],),
             );
         }
 
@@ -424,6 +484,6 @@ mod tests {
         // make sure that the entities' component was actually moved in the column
         let col = &a.columns[0];
         let comp = unsafe { col.get::<Vec2>(1) };
-        assert_eq!(comp.clone(), bundles[2].0);
+        assert_eq!(comp.clone(), bundles[2]);
     }
 }

lyra-ecs/src/main.rs

@@ -1,4 +1,3 @@
-//#![feature(cell_leak)]
 use crate::world::World;
 
 mod archetype;

lyra-ecs/src/query/borrow.rs

@@ -1,10 +1,10 @@
-use std::{marker::PhantomData, any::TypeId, ptr::NonNull};
+use std::{marker::PhantomData, any::TypeId, ptr::NonNull, cell::{Ref, RefCell, RefMut}};
 
 use super::{Fetch, Query, AsQuery, DefaultQuery};
 
 /// Fetcher for borrowing components from archetypes.
 pub struct FetchBorrow<'a, T> {
-    ptr: NonNull<u8>,
+    ptr: Option<Ref<'a, NonNull<u8>>>,
     size: usize,
     _phantom: PhantomData<&'a T>
 }
@@ -17,14 +17,15 @@ where
 
     fn dangling() -> Self {
         FetchBorrow {
-            ptr: NonNull::dangling(),
+            ptr: None,
             size: 0,
             _phantom: PhantomData::default()
         }
     }
 
     unsafe fn get_item(&mut self, entity: crate::world::ArchetypeEntityId) -> Self::Item {
-        let ptr = NonNull::new_unchecked(self.ptr.as_ptr()
+        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()
@@ -68,10 +69,10 @@ where
     unsafe fn fetch<'a>(&self, _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.data;
+        let col_data = col.borrow_ptr();
 
         FetchBorrow {
-            ptr: NonNull::new_unchecked(col_data.as_ptr()),
+            ptr: Some(col_data),
             size: col.info.layout.size(),
             _phantom: PhantomData,
         }
@@ -100,7 +101,7 @@ impl<T: 'static> DefaultQuery for &T {
 
 /// A fetcher for mutably borrowing components from archetypes.
 pub struct FetchBorrowMut<'a, T> {
-    ptr: NonNull<u8>,
+    ptr: Option<RefMut<'a, NonNull<u8>>>,
     size: usize,
     _phantom: PhantomData<&'a T>
 }
@@ -113,14 +114,15 @@ where
 
     fn dangling() -> Self {
         FetchBorrowMut {
-            ptr: NonNull::dangling(),
+            ptr: None,
             size: 0,
             _phantom: PhantomData::default()
         }
     }
 
     unsafe fn get_item(&mut self, entity: crate::world::ArchetypeEntityId) -> Self::Item {
-        let ptr = NonNull::new_unchecked(self.ptr.as_ptr()
+        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()
@@ -164,10 +166,10 @@ where
     unsafe fn fetch<'a>(&self, _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.data;
+        let col_data = col.borrow_mut_ptr();
 
         FetchBorrowMut {
-            ptr: NonNull::new_unchecked(col_data.as_ptr()),
+            ptr: Some(col_data),
             size: col.info.layout.size(),
             _phantom: PhantomData,
         }
@@ -196,11 +198,11 @@ impl<T: 'static> DefaultQuery for &mut T {
 
 #[cfg(test)]
 mod tests {
-    use std::any::TypeId;
+    use std::{any::TypeId, mem::size_of, marker::PhantomData};
 
-    use crate::{world::World, archetype::Archetype, query::View, tests::Vec2};
+    use crate::{world::{World, Entity, EntityId}, archetype::{Archetype, ArchetypeId}, query::View, tests::Vec2, bundle::Bundle};
 
-    use super::{QueryBorrow, QueryBorrowMut};
+    use super::{QueryBorrow, QueryBorrowMut, FetchBorrowMut};
 
     /// Creates a world with two entities, one at Vec(10, 50) and
     /// the other at Vec2(25, 30).
@@ -262,4 +264,29 @@ mod tests {
         }
         println!("They were modified!");
     }
+
+    #[test]
+    fn only_one_mutable_borrow() {
+        let info = (Vec2::new(0.0, 0.0),).info();
+        let mut a = Archetype::from_bundle_info(ArchetypeId(0), info);
+
+        for i in 0..10 {
+            a.add_entity(Entity {
+                id: EntityId(i),
+                generation: 0,
+            }, (Vec2::rand(),));
+        }
+
+        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()),
+            size: size_of::<Vec2>(),
+            _phantom: PhantomData,
+        };
+
+        assert!(col.try_borrow_mut_ptr().is_err());
+        drop(bmut);
+        assert!(col.try_borrow_mut_ptr().is_ok());
+    }
 }