From cedc29f39a42f8a2fb0198a85a167c07203c4d91 Mon Sep 17 00:00:00 2001
From: lisyarus <lisyarus@gmail.com>
Date: Tue, 21 May 2024 19:20:04 +0300
Subject: [PATCH] Allow creating ECS entities inside batch systems

---
 libs/ecs/include/psemek/ecs/container.hpp | 31 +++++++++++------------
 libs/ecs/source/container.cpp             | 17 +++++++++++++
 2 files changed, 32 insertions(+), 16 deletions(-)

diff --git a/libs/ecs/include/psemek/ecs/container.hpp b/libs/ecs/include/psemek/ecs/container.hpp
index f4680cb7..6c0b7c05 100644
--- a/libs/ecs/include/psemek/ecs/container.hpp
+++ b/libs/ecs/include/psemek/ecs/container.hpp
@@ -245,6 +245,12 @@ namespace psemek::ecs
 		 *    the function should use the entity handles array view to compute the number
 		 *    of visited entities.
 		 *
+		 *  The function can freely create or destroy entities, and or attach/detach
+		 *    components to existing entities. It is unspecified whether the function
+		 *    will or will not visit newly created entities during this `apply()` call.
+		 *    The function is guaranteed not to visit destroyed entities (unless it did
+		 *    so before the entity was destroyed).
+		 *
 		 *  An optional query cache can be supplied to speed up iteration. If the caller doesn't
 		 *    supply a query cache, a new query cache is created as if by calling `cache<Components...>()`.
 		 *    In any case, the (user-provided or newly-created) query cache is returned.
@@ -256,8 +262,6 @@ namespace psemek::ecs
 		 *    constness
 		 *  @warning If any two of the passed component types are equal, the call fails with
 		 *    a compilation error
-		 *  @warning If the function creates or destroyes entities, or attaches or detaches components,
-		 *    the behavior is undefined
 		 *  @warning If the function recursively calls `apply()` or `batch_apply()`, the behavior is undefined
 		 */
 		template <typename ... Components, typename Function>
@@ -393,6 +397,7 @@ namespace psemek::ecs
 		detail::table * insert_table(std::vector<std::unique_ptr<detail::column>> columns);
 		void do_destroy(handle const & entity);
 		void remove_row(detail::table & table, std::uint32_t row, util::span<detail::entity_data> entities);
+		void finilize_iteration(detail::table & table);
 	};
 
 	template <typename Component>
@@ -570,20 +575,7 @@ namespace psemek::ecs
 			}
 
 			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);
-
-			std::uint32_t row = entry.table->row_count();
-			for (auto const & handle : entry.table->flush_delayed())
-			{
-				auto & data = entities[handle.id];
-				data.table = entry.table;
-				data.row = row++;
-			}
+			finilize_iteration(*entry.table);
 		}
 
 		return cache;
@@ -603,12 +595,19 @@ namespace psemek::ecs
 
 		for (auto const & entry : cache->entries)
 		{
+			auto & iteration_data = entry.table->get_iteration_data();
+			iteration_data.emplace();
+			iteration_data->current_row = entry.table->row_count();
+
 			typename detail::filter_with<detail::static_apply_helper, std::tuple<Components...>>::type apply_helper(*this, entry.table->entity_handles());
 
 			for (std::size_t i = 0; i < cache->with_uuids.size(); ++i)
 				apply_helper.pointers[i] = entry.table->column(cache->with_uuids[i])->data();
 
 			apply_helper.batch_apply(function);
+
+			iteration_data.reset();
+			finilize_iteration(*entry.table);
 		}
 
 		return cache;
diff --git a/libs/ecs/source/container.cpp b/libs/ecs/source/container.cpp
index 067e8416..1e62ff73 100644
--- a/libs/ecs/source/container.cpp
+++ b/libs/ecs/source/container.cpp
@@ -93,4 +93,21 @@ namespace psemek::ecs
 		entities[swap_handle.id].row = row;
 	}
 
+	void container::finilize_iteration(detail::table & table)
+	{
+		auto remove_queue = 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(table, row, entities);
+
+		std::uint32_t row = table.row_count();
+		for (auto const & handle : table.flush_delayed())
+		{
+			auto & data = entities[handle.id];
+			data.table = &table;
+			data.row = row++;
+		}
+	}
+
 }