SystemBuilder: better handling of system creation with non-default kind while deferred

kind_id() checks for existing phases and removes them, but the addition of the default
OnUpdate phase has not been processed yet by that point (if deferred).

Special case for existing entities being passed in, because they effectively have the same
issue; the system builder cannot reliably know whether a kind is going to be set on the
existing entity in the next defer flush.

Author: Georgiy-Tugai

flecs_ecs/src/addons/system/system_builder.rs

@@ -13,6 +13,9 @@ where
     pub(crate) desc: sys::ecs_system_desc_t,
     term_builder: TermBuilder,
     world: WorldRef<'a>,
+    /// Skip setting default phase (`OnUpdate`) if `kind` was set,
+    /// or an existing entity was passed in (via [`Self::new_from_desc`])
+    kind_set: bool,
     _phantom: core::marker::PhantomData<&'a T>,
 }
 
@@ -26,23 +29,14 @@ where
             desc: Default::default(),
             term_builder: TermBuilder::default(),
             world: world.into(),
+            kind_set: false,
             _phantom: core::marker::PhantomData,
         };
 
         obj.desc.entity = unsafe { sys::ecs_entity_init(obj.world_ptr_mut(), &Default::default()) };
 
         T::populate(&mut obj);
 
-        #[cfg(feature = "flecs_pipeline")]
-        unsafe {
-            sys::ecs_add_id(
-                world.world_ptr_mut(),
-                obj.desc.entity,
-                ecs_dependson(ECS_ON_UPDATE),
-            );
-            sys::ecs_add_id(world.world_ptr_mut(), obj.desc.entity, ECS_ON_UPDATE);
-        }
-
         obj
     }
 
@@ -51,26 +45,20 @@ where
             desc,
             term_builder: TermBuilder::default(),
             world: world.into(),
+            kind_set: false,
             _phantom: core::marker::PhantomData,
         };
 
         if obj.desc.entity == 0 {
             obj.desc.entity =
                 unsafe { sys::ecs_entity_init(obj.world_ptr_mut(), &Default::default()) };
+        } else {
+            // Can't make assumptions about the kind on an existing entity.
+            obj.kind_set = true;
         }
 
         T::populate(&mut obj);
 
-        #[cfg(feature = "flecs_pipeline")]
-        unsafe {
-            sys::ecs_add_id(
-                world.world_ptr_mut(),
-                obj.desc.entity,
-                ecs_dependson(ECS_ON_UPDATE),
-            );
-            sys::ecs_add_id(world.world_ptr_mut(), obj.desc.entity, ECS_ON_UPDATE);
-        }
-
         obj
     }
 
@@ -82,6 +70,7 @@ where
             desc: Default::default(),
             term_builder: TermBuilder::default(),
             world: world.into(),
+            kind_set: false,
             _phantom: core::marker::PhantomData,
         };
 
@@ -95,15 +84,6 @@ where
 
         T::populate(&mut obj);
 
-        #[cfg(feature = "flecs_pipeline")]
-        unsafe {
-            sys::ecs_add_id(
-                world.world_ptr_mut(),
-                obj.desc.entity,
-                ecs_dependson(ECS_ON_UPDATE),
-            );
-            sys::ecs_add_id(world.world_ptr_mut(), obj.desc.entity, ECS_ON_UPDATE);
-        }
         obj
     }
 
@@ -134,6 +114,7 @@ where
             if phase != 0 {
                 sys::ecs_add_id(self.world_ptr_mut(), self.desc.entity, ecs_dependson(phase));
                 sys::ecs_add_id(self.world_ptr_mut(), self.desc.entity, phase);
+                self.kind_set = true;
             }
         };
         self
@@ -249,6 +230,22 @@ where
     /// * C++ API: `node_builder::build`
     #[doc(alias = "node_builder::build")]
     fn build(&mut self) -> Self::BuiltType {
+        #[cfg(feature = "flecs_pipeline")]
+        if !self.kind_set {
+            unsafe {
+                sys::ecs_add_id(
+                    self.world().world_ptr_mut(),
+                    self.desc.entity,
+                    ecs_dependson(ECS_ON_UPDATE),
+                );
+                sys::ecs_add_id(
+                    self.world().world_ptr_mut(),
+                    self.desc.entity,
+                    ECS_ON_UPDATE,
+                );
+            }
+        }
+
         let system = System::new(self.world(), self.desc);
         for s in self.term_builder.str_ptrs_to_free.iter_mut() {
             unsafe { core::mem::ManuallyDrop::drop(s) };

flecs_ecs/tests/flecs/system_test.rs

@@ -1652,6 +1652,26 @@ fn system_custom_pipeline_w_kind() {
     });
 }
 
+#[test]
+fn system_kind_while_deferred() {
+    let world = World::new();
+
+    let sys = world.defer(|| {
+        world
+            .system::<()>()
+            .kind_id(flecs::pipeline::OnValidate::ID)
+            .run(|_| {})
+    });
+
+    world.progress();
+
+    let sys = sys.entity_view(&world);
+    assert!(sys.has_id(flecs::pipeline::OnValidate::ID));
+    assert!(sys.has_first::<flecs::DependsOn>(flecs::pipeline::OnValidate::ID));
+    assert!(!sys.has_id(flecs::pipeline::OnUpdate::ID));
+    assert!(!sys.has_first::<flecs::DependsOn>(flecs::pipeline::OnUpdate::ID));
+}
+
 #[test]
 fn system_create_w_no_template_args() {
     let world = World::new();