diff --git a/benches/benchmarks.rs b/benches/benchmarks.rs index ef3555b..084706a 100644 --- a/benches/benchmarks.rs +++ b/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); diff --git a/src/lib.rs b/src/lib.rs index 830a9c3..0029e85 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -132,72 +132,157 @@ impl Spindex2d { } /// Find up to `k` closest points to `point`. - pub fn knn(&self, point: &Point2d, k: usize) -> Option> { + pub fn knn(&self, point: &Point2d, k: usize) -> Vec { todo!() } /// Search within a given `radius` from `point`. - pub fn radius(&self, point: &Point2d, radius: f64) -> Option> { + pub fn radius(&self, point: &Point2d, radius: f64) -> Vec { 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> { - todo!() + pub fn aabb(&self, point: &Point2d, span: f64) -> Vec { + 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 { - 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() - } + + 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 {