feat: set_missing

src/commands.rs

@@ -4,8 +4,9 @@ use alloc::{boxed::Box, format, vec::Vec};
 use anyhow::Context;
 
 use crate::{
-    buffer::MultiComponentBuffer, BatchSpawn, Component, ComponentDesc, ComponentValue, Entity,
-    EntityBuilder, World,
+    buffer::MultiComponentBuffer,
+    writer::{MissingDyn, SingleComponentWriter},
+    BatchSpawn, Component, ComponentDesc, ComponentValue, Entity, EntityBuilder, World,
 };
 
 type DeferFn = Box<dyn Fn(&mut World) -> anyhow::Result<()> + Send + Sync>;
@@ -25,6 +26,11 @@ enum Command {
         desc: ComponentDesc,
         offset: usize,
     },
+    SetMissing {
+        id: Entity,
+        desc: ComponentDesc,
+        offset: usize,
+    },
     /// Despawn an entity
     Despawn(Entity),
     /// Remove a component from an entity
@@ -55,6 +61,12 @@ impl fmt::Debug for Command {
                 .field("desc", desc)
                 .field("offset", offset)
                 .finish(),
+            Self::SetMissing { id, desc, offset } => f
+                .debug_struct("SetMissing")
+                .field("id", id)
+                .field("desc", desc)
+                .field("offset", offset)
+                .finish(),
             Self::Despawn(arg0) => f.debug_tuple("Despawn").field(arg0).finish(),
             Self::Remove {
                 id,
@@ -96,9 +108,7 @@ impl CommandBuffer {
         Self::default()
     }
 
-    /// Deferred set a component for `id`.
-    /// Unlike, [`World::set`] it does not return the old value as that is
-    /// not known at call time.
+    /// Set a component for `id`.
     pub fn set<T: ComponentValue>(
         &mut self,
         id: Entity,
@@ -115,6 +125,24 @@ impl CommandBuffer {
         self
     }
 
+    /// Set a component for `id` if it does not exist when the commandbuffer is applied.
+    ///
+    /// This avoid accidentally overwriting a component that was added by another system.
+    pub fn set_missing<T: ComponentValue>(
+        &mut self,
+        id: Entity,
+        component: Component<T>,
+        value: T,
+    ) -> &mut Self {
+        let offset = self.inserts.push(value);
+        self.commands.push(Command::SetMissing {
+            id,
+            desc: component.desc(),
+            offset,
+        });
+
+        self
+    }
     /// Deferred removal of a component for `id`.
     /// Unlike, [`World::remove`] it does not return the old value as that is
     /// not known at call time.
@@ -215,6 +243,13 @@ impl CommandBuffer {
                         .map_err(|v| v.into_anyhow())
                         .with_context(|| format!("Failed to set component {}", desc.name()))?;
                 },
+                Command::SetMissing { id, desc, offset } => unsafe {
+                    let value = self.inserts.take_dyn(offset);
+                    world
+                        .set_with_writer(id, SingleComponentWriter::new(desc, MissingDyn { value }))
+                        .map_err(|v| v.into_anyhow())
+                        .with_context(|| format!("Failed to set component {}", desc.name()))?;
+                },
                 Command::Despawn(id) => world
                     .despawn(id)
                     .map_err(|v| v.into_anyhow())
@@ -239,3 +274,42 @@ impl CommandBuffer {
         self.commands.clear()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::{component, FetchExt, Query};
+
+    use super::*;
+
+    #[test]
+    fn set_missing() {
+        use alloc::string::String;
+        use alloc::string::ToString;
+
+        component! {
+            a: String,
+        }
+
+        let mut world = World::new();
+        let mut cmd = CommandBuffer::new();
+
+        let mut query = Query::new((a().modified().satisfied(), a().cloned()));
+
+        let id = EntityBuilder::new().spawn(&mut world);
+
+        assert!(query.collect_vec(&world).is_empty());
+
+        cmd.set_missing(id, a(), "Foo".into())
+            .set_missing(id, a(), "Bar".into());
+
+        cmd.apply(&mut world).unwrap();
+
+        assert_eq!(query.collect_vec(&world), [(true, "Foo".to_string())]);
+        assert_eq!(query.collect_vec(&world), [(false, "Foo".to_string())]);
+
+        cmd.set_missing(id, a(), "Baz".into());
+        cmd.apply(&mut world).unwrap();
+
+        assert_eq!(query.collect_vec(&world), [(false, "Foo".to_string())]);
+    }
+}

src/entity_ref.rs

@@ -13,7 +13,7 @@ use crate::{
     error::MissingComponent,
     format::EntityFormatter,
     name,
-    writer::{EntityWriter, FnWriter, Replace, SingleComponentWriter, WriteDedup},
+    writer::{EntityWriter, FnWriter, Missing, Replace, SingleComponentWriter, WriteDedup},
     Component, ComponentKey, ComponentValue, Entity, RelationExt, World,
 };
 use crate::{RelationIter, RelationIterMut};
@@ -153,6 +153,17 @@ impl<'a> EntityRefMut<'a> {
         .left()
     }
 
+    /// Set a component for the entity only if it is missing.
+    ///
+    /// Does not disturb or generate a change event if the component is present
+    pub fn set_missing<T: ComponentValue>(&mut self, component: Component<T>, value: T) -> bool {
+        self.set_with_writer(SingleComponentWriter::new(
+            component.desc(),
+            Missing { value },
+        ))
+        .is_right()
+    }
+
     /// Set a component for the entity.
     ///
     /// Does not trigger a modification event if the value is the same
@@ -652,6 +663,35 @@ mod test {
         assert_eq!(query.collect_vec(&world), ["Bar"]);
     }
 
+    #[test]
+    fn set_missing() {
+        use alloc::string::String;
+        use alloc::string::ToString;
+
+        component! {
+            a: String,
+        }
+
+        let mut world = World::new();
+
+        let mut query = Query::new((a().modified().satisfied(), a().cloned()));
+
+        let id = EntityBuilder::new().spawn(&mut world);
+
+        assert!(query.collect_vec(&world).is_empty());
+
+        let mut entity = world.entity_mut(id).unwrap();
+        assert!(entity.set_missing(a(), "Foo".into()));
+
+        assert_eq!(query.collect_vec(&world), [(true, "Foo".to_string())]);
+        assert_eq!(query.collect_vec(&world), [(false, "Foo".to_string())]);
+
+        let mut entity = world.entity_mut(id).unwrap();
+        assert!(!entity.set_missing(a(), "Bar".into()));
+
+        assert_eq!(query.collect_vec(&world), [(false, "Foo".to_string())]);
+    }
+
     #[test]
     fn update_dedup() {
         use alloc::string::String;

src/writer.rs

@@ -220,6 +220,30 @@ impl<T: ComponentValue> ComponentPusher for Replace<T> {
     }
 }
 
+pub(crate) struct Missing<T: ComponentValue> {
+    pub(crate) value: T,
+}
+
+impl<T: ComponentValue> ComponentUpdater for Missing<T> {
+    type Updated = ();
+
+    unsafe fn update(self, _: &mut CellData, _: Slot, _: Entity, _: u32) {}
+}
+
+impl<T: ComponentValue> ComponentPusher for Missing<T> {
+    type Pushed = ();
+
+    unsafe fn push(mut self, data: &mut CellData, id: Entity, tick: u32) {
+        let slot = data.storage.len();
+
+        data.storage.extend(&mut self.value as *mut T as *mut u8, 1);
+
+        mem::forget(self.value);
+
+        data.set_added(&[id], Slice::single(slot), tick);
+    }
+}
+
 pub(crate) struct WriteDedup<T: ComponentValue> {
     pub(crate) value: T,
 }
@@ -290,6 +314,32 @@ impl ComponentPusher for ReplaceDyn {
     }
 }
 
+pub(crate) struct MissingDyn {
+    pub(crate) value: *mut u8,
+}
+
+impl ComponentUpdater for MissingDyn {
+    type Updated = ();
+
+    unsafe fn update(self, data: &mut CellData, _: Slot, _: Entity, _: u32) {
+        let desc = data.storage.desc();
+        unsafe {
+            desc.drop(self.value);
+        }
+    }
+}
+
+impl ComponentPusher for MissingDyn {
+    type Pushed = ();
+
+    unsafe fn push(self, data: &mut CellData, id: Entity, tick: u32) {
+        let slot = data.storage.len();
+        data.storage.extend(self.value, 1);
+
+        data.set_added(&[id], Slice::single(slot), tick);
+    }
+}
+
 pub(crate) struct Buffered<'b> {
     pub(crate) buffer: &'b mut ComponentBuffer,
 }