From 0ecee89421357be2d08a745410cc9e31fbfcdc3e Mon Sep 17 00:00:00 2001
From: Joe <Joe>
Date: Sat, 18 Oct 2025 14:22:28 -0700
Subject: [PATCH] add tweakable horizon and benchmark

---
 benches/benchmarks.rs | 33 ++++++++++++++++++++++++++++-----
 src/lib.rs            | 25 ++++++++++++++++---------
 2 files changed, 44 insertions(+), 14 deletions(-)

diff --git a/benches/benchmarks.rs b/benches/benchmarks.rs
index e9244ac..ef3555b 100644
--- a/benches/benchmarks.rs
+++ b/benches/benchmarks.rs
@@ -14,7 +14,7 @@ use criterion::{criterion_group, criterion_main};
 const SEED_1: &[u8; 32] = b"Gv0aHMtHkBGsUXNspGU9fLRuCWkZWHZx";
 // const SEED_2: &[u8; 32] = b"km7DO4GeaFZfTcDXVpnO7ZJlgUY7hZiS";
 
-const DEFAULT_NUM_POINTS: usize = 1_000_000;
+const DEFAULT_NUM_POINTS: usize = 100_000;
 
 fn new_index(c: &mut Criterion) {
     c.bench_function("new index", move |b| {
@@ -60,8 +60,8 @@ fn update_positions(c: &mut Criterion) {
     });
 }
 
-fn nearest(c: &mut Criterion) {
-    c.bench_function("nearst point", move |b| {
+fn nearest_default_horizon(c: &mut Criterion) {
+    c.bench_function("nearst point with default horizon", move |b| {
         let points: Vec<_> = create_random_points(DEFAULT_NUM_POINTS, SEED_1)
             .into_iter()
             .map(|[x, y]| Point2d { x, y })
@@ -71,11 +71,34 @@ fn nearest(c: &mut Criterion) {
 
         let mut i = 0;
         b.iter(|| {
-            index.nearest(&points[i]);
             i += 1;
+            index.nearest(&points[i - 1]).is_some()
         });
     });
 }
 
-criterion_group!(benches, new_index, update_positions, nearest);
+fn nearest_short_horizon(c: &mut Criterion) {
+    c.bench_function("nearst point with short horizon", 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).with_horizon(20);
+
+        let mut i = 0;
+        b.iter(|| {
+            i += 1;
+            index.nearest(&points[i - 1]).is_some()
+        });
+    });
+}
+
+criterion_group!(
+    benches,
+    new_index,
+    update_positions,
+    nearest_default_horizon,
+    nearest_short_horizon
+);
 criterion_main!(benches);
diff --git a/src/lib.rs b/src/lib.rs
index df8fd94..830a9c3 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -67,7 +67,7 @@ impl PVal {
 pub struct Spindex2d {
     pub xs: RwLock<Vec<PVal>>,
     pub ys: RwLock<Vec<PVal>>,
-    len: usize,
+    horizon: usize,
 }
 
 impl Spindex2d {
@@ -92,10 +92,11 @@ impl Spindex2d {
             });
         });
 
+        let horizon = (len / (10usize.pow(len.ilog10() / 2))).max(50);
         let index = Self {
             xs: xs.into(),
             ys: ys.into(),
-            len,
+            horizon,
         };
 
         index.sort();
@@ -103,6 +104,10 @@ impl Spindex2d {
         index
     }
 
+    pub fn with_horizon(self, horizon: usize) -> Self {
+        Self { horizon, ..self }
+    }
+
     /// This method assumes that the points have the same cardinality and order as the points used
     /// when constructing this index. May panic if the updated points argument is longer than the
     /// original points.
@@ -148,16 +153,18 @@ impl Spindex2d {
         let xs = self.xs.read().unwrap();
         let ys = self.ys.read().unwrap();
 
-        let mut nxs = HashMap::with_capacity(100);
-        let mut nys = HashMap::with_capacity(100);
+        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 nx = match xs.binary_search(&p) {
                     Ok(i) | Err(i) => i,
                 };
-                let start = nx.saturating_sub(50);
-                let end = (nx + (50)).min(self.len);
+                let start = nx.saturating_sub(horizon);
+                let end = (nx + horizon).min(len);
                 for p in xs[start..end].iter() {
                     nxs.insert(p.point_index(), p.val());
                 }
@@ -168,14 +175,14 @@ impl Spindex2d {
                 let ny = match ys.binary_search(&p) {
                     Ok(i) | Err(i) => i,
                 };
-                let start = ny.saturating_sub(50);
-                let end = (ny + 50).min(self.len);
+                let start = ny.saturating_sub(horizon);
+                let end = (ny + horizon).min(len);
                 for p in ys[start..end].iter() {
                     nys.insert(p.point_index(), p.val());
                 }
             });
         });
-        let mut points = Vec::with_capacity(100);
+        let mut points = Vec::with_capacity(horizon * 2);
         for (id, xval) in nxs.into_iter() {
             if let Some(&yval) = nys.get(&id) {
                 points.push(Point2d::new(xval, yval));