diff --git a/CHANGELOG.md b/CHANGELOG.md index 621b23c5..f5f5e95d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,9 @@ +## Unreleased + +### Added + +- Added `Bvh::traverse_indexed` to traverse only a part of the subtree, and to retrieve the node’s indices. + ## 0.26.0 ### Breaking changes diff --git a/src/partitioning/bvh/bvh_tests.rs b/src/partitioning/bvh/bvh_tests.rs index 4b2898b2..143d3170 100644 --- a/src/partitioning/bvh/bvh_tests.rs +++ b/src/partitioning/bvh/bvh_tests.rs @@ -1,6 +1,6 @@ use crate::bounding_volume::Aabb; use crate::math::{Real, Vector}; -use crate::partitioning::{Bvh, BvhBuildStrategy}; +use crate::partitioning::{Bvh, BvhBuildStrategy, BvhNode, BvhNodeIndex, TraversalAction}; fn make_test_aabb(i: usize) -> Aabb { Aabb::from_half_extents(Vector::splat(i as Real).into(), Vector::splat(1.0)) @@ -15,7 +15,7 @@ fn test_leaves_iteration() { let bvh = Bvh::from_leaves(BvhBuildStrategy::Binned, &leaves); // Only allow nodes with mins.x <= 3.0 (should only pass leaf 0) - let check = |node: &crate::partitioning::BvhNode| -> bool { node.mins.x <= 3.0 }; + let check = |node: &BvhNode| -> bool { node.mins.x <= 3.0 }; let mut found_invalid_leaf = false; for leaf_index in bvh.leaves(check) { @@ -31,6 +31,127 @@ fn test_leaves_iteration() { } } +#[test] +fn test_traverse_indexed() { + // Empty tree: callback must never fire, regardless of `subtree` being `None`. + let empty = Bvh::new(); + empty.traverse_indexed(None, |_, _| { + panic!("callback should not be called on an empty BVH"); + }); + + // Single-leaf tree exercises the partial-root branch when starting from the root. + let single = Bvh::from_leaves(BvhBuildStrategy::Binned, &[make_test_aabb(0)]); + let mut single_visited = std::vec::Vec::new(); + single.traverse_indexed(None, |node, idx| { + single_visited.push((idx, node.leaf_data())); + TraversalAction::Continue + }); + assert_eq!(single_visited.len(), 1); + assert_eq!(single_visited[0].0, BvhNodeIndex::left(0)); + assert_eq!(single_visited[0].1, Some(0)); + + // Multi-leaf tree: traversing from the root must visit every leaf, and the + // index passed to the callback must round-trip through `bvh.nodes`. + let leaves: std::vec::Vec<_> = (0..16).map(make_test_aabb).collect(); + let bvh = Bvh::from_leaves(BvhBuildStrategy::Binned, &leaves); + + let mut seen_leaves = std::vec::Vec::new(); + let mut traverse_indexed_calls = std::vec::Vec::new(); + bvh.traverse_indexed(None, |node, idx| { + // Every reported index must point to the same node we just received. + let by_idx: &BvhNode = &bvh.nodes[idx]; + assert!(core::ptr::eq(by_idx, node)); + + traverse_indexed_calls.push(idx); + if let Some(data) = node.leaf_data() { + seen_leaves.push(data); + } + TraversalAction::Continue + }); + seen_leaves.sort(); + assert_eq!(seen_leaves, (0..16).collect::>()); + + // `traverse_indexed(None, ...)` must visit exactly the same nodes (in the same + // order) as `traverse`. + let mut traverse_nodes: std::vec::Vec<*const BvhNode> = std::vec::Vec::new(); + bvh.traverse(|node| { + traverse_nodes.push(node as *const _); + TraversalAction::Continue + }); + let indexed_nodes: std::vec::Vec<*const BvhNode> = traverse_indexed_calls + .iter() + .map(|idx| &bvh.nodes[*idx] as *const _) + .collect(); + assert_eq!(traverse_nodes, indexed_nodes); + + // Starting from a specific subtree must only visit that subtree (the start + // node and its descendants), and every reported leaf must belong to it. + let subtree_root_idx = BvhNodeIndex::left(0); + let mut subtree_leaves = std::vec::Vec::new(); + let mut subtree_visited = std::vec::Vec::new(); + bvh.traverse_indexed(Some(subtree_root_idx), |node, idx| { + subtree_visited.push(idx); + if let Some(data) = node.leaf_data() { + subtree_leaves.push(data); + } + TraversalAction::Continue + }); + assert_eq!(subtree_visited[0], subtree_root_idx); + // The subtree's leaves must be a non-empty strict subset of the full set. + assert!(!subtree_leaves.is_empty()); + assert!(subtree_leaves.len() < 16); + for leaf in &subtree_leaves { + assert!(seen_leaves.contains(leaf)); + } + // Leaf count reported by the subtree's root must match the visited leaves. + assert_eq!( + bvh.nodes[subtree_root_idx].leaf_count() as usize, + subtree_leaves.len() + ); + + // Starting from a leaf node visits exactly that leaf. + let leaf_idx = *traverse_indexed_calls + .iter() + .find(|idx| bvh.nodes[**idx].is_leaf()) + .expect("the tree must contain at least one leaf"); + let mut leaf_only = std::vec::Vec::new(); + bvh.traverse_indexed(Some(leaf_idx), |node, idx| { + leaf_only.push((idx, node.leaf_data())); + TraversalAction::Continue + }); + assert_eq!(leaf_only.len(), 1); + assert_eq!(leaf_only[0].0, leaf_idx); + assert!(leaf_only[0].1.is_some()); + + // `Prune` at the start node must visit it once and stop. + let mut prune_visits = 0; + bvh.traverse_indexed(Some(BvhNodeIndex::left(0)), |_, _| { + prune_visits += 1; + TraversalAction::Prune + }); + assert_eq!(prune_visits, 1); + + // `EarlyExit` at the start node must visit it once and stop. + let mut exit_visits = 0; + bvh.traverse_indexed(Some(BvhNodeIndex::left(0)), |_, _| { + exit_visits += 1; + TraversalAction::EarlyExit + }); + assert_eq!(exit_visits, 1); + + // `EarlyExit` partway through must short-circuit the full traversal. + let mut early = 0; + bvh.traverse_indexed(None, |_, _| { + early += 1; + if early >= 3 { + TraversalAction::EarlyExit + } else { + TraversalAction::Continue + } + }); + assert_eq!(early, 3); +} + #[test] fn bvh_build_and_removal() { // Check various combination of building pattern and removal pattern. diff --git a/src/partitioning/bvh/bvh_traverse.rs b/src/partitioning/bvh/bvh_traverse.rs index 9252bbf7..18748336 100644 --- a/src/partitioning/bvh/bvh_traverse.rs +++ b/src/partitioning/bvh/bvh_traverse.rs @@ -1,6 +1,6 @@ use super::BvhNode; use crate::math::Real; -use crate::partitioning::Bvh; +use crate::partitioning::{Bvh, BvhNodeIndex}; use smallvec::SmallVec; const TRAVERSAL_STACK_SIZE: usize = 32; @@ -288,27 +288,54 @@ impl Bvh { /// - [`cast_ray`](Self::cast_ray) - Ray casting with best-first traversal /// - [`TraversalAction`] - Controls traversal flow pub fn traverse(&self, mut check_node: impl FnMut(&BvhNode) -> TraversalAction) { + self.traverse_indexed(None, move |node, _| check_node(node)); + } + + /// Similar to [`Self::traverse`] but starts the traversal with the node at index `subtree`, + /// and the `check_node` closure is given the node index. + /// + /// If `subtree` is `None` the traversal starts at the root of the tree. + pub fn traverse_indexed( + &self, + subtree: Option, + mut check_node: impl FnMut(&BvhNode, BvhNodeIndex) -> TraversalAction, + ) { let mut stack = Self::traversal_stack(); let mut curr_id = 0; - if self.nodes.is_empty() { - return; - } else if self.nodes[0].right.leaf_count() == 0 { - // Special case for partial root. - let _ = check_node(&self.nodes[0].left); - return; + if let Some(subtree) = subtree { + let to_check = &self.nodes[subtree]; + match check_node(to_check, subtree) { + TraversalAction::Continue => { + if to_check.is_leaf() { + return; + } + + curr_id = to_check.children; + } + TraversalAction::Prune | TraversalAction::EarlyExit => return, + } + } else { + // Start at the root. + if self.nodes.is_empty() { + return; + } else if self.nodes[0].right.leaf_count() == 0 { + // Special case for partial root. + let _ = check_node(&self.nodes[0].left, BvhNodeIndex::left(0)); + return; + } } loop { let node = &self.nodes[curr_id as usize]; let left = &node.left; let right = &node.right; - let go_left = match check_node(left) { + let go_left = match check_node(left, BvhNodeIndex::left(curr_id)) { TraversalAction::Continue => !left.is_leaf(), TraversalAction::Prune => false, TraversalAction::EarlyExit => return, }; - let go_right = match check_node(right) { + let go_right = match check_node(right, BvhNodeIndex::right(curr_id)) { TraversalAction::Continue => !right.is_leaf(), TraversalAction::Prune => false, TraversalAction::EarlyExit => return,