feat(bvh-region): add raycast query

crates/bvh-region/Cargo.toml

@@ -16,6 +16,7 @@ derive_more = { workspace = true }
 fastrand = { workspace = true }
 geometry = { workspace = true }
 glam = { workspace = true, features = ["serde"] }
+num-traits = { workspace = true }
 ordered-float = { workspace = true }
 plotters = { workspace = true, features = ["plotters-bitmap", "image"], optional = true }
 plotters-bitmap = { workspace = true, optional = true }

crates/bvh-region/src/query.rs

@@ -1,2 +1,3 @@
 mod closest;
 mod range;
+mod ray;

crates/bvh-region/src/query/closest.rs

@@ -2,6 +2,8 @@ use std::{cmp::Reverse, collections::BinaryHeap, fmt::Debug};
 
 use geometry::aabb::Aabb;
 use glam::Vec3;
+use num_traits::Zero;
+use ordered_float::NotNan;
 
 use crate::{Bvh, Node, utils::NodeOrd};
 
@@ -29,7 +31,12 @@ impl<T: Debug> Bvh<T> {
 
         // let mut stack: SmallVec<&BvhNode, 64> = SmallVec::new();
         let mut heap: BinaryHeap<_> = std::iter::once(on)
-            .map(|node| Reverse(NodeOrd { node, dist2: 0.0 }))
+            .map(|node| {
+                Reverse(NodeOrd {
+                    node,
+                    by: NotNan::zero(),
+                })
+            })
             .collect();
 
         while let Some(on) = heap.pop() {
@@ -43,9 +50,12 @@ impl<T: Debug> Bvh<T> {
             for child in on.children(self) {
                 match child {
                     Node::Internal(internal) => {
+                        let dist2 = internal.aabb.dist2(target);
+                        let dist2 = NotNan::new(dist2).unwrap();
+
                         heap.push(Reverse(NodeOrd {
                             node: internal,
-                            dist2: internal.aabb.dist2(target),
+                            by: dist2,
                         }));
                     }
                     Node::Leaf(leaf) => {

crates/bvh-region/src/query/ray.rs

@@ -0,0 +1,87 @@
+use std::{cmp::Reverse, collections::BinaryHeap, fmt::Debug};
+
+use geometry::{aabb::Aabb, ray::Ray};
+use ordered_float::NotNan;
+
+use crate::{Bvh, Node, utils::NodeOrd};
+
+impl<T: Debug> Bvh<T> {
+    /// Returns the closest element hit by the ray and the intersection distance (t) along the ray.
+    ///
+    /// If no element is hit, returns `None`.
+    #[allow(clippy::excessive_nesting)]
+    pub fn get_closest_ray(
+        &self,
+        ray: Ray,
+        get_aabb: impl Fn(&T) -> Aabb,
+    ) -> Option<(&T, NotNan<f32>)> {
+        let mut closest_t = NotNan::new(f32::INFINITY).unwrap();
+        let mut closest_elem = None;
+
+        let root = self.root();
+
+        match root {
+            Node::Leaf(elems) => {
+                // Only a leaf: check all elements directly.
+                for elem in elems {
+                    if let Some(t) = get_aabb(elem).intersect_ray(&ray) {
+                        if t < closest_t && t.into_inner() >= 0.0 {
+                            closest_t = t;
+                            closest_elem = Some(elem);
+                        }
+                    }
+                }
+            }
+            Node::Internal(internal) => {
+                let mut heap: BinaryHeap<_> = BinaryHeap::new();
+
+                // Check if the ray hits the root node's AABB
+                if let Some(t) = internal.aabb.intersect_ray(&ray) {
+                    if t.into_inner() >= 0.0 {
+                        heap.push(Reverse(NodeOrd {
+                            node: internal,
+                            by: t,
+                        }));
+                    }
+                }
+
+                while let Some(Reverse(current)) = heap.pop() {
+                    let node = current.node;
+                    let node_t = current.by;
+
+                    // If the node AABB is farther than any known intersection, prune
+                    if node_t > closest_t {
+                        continue;
+                    }
+
+                    for child in node.children(self) {
+                        match child {
+                            Node::Internal(child_node) => {
+                                if let Some(t) = child_node.aabb.intersect_ray(&ray) {
+                                    if t < closest_t && t.into_inner() >= 0.0 {
+                                        heap.push(Reverse(NodeOrd {
+                                            node: child_node,
+                                            by: t,
+                                        }));
+                                    }
+                                }
+                            }
+                            Node::Leaf(elems) => {
+                                for elem in elems {
+                                    if let Some(t) = get_aabb(elem).intersect_ray(&ray) {
+                                        if t < closest_t && t.into_inner() >= 0.0 {
+                                            closest_t = t;
+                                            closest_elem = Some(elem);
+                                        }
+                                    }
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
+        closest_elem.map(|elem| (elem, closest_t))
+    }
+}

crates/bvh-region/src/utils.rs

@@ -22,27 +22,27 @@ pub trait GetAabb<T>: Fn(&T) -> Aabb {}
 
 impl<T, F> GetAabb<T> for F where F: Fn(&T) -> Aabb {}
 
-#[derive(Constructor)]
-pub struct NodeOrd<'a> {
+#[derive(Constructor, Copy, Clone, Debug)]
+pub struct NodeOrd<'a, T> {
     pub node: &'a BvhNode,
-    pub dist2: f64,
+    pub by: T,
 }
 
-impl PartialEq<Self> for NodeOrd<'_> {
+impl<T: PartialEq> PartialEq<Self> for NodeOrd<'_, T> {
     fn eq(&self, other: &Self) -> bool {
-        self.dist2 == other.dist2
+        self.by == other.by
     }
 }
-impl PartialOrd for NodeOrd<'_> {
+impl<T: PartialOrd> PartialOrd for NodeOrd<'_, T> {
     fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
-        Some(self.cmp(other))
+        self.by.partial_cmp(&other.by)
     }
 }
 
-impl Eq for NodeOrd<'_> {}
+impl<T: Eq> Eq for NodeOrd<'_, T> {}
 
-impl Ord for NodeOrd<'_> {
+impl<T: Ord> Ord for NodeOrd<'_, T> {
     fn cmp(&self, other: &Self) -> std::cmp::Ordering {
-        self.dist2.partial_cmp(&other.dist2).unwrap()
+        self.by.cmp(&other.by)
     }
 }

crates/bvh-region/tests/simple.proptest-regressions

@@ -0,0 +1,7 @@
+# Seeds for failure cases proptest has generated in the past. It is
+# automatically read and these particular cases re-run before any
+# novel cases are generated.
+#
+# It is recommended to check this file in to source control so that
+# everyone who runs the test benefits from these saved cases.
+cc 6eda8ce9f67e5e83265d2d9b12b641fd177829ff6893611a8fba5fb06af97015 # shrinks to elements = [[-126.04, 0.00, -677.53] -> [0.00, 726.19, 953.13], [-694.23, 0.00, 0.00] -> [0.00, 627.63, 922.19], [0.00, 0.00, 0.00] -> [0.00, 0.00, 0.00], [0.00, 0.00, 0.00] -> [0.00, 0.00, 0.00], [0.00, 0.00, 0.00] -> [0.00, 0.00, 0.00], [0.00, 0.00, 0.00] -> [0.00, 0.00, 0.00], [0.00, 0.00, 0.00] -> [0.00, 0.00, 0.00], [0.00, 0.00, -53.49] -> [0.00, 0.00, 0.00], [0.00, 0.00, 0.00] -> [0.00, 0.00, 0.00], [0.00, 0.00, -947.84] -> [0.00, 0.00, 0.00], [0.00, 0.00, -290.94] -> [0.00, 0.00, 0.00], [0.00, 0.00, -486.91] -> [0.00, 0.00, 0.00], [0.00, 0.00, -418.65] -> [0.00, 0.00, 0.00], [0.00, 0.00, -866.01] -> [0.00, 0.00, 0.00], [0.00, 0.00, 0.00] -> [0.00, 0.00, 0.00], [0.00, 0.00, 0.00] -> [0.00, 0.00, 0.00], [0.00, 0.00, 0.00] -> [0.00, 929.05, 0.00]], origin = (-4.9637938, 474.28943, 904.7494), direction = (0.0, 866.90247, 0.0)

crates/bvh-region/tests/simple.rs

@@ -2,8 +2,12 @@ use std::collections::HashSet;
 
 use approx::assert_relative_eq;
 use bvh_region::Bvh;
-use geometry::aabb::{Aabb, OrderedAabb};
+use geometry::{
+    aabb::{Aabb, OrderedAabb},
+    ray::Ray,
+};
 use glam::Vec3;
+use ordered_float::NotNan;
 use proptest::prelude::*;
 
 const fn copied<T: Copy>(value: &T) -> T {
@@ -156,3 +160,78 @@ proptest! {
         prop_assert_eq!(&bvh_set, &brute_force_set, "Mismatch between BVH range and brute force collision sets: {:?} != {:?}", bvh_set, brute_force_set);
     }
 }
+
+/// Computes the closest intersection of `ray` with the list of `elements` via brute force.
+fn brute_force_closest_ray(elements: &[Aabb], ray: Ray) -> Option<(&Aabb, NotNan<f32>)> {
+    let mut closest_t = NotNan::new(f32::INFINITY).unwrap();
+    let mut closest_elem = None;
+
+    for aabb in elements {
+        if let Some(t) = aabb.intersect_ray(&ray) {
+            if t < closest_t && t.into_inner() >= 0.0 {
+                closest_t = t;
+                closest_elem = Some(aabb);
+            }
+        }
+    }
+
+    closest_elem.map(|e| (e, closest_t))
+}
+
+proptest! {
+    #[test]
+    fn test_get_closest_ray_correctness(
+        elements in prop::collection::vec(
+            // Generate random AABBs by picking two random points and making one the min and the other the max.
+            (-1000.0..1000.0f32, -1000.0..1000.0f32, -1000.0..1000.0f32, -1000.0..1000.0f32, -1000.0..1000.0f32, -1000.0..1000.0f32)
+                .prop_map(|(x1, y1, z1, x2, y2, z2)| {
+                    let min_x = x1.min(x2);
+                    let max_x = x1.max(x2);
+                    let min_y = y1.min(y2);
+                    let max_y = y1.max(y2);
+                    let min_z = z1.min(z2);
+                    let max_z = z1.max(z2);
+                    Aabb::from([min_x, min_y, min_z, max_x, max_y, max_z])
+                }),
+            1..50 // vary the number of elements
+        ),
+        origin in (-1000.0..1000.0f32, -1000.0..1000.0f32, -1000.0..1000.0f32),
+        direction in (-1000.0..1000.0f32, -1000.0..1000.0f32, -1000.0..1000.0f32)
+    ) {
+        // If the direction is zero, we skip this test case. A ray with zero direction doesn't make sense.
+        let dir_vec = Vec3::new(direction.0, direction.1, direction.2);
+        let zero_vec = Vec3::new(0.0, 0.0, 0.0);
+        if dir_vec == zero_vec {
+            return Ok(());
+        }
+
+        let ray = Ray::new(
+            Vec3::new(origin.0, origin.1, origin.2),
+            dir_vec
+        );
+
+        let bvh = Bvh::build(elements.clone(), copied);
+        let bvh_closest = bvh.get_closest_ray(ray, copied);
+        let brute_closest = brute_force_closest_ray(&elements, ray);
+
+        match (bvh_closest, brute_closest) {
+            (None, None) => {
+                // Both found no intersections.
+            },
+            (Some((.., bvh_t)), Some((.., brute_t))) => {
+                // Check that the chosen elements and intersection distances are close.
+                // Because multiple AABBs might have the exact same intersection distance, we can't always assert equality of the element references.
+                // But at minimum, we check that the distances are very close.
+                let diff = (bvh_t.into_inner() - brute_t.into_inner()).abs();
+                prop_assert!(diff < 1e-6, "Distances differ significantly; BVH: {}, brute force: {}", bvh_t, brute_t);
+
+                // If desired (and if you trust that intersections are unique), you could also check:
+                // prop_assert_eq!(bvh_elem, brute_elem);
+            },
+            (bvh_val, brute_val) => {
+                // Mismatch: One found an intersection, the other didn't.
+                prop_assert!(false, "Mismatch between BVH closest ray and brute force closest ray; BVH: {:?}, brute force: {:?}", bvh_val, brute_val);
+            }
+        }
+    }
+}