Add feature to traverse a bvh passing a context from parent to children

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## v0.13.7
 
+### Unreleased
+
+- Added new visitor logic to QBVH to allow parent nodes to pass a context to their children.
+
 ### Modified
 
 - The `point_in_poly2d` now handles arbitrary (convex and non-convex) polygons. The previous implementation

src/partitioning/qbvh/traversal.rs

@@ -2,7 +2,7 @@
 
 use crate::bounding_volume::{Aabb, SimdAabb};
 use crate::math::Real;
-use crate::partitioning::visitor::SimdSimultaneousVisitStatus;
+use crate::partitioning::visitor::{SimdSimultaneousVisitStatus, SimdVisitorWithContext};
 use crate::partitioning::{
     GenericQbvh, QbvhStorage, SimdBestFirstVisitStatus, SimdBestFirstVisitor,
     SimdSimultaneousVisitor, SimdVisitStatus, SimdVisitor,
@@ -109,6 +109,74 @@ impl<LeafData: IndexedData, Storage: QbvhStorage<LeafData>> GenericQbvh<LeafData
         true
     }
 
+    /// Performs a depth-first traversal on the BVH. Passes a context from the
+    /// parent to the children.
+    ///
+    /// # Return
+    ///
+    /// Returns `false` if the traversal exitted early, and `true` otherwise.
+    pub fn traverse_depth_first_with_context<Context: Clone>(
+        &self,
+        visitor: &mut impl SimdVisitorWithContext<LeafData, SimdAabb, Context>,
+        context: Context,
+    ) -> bool {
+        self.traverse_depth_first_node_with_stack_and_context(visitor, &mut Vec::new(), 0, context)
+    }
+
+    /// Performs a depth-first traversal on the BVH and propagates a context down,
+    /// from the root to each of its descendants. The context can be modified
+    /// during the query.
+    ///
+    /// # Return
+    ///
+    /// Returns `false` if the traversal exited early, and `true` otherwise.
+    pub fn traverse_depth_first_node_with_stack_and_context<Context: Clone>(
+        &self,
+        visitor: &mut impl SimdVisitorWithContext<LeafData, SimdAabb, Context>,
+        stack: &mut Vec<(u32, Context)>,
+        start_node: u32,
+        context: Context,
+    ) -> bool {
+        stack.clear();
+
+        if !self.nodes.is_empty() {
+            stack.push((start_node, context));
+        }
+        while let Some((entry, context)) = stack.pop() {
+            let node = &self.nodes[entry as usize];
+            let leaf_data = if node.is_leaf() {
+                Some(
+                    array![|ii| Some(&self.proxies.get_at(node.children[ii] as usize)?.data); SIMD_WIDTH],
+                )
+            } else {
+                None
+            };
+
+            let (visit_result, contexts) = visitor.visit(&node.simd_aabb, leaf_data, context);
+            match visit_result {
+                SimdVisitStatus::ExitEarly => {
+                    return false;
+                }
+                SimdVisitStatus::MaybeContinue(mask) => {
+                    let bitmask = mask.bitmask();
+
+                    for ii in 0..SIMD_WIDTH {
+                        if (bitmask & (1 << ii)) != 0 && !node.is_leaf() {
+                            // Internal node, visit the child.
+                            // Un fortunately, we have this check because invalid Aabbs
+                            // return a hit as well.
+                            if node.children[ii] as usize <= self.nodes.len() {
+                                stack.push((node.children[ii], contexts[ii].clone()));
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
+        true
+    }
+
     /// Performs a best-first-search on the BVH.
     ///
     /// Returns the content of the leaf with the smallest associated cost, and a result of

src/partitioning/visitor.rs

@@ -78,6 +78,37 @@ where
         (self)(bv, data)
     }
 }
+
+/// Trait implemented by visitor called during the traversal of a spatial partitioning data structure.
+pub trait SimdVisitorWithContext<LeafData, SimdBV, Context: Clone> {
+    /// Execute an operation on the content of a node of the spatial partitioning structure.
+    ///
+    /// Returns whether the traversal should continue on the node's children, if it should not continue
+    /// on those children, or if the whole traversal should be exited early. Also returns
+    /// a context, which may or may not be identical to the input context.
+    fn visit(
+        &mut self,
+        bv: &SimdBV,
+        data: Option<[Option<&LeafData>; SIMD_WIDTH]>,
+        context: Context,
+    ) -> (SimdVisitStatus, [Context; SIMD_WIDTH]);
+}
+
+impl<T, LeafData, SimdBV, Context: Clone> SimdVisitorWithContext<LeafData, SimdBV, Context> for T
+where
+    T: SimdVisitor<LeafData, SimdBV>,
+{
+    fn visit(
+        &mut self,
+        bv: &SimdBV,
+        data: Option<[Option<&LeafData>; SIMD_WIDTH]>,
+        context: Context,
+    ) -> (SimdVisitStatus, [Context; SIMD_WIDTH]) {
+        let contexts = std::array::from_fn(|_| context.clone());
+        (self.visit(bv, data), contexts)
+    }
+}
+
 /// Trait implemented by visitor called during a simultaneous spatial partitioning data structure tarversal.
 pub trait SimdSimultaneousVisitor<T1, T2, SimdBV> {
     /// Execute an operation on the content of two nodes, one from each structure.

src/utils/weighted_value.rs

@@ -1,5 +1,5 @@
 use crate::math::Real;
-use std::cmp::{Eq, Ord, Ordering, PartialEq, PartialOrd};
+use std::cmp::Ordering;
 
 #[derive(Copy, Clone)]
 pub struct WeightedValue<T> {