From db525525bd32af0b5ea147bc751a3a7d251c0cfb Mon Sep 17 00:00:00 2001
From: Georgiy Tugai <3786806+Georgiy-Tugai@users.noreply.github.com>
Date: Mon, 3 Mar 2025 10:49:25 +0100
Subject: [PATCH] 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.
---
 flecs_ecs/src/addons/system/system_builder.rs | 55 +++++++++----------
 flecs_ecs/tests/flecs/system_test.rs          | 20 +++++++
 2 files changed, 46 insertions(+), 29 deletions(-)

diff --git a/flecs_ecs/src/addons/system/system_builder.rs b/flecs_ecs/src/addons/system/system_builder.rs
index f4ca9c10..7bb3ea33 100644
--- a/flecs_ecs/src/addons/system/system_builder.rs
+++ b/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,6 +29,7 @@ where
             desc: Default::default(),
             term_builder: TermBuilder::default(),
             world: world.into(),
+            kind_set: false,
             _phantom: core::marker::PhantomData,
         };
 
@@ -33,16 +37,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
     }
 
@@ -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) };
diff --git a/flecs_ecs/tests/flecs/system_test.rs b/flecs_ecs/tests/flecs/system_test.rs
index b6674000..fd974d52 100644
--- a/flecs_ecs/tests/flecs/system_test.rs
+++ b/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();