From b57239fac29bfe81cf7e3b68677d9d9cab5a8edf Mon Sep 17 00:00:00 2001
From: Ukendio <ukendio@gmail.com>
Date: Sun, 11 Aug 2024 04:03:18 +0200
Subject: [PATCH] optimize hook calls and add tests

---
 src/init.luau   |  5 ++---
 test/tests.luau | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/src/init.luau b/src/init.luau
index 28ed521..bbed609 100644
--- a/src/init.luau
+++ b/src/init.luau
@@ -520,7 +520,7 @@ local function world_set(world: World, entity: i53, id: i53, data: unknown)
 	local tr = to.records[id]
 	to.columns[tr.column][record.row] = data
 
-	invoke_hook(world, EcsOnAdd, id, entity, data)
+	invoke_hook(world, EcsOnSet, id, entity, data)
 end
 
 local function world_component(world: World): i53
@@ -556,8 +556,6 @@ local function archetype_traverse_remove(world: World, id: i53, from: Archetype)
 end
 
 local function world_remove(world: World, entity: i53, id: i53)
-    invoke_hook(world, EcsOnRemove, id, entity)
-
 	local entity_index = world.entityIndex
 	local record = entity_index.sparse[entity]
 	local from = record.archetype
@@ -567,6 +565,7 @@ local function world_remove(world: World, entity: i53, id: i53)
 	local to = archetype_traverse_remove(world, id, from)
 
 	if from and not (from == to) then
+	    invoke_hook(world, EcsOnRemove, id, entity)
 		entity_move(entity_index, entity, record, to)
 	end
 end
diff --git a/test/tests.luau b/test/tests.luau
index 27aedc0..1cd29a4 100644
--- a/test/tests.luau
+++ b/test/tests.luau
@@ -172,6 +172,7 @@ TEST("world:remove()", function()
 end)
 
 TEST("world:add()", function()
+
     do CASE "idempotent"
         local world = jecs.World.new()
         local d = debug_world_inspect(world)
@@ -968,4 +969,39 @@ TEST("changetracker:track()", function()
 
 end)
 
+TEST("Hooks", function()
+    do CASE "OnAdd"
+        local world = jecs.World.new()
+        local Transform = world:component()
+        local e1 = world:entity()
+        world:set(Transform, jecs.OnAdd, function(entity)
+            CHECK(e1 == entity)
+        end)
+        world:add(e1, Transform)
+    end
+
+    do CASE "OnSet"
+        local world = jecs.World.new()
+        local Number = world:component()
+        local e1 = world:entity()
+        world:set(Number, jecs.OnSet, function(entity, data)
+            CHECK(e1 == entity)
+            CHECK(data == 1)
+        end)
+        world:set(e1, Number, 1)
+    end
+
+    do CASE "OnRemove"
+        local world = jecs.World.new()
+        local A = world:component()
+        local e1 = world:entity()
+        world:add(e1, A)
+        world:set(A, jecs.OnRemove, function(entity)
+            CHECK(e1 == entity)
+            CHECK(world:has(e1, A))
+        end)
+        world:remove(e1, A)
+    end
+
+end)
 FINISH()