diff --git a/src/unittest/CMakeLists.txt b/src/unittest/CMakeLists.txt index ec52ee6bf..128993a7f 100644 --- a/src/unittest/CMakeLists.txt +++ b/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 diff --git a/src/unittest/test_k_d_tree.cpp b/src/unittest/test_k_d_tree.cpp new file mode 100644 index 000000000..00ce2a571 --- /dev/null +++ b/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 +#include + +class TestKdTree : public TestBase +{ +public: + TestKdTree() { TestManager::registerTestModule(this); } + const char *getName() { return "TestKdTree"; } + + void runTests(IGameDef *gamedef); + + // TODO void cube(); + void randomInserts(); +}; + +template +class ObjectVector { +public: + using Point = std::array; + 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 + 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 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 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 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 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()); + } +} \ No newline at end of file diff --git a/src/util/k_d_tree.h b/src/util/k_d_tree.h new file mode 100644 index 000000000..3f2972059 --- /dev/null +++ b/src/util/k_d_tree.h @@ -0,0 +1,447 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +using Idx = std::size_t; + +template +class Points { +public: + using Point = std::array; + //! Empty + Points() : n(0), coords(nullptr) {} + //! Leaves coords uninitialized! + Points(Idx n) : n(n), coords(std::make_unique(Dim * n)) {} + //! Copying constructor + Points(Idx n, const std::array &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 coords; +}; + +template +class SortedIndices { +private: + // SortedIndices(Points &&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(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 &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 indices; +}; + +template +class SortedPoints { +public: + SortedPoints() : points(), indices() {} + + //! Single point + // TODO remove this + SortedPoints(const std::array &point) : points(1), indices(1) { + points.setPoint(0, point); + } + + //! Sort points + SortedPoints(Idx n, const std::array 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::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 points; + SortedIndices indices; +}; + +template +class KdTree { +public: + using Point = std::array; + + //! 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(1)) + , tree(std::make_unique(1)) + , deleted(1) + { + tree[0] = 0; + ids[0] = id; + } + + //! Build a tree + KdTree(Idx n, Id const *ids, std::array pts) + : items(n, pts) + , tree(std::make_unique(n)) + , ids(std::make_unique(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(cap()); + ids = std::make_unique(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(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 alivePoints() { + const auto newIndices = std::make_unique(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 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 + 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 + 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 &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 + 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 items; + std::unique_ptr ids; + std::unique_ptr 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 deleted; +}; + +template +class DynamicKdTrees { + using Tree = KdTree; +public: + using Point = typename Tree::Point; + void insert(const std::array &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 + 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(); + + // Collect all live points and corresponding IDs. + const auto ids = std::make_unique(n_entries); + Points pts; + Idx i = 0; + for (const auto tree : trees) { + if (tree.empty()) + continue; + tree.foreach([&](Idx _, std::array 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 point_ptrs; + std::vector 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 trees; + struct DelEntry { + uint8_t treeIdx; + Idx inTree; + }; + std::unordered_map del_entries; + Idx n_entries; + Idx deleted; +}; \ No newline at end of file