ecs: initial component change tracking

(changing entity archetypes currently breaks tracking)

crates/bevy_ecs/hecs/src/archetype.rs

@@ -50,19 +50,7 @@ pub struct Archetype {
 impl Archetype {
     #[allow(missing_docs)]
     pub fn new(types: Vec<TypeInfo>) -> Self {
-        debug_assert!(
-            types.windows(2).all(|x| x[0] < x[1]),
-            "type info unsorted or contains duplicates"
-        );
-        Self {
-            types,
-            state: HashMap::default(),
-            entities: Box::new([]),
-            len: 0,
-            data: UnsafeCell::new(NonNull::dangling()),
-            data_size: 0,
-            grow_size: 64,
-        }
+        Self::with_grow(types, 64)
     }
 
     #[allow(missing_docs)]
@@ -71,9 +59,13 @@ impl Archetype {
             types.windows(2).all(|x| x[0] < x[1]),
             "type info unsorted or contains duplicates"
         );
+        let mut state = HashMap::with_capacity(types.len());
+        for ty in &types {
+            state.insert(ty.id, TypeState::new());
+        }
         Self {
+            state,
             types,
-            state: HashMap::default(),
             entities: Box::new([]),
             len: 0,
             data: UnsafeCell::new(NonNull::dangling()),
@@ -97,7 +89,8 @@ impl Archetype {
         self.len = 0;
     }
 
-    pub(crate) fn has<T: Component>(&self) -> bool {
+    #[allow(missing_docs)]
+    pub fn has<T: Component>(&self) -> bool {
         self.has_dynamic(TypeId::of::<T>())
     }
 
@@ -105,6 +98,7 @@ impl Archetype {
         self.state.contains_key(&id)
     }
 
+    // TODO: this should be unsafe i think
     #[allow(missing_docs)]
     #[inline]
     pub fn get<T: Component>(&self) -> Option<NonNull<T>> {
@@ -116,6 +110,29 @@ impl Archetype {
         })
     }
 
+    // TODO: this should be unsafe i think
+    #[allow(missing_docs)]
+    #[inline]
+    pub fn get_with_modified<T: Component>(&self) -> Option<(NonNull<T>, NonNull<bool>)> {
+        let state = self.state.get(&TypeId::of::<T>())?;
+        Some(unsafe {
+            (
+                NonNull::new_unchecked(
+                    (*self.data.get()).as_ptr().add(state.offset).cast::<T>() as *mut T
+                ),
+                NonNull::new_unchecked(state.modified_entities.as_ptr() as *mut bool),
+            )
+        })
+    }
+
+    // TODO: this should be unsafe i think
+    #[allow(missing_docs)]
+    #[inline]
+    pub fn get_modified<T: Component>(&self) -> Option<NonNull<bool>> {
+        let state = self.state.get(&TypeId::of::<T>())?;
+        Some(unsafe { NonNull::new_unchecked(state.modified_entities.as_ptr() as *mut bool) })
+    }
+
     #[allow(missing_docs)]
     pub fn borrow<T: Component>(&self) {
         if self
@@ -212,6 +229,12 @@ impl Archetype {
         self.entities.len() as u32
     }
 
+    pub fn clear_trackers(&mut self) {
+        for type_state in self.state.values_mut() {
+            type_state.clear_trackers();
+        }
+    }
+
     fn grow(&mut self, increment: u32) {
         unsafe {
             let old_count = self.len as usize;
@@ -220,11 +243,17 @@ impl Archetype {
             new_entities[0..old_count].copy_from_slice(&self.entities[0..old_count]);
             self.entities = new_entities;
 
+            for type_state in self.state.values_mut() {
+                type_state.modified_entities.resize_with(count, || false);
+            }
+
             let old_data_size = mem::replace(&mut self.data_size, 0);
-            let mut state = HashMap::with_capacity(self.types.len());
+            let mut old_offsets = Vec::with_capacity(self.types.len());
             for ty in &self.types {
                 self.data_size = align(self.data_size, ty.layout.align());
-                state.insert(ty.id, TypeState::new(self.data_size));
+                let ty_state = self.state.get_mut(&ty.id).unwrap();
+                old_offsets.push(ty_state.offset);
+                ty_state.offset = self.data_size;
                 self.data_size += ty.layout.size() * count;
             }
             let new_data = if self.data_size == 0 {
@@ -240,9 +269,9 @@ impl Archetype {
                 .unwrap()
             };
             if old_data_size != 0 {
-                for ty in &self.types {
-                    let old_off = self.state.get(&ty.id).unwrap().offset;
-                    let new_off = state.get(&ty.id).unwrap().offset;
+                for (i, ty) in self.types.iter().enumerate() {
+                    let old_off = old_offsets[i];
+                    let new_off = self.state.get(&ty.id).unwrap().offset;
                     ptr::copy_nonoverlapping(
                         (*self.data.get()).as_ptr().add(old_off),
                         new_data.as_ptr().add(new_off),
@@ -252,7 +281,6 @@ impl Archetype {
             }
 
             self.data = UnsafeCell::new(new_data);
-            self.state = state;
         }
     }
 
@@ -266,6 +294,7 @@ impl Archetype {
                 .as_ptr();
             (ty.drop)(removed);
             if index != last {
+                // TODO: copy component tracker state here
                 ptr::copy_nonoverlapping(
                     self.get_dynamic(ty.id, ty.layout.size(), last)
                         .unwrap()
@@ -297,6 +326,7 @@ impl Archetype {
                 .unwrap()
                 .as_ptr();
             f(moved, ty.id(), ty.layout().size());
+            // TODO: copy component tracker state here
             if index != last {
                 ptr::copy_nonoverlapping(
                     self.get_dynamic(ty.id, ty.layout.size(), last)
@@ -352,13 +382,21 @@ impl Drop for Archetype {
 struct TypeState {
     offset: usize,
     borrow: AtomicBorrow,
+    modified_entities: Vec<bool>,
 }
 
 impl TypeState {
-    fn new(offset: usize) -> Self {
+    fn new() -> Self {
         Self {
-            offset,
+            offset: 0,
             borrow: AtomicBorrow::new(),
+            modified_entities: Vec::new(),
+        }
+    }
+
+    fn clear_trackers(&mut self) {
+        for modified in self.modified_entities.iter_mut() {
+            *modified = false;
         }
     }
 }

crates/bevy_ecs/hecs/src/query.rs

@@ -42,6 +42,11 @@ pub trait Fetch<'a>: Sized {
     /// Release dynamic borrows acquired by `borrow`
     fn release(archetype: &Archetype);
 
+    /// if this returns true, the current item will be skipped during iteration
+    unsafe fn should_skip(&self) -> bool {
+        false
+    }
+
     /// Access the next item in this archetype without bounds checking
     ///
     /// # Safety
@@ -519,11 +524,20 @@ struct ChunkIter<Q: Query> {
 
 impl<Q: Query> ChunkIter<Q> {
     unsafe fn next<'a, 'w>(&mut self) -> Option<<Q::Fetch as Fetch<'a>>::Item> {
-        if self.len == 0 {
-            return None;
+        loop {
+            if self.len == 0 {
+                return None;
+            }
+
+            self.len -= 1;
+            if self.fetch.should_skip() {
+                // we still need to progress the iterator
+                let _ = self.fetch.next();
+                continue;
+            }
+
+            break Some(self.fetch.next())
         }
-        self.len -= 1;
-        Some(self.fetch.next())
     }
 }
 

crates/bevy_ecs/hecs/src/world.rs

@@ -559,6 +559,13 @@ impl World {
     pub fn get_entity_location(&self, entity: Entity) -> Option<Location> {
         self.entities.get(entity).ok()
     }
+
+    /// Clears each entity's tracker state. For example, each entity's component "modified" state will be reset to `false`. 
+    pub fn clear_trackers(&mut self) {
+        for archetype in self.archetypes.iter_mut() {
+            archetype.clear_trackers();
+        }
+    }
 }
 
 unsafe impl Send for World {}

crates/bevy_ecs/src/lib.rs

@@ -2,12 +2,12 @@ pub use hecs::{Query as HecsQuery, *};
 mod resource;
 mod schedule;
 mod system;
-mod world_builder;
+mod world;
 
 pub use resource::*;
 pub use schedule::*;
 pub use system::{*, Query};
-pub use world_builder::*;
+pub use world::*;
 
 pub mod prelude {
     pub use crate::{
@@ -15,7 +15,7 @@ pub mod prelude {
         system::{
             Commands, IntoForEachSystem, IntoQuerySystem, IntoThreadLocalSystem, Query, System,
         },
-        world_builder::WorldBuilderSource,
+        world::{WorldBuilderSource, ComMut},
         Bundle, Component, Entity, Ref, RefMut, With, Without, World,
     };
 }

crates/bevy_ecs/src/resource/resource_query.rs

@@ -62,7 +62,7 @@ impl<'a, T: Component> Deref for Res<'a, T> {
     }
 }
 
-/// Unique borrow of an entity's component
+/// Unique borrow of a resource
 pub struct ResMut<'a, T: Component> {
     archetype: &'a Archetype,
     target: NonNull<T>,

crates/bevy_ecs/src/world/component.rs

@@ -0,0 +1,172 @@
+use crate::{Archetype, Component, HecsQuery};
+use hecs::{Access, Fetch};
+use std::{
+    marker::PhantomData,
+    ops::{Deref, DerefMut},
+    ptr::NonNull,
+};
+
+/// Unique borrow of an entity's component
+pub struct ComMut<'a, T: Component> {
+    value: &'a mut T,
+    modified: &'a mut bool,
+}
+
+unsafe impl<T: Component> Send for ComMut<'_, T> {}
+unsafe impl<T: Component> Sync for ComMut<'_, T> {}
+
+impl<'a, T: Component> Deref for ComMut<'a, T> {
+    type Target = T;
+    fn deref(&self) -> &T {
+        self.value
+    }
+}
+
+impl<'a, T: Component> DerefMut for ComMut<'a, T> {
+    fn deref_mut(&mut self) -> &mut T {
+        *self.modified = true;
+        self.value
+    }
+}
+
+impl<'a, T: Component> HecsQuery for ComMut<'a, T> {
+    type Fetch = FetchComMut<T>;
+}
+#[doc(hidden)]
+pub struct FetchComMut<T>(NonNull<T>, NonNull<bool>);
+
+impl<'a, T: Component> Fetch<'a> for FetchComMut<T> {
+    type Item = ComMut<'a, T>;
+
+    fn access(archetype: &Archetype) -> Option<Access> {
+        if archetype.has::<T>() {
+            Some(Access::Write)
+        } else {
+            None
+        }
+    }
+
+    fn borrow(archetype: &Archetype) {
+        archetype.borrow_mut::<T>();
+    }
+    unsafe fn get(archetype: &'a Archetype, offset: usize) -> Option<Self> {
+        archetype
+            .get_with_modified::<T>()
+            .map(|(components, modified)| {
+                Self(
+                    NonNull::new_unchecked(components.as_ptr().add(offset)),
+                    NonNull::new_unchecked(modified.as_ptr().add(offset)),
+                )
+            })
+    }
+    fn release(archetype: &Archetype) {
+        archetype.release_mut::<T>();
+    }
+
+    unsafe fn next(&mut self) -> ComMut<'a, T> {
+        let component = self.0.as_ptr();
+        let modified = self.1.as_ptr();
+        self.0 = NonNull::new_unchecked(component.add(1));
+        self.1 = NonNull::new_unchecked(modified.add(1));
+        ComMut {
+            value: &mut *component,
+            modified: &mut *modified,
+        }
+    }
+}
+
+pub struct Changed<T, Q>(PhantomData<(Q, fn(T))>);
+
+impl<T: Component, Q: HecsQuery> HecsQuery for Changed<T, Q> {
+    type Fetch = FetchChanged<T, Q::Fetch>;
+}
+
+#[doc(hidden)]
+pub struct FetchChanged<T, F>(F, PhantomData<fn(T)>, NonNull<bool>);
+
+impl<'a, T: Component, F: Fetch<'a>> Fetch<'a> for FetchChanged<T, F> {
+    type Item = F::Item;
+
+    fn access(archetype: &Archetype) -> Option<Access> {
+        if archetype.has::<T>() {
+            F::access(archetype)
+        } else {
+            None
+        }
+    }
+
+    fn borrow(archetype: &Archetype) {
+        F::borrow(archetype)
+    }
+    unsafe fn get(archetype: &'a Archetype, offset: usize) -> Option<Self> {
+        if !archetype.has::<T>() {
+            return None;
+        }
+        Some(Self(
+            F::get(archetype, offset)?,
+            PhantomData,
+            NonNull::new_unchecked(archetype.get_modified::<T>()?.as_ptr().add(offset)),
+        ))
+    }
+    fn release(archetype: &Archetype) {
+        F::release(archetype)
+    }
+
+    unsafe fn should_skip(&self) -> bool {
+        // skip if the current item wasn't changed
+        !*self.2.as_ref()
+    }
+
+    unsafe fn next(&mut self) -> F::Item {
+        self.2 = NonNull::new_unchecked(self.2.as_ptr().add(1));
+        self.0.next()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{Changed, ComMut};
+    use hecs::{Entity, World};
+
+    struct A(usize);
+    struct B;
+    struct C;
+
+    #[test]
+    fn modified_trackers() {
+        let mut world = World::default();
+        let e1 = world.spawn((A(0), B));
+        world.spawn((A(0), B));
+        let e3 = world.spawn((A(0), B));
+        world.spawn((A(0), B));
+
+        for (i, mut a) in world.query::<ComMut<A>>().iter().enumerate() {
+            if i % 2 == 0 {
+                a.0 += 1;
+            }
+        }
+
+        let changed_entities = world
+            .query::<Changed<A, Entity>>()
+            .iter()
+            .collect::<Vec<Entity>>();
+        assert_eq!(changed_entities, vec![e1, e3]);
+
+        // ensure changing an entity's archetypes also moves its modified state
+        world.insert(e1, (C,)).unwrap();
+        
+        let changed_entities = world
+            .query::<Changed<A, Entity>>()
+            .iter()
+            .collect::<Vec<Entity>>();
+        assert_eq!(changed_entities, vec![e1, e3]);
+
+        world.clear_trackers();
+
+        assert!(world
+            .query::<Changed<A, Entity>>()
+            .iter()
+            .collect::<Vec<Entity>>()
+            .is_empty());
+    }
+}

crates/bevy_ecs/src/world/mod.rs

@@ -0,0 +1,5 @@
+mod component;
+mod world_builder;
+
+pub use world_builder::*;
+pub use component::*;