Make the WQuadTree more generic and use it as the trimesh acceleration structure.

8 files changed, 146 insertions, 98 deletions

diff --git a/src/geometry/broad_phase.rs b/src/geometry/broad_phase.rs
index a43b7af..a1333d0 100644
--- a/src/geometry/broad_phase.rs
+++ b/src/geometry/broad_phase.rs
@@ -144,7 +144,7 @@ impl WAABBHierarchy {
                 if let Some(waabb2) = level.get(*i) {
                     // NOTE: using `intersect.bitmask()` and performing bit comparisons
                     // is much more efficient than testing if each intersect.extract(i) is true.
-                    let intersect = waabb1.intersects_lanewise(waabb2);
+                    let intersect = waabb1.intersects(waabb2);
                     let bitmask = intersect.bitmask();
 
                     for j in 0..SIMD_WIDTH {
diff --git a/src/geometry/collider_set.rs b/src/geometry/collider_set.rs
index 3a259ae..dbdd2a4 100644
--- a/src/geometry/collider_set.rs
+++ b/src/geometry/collider_set.rs
@@ -38,7 +38,7 @@ impl ColliderSet {
     }
 
     /// Iterate through all the colliders on this set.
-    pub fn iter(&self) -> impl Iterator<Item = (ColliderHandle, &Collider)> {
+    pub fn iter(&self) -> impl ExactSizeIterator<Item = (ColliderHandle, &Collider)> {
         self.colliders.iter()
     }
 
diff --git a/src/geometry/contact_generator/trimesh_shape_contact_generator.rs b/src/geometry/contact_generator/trimesh_shape_contact_generator.rs
index 78f26bb..d759ae3 100644
--- a/src/geometry/contact_generator/trimesh_shape_contact_generator.rs
+++ b/src/geometry/contact_generator/trimesh_shape_contact_generator.rs
@@ -5,7 +5,7 @@ use crate::geometry::{Collider, ContactManifold, Shape, Trimesh, WAABBHierarchyI
 use crate::ncollide::bounding_volume::{BoundingVolume, AABB};
 
 pub struct TrimeshShapeContactGeneratorWorkspace {
-    interferences: WAABBHierarchyIntersections,
+    interferences: Vec<usize>,
     local_aabb2: AABB<f32>,
     old_interferences: Vec<usize>,
     old_manifolds: Vec<ContactManifold>,
@@ -14,7 +14,7 @@ pub struct TrimeshShapeContactGeneratorWorkspace {
 impl TrimeshShapeContactGeneratorWorkspace {
     pub fn new() -> Self {
         Self {
-            interferences: WAABBHierarchyIntersections::new(),
+            interferences: Vec::new(),
             local_aabb2: AABB::new_invalid(),
             old_interferences: Vec::new(),
             old_manifolds: Vec::new(),
@@ -74,7 +74,7 @@ fn do_generate_contacts(
         let local_aabb2 = new_local_aabb2; // .loosened(ctxt.prediction_distance * 2.0); // FIXME: what would be the best value?
         std::mem::swap(
             &mut workspace.old_interferences,
-            workspace.interferences.computed_interferences_mut(),
+            &mut workspace.interferences,
         );
         std::mem::swap(&mut workspace.old_manifolds, &mut ctxt.pair.manifolds);
         ctxt.pair.manifolds.clear();
@@ -108,16 +108,14 @@ fn do_generate_contacts(
         //     workspace.old_manifolds.len()
         // );
 
-        trimesh1
-            .waabbs()
-            .compute_interferences_with(local_aabb2, &mut workspace.interferences);
+        workspace.interferences = trimesh1.waabbs().intersect_aabb(&local_aabb2);
         workspace.local_aabb2 = local_aabb2;
     }
 
     /*
      * Dispatch to the specific solver by keeping the previous manifold if we already had one.
      */
-    let new_interferences = workspace.interferences.computed_interferences();
+    let new_interferences = &workspace.interferences;
     let mut old_inter_it = workspace.old_interferences.drain(..).peekable();
     let mut old_manifolds_it = workspace.old_manifolds.drain(..);
 
diff --git a/src/geometry/proximity_detector/trimesh_shape_proximity_detector.rs b/src/geometry/proximity_detector/trimesh_shape_proximity_detector.rs
index 3dd7381..0a3ff44 100644
--- a/src/geometry/proximity_detector/trimesh_shape_proximity_detector.rs
+++ b/src/geometry/proximity_detector/trimesh_shape_proximity_detector.rs
@@ -5,7 +5,7 @@ use crate::geometry::{Collider, Proximity, Shape, Trimesh, WAABBHierarchyInterse
 use crate::ncollide::bounding_volume::{BoundingVolume, AABB};
 
 pub struct TrimeshShapeProximityDetectorWorkspace {
-    interferences: WAABBHierarchyIntersections,
+    interferences: Vec<usize>,
     local_aabb2: AABB<f32>,
     old_interferences: Vec<usize>,
 }
@@ -13,7 +13,7 @@ pub struct TrimeshShapeProximityDetectorWorkspace {
 impl TrimeshShapeProximityDetectorWorkspace {
     pub fn new() -> Self {
         Self {
-            interferences: WAABBHierarchyIntersections::new(),
+            interferences: Vec::new(),
             local_aabb2: AABB::new_invalid(),
             old_interferences: Vec::new(),
         }
@@ -67,19 +67,17 @@ fn do_detect_proximity(
         let local_aabb2 = new_local_aabb2; // .loosened(ctxt.prediction_distance * 2.0); // FIXME: what would be the best value?
         std::mem::swap(
             &mut workspace.old_interferences,
-            &mut workspace.interferences.computed_interferences_mut(),
+            &mut workspace.interferences,
         );
 
-        trimesh1
-            .waabbs()
-            .compute_interferences_with(local_aabb2, &mut workspace.interferences);
+        workspace.interferences = trimesh1.waabbs().intersect_aabb(&local_aabb2);
         workspace.local_aabb2 = local_aabb2;
     }
 
     /*
      * Dispatch to the specific solver by keeping the previous manifold if we already had one.
      */
-    let new_interferences = workspace.interferences.computed_interferences();
+    let new_interferences = &workspace.interferences;
     let mut old_inter_it = workspace.old_interferences.drain(..).peekable();
     let mut best_proximity = Proximity::Disjoint;
 
diff --git a/src/geometry/trimesh.rs b/src/geometry/trimesh.rs
index 62731b6..f9e6034 100644
--- a/src/geometry/trimesh.rs
+++ b/src/geometry/trimesh.rs
@@ -1,4 +1,4 @@
-use crate::geometry::{Triangle, WAABBHierarchy};
+use crate::geometry::{Triangle, WQuadtree};
 use crate::math::{Isometry, Point};
 use na::Point3;
 use ncollide::bounding_volume::{HasBoundingVolume, AABB};
@@ -7,7 +7,7 @@ use ncollide::bounding_volume::{HasBoundingVolume, AABB};
 #[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
 /// A triangle mesh.
 pub struct Trimesh {
-    waabb_tree: WAABBHierarchy,
+    wquadtree: WQuadtree<usize>,
     aabb: AABB<f32>,
     vertices: Vec<Point<f32>>,
     indices: Vec<Point3<u32>>,
@@ -25,41 +25,24 @@ impl Trimesh {
             "A triangle mesh must contain at least one triangle."
         );
 
-        // z-sort the indices.
-        //        indices.sort_unstable_by(|idx, jdx| {
-        //            let ti = Triangle::new(
-        //                vertices[idx[0] as usize],
-        //                vertices[idx[1] as usize],
-        //                vertices[idx[2] as usize],
-        //            );
-        //            let tj = Triangle::new(
-        //                vertices[jdx[0] as usize],
-        //                vertices[jdx[1] as usize],
-        //                vertices[jdx[2] as usize],
-        //            );
-        //            let center_i = (ti.a.coords + ti.b.coords + ti.c.coords) / 3.0;
-        //            let center_j = (tj.a.coords + tj.b.coords + tj.c.coords) / 3.0;
-        //            crate::geometry::z_cmp_floats(center_i.as_slice(), center_j.as_slice())
-        //                .unwrap_or(std::cmp::Ordering::Equal)
-        //        });
         let aabb = AABB::from_points(&vertices);
+        let data = indices.iter().enumerate().map(|(i, idx)| {
+            let aabb = Triangle::new(
+                vertices[idx[0] as usize],
+                vertices[idx[1] as usize],
+                vertices[idx[2] as usize],
+            )
+            .local_bounding_volume();
+            (i, aabb)
+        });
 
-        let aabbs: Vec<_> = indices
-            .iter()
-            .map(|idx| {
-                Triangle::new(
-                    vertices[idx[0] as usize],
-                    vertices[idx[1] as usize],
-                    vertices[idx[2] as usize],
-                )
-                .local_bounding_volume()
-            })
-            .collect();
-
-        let waabb_tree = WAABBHierarchy::new(&aabbs);
+        let mut wquadtree = WQuadtree::new();
+        // NOTE: we apply no dilation factor because we won't
+        // update this tree dynamically.
+        wquadtree.clear_and_rebuild(data, 0.0);
 
         Self {
-            waabb_tree,
+            wquadtree,
             aabb,
             vertices,
             indices,
@@ -71,8 +54,8 @@ impl Trimesh {
         self.aabb.transform_by(pos)
     }
 
-    pub(crate) fn waabbs(&self) -> &WAABBHierarchy {
-        &self.waabb_tree
+    pub(crate) fn waabbs(&self) -> &WQuadtree<usize> {
+        &self.wquadtree
     }
 
     /// The number of triangles forming this mesh.
diff --git a/src/geometry/waabb.rs b/src/geometry/waabb.rs
index 1706370..cc420d9 100644
--- a/src/geometry/waabb.rs
+++ b/src/geometry/waabb.rs
@@ -156,7 +156,7 @@ impl WAABB {
     }
 
     #[cfg(feature = "dim2")]
-    pub fn contains_lanewise(&self, other: &WAABB) -> SimdBool {
+    pub fn contains(&self, other: &WAABB) -> SimdBool {
         self.mins.x.simd_le(other.mins.x)
             & self.mins.y.simd_le(other.mins.y)
             & self.maxs.x.simd_ge(other.maxs.x)
@@ -164,7 +164,7 @@ impl WAABB {
     }
 
     #[cfg(feature = "dim3")]
-    pub fn contains_lanewise(&self, other: &WAABB) -> SimdBool {
+    pub fn contains(&self, other: &WAABB) -> SimdBool {
         self.mins.x.simd_le(other.mins.x)
             & self.mins.y.simd_le(other.mins.y)
             & self.mins.z.simd_le(other.mins.z)
@@ -174,7 +174,7 @@ impl WAABB {
     }
 
     #[cfg(feature = "dim2")]
-    pub fn intersects_lanewise(&self, other: &WAABB) -> SimdBool {
+    pub fn intersects(&self, other: &WAABB) -> SimdBool {
         self.mins.x.simd_le(other.maxs.x)
             & other.mins.x.simd_le(self.maxs.x)
             & self.mins.y.simd_le(other.maxs.y)
@@ -182,7 +182,7 @@ impl WAABB {
     }
 
     #[cfg(feature = "dim3")]
-    pub fn intersects_lanewise(&self, other: &WAABB) -> SimdBool {
+    pub fn intersects(&self, other: &WAABB) -> SimdBool {
         self.mins.x.simd_le(other.maxs.x)
             & other.mins.x.simd_le(self.maxs.x)
             & self.mins.y.simd_le(other.maxs.y)
diff --git a/src/geometry/wquadtree.rs b/src/geometry/wquadtree.rs
index 17ecffe..627ad9a 100644
--- a/src/geometry/wquadtree.rs
+++ b/src/geometry/wquadtree.rs
@@ -7,6 +7,31 @@ use simba::simd::{SimdBool, SimdValue};
 use std::collections::VecDeque;
 use std::ops::Range;
 
+pub trait IndexedData: Copy {
+    fn default() -> Self;
+    fn index(&self) -> usize;
+}
+
+impl IndexedData for usize {
+    fn default() -> Self {
+        u32::MAX as usize
+    }
+
+    fn index(&self) -> usize {
+        *self
+    }
+}
+
+impl IndexedData for ColliderHandle {
+    fn default() -> Self {
+        ColliderSet::invalid_handle()
+    }
+
+    fn index(&self) -> usize {
+        self.into_raw_parts().0
+    }
+}
+
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
 #[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
 struct NodeIndex {
@@ -41,38 +66,32 @@ struct WQuadtreeNode {
 
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
 #[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
-struct WQuadtreeProxy {
+struct WQuadtreeProxy<T> {
     node: NodeIndex,
-    handle: ColliderHandle, // The collider handle. TODO: only set the collider generation here?
+    data: T, // The collider data. TODO: only set the collider generation here?
 }
 
-impl WQuadtreeProxy {
+impl<T: IndexedData> WQuadtreeProxy<T> {
     fn invalid() -> Self {
         Self {
             node: NodeIndex::invalid(),
-            handle: ColliderSet::invalid_handle(),
+            data: T::default(),
         }
     }
 }
 
 #[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
-pub struct WQuadtree {
+#[derive(Clone, Debug)]
+pub struct WQuadtree<T> {
     nodes: Vec<WQuadtreeNode>,
     dirty_nodes: VecDeque<u32>,
-    proxies: Vec<WQuadtreeProxy>,
+    proxies: Vec<WQuadtreeProxy<T>>,
 }
 
-impl WQuadtree {
-    pub fn new() -> Self {
-        WQuadtree {
-            nodes: Vec::new(),
-            dirty_nodes: VecDeque::new(),
-            proxies: Vec::new(),
-        }
-    }
-
-    pub fn pre_update(&mut self, handle: ColliderHandle) {
-        let id = handle.into_raw_parts().0;
+// FIXME: this should be generic too.
+impl WQuadtree<ColliderHandle> {
+    pub fn pre_update(&mut self, data: ColliderHandle) {
+        let id = data.into_raw_parts().0;
         let node_id = self.proxies[id].node.index;
         let node = &mut self.nodes[node_id as usize];
         if !node.dirty {
@@ -86,7 +105,7 @@ impl WQuadtree {
         let dilation_factor = SimdFloat::splat(dilation_factor);
 
         while let Some(id) = self.dirty_nodes.pop_front() {
-            // NOTE: this will handle the case where we reach the root of the tree.
+            // NOTE: this will data the case where we reach the root of the tree.
             if let Some(node) = self.nodes.get(id as usize) {
                 // Compute the new WAABB.
                 let mut new_aabbs = [AABB::new_invalid(); SIMD_WIDTH];
@@ -94,7 +113,7 @@ impl WQuadtree {
                     if node.leaf {
                         // We are in a leaf: compute the colliders' AABBs.
                         if let Some(proxy) = self.proxies.get(*child_id as usize) {
-                            let collider = &colliders[proxy.handle];
+                            let collider = &colliders[proxy.data];
                             *new_aabb = collider.compute_aabb();
                         }
                     } else {
@@ -107,7 +126,7 @@ impl WQuadtree {
 
                 let node = &mut self.nodes[id as usize];
                 let new_waabb = WAABB::from(new_aabbs);
-                if !node.waabb.contains_lanewise(&new_waabb).all() {
+                if !node.waabb.contains(&new_waabb).all() {
                     node.waabb = new_waabb;
                     node.waabb.dilate_by_factor(dilation_factor);
                     self.dirty_nodes.push_back(node.parent.index);
@@ -116,31 +135,40 @@ impl WQuadtree {
             }
         }
     }
+}
+
+impl<T: IndexedData> WQuadtree<T> {
+    pub fn new() -> Self {
+        WQuadtree {
+            nodes: Vec::new(),
+            dirty_nodes: VecDeque::new(),
+            proxies: Vec::new(),
+        }
+    }
 
-    pub fn clear_and_rebuild(&mut self, colliders: &ColliderSet, dilation_factor: f32) {
+    pub fn clear_and_rebuild(
+        &mut self,
+        data: impl ExactSizeIterator<Item = (T, AABB)>,
+        dilation_factor: f32,
+    ) {
         self.nodes.clear();
         self.proxies.clear();
 
         // Create proxies.
-        let mut indices = Vec::with_capacity(colliders.len());
-        self.proxies = vec![WQuadtreeProxy::invalid(); colliders.len()];
+        let mut indices = Vec::with_capacity(data.len());
+        let mut aabbs = vec![AABB::new_invalid(); data.len()];
+        self.proxies = vec![WQuadtreeProxy::invalid(); data.len()];
 
-        for (handle, collider) in colliders.iter() {
-            let index = handle.into_raw_parts().0;
+        for (data, aabb) in data {
+            let index = data.index();
             if index >= self.proxies.len() {
                 self.proxies.resize(index + 1, WQuadtreeProxy::invalid());
+                aabbs.resize(index + 1, AABB::new_invalid());
             }
 
-            self.proxies[index].handle = handle;
-            indices.push(index);
-        }
-
-        // Compute AABBs.
-        let mut aabbs = vec![AABB::new_invalid(); self.proxies.len()];
-        for (handle, collider) in colliders.iter() {
-            let index = handle.into_raw_parts().0;
-            let aabb = collider.compute_aabb();
+            self.proxies[index].data = data;
             aabbs[index] = aabb;
+            indices.push(index);
         }
 
         // Build the tree recursively.
@@ -279,7 +307,47 @@ impl WQuadtree {
         (id, my_aabb)
     }
 
-    pub fn cast_ray(&self, ray: &Ray, max_toi: f32) -> Vec<ColliderHandle> {
+    // FIXME: implement a visitor pattern to merge intersect_aabb
+    // and intersect_ray into a single method.
+    pub fn intersect_aabb(&self, aabb: &AABB) -> Vec<T> {
+        let mut res = Vec::new();
+
+        if self.nodes.is_empty() {
+            return res;
+        }
+
+        // Special case for the root.
+        let mut stack = vec![0u32];
+        let waabb = WAABB::splat(*aabb);
+        while let Some(inode) = stack.pop() {
+            let node = self.nodes[inode as usize];
+            let intersections = node.waabb.intersects(&waabb);
+            let bitmask = intersections.bitmask();
+
+            for ii in 0..SIMD_WIDTH {
+                if (bitmask & (1 << ii)) != 0 {
+                    if node.leaf {
+                        // We found a leaf!
+                        // Unfortunately, invalid AABBs return a intersection as well.
+                        if let Some(proxy) = self.proxies.get(node.children[ii] as usize) {
+                            res.push(proxy.data);
+                        }
+                    } else {
+                        // Internal node, visit the child.
+                        // Unfortunately, we have this check because invalid AABBs
+                        // return a intersection as well.
+                        if node.children[ii] as usize <= self.nodes.len() {
+                            stack.push(node.children[ii]);
+                        }
+                    }
+                }
+            }
+        }
+
+        res
+    }
+
+    pub fn cast_ray(&self, ray: &Ray, max_toi: f32) -> Vec<T> {
         let mut res = Vec::new();
 
         if self.nodes.is_empty() {
@@ -301,7 +369,7 @@ impl WQuadtree {
                         // We found a leaf!
                         // Unfortunately, invalid AABBs return a hit as well.
                         if let Some(proxy) = self.proxies.get(node.children[ii] as usize) {
-                            res.push(proxy.handle);
+                            res.push(proxy.data);
                         }
                     } else {
                         // Internal node, visit the child.
@@ -324,14 +392,14 @@ struct WQuadtreeIncrementalBuilderStep {
     parent: NodeIndex,
 }
 
-struct WQuadtreeIncrementalBuilder {
-    quadtree: WQuadtree,
+struct WQuadtreeIncrementalBuilder<T> {
+    quadtree: WQuadtree<T>,
     to_insert: Vec<WQuadtreeIncrementalBuilderStep>,
     aabbs: Vec<AABB>,
     indices: Vec<usize>,
 }
 
-impl WQuadtreeIncrementalBuilder {
+impl<T: IndexedData> WQuadtreeIncrementalBuilder<T> {
     pub fn new() -> Self {
         Self {
             quadtree: WQuadtree::new(),
diff --git a/src/pipeline/query_pipeline.rs b/src/pipeline/query_pipeline.rs
index ce4e063..10aa7ed 100644
--- a/src/pipeline/query_pipeline.rs
+++ b/src/pipeline/query_pipeline.rs
@@ -8,7 +8,7 @@ use ncollide::bounding_volume::BoundingVolume;
 /// A pipeline for performing queries on all the colliders of a scene.
 #[cfg_attr(feature = "serde-serialize", derive(Serialize, Deserialize))]
 pub struct QueryPipeline {
-    quadtree: WQuadtree,
+    quadtree: WQuadtree<ColliderHandle>,
     tree_built: bool,
     dilation_factor: f32,
 }
@@ -32,9 +32,10 @@ impl QueryPipeline {
     /// Update the acceleration structure on the query pipeline.
     pub fn update(&mut self, bodies: &RigidBodySet, colliders: &ColliderSet) {
         if !self.tree_built {
-            self.quadtree
-                .clear_and_rebuild(colliders, self.dilation_factor);
-            // self.tree_built = true; // FIXME: uncomment this once we handle insertion/removals properly.
+            let data = colliders.iter().map(|(h, c)| (h, c.compute_aabb()));
+            self.quadtree.clear_and_rebuild(data, self.dilation_factor);
+            // FIXME: uncomment this once we handle insertion/removals properly.
+            // self.tree_built = true;
             return;
         }