add aabb search and benchmark

benches/benchmarks.rs

@@ -94,11 +94,32 @@ fn nearest_short_horizon(c: &mut Criterion) {
     });
 }
 
+fn aabb_search(c: &mut Criterion) {
+    c.bench_function(
+        "search an axis-aligned bounding box 10 units a side",
+        move |b| {
+            let points: Vec<_> = create_random_points(DEFAULT_NUM_POINTS, SEED_1)
+                .into_iter()
+                .map(|[x, y]| Point2d { x, y })
+                .collect();
+
+            let index = Spindex2d::new(&points);
+
+            let mut i = 0;
+            b.iter(|| {
+                i += 1;
+                !index.aabb(&points[i - 1], 10.0).is_empty()
+            });
+        },
+    );
+}
+
 criterion_group!(
     benches,
     new_index,
     update_positions,
     nearest_default_horizon,
-    nearest_short_horizon
+    nearest_short_horizon,
+    aabb_search
 );
 criterion_main!(benches);

src/lib.rs

@@ -132,73 +132,158 @@ impl Spindex2d {
     }
 
     /// Find up to `k` closest points to `point`.
-    pub fn knn(&self, point: &Point2d, k: usize) -> Option<Vec<Point2d>> {
+    pub fn knn(&self, point: &Point2d, k: usize) -> Vec<Point2d> {
         todo!()
     }
 
     /// Search within a given `radius` from `point`.
-    pub fn radius(&self, point: &Point2d, radius: f64) -> Option<Vec<Point2d>> {
+    pub fn radius(&self, point: &Point2d, radius: f64) -> Vec<Point2d> {
         todo!()
     }
 
     /// Search an axis-aligned bounding box whose sides are `span` away from the given point.
-    pub fn aabb(&self, point: &Point2d, span: f64) -> Option<Vec<Point2d>> {
+    pub fn aabb(&self, point: &Point2d, span: f64) -> Vec<Point2d> {
-        todo!()
+        let span = span / 2.0;
+
+        let mut nxs = HashMap::new();
+        let mut nys = HashMap::new();
+
+        thread::scope(|s| {
+            // x axis
+            s.spawn(|| {
+                let xs = self.xs.read().unwrap();
+                let p = PVal::new(0, point.x);
+                let nx = match xs.binary_search(&p) {
+                    Ok(i) | Err(i) => i,
+                };
+
+                let mut i = nx;
+                // first walk backwards
+                loop {
+                    let p = xs[i];
+                    // `p.val` is less than `point.x` because the xs are sorted
+                    let dist = point.x - p.val();
+                    if dist <= span {
+                        nxs.insert(p.point_index(), p.val());
+                    } else {
+                        break;
+                    }
+                    if i == 0 {
+                        break;
+                    }
+                    i -= 1;
+                }
+
+                // now forwards
+                i = nx;
+                while i < xs.len() {
+                    let p = xs[i];
+                    let dist = p.val() - point.x;
+                    if dist <= span {
+                        nxs.insert(p.point_index(), p.val());
+                    } else {
+                        break;
+                    }
+                    i += 1;
+                }
+            });
+
+            // y axis
+            s.spawn(|| {
+                let ys = self.ys.read().unwrap();
+                let p = PVal::new(0, point.y);
+                let ny = match ys.binary_search(&p) {
+                    Ok(i) | Err(i) => i,
+                };
+                let mut i = ny;
+                // backwards
+                loop {
+                    let p = ys[i];
+                    // `p.val` is less than `point.y` because the ys are sorted
+                    let dist = point.y - p.val();
+                    if dist <= span {
+                        nys.insert(p.point_index(), p.val());
+                    } else {
+                        break;
+                    }
+                    if i == 0 {
+                        break;
+                    }
+                    i -= 1;
+                }
+
+                // forwards
+                i = ny;
+                while i < ys.len() {
+                    let p = ys[i];
+                    let dist = p.val() - point.y;
+                    if dist <= span {
+                        nys.insert(p.point_index(), p.val());
+                    } else {
+                        break;
+                    }
+                    i += 1;
+                }
+            });
+        });
+
+        let mut neighbors = Vec::with_capacity(nxs.len());
+        for (id, x) in nxs.into_iter() {
+            if let Some(&y) = nys.get(&id) {
+                neighbors.push(Point2d::new(x, y));
+            }
+        }
+        neighbors
     }
 
     pub fn nearest(&self, point: &Point2d) -> Option<Point2d> {
-        let x = point.x;
-        let y = point.y;
-
-        let xs = self.xs.read().unwrap();
-        let ys = self.ys.read().unwrap();
-
-        let len = xs.len();
         let horizon = self.horizon;
         let mut nxs = HashMap::with_capacity(horizon * 2);
         let mut nys = HashMap::with_capacity(horizon * 2);
         thread::scope(|s| {
             s.spawn(|| {
-                let p = PVal::new(0, x);
+                let xs = self.xs.read().unwrap();
+                let p = PVal::new(0, point.x);
                 let nx = match xs.binary_search(&p) {
                     Ok(i) | Err(i) => i,
                 };
                 let start = nx.saturating_sub(horizon);
-                let end = (nx + horizon).min(len);
+                let end = (nx + horizon).min(xs.len());
                 for p in xs[start..end].iter() {
                     nxs.insert(p.point_index(), p.val());
                 }
             });
 
             s.spawn(|| {
-                let p = PVal::new(0, y);
+                let ys = self.ys.read().unwrap();
+                let p = PVal::new(0, point.y);
                 let ny = match ys.binary_search(&p) {
                     Ok(i) | Err(i) => i,
                 };
                 let start = ny.saturating_sub(horizon);
-                let end = (ny + horizon).min(len);
+                let end = (ny + horizon).min(ys.len());
                 for p in ys[start..end].iter() {
                     nys.insert(p.point_index(), p.val());
                 }
             });
         });
-        let mut points = Vec::with_capacity(horizon * 2);
+
+        let mut points = Vec::with_capacity(nxs.len());
+
         for (id, xval) in nxs.into_iter() {
             if let Some(&yval) = nys.get(&id) {
                 points.push(Point2d::new(xval, yval));
             }
         }
-        if points.is_empty() {
+
-            None
-        } else {
         points.sort_unstable_by(|a, b| {
             let adist = a.distance_squared(point);
             let bdist = b.distance_squared(point);
             adist.total_cmp(&bdist)
         });
+
         points.first().cloned()
     }
-    }
 
     pub fn points(&self) -> Vec<Point2d> {
         let mut points = vec![Point2d::default(); self.xs.read().unwrap().len()];