Support removing ecs entities while iterating over them using apply

libs/ecs/include/psemek/ecs/detail/table.hpp

@@ -10,6 +10,7 @@
 #include <type_traits>
 #include <vector>
 #include <array>
+#include <optional>
 
 namespace psemek::ecs::detail
 {
@@ -21,6 +22,11 @@ namespace psemek::ecs::detail
 			std::uint8_t * data = nullptr;
 		};
 
+		struct iteration_data
+		{
+			std::size_t current_row = 0;
+		};
+
 		util::span<util::uuid const> get_component_uuids() const
 		{
 			return component_uuids_;
@@ -31,6 +37,11 @@ namespace psemek::ecs::detail
 			return component_pointers_;
 		}
 
+		util::span<entity_handle const> get_entity_handles() const
+		{
+			return entity_handles_;
+		}
+
 		std::size_t component_count() const
 		{
 			return component_uuids_.size();
@@ -45,9 +56,19 @@ namespace psemek::ecs::detail
 		virtual void swap_rows(std::size_t row1, std::size_t row2) = 0;
 		virtual void pop_row() = 0;
 
-		util::span<entity_handle const> get_entity_handles() const
+		std::optional<iteration_data> & get_iteration_data()
 		{
-			return entity_handles_;
+			return iteration_data_;
+		}
+
+		void push_remove(std::uint32_t row)
+		{
+			remove_queue_.push_back(row);
+		}
+
+		std::vector<std::uint32_t> grab_remove_queue()
+		{
+			return std::move(remove_queue_);
 		}
 
 		virtual ~table() = default;
@@ -58,6 +79,10 @@ namespace psemek::ecs::detail
 		std::size_t row_count_ = 0;
 
 		std::vector<entity_handle> entity_handles_;
+
+		std::optional<iteration_data> iteration_data_;
+
+		std::vector<std::uint32_t> remove_queue_;
 	};
 
 	template <typename ... Components>

libs/ecs/include/psemek/ecs/entity_container.hpp

@@ -8,6 +8,7 @@
 #include <psemek/ecs/detail/all_different_types.hpp>
 #include <psemek/ecs/entity_accessor.hpp>
 #include <psemek/util/span.hpp>
+#include <psemek/util/range.hpp>
 
 namespace psemek::ecs
 {
@@ -60,8 +61,10 @@ namespace psemek::ecs
 		//    void(entity_handle, components...)
 		//    void(entity_container, components...)
 		//    void(entity_container, entity_handle, components...)
+		// The function can create or destroy any entities.
 		// An optional query cache can be supplied to speed up the call
 		// UB if the cache wasn't created with the exact same component sequence
+		// UB if accessing the entitie's component after destroying it
 		template <typename ... Components, typename Function>
 		void apply(Function && function, query_cache cache = {});
 
@@ -75,6 +78,7 @@ namespace psemek::ecs
 		//    void(entity_container, span<entity_handle const>, span<components>...)
 		// An optional query cache can be supplied to speed up the call
 		// UB if the cache wasn't created with the exact same component sequence
+		// UB if the function tries to create or destroy entities
 		template <typename ... Components, typename Function>
 		void batch_apply(Function && function, query_cache cache = {});
 
@@ -83,6 +87,8 @@ namespace psemek::ecs
 		mutable detail::component_index component_index_;
 		detail::table_container table_container_;
 		detail::query_cache_container query_cache_container_;
+
+		void remove_row(detail::table & table, std::uint32_t row, util::span<detail::entity_data> entities);
 	};
 
 	template <typename ... Components>
@@ -137,16 +143,27 @@ namespace psemek::ecs
 
 		for (auto const & entry : cache->tables)
 		{
+			auto & iteration_data = entry.table->get_iteration_data();
+			iteration_data.emplace();
 			detail::static_apply_helper<Components...> apply_helper(*this, entry.table->get_entity_handles());
 
 			for (std::size_t i = 0; i < sizeof...(Components); ++i)
 				apply_helper.pointers[i] = entry.table->get_component_pointers()[entry.column_ids[i]];
 
-			for (std::size_t i = 0; i < apply_helper.size(); ++i)
+			for (std::size_t i = 0; i < entry.table->row_count(); ++i)
 			{
+				iteration_data->current_row = i;
 				apply_helper.apply(function);
 				apply_helper.advance();
 			}
+
+			iteration_data.reset();
+
+			auto remove_queue = entry.table->grab_remove_queue();
+			std::sort(remove_queue.begin(), remove_queue.end());
+			auto entities = entity_list_.get_entities();
+			for (auto row : util::reversed(remove_queue))
+				remove_row(*entry.table, row, entities);
 		}
 	}
 

libs/ecs/source/entity_container.cpp

@@ -10,16 +10,20 @@ namespace psemek::ecs
 
 	void entity_container::destroy(entity_handle const & entity)
 	{
-		// Swap with the last row in that table
 		auto entities = entity_list_.get_entities();
 		auto & data = entities[entity.id];
-		auto table_entity_handles = data.table->get_entity_handles();
-		data.table->swap_rows(data.row, table_entity_handles.size() - 1);
-		data.table->pop_row();
-		auto swap_handle = table_entity_handles[data.row];
-		auto & swap_data = entities[swap_handle.id];
-		swap_data.row = data.row;
-		entity_list_.destroy(entity.id);
+		auto & iteration_data = data.table->get_iteration_data();
+
+		if (!iteration_data || iteration_data->current_row < data.row)
+		{
+			remove_row(*data.table, data.row, entities);
+			entity_list_.destroy(entity.id);
+		}
+		else
+		{
+			data.table->push_remove(data.row);
+			entity_list_.destroy(entity.id);
+		}
 	}
 
 	entity_accessor entity_container::get(entity_handle const & entity)
@@ -28,4 +32,14 @@ namespace psemek::ecs
 		return {data.table, data.row};
 	}
 
+	void entity_container::remove_row(detail::table & table, std::uint32_t row, util::span<detail::entity_data> entities)
+	{
+		// Swap with the last row in that table
+		auto table_entity_handles = table.get_entity_handles();
+		table.swap_rows(row, table_entity_handles.size() - 1);
+		table.pop_row();
+		auto swap_handle = table_entity_handles[row];
+		entities[swap_handle.id].row = row;
+	}
+
 }

libs/ecs/tests/apply.cpp

@@ -169,3 +169,73 @@ test_case(ecs_apply_batch_components)
 	expect_equal(count, expected_count);
 	expect_equal(sum, expected_sum);
 }
+
+test_case(ecs_apply_remove_forward)
+{
+	entity_container container;
+
+	std::vector<entity_handle> handles;
+	int const count = 1024 * 1024;
+	for (int i = 0; i < count; ++i)
+		handles.push_back(container.create());
+
+	int call_count = 0;
+	container.apply([&]{
+		++call_count;
+		if (call_count == count)
+			for (auto h : handles)
+				container.destroy(h);
+	});
+
+	for (auto h : handles)
+		expect(!container.alive(h));
+}
+
+test_case(ecs_apply_remove_reversed)
+{
+	entity_container container;
+
+	std::vector<entity_handle> handles;
+	int const count = 1024 * 1024;
+	for (int i = 0; i < count; ++i)
+		handles.push_back(container.create());
+	std::reverse(handles.begin(), handles.end());
+
+	int call_count = 0;
+	container.apply([&]{
+		++call_count;
+		if (call_count == count)
+			for (auto h : handles)
+				container.destroy(h);
+	});
+
+	expect_equal(call_count, call_count);
+
+	for (auto h : handles)
+		expect(!container.alive(h));
+}
+
+test_case(ecs_apply_remove_random)
+{
+	entity_container container;
+	random::generator rng;
+
+	std::vector<entity_handle> handles;
+	int const count = 1024 * 1024;
+	for (int i = 0; i < count; ++i)
+		handles.push_back(container.create());
+	std::shuffle(handles.begin(), handles.end(), rng);
+
+	int call_count = 0;
+	container.apply([&]{
+		++call_count;
+		if (call_count == count)
+			for (auto h : handles)
+				container.destroy(h);
+	});
+
+	expect_equal(call_count, call_count);
+
+	for (auto h : handles)
+		expect(!container.alive(h));
+}