From 04c2e59ada7da892e4fa33049cab2f46684a67a4 Mon Sep 17 00:00:00 2001
From: light7734 <light7734@tuta.io>
Date: Sun, 21 Sep 2025 08:41:56 +0330
Subject: [PATCH] fix(ecs): sparse_set not properly removing elements

---
 modules/ecs/private/sparse_set.test.cpp | 27 +++++++++++++++++++++++-
 modules/ecs/public/registry.hpp         | 15 +++++++++----
 modules/ecs/public/sparse_set.hpp       | 28 ++++++++++++++++++++-----
 3 files changed, 60 insertions(+), 10 deletions(-)

diff --git a/modules/ecs/private/sparse_set.test.cpp b/modules/ecs/private/sparse_set.test.cpp
index 1465eec..d92e9df 100644
--- a/modules/ecs/private/sparse_set.test.cpp
+++ b/modules/ecs/private/sparse_set.test.cpp
@@ -76,8 +76,33 @@ Suite element_raii = [] {
 			set.remove(idx);
 
 			expect_eq(set.get_size(), 10'000 - (idx + 1));
-			expect_throw([&] { std::ignore = set.at(idx); });
 			expect_false(set.contains(idx));
+			expect_throw([&] { std::ignore = set.at(idx); });
+		}
+	};
+
+	Case { "removed elements won't be iterated again" } = [] {
+		auto set = Set {};
+
+		for (auto idx : std::views::iota(0, 10'000))
+		{
+			set.insert(idx, idx);
+		}
+
+		set.remove(0);
+		set.remove(32);
+		set.remove(69);
+		set.remove(420);
+		set.remove(9'999);
+
+		for (auto &[identifier, value] : set)
+		{
+			expect_eq(identifier, value);
+			expect_ne(value, 0);
+			expect_ne(value, 32);
+			expect_ne(value, 69);
+			expect_ne(value, 420);
+			expect_ne(value, 9'999);
 		}
 	};
 };
diff --git a/modules/ecs/public/registry.hpp b/modules/ecs/public/registry.hpp
index 1377cfd..367e3bc 100644
--- a/modules/ecs/public/registry.hpp
+++ b/modules/ecs/public/registry.hpp
@@ -8,9 +8,16 @@ using Entity = uint32_t;
 
 /** A registry of components, the heart of an ECS architecture.
  *
- * @todo(Light): optimize multi-component views
- * @todo(Light): support more than 2-component views
- * @todo(Light): handle edge cases or specify the undefined behaviors
+ * @todo(Light): Implement grouping
+ * @todo(Light): Implement identifier recycling
+ * @todo(Light): Optimize views/each
+ * @todo(Light): Support >2 component views
+ * @todo(Light): Handle more edge cases or specify the undefined behaviors
+ *
+ * @ref https://skypjack.github.io/
+ * @ref https://github.com/alecthomas/entityx
+ * @ref https://github.com/skypjack/entt
+ * @ref https://github.com/SanderMertens/flecs
  */
 class Registry
 {
@@ -93,7 +100,7 @@ public:
 	}
 
 	template<typename Component_T>
-	auto view() -> SparseSet<Component_T, Entity>&
+	auto view() -> SparseSet<Component_T, Entity> &
 	{
 		return get_derived_set<Component_T>();
 	};
diff --git a/modules/ecs/public/sparse_set.hpp b/modules/ecs/public/sparse_set.hpp
index abaf441..496dc26 100644
--- a/modules/ecs/public/sparse_set.hpp
+++ b/modules/ecs/public/sparse_set.hpp
@@ -25,10 +25,6 @@ public:
 	virtual void remove(Identifier_T identifier) = 0;
 };
 
-/**
- *
- * @todo(Light): implement identifier recycling.
- */
 template<typename Value_T, typename Identifier_T = uint32_t>
 class SparseSet: public TypeErasedSparseSet<Identifier_T>
 {
@@ -77,7 +73,29 @@ public:
 		auto &idx = m_sparse[identifier];
 		auto &[entity, component] = m_dense[idx];
 
-		idx = null_identifier;
+		auto &[last_entity, last_component] = m_dense.back();
+		auto &last_idx = m_sparse[last_entity];
+
+		// removed entity is in dense's back, just pop and invalidate sparse[identifier]
+		if (entity == last_entity)
+		{
+			idx = null_identifier;
+			m_dense.pop_back();
+		}
+		else
+		{
+			// swap dense's 'back' to 'removed'
+			std::swap(component, last_component);
+			entity = last_entity;
+
+			// make sparse to point to new idx
+			last_idx = idx;
+
+			// pop dense and invalidate sparse[identifier]
+			idx = null_identifier;
+			m_dense.pop_back();
+		}
+
 		++m_dead_count;
 		--m_alive_count;
 	}