Dynamic k-d-tree firstest draft

src/unittest/CMakeLists.txt

@@ -10,6 +10,7 @@ set (UNITTEST_SRCS
 	${CMAKE_CURRENT_SOURCE_DIR}/test_connection.cpp
 	${CMAKE_CURRENT_SOURCE_DIR}/test_craft.cpp
 	${CMAKE_CURRENT_SOURCE_DIR}/test_datastructures.cpp
+	${CMAKE_CURRENT_SOURCE_DIR}/test_k_d_tree.cpp
 	${CMAKE_CURRENT_SOURCE_DIR}/test_filesys.cpp
 	${CMAKE_CURRENT_SOURCE_DIR}/test_inventory.cpp
 	${CMAKE_CURRENT_SOURCE_DIR}/test_irrptr.cpp

src/unittest/test_k_d_tree.cpp

@@ -0,0 +1,107 @@
+/*
+Minetest
+Copyright (C) 2024 sfan5
+
+This program is free software; you can redistribute it and/or modify
+it under the terms of the GNU Lesser General Public License as published by
+the Free Software Foundation; either version 2.1 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU Lesser General Public License for more details.
+
+You should have received a copy of the GNU Lesser General Public License along
+with this program; if not, write to the Free Software Foundation, Inc.,
+51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+*/
+
+#include "noise.h"
+#include "test.h"
+
+#include "util/k_d_tree.h"
+#include <algorithm>
+#include <unordered_set>
+
+class TestKdTree : public TestBase
+{
+public:
+	TestKdTree() { TestManager::registerTestModule(this); }
+	const char *getName() { return "TestKdTree"; }
+
+	void runTests(IGameDef *gamedef);
+
+	// TODO void cube();
+	void randomInserts();
+};
+
+template<uint8_t Dim, typename Component, typename Id>
+class ObjectVector {
+public:
+	using Point = std::array<Component, Dim>;
+	void insert(const Point &p, Id id) {
+		entries.push_back(Entry{p, id});
+	}
+	void remove(Id id) {
+		entries.erase(std::find_if(entries.begin(), entries.end(), [&](const auto &e) {
+			return e.id == id;
+		}));
+	}
+	template<typename F>
+	void rangeQuery(const Point &min, const Point &max, const F &cb) {
+		for (const auto &e : entries) {
+			for (uint8_t d = 0; d < Dim; ++d)
+				if (e.point[d] < min[d] || e.point[d] > max[d])
+					goto next;
+			cb(e.point, e.id); // TODO check
+			next: {}
+		}
+	}
+private:
+	struct Entry {
+		Point point;
+		Id id;
+	};
+	std::vector<Entry> entries;
+};
+
+static TestKdTree g_test_instance;
+
+void TestKdTree::runTests(IGameDef *gamedef)
+{
+	rawstream << "-------- k-d-tree" << std::endl;
+	TEST(randomInserts);
+}
+
+void TestKdTree::randomInserts() {
+	PseudoRandom pr(814538);
+
+	ObjectVector<3, f32, u16> objvec;
+	DynamicKdTrees<3, f32, u16> kds;
+	for (u16 id = 1; id < 1000; ++id) {
+		std::array<f32, 3> point;
+		for (uint8_t d = 0; d < 3; ++d)
+			point[d] = pr.range(-1000, 1000);
+		objvec.insert(point, id);
+		kds.insert(point, id);
+	}
+
+	for (int i = 0; i < 1000; ++i) {
+		std::array<f32, 3> min, max;
+		for (uint8_t d = 0; d < 3; ++d) {
+			min[d] = pr.range(-1500, 1500);
+			max[d] = min[d] + pr.range(1, 2500);
+		}
+		std::unordered_set<u16> expected_ids;
+		objvec.rangeQuery(min, max, [&](auto _, u16 id) {
+			UASSERT(expected_ids.count(id) == 0);
+			expected_ids.insert(id);
+		});
+		kds.rangeQuery(min, max, [&](auto point, u16 id) {
+			UASSERT(expected_ids.count(id) == 1);
+			expected_ids.erase(id);
+		});
+		UASSERT(expected_ids.empty());
+	}
+}

src/util/k_d_tree.h

@@ -0,0 +1,447 @@
+#pragma once
+
+#include <algorithm>
+#include <cstdint>
+#include <unordered_map>
+#include <vector>
+#include <memory>
+#include <iostream>
+
+using Idx = std::size_t;
+
+template<uint8_t Dim, typename Component>
+class Points {
+public:
+	using Point = std::array<Component, Dim>;
+	//! Empty
+	Points() : n(0), coords(nullptr) {}
+	//! Leaves coords uninitialized!
+	Points(Idx n) : n(n), coords(std::make_unique<Component[]>(Dim * n)) {}
+	//! Copying constructor
+	Points(Idx n, const std::array<Component const *, Dim> &coords) : Points(n) {
+		for (uint8_t d = 0; d < Dim; ++d)
+			std::copy(coords[d], coords[d] + n, begin(d));
+	}
+	Idx size() const {
+		return n;
+	}
+	void assign(Idx start, const Points &from) {
+		for (uint8_t d = 0; d < Dim; ++d)
+			std::copy(from.begin(d), from.end(d), begin(d) + start);
+	}
+	Point getPoint(Idx i) const {
+		Point point;
+		for (uint8_t d = 0; d < Dim; ++d)
+			point[d] = begin(d)[i];
+		return point;
+	}
+	void setPoint(Idx i, const Point &point) {
+		for (uint8_t d = 0; d < Dim; ++d)
+			begin(d)[i] = point[d];
+	}
+	// HACK interior mutability...
+	Component *begin(uint8_t d) const {
+	 	return coords.get() + d * n;
+	}
+	Component *end(uint8_t d) const {
+	 	return begin(d) + n;
+	}
+private:
+	Idx n;
+	std::unique_ptr<Component[]> coords;
+};
+
+template<uint8_t Dim>
+class SortedIndices {
+private:
+	// SortedIndices(Points<Dim, Idx> &&indices) : indices(indices) {}
+public:
+	//! empty
+	SortedIndices() : indices() {}
+
+	//! uninitialized indices
+	static SortedIndices newUninitialized(Idx n) {
+		return SortedIndices(n); // HACK can't be arsed to fix rn Points<Dim, Idx>(n));
+	}
+
+	// Identity permutation on all axes
+	SortedIndices(Idx n) : indices(n) {
+		for (uint8_t d = 0; d < Dim; ++d)
+			for (Idx i = 0; i < n; ++i)
+				indices.begin(d)[i] = i;
+	}
+
+	Idx size() const {
+		return indices.size();
+	}
+
+	bool empty() const {
+		return size() == 0;
+	}
+
+	struct SplitResult {
+		SortedIndices left, right;
+		Idx pivot;
+	};
+
+	//! Splits the sorted indices in the middle along the specified axis,
+	//! partitioning them into left (<=), the pivot, and right (>=).
+	SplitResult split(uint8_t axis, std::vector<bool> &markers) const {
+		const auto begin = indices.begin(axis);
+		Idx left_n = indices.size() / 2;
+		const auto mid = begin + left_n;
+
+		// Mark all points to be partitioned left
+		for (auto it = begin; it != mid; ++it)
+			markers[*it] = true;
+
+		SortedIndices left(left_n);
+		std::copy(begin, mid, left.indices.begin(axis));
+		SortedIndices right(indices.size() - left_n - 1);
+		std::copy(mid + 1, indices.end(axis), right.indices.begin(axis));
+
+		for (uint8_t d = 0; d < Dim; ++d) {
+			if (d == axis)
+				continue;
+			auto left_ptr = left.indices.begin(d);
+			auto right_ptr = right.indices.begin(d);
+			for (auto it = indices.begin(d); it != indices.end(d); ++it) {
+				if (*it != *mid) { // ignore pivot
+					*(markers[*it] ? left_ptr++ : right_ptr++) = *it;
+				}
+			}
+		}
+
+		// Unmark points, since we want to reuse the storage for markers
+		for (auto it = begin; it != mid; ++it)
+			markers[*it] = false;
+
+		return SplitResult{std::move(left), std::move(right), *mid};
+	}
+
+	Idx *begin(uint8_t d) const {
+		return indices.begin(d);
+	}
+
+	Idx *end(uint8_t d) const {
+		return indices.end(d);
+	}
+private:
+	Points<Dim, Idx> indices;
+};
+
+template<uint8_t Dim, class Component>
+class SortedPoints {
+public:
+	SortedPoints() : points(), indices() {}
+
+	//! Single point
+	// TODO remove this
+	SortedPoints(const std::array<Component, Dim> &point) : points(1), indices(1) {
+		points.setPoint(0, point);
+	}
+
+	//! Sort points
+	SortedPoints(Idx n, const std::array<Component const *, Dim> ptrs)
+		: points(n, ptrs), indices(n)
+	{
+		for (uint8_t d = 0; d < Dim; ++d) {
+			const auto coord = points.begin(d);
+			std::sort(indices.begin(d), indices.end(d), [&](auto i, auto j) {
+				return coord[i] < coord[j];
+			});
+		}
+	}
+
+	//! Merge two sets of sorted points
+	SortedPoints(const SortedPoints &a, const SortedPoints &b)
+		: points(a.size() + b.size())
+	{
+		const auto n = points.size();
+		indices = SortedIndices<Dim>::newUninitialized(n);
+		for (uint8_t d = 0; d < Dim; ++d) {
+			points.assign(0, a.points);
+			points.assign(a.points.size(), b.points);
+			const auto coord = points.begin(d);
+			auto a_ptr = a.indices.begin(d);
+			auto b_ptr = b.indices.begin(d);
+			auto dst_ptr = indices.begin(d);
+			while (a_ptr != a.indices.end(d) && b_ptr != b.indices.end(d)) {
+				const auto i = *a_ptr;
+				const auto j = *b_ptr + a.size();
+				if (coord[i] <= coord[j]) {
+					*(dst_ptr++) = i;
+					++a_ptr;
+				} else {
+					*(dst_ptr++) = j;
+					++b_ptr;
+				}
+			}
+			while (a_ptr != a.indices.end(d))
+				*(dst_ptr++) = *(a_ptr++);
+			while (b_ptr != b.indices.end(d))
+				*(dst_ptr++) = a.size() + *(b_ptr++);
+		}
+	}
+
+	Idx size() const {
+		// technically redundant with indices.size(),
+		// but that is irrelevant
+		return points.size();
+	}
+	
+// HACK private:
+	Points<Dim, Component> points;
+	SortedIndices<Dim> indices;
+};
+
+template<uint8_t Dim, class Component, class Id>
+class KdTree {
+public:
+	using Point = std::array<Component, Dim>;
+
+	//! Empty tree
+	KdTree()
+		: items()
+		, ids(nullptr)
+		, tree(nullptr)
+		, deleted()
+	{}
+
+	//! Build a tree containing just a single point
+	// TODO this will probably be obsolete soon (TM)
+	KdTree(const Point &point, const Id &id)
+		: items(point)
+		, ids(std::make_unique<Id[]>(1))
+		, tree(std::make_unique<Idx[]>(1))
+		, deleted(1)
+	{
+		tree[0] = 0;
+		ids[0] = id;
+	}
+
+	//! Build a tree
+	KdTree(Idx n, Id const *ids, std::array<Component const *, Dim> pts)
+		: items(n, pts)
+		, tree(std::make_unique<Id[]>(n))
+		, ids(std::make_unique<Id[]>(n))
+		, deleted(n)
+	{
+		std::copy(ids, ids + n, this->ids.get());
+		init(0, items);
+	}
+
+	//! Merge two trees. Both trees are assumed to have a power of two size.
+	KdTree(const KdTree &a, const KdTree &b)
+		: items(a.items, b.items)
+	{
+		tree = std::make_unique<Idx[]>(cap());
+		ids = std::make_unique<Id[]>(cap());
+		std::copy(a.ids.get(), a.ids.get() + a.cap(), ids.get());
+		std::copy(b.ids.get(), b.ids.get() + b.cap(), ids.get() + a.cap());
+		deleted = std::vector<bool>(cap());
+		std::copy(a.deleted.begin(), a.deleted.end(), deleted.begin());
+		std::copy(b.deleted.begin(), b.deleted.end(), deleted.begin() + a.items.size());
+		init(0, 0, items.indices);
+	}
+
+	// TODO remove dead code
+	/* SortedPoints<Dim, Component> alivePoints() {
+		const auto newIndices = std::make_unique<Idx[]>(items.n);
+		Idx j = 0;
+		for (Idx i = 0; i < items.n; ++i)
+			newIndices[i] = deleted[i] ? j : j++;
+		Idx aliveN = std::count(deleted.begin(), deleted.end(), false);
+		SortedPoints<Dim, Component> res(aliveN);
+		for (uint8_t d = 0; d < Dim; ++d) {
+			const auto pts = &res.points[d * aliveN];
+			const auto indices = &res.indices[d * aliveN];
+			const auto items_pts = &items.points[d * items.n];
+			const auto items_indices = &items.indices[d * aliveN];
+			Idx j = 0;
+			for (Idx i = 0; i < items.n; ++i) {
+				if (deleted[i]) continue;
+				pts[j++] = items_pts[i];
+				indices[newIndices[i]] = items_indices[i];
+			}
+		}
+		return res;
+	} */
+
+	template<typename F>
+	void rangeQuery(const Point &min, const Point &max,
+			const F &cb) const {
+		if (empty())
+			return;
+		rangeQuery(0, 0, min, max, cb);
+	}
+
+	void remove(Idx internalIdx) {
+		deleted[internalIdx] = true;
+	}
+
+	template<class F>
+	void foreach(F cb) {
+		for (Idx i = 0; i < cap(); ++i) {
+			if (!deleted[i]) {
+				cb(i, items.points.getPoint(i), ids[i]);
+			}
+		}
+	}
+
+	//! Capacity, not size, since some items may be marked as deleted
+	Idx cap() const {
+		return items.size();
+	}
+
+	//! "Empty" as in "never had anything"
+	bool empty() const {
+		return cap() == 0;
+	}
+
+private:
+	void init(Idx root, uint8_t axis, const SortedIndices<Dim> &sorted) {
+		// HACK abuse deleted marks as left/right marks
+		const auto split = sorted.split(axis, deleted);
+		tree[root] = split.pivot;
+		const auto next_axis = (axis + 1) % Dim;
+		if (!split.left.empty())
+			init(2 * root + 1, next_axis, split.left);
+		if (!split.right.empty())
+			init(2 * root + 2, next_axis, split.right);
+	}
+
+	template<typename F>
+	void rangeQuery(Idx root, uint8_t split,
+			const Point &min, const Point &max,
+			const F &cb) const {
+		if (root >= cap()) // TODO what if we overflow earlier?
+			return;
+		const auto ptid = tree[root];
+		const auto coord = items.points.begin(split)[ptid];
+		const auto leftChild = 2*root + 1;
+		const auto rightChild = 2*root + 2;
+		const auto nextSplit = (split + 1) % Dim;
+		if (min[split] > coord) {
+			rangeQuery(rightChild, nextSplit, min, max, cb);
+		} else if (max[split] < coord) {
+			rangeQuery(leftChild, nextSplit, min, max, cb);
+		} else {
+			rangeQuery(rightChild, nextSplit, min, max, cb);
+			rangeQuery(leftChild, nextSplit, min, max, cb);
+			if (deleted[root])
+				return;
+			const auto point = items.points.getPoint(ptid);
+			for (uint8_t d = 0; d < Dim; ++d)
+				if (point[d] < min[d] || point[d] > max[d])
+					return;
+			cb(items.points.getPoint(ptid), ids[ptid]);
+		}
+	}
+	SortedPoints<Dim, Component> items;
+	std::unique_ptr<Id[]> ids;
+	std::unique_ptr<Idx[]> tree;
+	// vector because this has the template specialization we want
+	// and i'm too lazy to implement bitsets myself right now
+	// just to shave off 16 redundant bytes (len + cap)
+	std::vector<bool> deleted;
+};
+
+template<uint8_t Dim, class Component, class Id>
+class DynamicKdTrees {
+	using Tree = KdTree<Dim, Component, Id>;
+public:
+	using Point = typename Tree::Point;
+	void insert(const std::array<Component, Dim> &point, const Id id) {
+		Tree tree(point, id);
+		for (uint8_t treeIdx = 0;; ++treeIdx) {
+			if (treeIdx >= trees.size()) {
+				tree.foreach([&](Idx objIdx, auto _, Id id) {
+					del_entries[id] = {treeIdx, objIdx};
+				});
+				trees.push_back(std::move(tree));
+				break;
+			}
+			if (trees[treeIdx].empty()) {
+				// TODO deduplicate
+				tree.foreach([&](Idx objIdx, auto _, Id id) {
+					del_entries[id] = {treeIdx, objIdx};
+				});
+				trees[treeIdx] = std::move(tree);
+				break;
+			}
+			tree = Tree(tree, trees[treeIdx]);
+			trees[treeIdx] = std::move(Tree());
+		}
+		++n_entries;
+	}
+	void remove(const Id id) {
+		const auto del_entry = del_entries.at(id);
+		trees.at(del_entry.treeIdx).remove(del_entry.inTree);
+		del_entries.erase(del_entry);
+		++deleted;
+		if (deleted > n_entries/2) // we want to shift out the one!
+			compactify();
+	}
+	template<typename F>
+	void rangeQuery(const Point &min, const Point &max,
+			const F &cb) const {
+		for (const auto &tree : trees)
+			tree.rangeQuery(min, max, cb);
+	}
+private:
+	void compactify() {
+		n_entries -= deleted;
+		del_entries = std::unordered_map<Id, DelEntry>();
+
+		// Collect all live points and corresponding IDs.
+		const auto ids = std::make_unique<Id[]>(n_entries);
+		Points<Dim, Component> pts;
+		Idx i = 0;
+		for (const auto tree : trees) {
+			if (tree.empty())
+				continue;
+			tree.foreach([&](Idx _, std::array<Component, Dim> point, Id id) {
+				for (uint8_t d = 0; d < Dim; ++d)
+					pts.get(d)[i] = point[d];
+				ids[i] = id;
+				++i;
+			});
+		}
+
+		// Construct a new forest.
+		// The pattern will be the current pattern, shifted down by one.
+		auto id_ptr = ids.get();
+		std::array<Component const *, Dim> point_ptrs;
+		std::vector<Tree> newTrees(trees.size() - 1);
+		Idx n = 1;
+		for (uint8_t d = 0; d < Dim; ++d)
+			point_ptrs[d] = pts.begin(d);
+		for (uint8_t i = 0; i < newTrees.size(); ++i, n *= 2) {
+			if (trees[i+1].empty())
+				continue;
+
+			// TODO maybe optimize from log² -> log?
+			// This can be achieved by doing a sorted merge of live points,
+			// then doing a radix sort.
+			Tree tree(n, id_ptr, point_ptrs);
+			id_ptr += n;
+			for (uint8_t d = 0; d < Dim; ++d)
+				point_ptrs[d] += n;
+			tree.foreach([&](Idx i, auto _, Id id) {
+				del_entries[id] = {n, i};
+			});
+			newTrees[i] = std::move(tree);
+		}
+		trees = std::move(newTrees);
+	}
+	// could use an array here since we've got a good bound on the size ahead of time but meh
+	std::vector<Tree> trees;
+	struct DelEntry {
+		uint8_t treeIdx;
+		Idx inTree;
+	};
+	std::unordered_map<Id, DelEntry> del_entries;
+	Idx n_entries;
+	Idx deleted;
+};