Use global borrow flags to avoid per-archetype borrow checking overhead.

This reduces the number of dynamic borrow checks to one per component instead of one per archetype, thereby avoiding dynamic allocations for the reference wrappers. However, it also means that e.g. World::get_mut can now fail due overlapping borrows for entities in different archetypes which would have worked before.
adamreichold · Jan 15, 2022 · f20203b · f20203b
1 parent e74d5e2
commit f20203b
Show file tree

Hide file tree

Showing 7 changed files with 337 additions and 268 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 name = "rs-ecs"
 description = "reasonably simple entity component system"
-version = "0.6.0"
+version = "0.6.1"
 edition = "2018"
 rust-version = "1.51"
 authors = ["Adam Reichold <[email protected]>"]

diff --git a/src/archetype.rs b/src/archetype.rs
@@ -1,9 +1,7 @@
 use std::alloc::{alloc, dealloc, handle_alloc_error, Layout};
-use std::any::{type_name, TypeId};
-use std::cell::UnsafeCell;
+use std::any::TypeId;
 use std::cmp::Reverse;
 use std::mem::{align_of, needs_drop};
-use std::ops::{Deref, DerefMut};
 use std::ptr::{copy_nonoverlapping, drop_in_place};
 use std::slice::from_raw_parts_mut;
 
@@ -233,26 +231,6 @@ impl Archetype {
         ty
     }
 
-    pub unsafe fn borrow<C>(&self, ty: u16) -> Ref<'_>
-    where
-        C: 'static,
-    {
-        let ty = self.type_metadata::<C>(ty);
-
-        Ref::new(&ty.borrow)
-            .unwrap_or_else(|| panic!("Component {} already borrowed", type_name::<C>()))
-    }
-
-    pub unsafe fn borrow_mut<C>(&self, ty: u16) -> RefMut<'_>
-    where
-        C: 'static,
-    {
-        let ty = self.type_metadata::<C>(ty);
-
-        RefMut::new(&ty.borrow)
-            .unwrap_or_else(|| panic!("Component {} already borrowed", type_name::<C>()))
-    }
-
     pub unsafe fn base_pointer<C>(&self, ty: u16) -> *mut C
     where
         C: 'static,
@@ -275,166 +253,36 @@ impl Archetype {
 }
 
 impl Archetype {
-    pub unsafe fn drop<C>(&mut self, idx: u32)
-    where
-        C: 'static,
-    {
-        if needs_drop::<C>() {
-            if let Some(ty) = self.find::<C>() {
-                let ty = self.types.get_unchecked(ty as usize);
-
-                let ptr = ty.base_pointer.add(ty.layout.size() * idx as usize);
-
-                (ty.drop)(ptr, 1);
-            }
-        }
-    }
-}
-
-impl Archetype {
-    pub unsafe fn get_raw<C>(&mut self, idx: u32) -> *mut C
+    pub unsafe fn get<C>(&self, idx: u32) -> *mut C
     where
         C: 'static,
     {
         let ty = self.find::<C>().unwrap();
         self.pointer::<C>(ty, idx)
     }
 
-    pub unsafe fn get<C>(&self, idx: u32) -> Option<Comp<'_, C>>
-    where
-        C: 'static,
-    {
-        let ty = self.find::<C>()?;
-        let _ref = self.borrow::<C>(ty);
-        let ptr = self.pointer::<C>(ty, idx);
-
-        let val = &*ptr;
-
-        Some(Comp { _ref, val })
-    }
-
-    pub unsafe fn get_mut<C>(&self, idx: u32) -> Option<CompMut<'_, C>>
+    pub unsafe fn drop<C>(&mut self, idx: u32)
     where
         C: 'static,
     {
-        let ty = self.find::<C>()?;
-        let _ref = self.borrow_mut::<C>(ty);
-        let ptr = self.pointer::<C>(ty, idx);
-
-        let val = &mut *ptr;
-
-        Some(CompMut { _ref, val })
-    }
-}
-
-/// An immutable borrow of a component.
-pub struct Comp<'a, C> {
-    _ref: Ref<'a>,
-    val: &'a C,
-}
-
-impl<C> Deref for Comp<'_, C> {
-    type Target = C;
-
-    fn deref(&self) -> &C {
-        self.val
-    }
-}
-
-/// A mutable borrow of a component.
-pub struct CompMut<'a, C> {
-    _ref: RefMut<'a>,
-    val: &'a mut C,
-}
-
-impl<C> Deref for CompMut<'_, C> {
-    type Target = C;
-
-    fn deref(&self) -> &C {
-        self.val
-    }
-}
-
-impl<C> DerefMut for CompMut<'_, C> {
-    fn deref_mut(&mut self) -> &mut C {
-        self.val
-    }
-}
-
-pub struct Ref<'a>(&'a UnsafeCell<isize>);
-
-impl<'a> Ref<'a> {
-    fn new(borrow: &'a UnsafeCell<isize>) -> Option<Self> {
-        let readers = unsafe { &mut *borrow.get() };
-
-        let new_readers = readers.wrapping_add(1);
-
-        if new_readers <= 0 {
-            cold();
-            return None;
-        }
-
-        *readers = new_readers;
-
-        Some(Self(borrow))
-    }
-}
-
-impl Drop for Ref<'_> {
-    fn drop(&mut self) {
-        unsafe {
-            *self.0.get() -= 1;
-        }
-    }
-}
-
-pub struct RefMut<'a>(&'a UnsafeCell<isize>);
-
-impl<'a> RefMut<'a> {
-    fn new(borrow: &'a UnsafeCell<isize>) -> Option<Self> {
-        let writers = unsafe { &mut *borrow.get() };
-
-        if *writers != 0 {
-            cold();
-            return None;
-        }
-
-        *writers = -1;
+        if needs_drop::<C>() {
+            if let Some(ty) = self.find::<C>() {
+                let ty = self.types.get_unchecked(ty as usize);
 
-        Some(Self(borrow))
-    }
-}
+                let ptr = ty.base_pointer.add(ty.layout.size() * idx as usize);
 
-impl Drop for RefMut<'_> {
-    fn drop(&mut self) {
-        unsafe {
-            *self.0.get() = 0;
+                (ty.drop)(ptr, 1);
+            }
         }
     }
 }
 
-#[cold]
-#[inline(always)]
-fn cold() {}
-
+#[derive(Clone, Copy)]
 struct TypeMetadata {
     id: TypeId,
     layout: Layout,
     drop: unsafe fn(*mut u8, usize),
     base_pointer: *mut u8,
-    borrow: UnsafeCell<isize>,
-}
-
-impl Clone for TypeMetadata {
-    fn clone(&self) -> Self {
-        Self {
-            id: self.id,
-            layout: self.layout,
-            drop: self.drop,
-            base_pointer: self.layout.align() as *mut u8,
-            borrow: Default::default(),
-        }
-    }
 }
 
 impl TypeMetadata {
@@ -451,7 +299,6 @@ impl TypeMetadata {
             layout: Layout::new::<C>(),
             drop: drop::<C>,
             base_pointer: align_of::<C>() as *mut u8,
-            borrow: Default::default(),
         }
     }
 }
@@ -460,6 +307,10 @@ impl TypeMetadata {
 pub struct TypeMetadataSet(Vec<TypeMetadata>);
 
 impl TypeMetadataSet {
+    pub fn ids(&self) -> impl ExactSizeIterator<Item = TypeId> + '_ {
+        self.0.iter().map(|ty| ty.id)
+    }
+
     pub fn insert<C>(&mut self)
     where
         C: 'static,
@@ -554,7 +405,7 @@ mod tests {
                 let ent = archetype.alloc();
 
                 archetype
-                    .get_raw::<CountDrops>(ent)
+                    .get::<CountDrops>(ent)
                     .write(CountDrops(drops.clone()));
             }
         }

diff --git a/src/borrow_flags.rs b/src/borrow_flags.rs
@@ -0,0 +1,116 @@
+use std::any::{type_name, TypeId};
+use std::cell::UnsafeCell;
+
+use crate::archetype::TypeMetadataSet;
+
+#[derive(Default)]
+pub struct BorrowFlags(Vec<(TypeId, UnsafeCell<isize>)>);
+
+impl BorrowFlags {
+    pub fn insert(&mut self, types: &TypeMetadataSet) {
+        for id in types.ids() {
+            if let Err(idx) = self.0.binary_search_by_key(&id, |(id, _)| *id) {
+                assert!(self.0.len() < u16::MAX as usize);
+
+                self.0.insert(idx, (id, UnsafeCell::new(0)));
+            }
+        }
+    }
+
+    pub fn find<C>(&self) -> Option<u16>
+    where
+        C: 'static,
+    {
+        self.0
+            .binary_search_by_key(&TypeId::of::<C>(), |(id, _)| *id)
+            .map(|idx| idx as u16)
+            .ok()
+    }
+
+    unsafe fn flag<C>(&self, ty: u16) -> &UnsafeCell<isize>
+    where
+        C: 'static,
+    {
+        let ty = ty as usize;
+        debug_assert!(ty < self.0.len());
+        let (id, flag) = self.0.get_unchecked(ty);
+        debug_assert_eq!(id, &TypeId::of::<C>());
+
+        flag
+    }
+
+    pub unsafe fn borrow<C>(&self, ty: u16) -> Ref<'_>
+    where
+        C: 'static,
+    {
+        let flag = self.flag::<C>(ty);
+
+        Ref::new(flag).unwrap_or_else(|| panic!("Component {} already borrowed", type_name::<C>()))
+    }
+
+    pub unsafe fn borrow_mut<C>(&self, ty: u16) -> RefMut<'_>
+    where
+        C: 'static,
+    {
+        let flag = self.flag::<C>(ty);
+
+        RefMut::new(flag)
+            .unwrap_or_else(|| panic!("Component {} already borrowed", type_name::<C>()))
+    }
+}
+
+pub struct Ref<'a>(&'a UnsafeCell<isize>);
+
+impl<'a> Ref<'a> {
+    pub(crate) fn new(borrow: &'a UnsafeCell<isize>) -> Option<Self> {
+        let readers = unsafe { &mut *borrow.get() };
+
+        let new_readers = readers.wrapping_add(1);
+
+        if new_readers <= 0 {
+            cold();
+            return None;
+        }
+
+        *readers = new_readers;
+
+        Some(Self(borrow))
+    }
+}
+
+impl Drop for Ref<'_> {
+    fn drop(&mut self) {
+        unsafe {
+            *self.0.get() -= 1;
+        }
+    }
+}
+
+pub struct RefMut<'a>(&'a UnsafeCell<isize>);
+
+impl<'a> RefMut<'a> {
+    pub(crate) fn new(borrow: &'a UnsafeCell<isize>) -> Option<Self> {
+        let writers = unsafe { &mut *borrow.get() };
+
+        if *writers != 0 {
+            cold();
+            return None;
+        }
+
+        *writers = -1;
+
+        Some(Self(borrow))
+    }
+}
+
+impl Drop for RefMut<'_> {
+    fn drop(&mut self) {
+        unsafe {
+            *self.0.get() = 0;
+        }
+    }
+}
+
+#[cold]
+#[inline(always)]
+fn cold() {}
diff --git a/src/lib.rs b/src/lib.rs
@@ -28,15 +28,15 @@
 #![warn(missing_docs)]
 
 mod archetype;
+mod borrow_flags;
 mod query;
 mod resources;
 mod world;
 
 pub use crate::{
-    archetype::{Comp, CompMut},
     query::{Matches, Query, QueryIter, QueryMap, QueryRef, QuerySpec, With, Without},
     resources::{Res, ResMut, Resources},
-    world::{Entity, World},
+    world::{Comp, CompMut, Entity, World},
 };
 
 #[cfg(feature = "rayon")]