From 31e364afdf7a0926632d476a5df7c5a985ba6a3e Mon Sep 17 00:00:00 2001
From: lolmanurfunny <77128366+lolmanurfunny@users.noreply.github.com>
Date: Wed, 2 Apr 2025 02:21:47 -0400
Subject: [PATCH] invoke OnRemove hooks when overwritten

---
 jecs.luau       | 22 ++++++++++++++---
 test/tests.luau | 65 +++++++++++++++++++++++++++++++++++++++++--------
 2 files changed, 74 insertions(+), 13 deletions(-)

diff --git a/jecs.luau b/jecs.luau
index 377f123..7d8a840 100644
--- a/jecs.luau
+++ b/jecs.luau
@@ -923,7 +923,19 @@ local function world_add(
 	if from == to then
 		return
 	end
+
+	local idr = world.component_index[id]
+	local idr_hooks = idr.hooks
+	local on_add = idr_hooks.on_add
+
 	if from then
+		if ECS_IS_PAIR(id) and bit32.band(idr.flags, ECS_ID_EXCLUSIVE) ~= 0 then
+			local on_remove = idr_hooks.on_remove
+			if on_remove then
+				on_remove(entity)
+			end
+		end
+
 		entity_move(entity_index, entity, record, to)
 	else
 		if #to.types > 0 then
@@ -931,9 +943,6 @@ local function world_add(
 		end
 	end
 
-	local idr = world.component_index[id]
-	local on_add = idr.hooks.on_add
-
 	if on_add then
 		on_add(entity)
 	end
@@ -966,6 +975,13 @@ local function world_set(world: ecs_world_t, entity: i53, id: i53, data: unknown
 	end
 
 	if from then
+		if ECS_IS_PAIR(id) and bit32.band(idr.flags, ECS_ID_EXCLUSIVE) ~= 0 then
+			local on_remove = idr_hooks.on_remove
+			if on_remove then
+				on_remove(entity)
+			end
+		end
+
 		-- If there was a previous archetype, then the entity needs to move the archetype
 		entity_move(entity_index, entity, record, to)
 	else
diff --git a/test/tests.luau b/test/tests.luau
index a938648..b0647aa 100644
--- a/test/tests.luau
+++ b/test/tests.luau
@@ -1949,22 +1949,67 @@ TEST("world:delete() invokes OnRemove hook", function()
 end)
 
 TEST("exclusive relationships", function()
-	local world = world_new()
-	local pair = jecs.pair
+	do CASE "enforce exclusivity"
+		local world = world_new()
+		local pair = jecs.pair
 
-	local child = world:entity()
+		local e = world:entity()
 
-	for _ = 1, 10 do
+		for _ = 1, 10 do
+			local A = world:entity()
+			local B = world:entity()
+
+			world:add(e, pair(world:entity(), e)) -- noise
+			world:add(e, pair(ChildOf, A))
+			world:add(e, pair(ChildOf, B))
+
+			CHECK(not world:has(e, pair(ChildOf, A)))
+			CHECK(world:has(e, pair(ChildOf, B)))
+		end
+	end
+	do CASE "invoke OnRemove when overwritten using world:add()"
+		local world = world_new()
+		local pair = jecs.pair
+
+		local relation = world:component()
+		world:add(relation, jecs.Exclusive)
+
+		local e = world:entity()
 		local A = world:entity()
 		local B = world:entity()
 
-		world:add(child, pair(world:entity(), child)) -- noise
-		world:add(child, pair(ChildOf, A))
-		world:add(child, pair(child, world:entity())) -- noise
-		world:add(child, pair(ChildOf, B))
+		local called = false
+		world:set(relation, jecs.OnRemove, function(e)
+			called = true
+			CHECK(world:target(e, relation) == A)
+		end)
 
-		CHECK(world:has(child, pair(ChildOf, B)))
-		CHECK(not world:has(child, pair(ChildOf, A)))
+		world:add(e, pair(relation, A))
+		world:add(e, pair(relation, B))
+
+		CHECK(called)
+	end
+	do CASE "invoke OnRemove when overwritten using world:set()"
+		local world = world_new()
+		local pair = jecs.pair
+
+		local relation = world:component()
+		world:add(relation, jecs.Exclusive)
+
+		local e = world:entity()
+		local A = world:entity()
+		local B = world:entity()
+
+		local called = false
+		world:set(relation, jecs.OnRemove, function(e)
+			called = true
+			CHECK(world:target(e, relation) == A)
+		end)
+
+		world:set(e, pair(relation, A), nil)
+		world:set(e, pair(relation, B), nil)
+
+		CHECK(called)
 	end
 end)