From bd1393c505c721f7face1af80e194f47adec6419 Mon Sep 17 00:00:00 2001
From: lisyarus <lisyarus@gmail.com>
Date: Mon, 3 Jun 2024 13:59:50 +0300
Subject: [PATCH] Improve working with ecs caches, support retrieving ecs cache
 & table count, and fix ecs query container memory leak

---
 libs/ecs/include/psemek/ecs/container.hpp     | 48 +++++++----------
 .../include/psemek/ecs/detail/query_cache.hpp |  6 ++-
 .../ecs/detail/query_cache_container.hpp      | 53 +++++++++++--------
 .../psemek/ecs/detail/table_container.hpp     |  9 +++-
 libs/ecs/source/container.cpp                 | 10 ++++
 libs/ecs/source/detail/query_cache.cpp        |  8 +--
 6 files changed, 73 insertions(+), 61 deletions(-)

diff --git a/libs/ecs/include/psemek/ecs/container.hpp b/libs/ecs/include/psemek/ecs/container.hpp
index 04b7232f..ca715d4d 100644
--- a/libs/ecs/include/psemek/ecs/container.hpp
+++ b/libs/ecs/include/psemek/ecs/container.hpp
@@ -381,6 +381,10 @@ namespace psemek::ecs
 		template <typename ... Components>
 		std::size_t memory_usage();
 
+		std::size_t cache_count();
+
+		std::size_t table_count();
+
 	private:
 		detail::entity_list entity_list_;
 		detail::table_container table_container_;
@@ -391,11 +395,6 @@ namespace psemek::ecs
 		util::object_pool<std::vector<util::uuid>> uuid_list_pool_;
 		util::object_pool<util::hash_set<util::uuid>> uuid_set_pool_;
 
-		std::vector<query_cache> callback_caches_;
-
-		template <typename ... Components, typename Constructor, typename Destructor>
-		query_cache cache_impl(Constructor && constructor, Destructor && destructor);
-
 		detail::table * insert_table(std::vector<std::unique_ptr<detail::column>> columns);
 		void do_destroy(handle entity);
 		void remove_row(detail::table & table, std::uint32_t row, util::span<detail::entity_data> entities);
@@ -563,7 +562,10 @@ namespace psemek::ecs
 	template <typename ... Components>
 	query_cache container::cache()
 	{
-		return cache_impl<Components...>(nullptr, nullptr);
+		typename detail::filter_with   <detail::component_uuid_helper, std::tuple<std::remove_cvref_t<Components>...>>::type with_uuids;
+		typename detail::filter_without<detail::component_uuid_helper, std::tuple<std::remove_cvref_t<Components>...>>::type without_uuids;
+
+		return query_cache_container_.create(with_uuids.get(), without_uuids.get(), table_container_);
 	}
 
 	template <typename ... Components, typename Function>
@@ -642,7 +644,7 @@ namespace psemek::ecs
 
 		static_assert(invocable_type::value, "function is not invocable with these components");
 
-		auto constructor = [function = std::move(function)](std::vector<std::uint32_t> const & column_indices) -> detail::table_callback {
+		auto constructor_factory = [function = std::move(function)](std::vector<std::uint32_t> const & column_indices) -> detail::table_callback {
 			return [function, column_indices](container & container, detail::table & table, std::uint32_t row, util::hash_set<util::uuid> const & attached_components, bool force){
 
 				bool const invoke = force || (detail::contains_helper<Components>::contains(attached_components) || ...);
@@ -659,9 +661,11 @@ namespace psemek::ecs
 			};
 		};
 
-		query_cache cache = this->cache_impl<Components...>(std::move(constructor), nullptr);
+		query_cache cache = this->cache<Components...>();
 
-		callback_caches_.push_back(cache);
+		cache->constructor_factories.push_back(constructor_factory);
+		for (auto const & entry : cache->entries)
+			entry.table->add_constructor(constructor_factory(entry.columns_indices));
 	}
 
 	template <typename ... Components, typename Function>
@@ -673,7 +677,7 @@ namespace psemek::ecs
 
 		static_assert(invocable_type::value, "function is not invocable with these components");
 
-		auto destructor = [function = std::move(function)](std::vector<std::uint32_t> const & column_indices) -> detail::table_callback {
+		auto destructor_factory = [function = std::move(function)](std::vector<std::uint32_t> const & column_indices) -> detail::table_callback {
 			return [function, column_indices](container & container, detail::table & table, std::uint32_t row, util::hash_set<util::uuid> const & detached_components, bool force){
 
 				bool const invoke = force || (detail::contains_helper<Components>::contains(detached_components) || ...);
@@ -690,9 +694,11 @@ namespace psemek::ecs
 			};
 		};
 
-		query_cache cache = this->cache_impl<Components...>(nullptr, std::move(destructor));
+		query_cache cache = this->cache<Components...>();
 
-		callback_caches_.push_back(cache);
+		cache->destructor_factories.push_back(destructor_factory);
+		for (auto const & entry : cache->entries)
+			entry.table->add_destructor(destructor_factory(entry.columns_indices));
 	}
 
 	template <typename Index, typename ... Args>
@@ -701,24 +707,6 @@ namespace psemek::ecs
 		return index_container_.get<Index>(*this, std::forward<Args>(args)...);
 	}
 
-	template <typename ... Components, typename Constructor, typename Destructor>
-	query_cache container::cache_impl(Constructor && constructor, Destructor && destructor)
-	{
-		typename detail::filter_with   <detail::component_uuid_helper, std::tuple<std::remove_cvref_t<Components>...>>::type with_uuids;
-		typename detail::filter_without<detail::component_uuid_helper, std::tuple<std::remove_cvref_t<Components>...>>::type without_uuids;
-
-		auto result = query_cache_container_.create(with_uuids.get(), without_uuids.get());
-
-		result->constructor_factory = std::move(constructor);
-		result->destructor_factory = std::move(destructor);
-
-		table_container_.apply([&](detail::table & table){
-			result->add(&table);
-		}, with_uuids.get(), without_uuids.get());
-
-		return result;
-	}
-
 	template <typename ... Components>
 	std::size_t container::memory_usage()
 	{
diff --git a/libs/ecs/include/psemek/ecs/detail/query_cache.hpp b/libs/ecs/include/psemek/ecs/detail/query_cache.hpp
index e4316f3f..2875801c 100644
--- a/libs/ecs/include/psemek/ecs/detail/query_cache.hpp
+++ b/libs/ecs/include/psemek/ecs/detail/query_cache.hpp
@@ -24,8 +24,10 @@ namespace psemek::ecs::detail
 		std::vector<util::uuid> without_uuids;
 		std::list<query_cache_entry> entries;
 
-		util::function<table_callback(std::vector<std::uint32_t> const &)> constructor_factory;
-		util::function<table_callback(std::vector<std::uint32_t> const &)> destructor_factory;
+		using callback_factory = util::function<table_callback(std::vector<std::uint32_t> const &)>;
+
+		std::vector<callback_factory> constructor_factories;
+		std::vector<callback_factory> destructor_factories;
 
 		void add(table * table);
 	};
diff --git a/libs/ecs/include/psemek/ecs/detail/query_cache_container.hpp b/libs/ecs/include/psemek/ecs/detail/query_cache_container.hpp
index af9c232d..c784767a 100644
--- a/libs/ecs/include/psemek/ecs/detail/query_cache_container.hpp
+++ b/libs/ecs/include/psemek/ecs/detail/query_cache_container.hpp
@@ -2,6 +2,7 @@
 
 #include <psemek/ecs/detail/query_cache.hpp>
 #include <psemek/ecs/detail/component_hash.hpp>
+#include <psemek/ecs/detail/table_container.hpp>
 #include <psemek/util/hash_table.hpp>
 
 #include <memory>
@@ -9,12 +10,12 @@
 namespace psemek::ecs::detail
 {
 
-	struct query_cache_set
+	struct query_cache_data
 	{
 		std::size_t hash;
 		util::hash_set<util::uuid> with_uuids;
 		util::hash_set<util::uuid> without_uuids;
-		std::vector<std::weak_ptr<query_cache>> caches;
+		std::shared_ptr<query_cache> cache;
 	};
 
 	struct query_cache_hash
@@ -24,7 +25,7 @@ namespace psemek::ecs::detail
 			return component_hash(uuids.first, uuids.second);
 		}
 
-		std::size_t operator()(std::unique_ptr<query_cache_set> const & set) const
+		std::size_t operator()(std::unique_ptr<query_cache_data> const & set) const
 		{
 			return set->hash;
 		}
@@ -32,7 +33,7 @@ namespace psemek::ecs::detail
 
 	struct query_cache_equal
 	{
-		bool operator()(std::pair<util::span<util::uuid const>, util::span<util::uuid const>> const & uuids, std::unique_ptr<query_cache_set> const & set) const
+		bool operator()(std::pair<util::span<util::uuid const>, util::span<util::uuid const>> const & uuids, std::unique_ptr<query_cache_data> const & set) const
 		{
 			if (uuids.first.size() != set->with_uuids.size())
 				return false;
@@ -51,7 +52,7 @@ namespace psemek::ecs::detail
 			return true;
 		}
 
-		bool operator()(std::unique_ptr<query_cache_set> const & set1, std::unique_ptr<query_cache_set> const & set2) const
+		bool operator()(std::unique_ptr<query_cache_data> const & set1, std::unique_ptr<query_cache_data> const & set2) const
 		{
 			return set1.get() == set2.get();
 		}
@@ -60,41 +61,49 @@ namespace psemek::ecs::detail
 	// TODO: store query caches in a bitmask trie balanced by subtree size
 	struct query_cache_container
 	{
-		std::shared_ptr<query_cache> create(util::span<util::uuid const> with_uuids, util::span<util::uuid const> without_uuids)
+		std::shared_ptr<query_cache> create(util::span<util::uuid const> with_uuids, util::span<util::uuid const> without_uuids, table_container & table_container)
 		{
-			auto result = std::make_shared<query_cache>();
-			result->with_uuids.assign(with_uuids.begin(), with_uuids.end());
-			result->without_uuids.assign(without_uuids.begin(), without_uuids.end());
 			auto it = caches_.find(std::pair{with_uuids, without_uuids});
 			if (it == caches_.end())
 			{
-				auto value = std::make_unique<query_cache_set>();
+				auto value = std::make_unique<query_cache_data>();
+				value->cache = std::make_shared<query_cache>();
 				for (auto const & uuid : with_uuids)
+				{
 					value->with_uuids.insert(uuid);
+					value->cache->with_uuids.push_back(uuid);
+				}
 				for (auto const & uuid : without_uuids)
+				{
 					value->without_uuids.insert(uuid);
+					value->cache->without_uuids.push_back(uuid);
+				}
 				value->hash = component_hash(with_uuids, without_uuids);
+
+				table_container.apply([&](table & table){
+					value->cache->add(&table);
+				}, with_uuids, without_uuids);
+
 				it = caches_.insert(std::move(value)).first;
 			}
-			it->get()->caches.push_back(result);
-			return result;
+			return it->get()->cache;
 		}
 
 		template <typename Function, typename ContainsUUID>
 		void apply(Function && function, ContainsUUID && contains_uuid)
 		{
-			for (auto & cache_set : caches_)
+			for (auto & cache_data : caches_)
 			{
 				bool good = true;
 
-				for (auto const & uuid : cache_set->with_uuids)
+				for (auto const & uuid : cache_data->with_uuids)
 					if (!contains_uuid(uuid))
 					{
 						good = false;
 						break;
 					}
 
-				for (auto const & uuid : cache_set->without_uuids)
+				for (auto const & uuid : cache_data->without_uuids)
 					if (contains_uuid(uuid))
 					{
 						good = false;
@@ -104,19 +113,17 @@ namespace psemek::ecs::detail
 				if (!good)
 					continue;
 
-				filter(cache_set->caches);
-				for (auto & cache : cache_set->caches)
-					function(*cache.lock());
+				function(*(cache_data->cache));
 			}
 		}
 
-	private:
-		util::hash_set<std::unique_ptr<query_cache_set>, query_cache_hash, query_cache_equal> caches_;
-
-		static void filter(std::vector<std::weak_ptr<query_cache>> & caches)
+		std::size_t cache_count() const
 		{
-			caches.erase(std::remove_if(caches.begin(), caches.end(), [](std::weak_ptr<query_cache> const & weak){ return !weak.lock(); }), caches.end());
+			return caches_.size();
 		}
+
+	private:
+		util::hash_set<std::unique_ptr<query_cache_data>, query_cache_hash, query_cache_equal> caches_;
 	};
 
 }
diff --git a/libs/ecs/include/psemek/ecs/detail/table_container.hpp b/libs/ecs/include/psemek/ecs/detail/table_container.hpp
index 1d9474b4..f10ae2e0 100644
--- a/libs/ecs/include/psemek/ecs/detail/table_container.hpp
+++ b/libs/ecs/include/psemek/ecs/detail/table_container.hpp
@@ -4,8 +4,6 @@
 #include <psemek/ecs/detail/table.hpp>
 #include <psemek/util/hash_table.hpp>
 
-#include <vector>
-
 namespace psemek::ecs::detail
 {
 
@@ -49,6 +47,8 @@ namespace psemek::ecs::detail
 		template <typename Function>
 		void apply(Function && function, util::span<util::uuid const> with_uuids, util::span<util::uuid const> without_uuids);
 
+		std::size_t table_count() const;
+
 	private:
 		util::hash_set<std::unique_ptr<table>, table_hashset_hash, table_hashset_equal> tables_;
 	};
@@ -93,4 +93,9 @@ namespace psemek::ecs::detail
 		}
 	}
 
+	inline std::size_t table_container::table_count() const
+	{
+		return tables_.size();
+	}
+
 }
diff --git a/libs/ecs/source/container.cpp b/libs/ecs/source/container.cpp
index c97953de..96fd665a 100644
--- a/libs/ecs/source/container.cpp
+++ b/libs/ecs/source/container.cpp
@@ -116,4 +116,14 @@ namespace psemek::ecs
 		}
 	}
 
+	std::size_t container::cache_count()
+	{
+		return query_cache_container_.cache_count();
+	}
+
+	std::size_t container::table_count()
+	{
+		return table_container_.table_count();
+	}
+
 }
diff --git a/libs/ecs/source/detail/query_cache.cpp b/libs/ecs/source/detail/query_cache.cpp
index 6c717937..1df6a577 100644
--- a/libs/ecs/source/detail/query_cache.cpp
+++ b/libs/ecs/source/detail/query_cache.cpp
@@ -12,11 +12,11 @@ namespace psemek::ecs::detail
 		for (auto const & uuid : with_uuids)
 			entry.columns_indices.push_back(*table->column_index(uuid));
 
-		if (constructor_factory)
-			table->add_constructor(constructor_factory(entry.columns_indices));
+		for (auto const & factory : constructor_factories)
+			entry.table->add_constructor(factory(entry.columns_indices));
 
-		if (destructor_factory)
-			table->add_destructor(destructor_factory(entry.columns_indices));
+		for (auto const & factory : destructor_factories)
+			entry.table->add_destructor(factory(entry.columns_indices));
 	}
 
 }