From 5d6c2e8d425e8778d63562ee5e0229a5baf86fdc Mon Sep 17 00:00:00 2001
From: Ukendio <ukendio@gmail.com>
Date: Sat, 21 Sep 2024 20:53:12 +0200
Subject: [PATCH] Fix unit tests and debug mode to enforce OnRemove rules

---
 src/init.luau   | 31 ++++++++++++++++++++++++++++++-
 test/tests.luau |  4 ++--
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/src/init.luau b/src/init.luau
index 98a7885..aaa8161 100644
--- a/src/init.luau
+++ b/src/init.luau
@@ -1613,7 +1613,7 @@ if _G.__JECS_DEBUG then
         throw(msg)
     end
 
-    local function get_name(world, id): string | number
+    local function get_name(world, id): string
         local name: string | nil
         if ECS_IS_PAIR(id) then
             name = `pair({get_name(world, ECS_ENTITY_T_HI(id))}, {get_name(world, ECS_ENTITY_T_LO(id))})`
@@ -1634,6 +1634,14 @@ if _G.__JECS_DEBUG then
         return not world_has_one_inline(world, ECS_ENTITY_T_HI(id), EcsComponent)
     end
 
+    local original_invoke_hook = invoke_hook
+    local invoked_hook = false
+    invoke_hook = function(...)
+        invoked_hook = true
+        original_invoke_hook(...)
+        invoked_hook = false
+    end
+
     World.query = function(world: World, ...)
         ASSERT((...), "Requires at least a single component")
         return world_query(world, ...)
@@ -1658,6 +1666,19 @@ if _G.__JECS_DEBUG then
             return
         end
 
+
+
+        if world_has_one_inline(world, entity, id) then
+            if invoked_hook then
+                local file, line = debug.info(2, "sl")
+                local hook_fn = `{file}::{line}`
+                throw(
+                ([[cannot call world:set within the
+                This world:set function invokes a structural change because %s doesn't exist on the entity.\n
+                call world:add when the hook %s is in process]]):format(get_name(world, id), hook_fn))
+            end
+        end
+
         world_set(world, entity, id, value)
     end
 
@@ -1670,6 +1691,10 @@ if _G.__JECS_DEBUG then
             return
         end
 
+        if invoked_hook then
+            local hook_fn = debug.info(2, "sl")
+            throw(`Cannot call world:add when the hook {hook_fn} is in process`)
+        end
         world_add(world, entity, id)
     end
 
@@ -1704,6 +1729,10 @@ if _G.__JECS_DEBUG then
         end
         return world_target(world, entity, relation, index)
     end
+
+    World.remove = function()
+
+    end
 end
 
 function World.new()
diff --git a/test/tests.luau b/test/tests.luau
index 3c6f8c5..503ebd8 100644
--- a/test/tests.luau
+++ b/test/tests.luau
@@ -1328,14 +1328,14 @@ TEST("Hooks", function()
 
             world:set(A, jecs.OnRemove, function(entity)
                 world:set(entity, B, true)
-                CHECK(not world:get(entity, A))
+                CHECK(world:get(entity, A))
                 CHECK(world:get(entity, B))
             end)
 
             world:set(e, A, true)
             world:remove(e, A)
             CHECK(not world:get(e, A))
-            CHECK(world:get(e, B))
+            CHECK(not world:get(e, B))
         end
     end