Add all-nearest-neighbors (a1NN) iterator for tree-to-tree lookup

rstar-benches/benches/benchmarks.rs

@@ -16,6 +16,7 @@ use criterion::{Bencher, Criterion, Fun};
 const SEED_1: &[u8; 32] = b"Gv0aHMtHkBGsUXNspGU9fLRuCWkZWHZx";
 const SEED_2: &[u8; 32] = b"km7DO4GeaFZfTcDXVpnO7ZJlgUY7hZiS";
 
+#[derive(Clone)]
 struct Params;
 
 impl RTreeParams for Params {
@@ -92,6 +93,28 @@ fn tree_creation_quality(c: &mut Criterion) {
     });
 }
 
+fn all_to_all_neighbors(c: &mut Criterion) {
+    const SIZE: usize = 1_000;
+    let points: Vec<_> = create_random_points(SIZE, SEED_1);
+    let tree_target = RTree::<_, Params>::bulk_load_with_params(points.clone());
+    let query_points = create_random_points(SIZE, SEED_2);
+    let tree_query = RTree::<_, Params>::bulk_load_with_params(query_points.clone());
+
+    let tree_target_cloned = tree_target.clone();
+    c.bench_function("all to all tree lookup", move |b| {
+        b.iter(|| {
+            tree_target.all_nearest_neighbors(&tree_query).count();
+        });
+    })
+    .bench_function("all to all point lookup", move |b| {
+        b.iter(|| {
+            for query_point in &query_points {
+                tree_target_cloned.nearest_neighbor(&query_point).unwrap();
+            }
+        });
+    });
+}
+
 fn locate_successful(c: &mut Criterion) {
     let points: Vec<_> = create_random_points(100_000, SEED_1);
     let query_point = points[500];
@@ -115,6 +138,7 @@ criterion_group!(
     bulk_load_baseline,
     bulk_load_comparison,
     tree_creation_quality,
+    all_to_all_neighbors,
     locate_successful,
     locate_unsuccessful
 );

rstar/src/aabb.rs

@@ -1,4 +1,4 @@
-use crate::point::{max_inline, Point, PointExt};
+use crate::point::{max_inline, min_inline, Point, PointExt};
 use crate::{Envelope, RTreeObject};
 use num_traits::{Bounded, One, Zero};
 
@@ -94,6 +94,64 @@ where
             self.min_point(point).sub(point).length_2()
         }
     }
+
+    /// Return an iterator over each corner vertex of the AABB.
+    ///
+    /// # Example
+    /// ```
+    /// use rstar::AABB;
+    ///
+    /// let aabb = AABB::from_corners([1.0, 2.0], [3.0, 4.0]);
+    /// let mut corners = aabb.iter_corners();
+    /// assert_eq!(corners.next(), Some([1.0, 2.0]));
+    /// assert_eq!(corners.next(), Some([3.0, 2.0]));
+    /// assert_eq!(corners.next(), Some([1.0, 4.0]));
+    /// assert_eq!(corners.next(), Some([3.0, 4.0]));
+    /// assert_eq!(corners.next(), None);
+    /// ```
+    pub fn iter_corners(&self) -> impl Iterator<Item=P> {
+        CornerIterator::new(self)
+    }
+}
+
+struct CornerIterator<P: Point> {
+    lower: P,
+    upper: P,
+    idx: usize,
+}
+
+impl<P> CornerIterator<P>
+where
+    P: Point,
+{
+    fn new(aabb: &AABB<P>) -> Self {
+        Self {
+            lower: aabb.lower,
+            upper: aabb.upper,
+            idx: 0,
+        }
+    }
+}
+
+impl<P> Iterator for CornerIterator<P>
+where
+    P: Point,
+{
+    type Item = P;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.idx & (1 << P::DIMENSIONS) != 0 {
+            None
+        } else {
+            let corner = P::generate(|i| if self.idx & (1 << i) != 0 {
+                self.upper.nth(i)
+            } else {
+                self.lower.nth(i)
+            });
+            self.idx += 1;
+            Some(corner)
+        }
+    }
 }
 
 impl<P> Envelope for AABB<P>
@@ -144,6 +202,12 @@ where
         self.distance_2(point)
     }
 
+    fn min_dist_2(&self, other: &Self) -> P::Scalar {
+        let l = self.min_point(&other.lower);
+        let u = self.min_point(&other.upper);
+        min_inline(other.distance_2(&l), other.distance_2(&u))
+    }
+
     fn min_max_dist_2(&self, point: &P) -> <P as Point>::Scalar {
         let l = self.lower.sub(point);
         let u = self.upper.sub(point);
@@ -169,6 +233,13 @@ where
         result - max_diff
     }
 
+    fn max_min_max_dist_2(&self, other: &Self) -> P::Scalar {
+        self.iter_corners()
+            .map(|corner| other.min_max_dist_2(&corner))
+            .max_by(|a, b| a.partial_cmp(b).unwrap_or_else(|| std::cmp::Ordering::Equal))
+            .unwrap_or_else(P::Scalar::zero)
+    }
+
     fn center(&self) -> Self::Point {
         let one = <Self::Point as Point>::Scalar::one();
         let two = one + one;