C++ implementation of a sparse voxel octree

svo.hpp

#ifndef SVO_HPP
#define SVO_HPP

#include "common_types.hpp"
#include "vec.hpp"
#include "vec_math.hpp"
#include "util_bits.hpp"
#include "util_common.hpp"

#include <algorithm>

namespace mve {

struct SVONode {
    virtual ~SVONode() = 0;
};

struct SVOBranch : public SVONode {
    std::array<uptr<SVONode>, 8> children;

    ~SVOBranch() final;
};

struct SVOLeaf : public SVONode {
    std::array<rgb32_t, 8> data{};

    ~SVOLeaf() final;
};

/** Sparse Voxel Octree. */
class SVO
{
private:
    using i32 = int32_t;
    using u32 = uint32_t;
    using u64 = uint64_t;

    uptr<SVOBranch> root = std::make_unique<SVOBranch>();
    size_t depth = 1;

public:
    SVO() = default;

    void insert(Vec3i32 pos, rgb32_t color) {
        ensureSpace(pos);
        auto octreeNodeIndex = indexOf(pos);
        insert(octreeNodeIndex, color);
    }

    Vec3i32 minIncl() const {
        return Vec3i32::filledWith(-(1 << depth));
    }


    Vec3i32 maxIncl() const {
        return Vec3i32::filledWith((1 << depth) - 1);
    }

    Vec3i32 minExcl() const {
        return Vec3i32::filledWith(-(1 << depth) - 1);
    }


    Vec3i32 maxExcl() const {
        return Vec3i32::filledWith(1 << depth);
    }

    rgb32_t& operator[](Vec3i32 pos) {
        ensureSpace(pos);
        auto octreeNodeIndex = indexOf(pos);
        return findOrCreate(octreeNodeIndex);
    }

    rgb32_t& at(Vec3i32 pos) {
        u32 lim = boundsTest(pos);
        ALWAYS_ASSERT(lim == 0);
        auto octreeNodeIndex = indexOf(pos);
        auto *result = find(octreeNodeIndex);
        ALWAYS_ASSERT(result != nullptr);
        return *result;
    }

    const rgb32_t& at(Vec3i32 pos) const {
        u32 lim = boundsTest(pos);
        ALWAYS_ASSERT(lim == 0);
        auto octreeNodeIndex = indexOf(pos);
        auto *result = find(octreeNodeIndex);
        ALWAYS_ASSERT(result != nullptr);
        return *result;
    }

private:
    rgb32_t& findOrCreate(u64 octreeNodeIndex) {
        SVONode *node = root.get();
        for (size_t s = depth * 3; s != size_t(-3); s -= 3) {
            u32 octDigit = (octreeNodeIndex >> s) & 0b111;
            if (s != 0) {
                auto *branch = downcast<SVOBranch*>(node);
                if (branch->children[octDigit] != nullptr) {
                    node = branch->children[octDigit].get();
                }
                else {
                    auto *child = s == 3 ? static_cast<SVONode*>(new SVOLeaf)
                                         : static_cast<SVONode*>(new SVOBranch);
                    node = child;
                    branch->children[octDigit] = std::unique_ptr<SVONode>{child};
                }
            }
            else {
                auto *leaf = downcast<SVOLeaf*>(node);
                return leaf->data[octDigit];
            }
        }
        DEBUG_ASSERT_UNREACHABLE();
    }

    rgb32_t* find(u64 octreeNodeIndex) const {
        SVONode *node = root.get();
        for (size_t s = depth * 3; s != size_t(-3); s -= 3) {
            u32 octDigit = (octreeNodeIndex >> s) & 0b111;
            if (s != 0) {
                auto *branch = downcast<SVOBranch*>(node);
                if (branch->children[octDigit] == nullptr) {
                    return nullptr;
                }
                else {
                    node = branch->children[octDigit].get();
                }
            }
            else {
                auto *leaf = downcast<SVOLeaf*>(node);
                return &leaf->data[octDigit];
            }
        }
        DEBUG_ASSERT_UNREACHABLE();
    }

    u64 indexOf(Vec3i32 pos) const {
        Vec3u32 uPos = static_vec_cast<u32>(pos - minIncl());
        return ileave3(uPos[0], uPos[1], uPos[2]);
    }

    void ensureSpace(Vec3i32 pos) {
        if (u32 lim = boundsTest(pos); lim != 0) {
            grow(lim);
        }
    }

    void insert(u64 octreeNodeIndex, u32 color) {
        findOrCreate(octreeNodeIndex) = color;
    }

    void grow(u32 lim) {
        for (size_t size = 1 << depth; size <= lim; depth <<= 1, size = 1 << depth) {
            growOnce();
        }
    }

    void growOnce() {
        for (size_t i = 0; i < 8; ++i) {
            if (root->children[i] == nullptr) {
                continue;
            }
            auto parent = std::make_unique<SVOBranch>();
            parent->children[~i & 0b111] = std::move(root->children[i]);
            root->children[i] = std::move(parent);
        }
    }

    /**
     * Tests whether the input vector lies within this octree. The result is an unsigned integer which indicates how
     * much the octree has to be enlarged to fit the vector.
     * The dimensions of the octree in all directions need to be > this integer.
     * @param v the input position
     * @return 0 if the test passes or a maximum-like coordinate if the test fails
     */
    uint32_t boundsTest(Vec3i32 v) const
    {
        constexpr auto absForBoundsTest = [](int32_t x) -> uint32_t {
            return static_cast<uint32_t>(x < 0 ? -x - 1 : x);
        };
        static_assert (absForBoundsTest(-5) == 4);
        u32 max = absForBoundsTest(v[0]) | absForBoundsTest(v[1]) | absForBoundsTest(v[2]);
        return max >= (1u << depth) ? max : 0;
    }
};

}

#endif // SVO_HPP