Merge pull request #80 from madmann91/rework_bvh_index

Rework node index data type

src/bvh/v2/bvh.h

@@ -8,6 +8,7 @@
 #include <vector>
 #include <stack>
 #include <utility>
+#include <algorithm>
 
 namespace bvh::v2 {
 
@@ -28,8 +29,28 @@ struct Bvh {
     bool operator == (const Bvh& other) const = default;
     bool operator != (const Bvh& other) const = default;
 
+    /// Returns whether the node located at the given index is the left child of its parent.
+    static BVH_ALWAYS_INLINE bool is_left_sibling(size_t node_id) { return node_id % 2 == 1; }
+
+    /// Returns the index of a sibling of a node.
+    static BVH_ALWAYS_INLINE size_t get_sibling_id(size_t node_id) {
+        return is_left_sibling(node_id) ? node_id + 1 : node_id - 1;
+    }
+
+    /// Returns the index of the left sibling of the node. This effectively returns the given index
+    /// unchanged if the node is the left sibling, or the other sibling index otherwise.
+    static BVH_ALWAYS_INLINE size_t get_left_sibling_id(size_t node_id) {
+        return is_left_sibling(node_id) ? node_id : node_id - 1;
+    }
+
+    /// Returns the index of the right sibling of the node. This effectively returns the given index
+    /// unchanged if the node is the right sibling, or the other sibling index otherwise.
+    static BVH_ALWAYS_INLINE size_t get_right_sibling_id(size_t node_id) {
+        return is_left_sibling(node_id) ? node_id + 1 : node_id;
+    }
+
     /// Returns the root node of this BVH.
-    const Node& get_root() const { return nodes[0]; }
+    BVH_ALWAYS_INLINE const Node& get_root() const { return nodes[0]; }
 
     /// Extracts the BVH rooted at the given node index.
     inline Bvh extract_bvh(size_t root_id) const;
@@ -79,18 +100,17 @@ Bvh<Node> Bvh<Node>::extract_bvh(size_t root_id) const {
         auto& dst_node = bvh.nodes[dst_id];
         dst_node = src_node;
         if (src_node.is_leaf()) {
-            dst_node.index.first_id = static_cast<typename Index::Type>(bvh.prim_ids.size());
+            dst_node.index.set_first_id(bvh.prim_ids.size());
             std::copy_n(
-                prim_ids.begin() + src_node.index.first_id,
-                src_node.index.prim_count,
+                prim_ids.begin() + src_node.index.first_id(),
+                src_node.index.prim_count(),
                 std::back_inserter(bvh.prim_ids));
         } else {
-            size_t first_id = bvh.nodes.size();
-            dst_node.index.first_id = static_cast<typename Index::Type>(first_id);
+            dst_node.index.set_first_id(bvh.nodes.size());
             bvh.nodes.emplace_back();
             bvh.nodes.emplace_back();
-            stack.emplace(src_node.index.first_id + 0, first_id + 0);
-            stack.emplace(src_node.index.first_id + 1, first_id + 1);
+            stack.emplace(src_node.index.first_id() + 0, dst_node.index.first_id() + 0);
+            stack.emplace(src_node.index.first_id() + 1, dst_node.index.first_id() + 1);
         }
     }
     return bvh;
@@ -104,9 +124,9 @@ void Bvh<Node>::traverse_top_down(Index start, Stack& stack, LeafFn&& leaf_fn, I
 restart:
     while (!stack.is_empty()) {
         auto top = stack.pop();
-        while (top.prim_count == 0) {
-            auto& left  = nodes[top.first_id];
-            auto& right = nodes[top.first_id + 1];
+        while (top.prim_count() == 0) {
+            auto& left  = nodes[top.first_id()];
+            auto& right = nodes[top.first_id() + 1];
             auto [hit_left, hit_right, should_swap] = inner_fn(left, right);
 
             if (hit_left) {
@@ -124,7 +144,7 @@ void Bvh<Node>::traverse_top_down(Index start, Stack& stack, LeafFn&& leaf_fn, I
                 goto restart;
         }
 
-        [[maybe_unused]] auto was_hit = leaf_fn(top.first_id, top.first_id + top.prim_count);
+        [[maybe_unused]] auto was_hit = leaf_fn(top.first_id(), top.first_id() + top.prim_count());
         if constexpr (IsAnyHit) {
             if (was_hit) return;
         }
@@ -163,8 +183,8 @@ void Bvh<Node>::traverse_bottom_up(LeafFn&& leaf_fn, InnerFn&& inner_fn) {
     for (size_t i = 0; i < nodes.size(); ++i) {
         if (nodes[i].is_leaf())
             continue;
-        parents[nodes[i].index.first_id] = i;
-        parents[nodes[i].index.first_id + 1] = i;
+        parents[nodes[i].index.first_id()] = i;
+        parents[nodes[i].index.first_id() + 1] = i;
     }
     std::vector<bool> seen(nodes.size(), false);
     for (size_t i = nodes.size(); i-- > 0;) {
@@ -175,7 +195,7 @@ void Bvh<Node>::traverse_bottom_up(LeafFn&& leaf_fn, InnerFn&& inner_fn) {
         size_t j = parents[i];
         while (j >= 0) {
             auto& node = nodes[j];
-            if (seen[j] || !seen[node.index.first_id] || !seen[node.index.first_id + 1])
+            if (seen[j] || !seen[node.index.first_id()] || !seen[node.index.first_id() + 1])
                 break;
             inner_fn(nodes[j]);
             seen[j] = true;
@@ -188,8 +208,8 @@ template <typename Node>
 template <typename LeafFn>
 void Bvh<Node>::refit(LeafFn&& leaf_fn) {
     traverse_bottom_up(leaf_fn, [&] (Node& node) {
-        const auto& left  = nodes[node.index.first_id];
-        const auto& right = nodes[node.index.first_id + 1];
+        const auto& left  = nodes[node.index.first_id()];
+        const auto& right = nodes[node.index.first_id() + 1];
         node.set_bbox(left.get_bbox().extend(right.get_bbox()));
     });
 }

src/bvh/v2/index.h

@@ -0,0 +1,83 @@
+#ifndef BVH_V2_INDEX_H
+#define BVH_V2_INDEX_H
+
+#include "bvh/v2/utils.h"
+
+#include <cassert>
+#include <cstddef>
+
+namespace bvh::v2 {
+
+/// Packed index data structure. This index can either refer to a range of primitives for a BVH
+/// leaf, or to the children of a BVH node. In either case, the index corresponds to a contiguous
+/// range, which means that:
+///
+/// - For leaves, primitives in a BVH node should be accessed via:
+///
+///     size_t begin = index.first_id();
+///     size_t end   = begin + index.prim_count();
+///     for (size_t i = begin; i < end; ++i) {
+///         size_t prim_id = bvh.prim_ids[i];
+///         // ...
+///     }
+///
+///   Note that for efficiency, reordering the original data to avoid the indirection via
+///   `bvh.prim_ids` is preferable.
+///
+/// - For inner nodes, children should be accessed via:
+///
+///   auto& left_child = bvh.nodes[index.first_id()];
+///   auto& right_child = bvh.nodes[index.first_id() + 1];
+///
+template <size_t Bits, size_t PrimCountBits>
+struct Index {
+    using Type = UnsignedIntType<Bits>;
+
+    static constexpr size_t bits = Bits;
+    static constexpr size_t prim_count_bits = PrimCountBits;
+    static constexpr Type max_prim_count = make_bitmask<Type>(prim_count_bits);
+
+    static_assert(PrimCountBits < Bits);
+
+    Type value;
+
+    Index() = default;
+    explicit Index(Type value) : value(value) {}
+
+    bool operator == (const Index&) const = default;
+    bool operator != (const Index&) const = default;
+
+    BVH_ALWAYS_INLINE Type first_id() const { return value >> prim_count_bits; }
+    BVH_ALWAYS_INLINE Type prim_count() const { return value & max_prim_count; }
+    BVH_ALWAYS_INLINE bool is_leaf() const { return prim_count() != 0; }
+    BVH_ALWAYS_INLINE bool is_inner() const { return !is_leaf(); }
+
+    BVH_ALWAYS_INLINE void set_first_id(Type first_id) {
+        *this = Index { first_id, prim_count() };
+    }
+
+    BVH_ALWAYS_INLINE void set_prim_count(Type prim_count) {
+        *this = Index { first_id(), prim_count };
+    }
+
+    static BVH_ALWAYS_INLINE Index make_leaf(Type first_prim, Type prim_count) {
+        assert(prim_count != 0);
+        return Index { first_prim, prim_count };
+    }
+
+    static BVH_ALWAYS_INLINE Index make_inner(Type first_child) {
+        return Index { first_child, 0 };
+    }
+
+private:
+    explicit Index(Type first_id, Type prim_count)
+        : value((first_id << prim_count_bits) | (prim_count & max_prim_count))
+    {
+        assert(first_id <= make_bitmask<Type>(Bits - PrimCountBits));
+        assert(prim_count <= max_prim_count);
+    }
+};
+
+} // namespace bvh::v2
+
+#endif

src/bvh/v2/mini_tree_builder.h

@@ -225,8 +225,8 @@ class MiniTreeBuilder {
                 if (node.get_bbox().get_half_area() < threshold || node.is_leaf()) {
                     pruned_roots.emplace_back(i, node_id);
                 } else {
-                    stack.push(node.index.first_id);
-                    stack.push(node.index.first_id + 1);
+                    stack.push(node.index.first_id());
+                    stack.push(node.index.first_id() + 1);
                 }
             }
         }
@@ -274,16 +274,16 @@ class MiniTreeBuilder {
         // Helper function to copy and fix the child/primitive index of a node
         auto copy_node = [&] (size_t i, Node& dst_node, const Node& src_node) {
             dst_node = src_node;
-            dst_node.index.first_id += static_cast<typename Node::Index::Type>(
-                src_node.is_leaf() ? prim_offsets[i] : node_offsets[i]);
+            dst_node.index.set_first_id(dst_node.index.first_id() +
+                (src_node.is_leaf() ? prim_offsets[i] : node_offsets[i]));
         };
 
         // Make the leaves of the top BVH point to the right internal nodes
         for (auto& node : bvh.nodes) {
             if (!node.is_leaf())
                 continue;
-            assert(node.index.prim_count == 1);
-            size_t tree_id = bvh.prim_ids[node.index.first_id];
+            assert(node.index.prim_count() == 1);
+            size_t tree_id = bvh.prim_ids[node.index.first_id()];
             copy_node(tree_id, node, mini_trees[tree_id].get_root());
         }
 

src/bvh/v2/node.h

@@ -1,7 +1,7 @@
 #ifndef BVH_V2_NODE_H
 #define BVH_V2_NODE_H
 
-#include "bvh/v2/utils.h"
+#include "bvh/v2/index.h"
 #include "bvh/v2/vec.h"
 #include "bvh/v2/bbox.h"
 #include "bvh/v2/ray.h"
@@ -13,64 +13,35 @@
 
 namespace bvh::v2 {
 
+/// Binary BVH node, containing its bounds and an index into its children or the primitives it
+/// contains. By definition, inner BVH nodes do not contain primitives; only leaves do.
 template <
     typename T,
     size_t Dim,
     size_t IndexBits = sizeof(T) * CHAR_BIT,
     size_t PrimCountBits = 4>
 struct Node {
     using Scalar = T;
+    using Index = bvh::v2::Index<IndexBits, PrimCountBits>;
+
     static constexpr size_t dimension = Dim;
     static constexpr size_t prim_count_bits = PrimCountBits;
     static constexpr size_t index_bits = IndexBits;
-    static constexpr size_t max_prim_count = make_bitmask<size_t>(prim_count_bits);
 
+    /// Bounds of the node, laid out in memory as `[min_x, max_x, min_y, max_y, ...]`. Users should
+    /// not really depend on a specific order and instead use `get_bbox()` and extract the `min`
+    /// and/or `max` components accordingly.
     std::array<T, Dim * 2> bounds;
-    struct Index {
-        using Type = UnsignedIntType<IndexBits>;
-        Type first_id   : std::numeric_limits<Type>::digits - prim_count_bits;
-        Type prim_count : prim_count_bits;
-
-        BVH_ALWAYS_INLINE bool operator == (const Index& other) const {
-            return first_id == other.first_id && prim_count == other.prim_count;
-        }
 
-        bool operator != (const Index&) const = default;
-    } index;
-
-    static_assert(sizeof(Index) == sizeof(typename Index::Type));
+    /// Index to the children of an inner node, or to the primitives for a leaf node.
+    Index index;
 
     Node() = default;
 
     bool operator == (const Node&) const = default;
     bool operator != (const Node&) const = default;
 
-    BVH_ALWAYS_INLINE bool is_leaf() const { return index.prim_count != 0; }
-    static BVH_ALWAYS_INLINE bool is_left_sibling(size_t node_id) { return node_id % 2 == 1; }
-
-    static BVH_ALWAYS_INLINE size_t get_sibling_id(size_t node_id) {
-        return is_left_sibling(node_id) ? node_id + 1 : node_id - 1;
-    }
-
-    static BVH_ALWAYS_INLINE size_t get_left_sibling_id(size_t node_id) {
-        return is_left_sibling(node_id) ? node_id : node_id - 1;
-    }
-
-    static BVH_ALWAYS_INLINE size_t get_right_sibling_id(size_t node_id) {
-        return is_left_sibling(node_id) ? node_id + 1 : node_id;
-    }
-
-    BVH_ALWAYS_INLINE void make_leaf(size_t first_prim, size_t prim_count) {
-        assert(prim_count != 0);
-        assert(prim_count <= max_prim_count);
-        index.prim_count = static_cast<typename Index::Type>(prim_count);
-        index.first_id = static_cast<typename Index::Type>(first_prim);
-    }
-
-    BVH_ALWAYS_INLINE void make_inner(size_t first_child) {
-        index.prim_count = 0;
-        index.first_id = static_cast<typename Index::Type>(first_child);
-    }
+    BVH_ALWAYS_INLINE bool is_leaf() const { return index.is_leaf(); }
 
     BVH_ALWAYS_INLINE BBox<T, Dim> get_bbox() const {
         return BBox<T, Dim>(
@@ -119,16 +90,14 @@ struct Node {
     BVH_ALWAYS_INLINE void serialize(OutputStream& stream) const {
         for (auto&& bound : bounds)
             stream.write(bound);
-        stream.write(static_cast<size_t>(index.first_id));
-        stream.write(static_cast<size_t>(index.prim_count));
+        stream.write(index.value);
     }
 
-    static inline Node deserialize(InputStream& stream) {
+    static BVH_ALWAYS_INLINE Node deserialize(InputStream& stream) {
         Node node;
         for (auto& bound : node.bounds)
             bound = stream.read<T>();
-        node.index.first_id = stream.read<size_t>();
-        node.index.prim_count = stream.read<size_t>();
+        node.index = Index(stream.read<typename Index::Type>());
         return node;
     }