bvh/v2/bvh.h

0001 #ifndef BVH_V2_BVH_H
0002 #define BVH_V2_BVH_H
0003
0004 #include "bvh/v2/node.h"
0005
0006 #include <cstddef>
0007 #include <iterator>
0008 #include <vector>
0009 #include <stack>
0010 #include <utility>
0011 #include <tuple>
0012 #include <algorithm>
0013
0014 namespace bvh::v2 {
0015
0016 template <typename Node>
0017 struct Bvh {
0018     using Index = typename Node::Index;
0019     using Scalar = typename Node::Scalar;
0020     using Ray = bvh::v2::Ray<Scalar, Node::dimension>;
0021
0022     std::vector<Node> nodes;
0023     std::vector<size_t> prim_ids;
0024
0025     Bvh() = default;
0026     Bvh(Bvh&&) = default;
0027
0028     Bvh& operator = (Bvh&&) = default;
0029
0030     //bool operator == (const Bvh& other) const = default;
0031     //bool operator != (const Bvh& other) const = default;
0032     bool operator == (const Bvh& other) const {
0033         return other.nodes == nodes && other.prim_ids == prim_ids;
0034     }
0035     bool operator != (const Bvh& other) const {
0036         return other.nodes != nodes || other.prim_ids != prim_ids;
0037     }
0038
0039     /// Returns whether the node located at the given index is the left child of its parent.
0040     static BVH_ALWAYS_INLINE bool is_left_sibling(size_t node_id) { return node_id % 2 == 1; }
0041
0042     /// Returns the index of a sibling of a node.
0043     static BVH_ALWAYS_INLINE size_t get_sibling_id(size_t node_id) {
0044         return is_left_sibling(node_id) ? node_id + 1 : node_id - 1;
0045     }
0046
0047     /// Returns the index of the left sibling of the node. This effectively returns the given index
0048     /// unchanged if the node is the left sibling, or the other sibling index otherwise.
0049     static BVH_ALWAYS_INLINE size_t get_left_sibling_id(size_t node_id) {
0050         return is_left_sibling(node_id) ? node_id : node_id - 1;
0051     }
0052
0053     /// Returns the index of the right sibling of the node. This effectively returns the given index
0054     /// unchanged if the node is the right sibling, or the other sibling index otherwise.
0055     static BVH_ALWAYS_INLINE size_t get_right_sibling_id(size_t node_id) {
0056         return is_left_sibling(node_id) ? node_id + 1 : node_id;
0057     }
0058
0059     /// Returns the root node of this BVH.
0060     BVH_ALWAYS_INLINE const Node& get_root() const { return nodes[0]; }
0061
0062     /// Extracts the BVH rooted at the given node index.
0063     inline Bvh extract_bvh(size_t root_id) const;
0064
0065     /// Traverses the BVH from the given index in `start` using the provided stack. Every leaf
0066     /// encountered on the way is processed using the given `LeafFn` function, and every pair of
0067     /// nodes is processed with the function in `HitFn`, which returns a triplet of booleans
0068     /// indicating whether the first child should be processed, whether the second child should be
0069     /// processed, and whether to traverse the second child first instead of the other way around.
0070     template <bool IsAnyHit, typename Stack, typename LeafFn, typename InnerFn>
0071     inline void traverse_top_down(Index start, Stack&, LeafFn&&, InnerFn&&) const;
0072
0073     /// Intersects the BVH with a single ray, using the given function to intersect the contents
0074     /// of a leaf. The algorithm starts at the node index `start` and uses the given stack object.
0075     /// When `IsAnyHit` is true, the function stops at the first intersection (useful for shadow
0076     /// rays), otherwise it finds the closest intersection. When `IsRobust` is true, a slower but
0077     /// numerically robust ray-box test is used, otherwise a fast, but less precise test is used.
0078     template <bool IsAnyHit, bool IsRobust, typename Stack, typename LeafFn, typename InnerFn = IgnoreArgs>
0079     inline void intersect(const Ray& ray, Index start, Stack&, LeafFn&&, InnerFn&& = {}) const;
0080
0081     /// Traverses this BVH from the bottom to the top, using the given function objects to process
0082     /// leaves and inner nodes.
0083     template <typename LeafFn = IgnoreArgs, typename InnerFn = IgnoreArgs>
0084     inline void traverse_bottom_up(LeafFn&& = {}, InnerFn&& = {});
0085
0086     /// Refits the BVH, using the given function object to recompute the bounding box of the leaves.
0087     template <typename LeafFn = IgnoreArgs>
0088     inline void refit(LeafFn&& = {});
0089
0090     inline void serialize(OutputStream&) const;
0091     static inline Bvh deserialize(InputStream&);
0092 };
0093
0094 template <typename Node>
0095 Bvh<Node> Bvh<Node>::extract_bvh(size_t root_id) const {
0096     assert(root_id != 0);
0097
0098     Bvh bvh;
0099     bvh.nodes.emplace_back();
0100
0101     std::stack<std::pair<size_t, size_t>> stack;
0102     stack.emplace(root_id, 0);
0103     while (!stack.empty()) {
0104         auto [src_id, dst_id] = stack.top();
0105         stack.pop();
0106         const auto& src_node = nodes[src_id];
0107         auto& dst_node = bvh.nodes[dst_id];
0108         dst_node = src_node;
0109         if (src_node.is_leaf()) {
0110             dst_node.index.set_first_id(bvh.prim_ids.size());
0111             std::copy_n(
0112                 prim_ids.begin() + src_node.index.first_id(),
0113                 src_node.index.prim_count(),
0114                 std::back_inserter(bvh.prim_ids));
0115         } else {
0116             dst_node.index.set_first_id(bvh.nodes.size());
0117             stack.emplace(src_node.index.first_id() + 0, bvh.nodes.size() + 0);
0118             stack.emplace(src_node.index.first_id() + 1, bvh.nodes.size() + 1);
0119             // Note: This may invalidate `dst_node` so has to happen after any access to it.
0120             bvh.nodes.emplace_back();
0121             bvh.nodes.emplace_back();
0122         }
0123     }
0124     return bvh;
0125 }
0126
0127 template <typename Node>
0128 template <bool IsAnyHit, typename Stack, typename LeafFn, typename InnerFn>
0129 void Bvh<Node>::traverse_top_down(Index start, Stack& stack, LeafFn&& leaf_fn, InnerFn&& inner_fn) const
0130 {
0131     stack.push(start);
0132 restart:
0133     while (!stack.is_empty()) {
0134         auto top = stack.pop();
0135         while (top.prim_count() == 0) {
0136             auto& left  = nodes[top.first_id()];
0137             auto& right = nodes[top.first_id() + 1];
0138             auto [hit_left, hit_right, should_swap] = inner_fn(left, right);
0139
0140             if (hit_left) {
0141                 auto near_index = left.index;
0142                 if (hit_right) {
0143                     auto far_index = right.index;
0144                     if (should_swap)
0145                         std::swap(near_index, far_index);
0146                     stack.push(far_index);
0147                 }
0148                 top = near_index;
0149             } else if (hit_right) {
0150                 top = right.index;
0151             }
0152             else [[unlikely]] {
0153                 goto restart;
0154             }
0155         }
0156
0157         [[maybe_unused]] auto was_hit = leaf_fn(top.first_id(), top.first_id() + top.prim_count());
0158         if constexpr (IsAnyHit) {
0159             if (was_hit) return;
0160         }
0161     }
0162 }
0163
0164 template <typename Node>
0165 template <bool IsAnyHit, bool IsRobust, typename Stack, typename LeafFn, typename InnerFn>
0166 void Bvh<Node>::intersect(const Ray& ray, Index start, Stack& stack, LeafFn&& leaf_fn, InnerFn&& inner_fn) const {
0167     auto inv_dir = ray.template get_inv_dir<!IsRobust>();
0168     auto inv_dir_pad_or_inv_org = IsRobust ? ray.pad_inv_dir(inv_dir) : -inv_dir * ray.org;
0169     auto octant = ray.get_octant();
0170
0171     traverse_top_down<IsAnyHit>(start, stack, leaf_fn, [&] (const Node& left, const Node& right) {
0172         inner_fn(left, right);
0173         std::pair<Scalar, Scalar> intr_left, intr_right;
0174         if constexpr (IsRobust) {
0175             intr_left  = left .intersect_robust(ray, inv_dir, inv_dir_pad_or_inv_org, octant);
0176             intr_right = right.intersect_robust(ray, inv_dir, inv_dir_pad_or_inv_org, octant);
0177         } else {
0178             intr_left  = left .intersect_fast(ray, inv_dir, inv_dir_pad_or_inv_org, octant);
0179             intr_right = right.intersect_fast(ray, inv_dir, inv_dir_pad_or_inv_org, octant);
0180         }
0181         return std::make_tuple(
0182             intr_left.first <= intr_left.second,
0183             intr_right.first <= intr_right.second,
0184             !IsAnyHit && intr_left.first > intr_right.first);
0185     });
0186 }
0187
0188 template <typename Node>
0189 template <typename LeafFn, typename InnerFn>
0190 void Bvh<Node>::traverse_bottom_up(LeafFn&& leaf_fn, InnerFn&& inner_fn) {
0191     std::vector<size_t> parents(nodes.size(), 0);
0192     for (size_t i = 0; i < nodes.size(); ++i) {
0193         if (nodes[i].is_leaf())
0194             continue;
0195         parents[nodes[i].index.first_id()] = i;
0196         parents[nodes[i].index.first_id() + 1] = i;
0197     }
0198     std::vector<bool> seen(nodes.size(), false);
0199     for (size_t i = nodes.size(); i-- > 0;) {
0200         if (!nodes[i].is_leaf())
0201             continue;
0202         leaf_fn(nodes[i]);
0203         seen[i] = true;
0204         for (size_t j = parents[i];; j = parents[j]) {
0205             auto& node = nodes[j];
0206             if (seen[j] || !seen[node.index.first_id()] || !seen[node.index.first_id() + 1])
0207                 break;
0208             inner_fn(nodes[j]);
0209             seen[j] = true;
0210         }
0211     }
0212 }
0213
0214 template <typename Node>
0215 template <typename LeafFn>
0216 void Bvh<Node>::refit(LeafFn&& leaf_fn) {
0217     traverse_bottom_up(leaf_fn, [&] (Node& node) {
0218         const auto& left  = nodes[node.index.first_id()];
0219         const auto& right = nodes[node.index.first_id() + 1];
0220         node.set_bbox(left.get_bbox().extend(right.get_bbox()));
0221     });
0222 }
0223
0224 template <typename Node>
0225 void Bvh<Node>::serialize(OutputStream& stream) const {
0226     stream.write(nodes.size());
0227     stream.write(prim_ids.size());
0228     for (auto&& node : nodes)
0229         node.serialize(stream);
0230     for (auto&& prim_id : prim_ids)
0231         stream.write(prim_id);
0232 }
0233
0234 template <typename Node>
0235 Bvh<Node> Bvh<Node>::deserialize(InputStream& stream) {
0236     Bvh bvh;
0237     bvh.nodes.resize(stream.read<size_t>());
0238     bvh.prim_ids.resize(stream.read<size_t>());
0239     for (auto& node : bvh.nodes)
0240         node = Node::deserialize(stream);
0241     for (auto& prim_id : bvh.prim_ids)
0242         prim_id = stream.read<size_t>();
0243     return bvh;
0244 }
0245
0246 } // namespace bvh::v2
0247
0248 #endif