Fix memory leak

ecs/src/commands/command.rs

@@ -1,4 +1,4 @@
-use std::{alloc::Layout, any::TypeId, borrow::Cow};
+use std::{alloc::Layout, any::TypeId, borrow::Cow, mem, ptr};
 
 use crate::{Entity, World};
 
@@ -8,17 +8,17 @@ pub enum Command {
     AddComponent {
         entity: Entity,
         type_id: TypeId,
-        component: *mut u8, // TODO: watch out for memory leak?
+        component: *mut u8, // Is set to null after ownership is transferred to the world
         name: Cow<'static, str>,
         layout: Layout,
         drop: unsafe fn(*mut u8),
     },
 }
 
 impl Command {
-    pub(super) fn execute(self, world: &mut World) {
-        match self {
-            Command::Despawn(entity) => {
+    pub(super) fn execute(mut self, world: &mut World) {
+        match &mut self {
+            &mut Command::Despawn(entity) => {
                 world.despawn(entity);
             }
             Command::AddComponent {
@@ -29,19 +29,34 @@ impl Command {
                 layout,
                 drop,
             } => {
+                let name = mem::take(name);
                 let comp_id = match world
                     .component_registry_mut()
-                    .component_id_from_type_id(type_id)
+                    .component_id_from_type_id(*type_id)
                 {
                     Some(id) => id,
                     None => unsafe {
                         world
                             .component_registry_mut()
-                            .register_raw(type_id, name, layout, drop)
+                            .register_raw(*type_id, name, *layout, *drop)
                     },
                 };
-                unsafe { world.add_raw(entity, component, comp_id) };
+                unsafe { world.add_raw(*entity, *component, comp_id) };
+                *component = ptr::null_mut();
             }
         }
     }
 }
+
+impl Drop for Command {
+    fn drop(&mut self) {
+        match self {
+            &mut Command::AddComponent {
+                component, drop, ..
+            } if !component.is_null() => unsafe {
+                drop(component);
+            },
+            _ => {}
+        }
+    }
+}

ecs/src/lib.rs

@@ -755,6 +755,22 @@ mod tests {
         });
     }
 
+    #[test]
+    fn drop_command_buffer_while_it_owns_components() {
+        let world = World::default();
+        let counter = Rc::new(Cell::new(0));
+        {
+            let mut command_buffer = CommandBuffer::new();
+            let mut commands = Commands::new(&mut command_buffer, world.entities());
+
+            let e1 = commands.spawn();
+            commands.add(e1, Counter::named(counter.clone(), "a"));
+            assert_eq!(counter.get(), 1);
+        }
+
+        assert_eq!(counter.get(), 0);
+    }
+
     #[derive(Debug)]
     struct Counter(Rc<Cell<usize>>, &'static str);
     impl Counter {