MeshPager.mpp

export module CppUtils.FileSystem.MeshPager;

import std;
import CppUtils.Memory;
import CppUtils.Container.MeshNetwork;
import CppUtils.Container.SafeShared;
import CppUtils.FileSystem.IndexedStorage;
import CppUtils.Type.Serializer;
import CppUtils.Thread.SharedPtr;

export namespace CppUtils::FileSystem
{
	template<class Key, class Value>
	struct MeshSerializedNode final
	{
		Value value;
		std::unordered_map<Key, std::vector<std::size_t>> branches;
		bool isRoot = false;
	};

	template<class Key, class Value>
	class MeshPager final
	{
	public:
		struct Policy final
		{
			MeshPager* pager = nullptr;

			auto resolve(std::size_t targetId) const -> Container::MeshNodePtr<Key, Value, Policy>
			{
				if (pager)
					return pager->loadNode(targetId);
				return {};
			}
		};

		using MeshNodePtr = Container::MeshNodePtr<Key, Value, Policy>;

	private:
		using SerializedNode = MeshSerializedNode<Key, Value>;

		std::filesystem::path m_directory;
		FileSystem::IndexedStorage m_storage;
		std::unordered_map<std::size_t, MeshNodePtr> m_loadedNodes;
		std::unordered_map<std::size_t, std::unordered_map<Key, std::unordered_set<std::size_t>>> m_knownLinks;
		std::size_t m_nextId = 1;

		std::vector<std::pair<std::size_t, std::size_t>> m_activeFocuses;
		mutable std::recursive_mutex m_mutex;

	public:
		class ScopedFocus final
		{
		private:
			MeshPager& m_pager;
			std::size_t m_nodeId = 0;
			std::size_t m_distance = 0;

		public:
			ScopedFocus(MeshPager& pager, std::size_t nodeId, std::size_t distance):
				m_pager{pager},
				m_nodeId{nodeId},
				m_distance{distance}
			{
				m_pager.addFocus(m_nodeId, m_distance);
			}

			~ScopedFocus()
			{
				if (m_nodeId != 0)
					m_pager.removeFocus(m_nodeId, m_distance);
			}

			ScopedFocus(const ScopedFocus&) = delete;
			ScopedFocus& operator=(const ScopedFocus&) = delete;

			ScopedFocus(ScopedFocus&& other) noexcept:
				m_pager{other.m_pager},
				m_nodeId{std::exchange(other.m_nodeId, 0uz)},
				m_distance{std::exchange(other.m_distance, 0uz)}
			{}

			ScopedFocus& operator=(ScopedFocus&& other) noexcept = delete;
		};

		explicit MeshPager(std::filesystem::path directory, std::optional<std::size_t> maxChunkSize = std::nullopt):
			m_directory{std::move(directory)},
			m_storage{m_directory, [&] {
						  using namespace CppUtils::Memory::Literals;
						  return maxChunkSize.value_or(64_MiB);
					  }()}
		{
			auto indexAccessor = m_storage.getIndex();
			auto keys = indexAccessor.value() | std::views::keys;
			if (auto it = std::ranges::max_element(keys); it != std::ranges::end(keys))
				m_nextId = *it + 1;
		}

		~MeshPager()
		{
			auto lock = std::unique_lock{m_mutex};
			auto nodeIds = std::vector<std::size_t>{};
			for (const auto& [nodeId, _] : m_loadedNodes)
				nodeIds.push_back(nodeId);
			for (auto nodeId : nodeIds)
				unloadNode(nodeId);
		}

		auto track(MeshNodePtr node) -> std::size_t
		{
			auto lock = std::unique_lock{m_mutex};
			if (not node)
				return 0;
			if (node.getId() == 0)
				node.setId(m_nextId++);
			auto nodeId = node.getId();
			node.setPolicy(Policy{this});
			m_loadedNodes[nodeId] = node;
			updateLinks(node);
			return nodeId;
		}

		[[nodiscard]] auto acquireFocus(MeshNodePtr node, std::size_t distance = 1) -> ScopedFocus
		{
			auto nodeId = track(node);
			return ScopedFocus{*this, nodeId, distance};
		}

		auto updateFocuses() -> void
		{
			auto lock = std::unique_lock{m_mutex};
			struct FocusState final
			{
				std::size_t nodeId;
				std::size_t distance;
				std::size_t maxDistance;
			};

			auto nodesToKeep = std::unordered_set<std::size_t>{};
			auto queue = std::queue<FocusState>{};

			for (const auto& [focusId, maxDistance] : m_activeFocuses)
				if (focusId != 0 and nodesToKeep.insert(focusId).second)
					queue.push({focusId, 0uz, maxDistance});

			while (not std::ranges::empty(queue))
			{
				auto state = queue.front();
				queue.pop();
				if (state.distance >= state.maxDistance)
					continue;

				if (auto it = m_loadedNodes.find(state.nodeId); it != std::ranges::end(m_loadedNodes))
					updateLinks(it->second);

				for (const auto& [key, neighbors] : m_knownLinks[state.nodeId])
					for (auto neighborId : neighbors)
					{
						if (not m_loadedNodes.contains(neighborId))
							loadNode(neighborId);

						if (m_loadedNodes.contains(neighborId) and not nodesToKeep.contains(neighborId))
						{
							nodesToKeep.insert(neighborId);
							queue.push({neighborId, state.distance + 1, state.maxDistance});
						}
					}
			}

			auto idsToUnload = std::vector<std::size_t>{};
			auto idsToRemove = std::vector<std::size_t>{};
			for (const auto& [nodeId, node] : m_loadedNodes)
				if (not nodesToKeep.contains(nodeId))
				{
					if (node.getDistanceFromRoot() == std::numeric_limits<std::size_t>::max())
						idsToRemove.push_back(nodeId);
					else
						idsToUnload.push_back(nodeId);
				}

			for (auto nodeId : idsToRemove)
				removeNode(nodeId);

			for (auto nodeId : idsToUnload)
				unloadNode(nodeId);

			for (auto id : nodesToKeep)
			{
				auto node = m_loadedNodes[id];
				for (const auto& [key, neighbors] : m_knownLinks[id])
				{
					auto branch = node[key];
					for (auto neighborId : neighbors)
						if (auto it = m_loadedNodes.find(neighborId);
							it != std::ranges::end(m_loadedNodes) and not branch.contains(it->second) and node.getDistanceFromRoot() != std::numeric_limits<std::size_t>::max())
							branch >> it->second;
				}
			}
		}

	private:
		auto addFocus(std::size_t nodeId, std::size_t distance) -> void
		{
			auto lock = std::unique_lock{m_mutex};
			m_activeFocuses.push_back({nodeId, distance});
			updateFocuses();
		}

		auto removeFocus(std::size_t nodeId, std::size_t distance) -> void
		{
			auto lock = std::unique_lock{m_mutex};
			if (auto it = std::ranges::find(m_activeFocuses, std::make_pair(nodeId, distance));
				it != std::ranges::end(m_activeFocuses))
			{
				m_activeFocuses.erase(it);
				updateFocuses();
			}
		}

	public:
		auto updateLinks(MeshNodePtr node) -> void
		{
			auto lock = std::unique_lock{m_mutex};
			auto nodeId = node.getId();
			if (nodeId == 0)
				return;

			auto newLinks = std::unordered_map<Key, std::unordered_set<std::size_t>>{};

			auto nodeAccessor = node->sharedAccess();
			auto branchesAccessor = nodeAccessor->value.branches.sharedAccess();
			for (const auto& [key, branchLinks] : branchesAccessor.value())
			{
				auto& ids = newLinks[key];
				auto vectorAccessor = branchLinks->uniqueAccess();
				for (auto& link : vectorAccessor.value())
				{
					if (link.targetId == 0)
					{
						if (auto shared = link.cache.lock())
						{
							auto neighbor = MeshNodePtr{shared};
							auto neighborId = neighbor.getId();
							if (neighborId == 0)
							{
								neighbor.setId(m_nextId++);
								neighborId = neighbor.getId();
								neighbor.setPolicy(Policy{this});
								m_loadedNodes[neighborId] = neighbor;
							}
							link.targetId = neighborId;
						}
					}
					if (link.targetId != 0)
						ids.insert(link.targetId);
				}
			}
			if (std::ranges::empty(newLinks))
				m_knownLinks.erase(nodeId);
			else
				m_knownLinks[nodeId] = std::move(newLinks);
		}

		auto loadNode(std::size_t nodeId) -> MeshNodePtr
		{
			auto lock = std::unique_lock{m_mutex};
			if (auto it = m_loadedNodes.find(nodeId); it != std::ranges::end(m_loadedNodes))
				return it->second;

			if (auto data = m_storage.load<SerializedNode>(nodeId))
			{
				auto node = data->isRoot ? MeshNodePtr::makeRoot(data->value) : MeshNodePtr::make(data->value);
				node.setId(nodeId);
				node.setPolicy(Policy{this});
				m_loadedNodes[nodeId] = node;
				for (const auto& [branchKey, neighborIds] : data->branches)
				{
					using Link = MeshNodePtr::Link;
					m_knownLinks[nodeId][branchKey] = std::unordered_set<std::size_t>(std::ranges::begin(neighborIds), std::ranges::end(neighborIds));
					auto branchLinks = Container::makeSafeShared<std::vector<Link>>();
					{
						auto links = branchLinks->uniqueAccess();
						for (auto neighborId : neighborIds)
							links->push_back(Link{neighborId, {}});
					}
					node->uniqueAccess()->value.branches.uniqueAccess().value()[branchKey] = std::move(branchLinks);
				}
				return node;
			}
			return MeshNodePtr{};
		}

		auto unloadNode(std::size_t nodeId) -> void
		{
			auto lock = std::unique_lock{m_mutex};
			auto it = m_loadedNodes.find(nodeId);
			if (it == std::ranges::end(m_loadedNodes))
				return;

			updateLinks(it->second);

			auto data = SerializedNode{};
			data.value = it->second.getValue().value();
			for (const auto& [key, ids] : m_knownLinks[nodeId])
				data.branches[key] = std::vector<std::size_t>(std::ranges::begin(ids), std::ranges::end(ids));
			data.isRoot = (it->second.getDistanceFromRoot() == 0);

			m_storage.store(nodeId, data);
			it->second.setPolicy(Policy{nullptr});
			m_loadedNodes.erase(it);
		}

		auto removeNode(std::size_t nodeId) -> void
		{
			auto lock = std::unique_lock{m_mutex};
			if (nodeId == 0)
				return;

			if (auto node = loadNode(nodeId))
			{
				node.clearBranches();
				node.setPolicy(Policy{nullptr});
				m_loadedNodes.erase(nodeId);
			}
			m_knownLinks.erase(nodeId);
			for (auto& [knownId, branchMap] : m_knownLinks)
				for (auto& [key, neighbors] : branchMap)
					neighbors.erase(nodeId);

			m_storage.remove(nodeId);

			updateFocuses();
		}
	};
}

namespace CppUtils::Type::Binary
{
	template<class Key, class Value>
	struct Serializer<CppUtils::FileSystem::MeshSerializedNode<Key, Value>> final
	{
		using SerializedNode = CppUtils::FileSystem::MeshSerializedNode<Key, Value>;

		static inline auto serialize(const SerializedNode& node, std::vector<std::byte>& buffer) -> void
		{
			Type::serialize(node.value, buffer);
			Type::serialize(node.branches, buffer);
			Type::serialize(node.isRoot, buffer);
		}

		[[nodiscard]] static inline auto deserialize(std::span<const std::byte>& view) -> std::expected<SerializedNode, std::string_view>
		{
			auto value = Type::deserialize<Value>(view);
			if (not value)
				return std::unexpected{value.error()};

			auto branches = Type::deserialize<std::unordered_map<Key, std::vector<std::size_t>>>(view);
			if (not branches)
				return std::unexpected{branches.error()};

			auto isRoot = Type::deserialize<bool>(view);
			if (not isRoot)
				return std::unexpected{isRoot.error()};

			auto node = SerializedNode{};
			node.value = std::move(value.value());
			node.branches = std::move(branches.value());
			node.isRoot = isRoot.value();
			return node;
		}
	};
}