use a static thread-safe arena allocator for storing temporary AABB search data

Cargo.lock

@@ -11,62 +11,24 @@ dependencies = [
  "memchr",
 ]
 
+[[package]]
+name = "allocator-api2"
+version = "0.2.21"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923"
+
 [[package]]
 name = "anes"
 version = "0.1.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299"
 
-[[package]]
-name = "anstream"
-version = "0.6.21"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "43d5b281e737544384e969a5ccad3f1cdd24b48086a0fc1b2a5262a26b8f4f4a"
-dependencies = [
- "anstyle",
- "anstyle-parse",
- "anstyle-query",
- "anstyle-wincon",
- "colorchoice",
- "is_terminal_polyfill",
- "utf8parse",
-]
-
 [[package]]
 name = "anstyle"
 version = "1.0.13"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "5192cca8006f1fd4f7237516f40fa183bb07f8fbdfedaa0036de5ea9b0b45e78"
 
-[[package]]
-name = "anstyle-parse"
-version = "0.2.7"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4e7644824f0aa2c7b9384579234ef10eb7efb6a0deb83f9630a49594dd9c15c2"
-dependencies = [
- "utf8parse",
-]
-
-[[package]]
-name = "anstyle-query"
-version = "1.1.4"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9e231f6134f61b71076a3eab506c379d4f36122f2af15a9ff04415ea4c3339e2"
-dependencies = [
- "windows-sys 0.60.2",
-]
-
-[[package]]
-name = "anstyle-wincon"
-version = "3.0.10"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3e0633414522a32ffaac8ac6cc8f748e090c5717661fddeea04219e2344f5f2a"
-dependencies = [
- "anstyle",
- "once_cell_polyfill",
- "windows-sys 0.60.2",
-]
-
 [[package]]
 name = "async-task"
 version = "4.7.1"
@@ -89,6 +51,22 @@ version = "1.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8"
 
+[[package]]
+name = "bitflags"
+version = "2.9.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2261d10cca569e4643e526d8dc2e62e433cc8aba21ab764233731f8d369bf394"
+
+[[package]]
+name = "blink-alloc"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e669f146bb8b2327006ed94c69cf78c8ec81c100192564654230a40b4f091d82"
+dependencies = [
+ "allocator-api2",
+ "parking_lot",
+]
+
 [[package]]
 name = "bumpalo"
 version = "3.19.0"
@@ -149,10 +127,8 @@ version = "4.5.49"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0025e98baa12e766c67ba13ff4695a887a1eba19569aad00a472546795bd6730"
 dependencies = [
- "anstream",
  "anstyle",
  "clap_lex",
- "strsim",
 ]
 
 [[package]]
@@ -161,12 +137,6 @@ version = "0.7.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a1d728cc89cf3aee9ff92b05e62b19ee65a02b5702cff7d5a377e32c6ae29d8d"
 
-[[package]]
-name = "colorchoice"
-version = "1.0.4"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b05b61dc5112cbb17e4b6cd61790d9845d13888356391624cbe7e41efeac1e75"
-
 [[package]]
 name = "criterion"
 version = "0.7.0"
@@ -237,6 +207,18 @@ version = "1.15.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719"
 
+[[package]]
+name = "equivalent"
+version = "1.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f"
+
+[[package]]
+name = "foldhash"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb"
+
 [[package]]
 name = "forte"
 version = "1.0.0-alpha.4"
@@ -273,10 +255,16 @@ dependencies = [
 ]
 
 [[package]]
-name = "is_terminal_polyfill"
+name = "hashbrown"
-version = "1.70.1"
+version = "0.16.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7943c866cc5cd64cbc25b2e01621d07fa8eb2a1a23160ee81ce38704e97b8ecf"
+checksum = "5419bdc4f6a9207fbeba6d11b604d481addf78ecd10c11ad51e76c2f6482748d"
+dependencies = [
+ "allocator-api2",
+ "equivalent",
+ "foldhash",
+ "rustc-std-workspace-alloc",
+]
 
 [[package]]
 name = "itertools"
@@ -315,6 +303,15 @@ version = "0.2.177"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2874a2af47a2325c2001a6e6fad9b16a53b802102b528163885171cf92b15976"
 
+[[package]]
+name = "lock_api"
+version = "0.4.14"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965"
+dependencies = [
+ "scopeguard",
+]
+
 [[package]]
 name = "log"
 version = "0.4.28"
@@ -351,18 +348,35 @@ version = "1.21.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d"
 
-[[package]]
-name = "once_cell_polyfill"
-version = "1.70.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a4895175b425cb1f87721b59f0f286c2092bd4af812243672510e1ac53e2e0ad"
-
 [[package]]
 name = "oorandom"
 version = "11.1.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e"
 
+[[package]]
+name = "parking_lot"
+version = "0.12.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a"
+dependencies = [
+ "lock_api",
+ "parking_lot_core",
+]
+
+[[package]]
+name = "parking_lot_core"
+version = "0.9.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1"
+dependencies = [
+ "cfg-if",
+ "libc",
+ "redox_syscall",
+ "smallvec",
+ "windows-link",
+]
+
 [[package]]
 name = "pin-project-lite"
 version = "0.2.16"
@@ -488,6 +502,15 @@ dependencies = [
  "crossbeam-utils",
 ]
 
+[[package]]
+name = "redox_syscall"
+version = "0.5.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d"
+dependencies = [
+ "bitflags",
+]
+
 [[package]]
 name = "regex"
 version = "1.12.2"
@@ -517,6 +540,12 @@ version = "0.8.8"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7a2d987857b319362043e95f5353c0535c1f58eec5336fdfcf626430af7def58"
 
+[[package]]
+name = "rustc-std-workspace-alloc"
+version = "1.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f9d441c3b2ebf55cebf796bfdc265d67fa09db17b7bb6bd4be75c509e1e8fec3"
+
 [[package]]
 name = "rustversion"
 version = "1.0.22"
@@ -538,6 +567,12 @@ dependencies = [
  "winapi-util",
 ]
 
+[[package]]
+name = "scopeguard"
+version = "1.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49"
+
 [[package]]
 name = "serde"
 version = "1.0.228"
@@ -600,19 +635,14 @@ checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03"
 name = "spacedog"
 version = "0.1.0"
 dependencies = [
- "clap",
+ "blink-alloc",
  "criterion",
  "forte",
+ "hashbrown",
  "rand",
  "rand_hc",
 ]
 
-[[package]]
-name = "strsim"
-version = "0.11.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f"
-
 [[package]]
 name = "syn"
 version = "2.0.106"
@@ -706,12 +736,6 @@ version = "1.0.19"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f63a545481291138910575129486daeaf8ac54aee4387fe7906919f7830c7d9d"
 
-[[package]]
-name = "utf8parse"
-version = "0.2.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821"
-
 [[package]]
 name = "valuable"
 version = "0.1.1"

Cargo.toml

@@ -3,17 +3,16 @@ name = "spacedog"
 version = "0.1.0"
 edition = "2024"
 
+[dependencies]
+blink-alloc = { version = "0.3", features = ["alloc", "nightly", "sync"] }
+forte = "1.0.0-alpha.4"
+hashbrown = { version = "0.16", features = ["alloc", "nightly"] }
+
 [dev-dependencies]
-clap = "4.5.49"
+criterion = "0.7"
-criterion = "0.7.0"
+rand = "0.9"
-rand = "0.9.2"
+rand_hc = "0.4"
-rand_hc = "0.4.0"
 
 [[bench]]
 name = "benchmarks"
 harness = false
-
-[dependencies]
-forte = "1.0.0-alpha.4"
-rand = "0.9.2"
-rand_hc = "0.4.0"

benches/benchmarks.rs

@@ -3,25 +3,16 @@ use spacedog::*;
 use rand::{Rng, SeedableRng};
 use rand_hc::Hc128Rng;
 
-fn create_random_points(num_points: usize, seed: &[u8; 32]) -> Vec<[f64; 2]> {
-    let mut rng = Hc128Rng::from_seed(*seed);
-    (0..num_points).map(|_| rng.random()).collect()
-}
-
 use criterion::Criterion;
 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 = 100_000;
 
 fn new_index(c: &mut Criterion) {
     c.bench_function("new index", move |b| {
-        let points: Vec<_> = create_random_points(DEFAULT_NUM_POINTS, SEED_1)
+        let points: Vec<_> = create_random_points(DEFAULT_NUM_POINTS, SEED_1);
-            .into_iter()
-            .map(|[x, y]| Point2d { x, y })
-            .collect();
 
         b.iter(|| {
             Spindex2d::new(&points);
@@ -31,19 +22,16 @@ fn new_index(c: &mut Criterion) {
 
 fn update_positions(c: &mut Criterion) {
     c.bench_function("update positions", move |b| {
-        let points: Vec<_> = create_random_points(DEFAULT_NUM_POINTS, SEED_1)
+        let points: Vec<_> = create_random_points(DEFAULT_NUM_POINTS, SEED_1);
-            .into_iter()
-            .map(|[x, y]| Point2d { x, y })
-            .collect();
 
         let mut rng = Hc128Rng::from_seed(*SEED_1);
 
         let jitters: Vec<_> = points
             .iter()
             .map(|p| {
-                let x = p.x + rng.random_range(-5.0..=5.0);
+                let x = p.x + rng.random_range(-0.5..=0.5);
-                let y = p.y + rng.random_range(-5.0..=5.0);
+                let y = p.y + rng.random_range(-0.5..=0.5);
-                Point2d { x, y }
+                Point2d::new(x, y)
             })
             .collect();
 
@@ -61,11 +49,8 @@ fn update_positions(c: &mut Criterion) {
 }
 
 fn nearest_default_horizon(c: &mut Criterion) {
-    c.bench_function("nearst point with default horizon", move |b| {
+    c.bench_function("nearest default", move |b| {
-        let points: Vec<_> = create_random_points(DEFAULT_NUM_POINTS, SEED_1)
+        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);
 
@@ -78,41 +63,46 @@ fn nearest_default_horizon(c: &mut Criterion) {
 }
 
 fn nearest_short_horizon(c: &mut Criterion) {
-    c.bench_function("nearst point with short horizon", move |b| {
+    c.bench_function("nearest short", move |b| {
-        let points: Vec<_> = create_random_points(DEFAULT_NUM_POINTS, SEED_1)
+        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 end = points.len() - 1;
+        let len = points.len();
 
         let mut i = 0;
         b.iter(|| {
-            i += 1;
+            i = (i + 1) % len;
-            index.nearest(&points[(i - 1).min(end)]).is_some()
+            index.nearest(&points[i]).is_some()
         });
     });
 }
 
 fn aabb_search(c: &mut Criterion) {
-    c.bench_function(
+    c.bench_function("aabb", move |b| {
-        "search an axis-aligned bounding box 10 units a side",
+        let points: Vec<_> = create_random_points(DEFAULT_NUM_POINTS, SEED_1);
-        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 index = Spindex2d::new(&points);
 
-            let mut i = 0;
+        let mut i = 0;
-            b.iter(|| {
+        b.iter(|| {
-                i += 1;
+            i += 1;
-                !index.aabb(&points[i - 1], 10.0).is_empty()
+            !index.aabb(&points[i - 1], 0.05).is_empty()
-            });
+        });
-        },
+    });
-    );
+}
+
+fn radius_search(c: &mut Criterion) {
+    c.bench_function("radius 0.1", move |b| {
+        let points: Vec<_> = create_random_points(DEFAULT_NUM_POINTS, SEED_1);
+
+        let index = Spindex2d::new(&points);
+
+        let mut i = 0;
+        b.iter(|| {
+            i += 1;
+            !index.radius(&points[i - 1], 0.1).is_empty()
+        });
+    });
 }
 
 criterion_group!(
@@ -121,6 +111,14 @@ criterion_group!(
     update_positions,
     nearest_default_horizon,
     nearest_short_horizon,
-    aabb_search
+    aabb_search,
+    radius_search
 );
 criterion_main!(benches);
+
+fn create_random_points(num_points: usize, seed: &[u8; 32]) -> Vec<Point2d> {
+    let mut rng = Hc128Rng::from_seed(*seed);
+    (0..num_points)
+        .map(|_| rng.random::<[f64; 2]>().into())
+        .collect()
+}

src/lib.rs

@@ -1,4 +1,9 @@
-use std::{collections::HashMap, sync::RwLock, thread};
+use std::{
+    collections::HashMap,
+    sync::{Mutex, RwLock},
+};
+
+use blink_alloc::SyncBlinkAlloc;
 
 static THREAD_POOL: forte::ThreadPool = forte::ThreadPool::new();
 
@@ -8,6 +13,13 @@ pub struct Point2d {
     pub y: f64,
 }
 
+impl From<[f64; 2]> for Point2d {
+    fn from(value: [f64; 2]) -> Self {
+        let [x, y] = value;
+        Point2d::new(x, y)
+    }
+}
+
 impl Point2d {
     pub fn new(x: f64, y: f64) -> Self {
         Self { x, y }
@@ -69,6 +81,7 @@ pub struct Spindex2d {
     pub xs: RwLock<Vec<PVal>>,
     pub ys: RwLock<Vec<PVal>>,
     horizon: usize,
+    alloc: Mutex<blink_alloc::SyncBlinkAlloc>,
 }
 
 impl Spindex2d {
@@ -100,6 +113,7 @@ impl Spindex2d {
             xs: xs.into(),
             ys: ys.into(),
             horizon,
+            alloc: SyncBlinkAlloc::new().into(),
         };
 
         index.sort();
@@ -141,91 +155,101 @@ impl Spindex2d {
 
     /// Search within a given `radius` from `point`.
     pub fn radius(&self, point: &Point2d, radius: f64) -> Vec<Point2d> {
-        todo!()
+        let d2 = radius.powi(2);
+        self.aabb(point, radius * 2.0)
+            .into_iter()
+            .filter(|v| v.distance_squared(point) <= d2)
+            .collect()
     }
 
     /// Search an axis-aligned bounding box whose sides are `span` away from the given point.
     pub fn aabb(&self, point: &Point2d, span: f64) -> Vec<Point2d> {
         let span = span / 2.0;
-        let sz = { self.xs.read().unwrap().len() / 2 };
-        let mut nxs = HashMap::with_capacity(sz);
-        let mut nys = HashMap::with_capacity(sz);
 
-        THREAD_POOL.scope(|s| {
+        let neighbors = {
-            // x axis
+            let alloc = self.alloc.lock().unwrap();
-            s.spawn(|_| {
+            let mut nys = hashbrown::HashMap::new_in(alloc.local());
-                let xs = self.xs.read().unwrap();
+            let mut nxs = hashbrown::HashMap::new_in(alloc.local());
-                let p = PVal::new(0, point.x);
-                let nx = match xs.binary_search(&p) {
-                    Ok(i) | Err(i) => i,
-                };
 
-                let mut i = nx;
+            THREAD_POOL.scope(|s| {
-                // first walk backwards
+                // x axis
-                while i > 0 {
+                s.spawn(|_| {
-                    i -= 1;
+                    let xs = self.xs.read().unwrap();
-                    let p = xs[i];
+                    let p = PVal::new(0, point.x);
-                    // `p.val` is less than `point.x` because the xs are sorted
+                    let nx = match xs.binary_search(&p) {
-                    let dist = point.x - p.val();
+                        Ok(i) | Err(i) => i,
-                    if dist <= span {
+                    };
-                        nxs.insert(p.point_index(), p.val());
+
+                    let mut i = nx;
+                    // first walk backwards
+                    //let mut nxs = nxs.lock().unwrap();
+                    while i > 0 {
+                        i -= 1;
+                        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());
+                        }
                     }
-                }
 
-                // now forwards
+                    // now forwards
-                i = nx;
+                    i = nx;
-                while i < xs.len() {
+                    while i < xs.len() {
-                    let p = xs[i];
+                        let p = xs[i];
-                    let dist = p.val() - point.x;
+                        let dist = p.val() - point.x;
-                    if dist <= span {
+                        if dist <= span {
-                        nxs.insert(p.point_index(), p.val());
+                            nxs.insert(p.point_index(), p.val());
-                    } else {
+                        } else {
-                        break;
+                            break;
+                        }
+                        i += 1;
                     }
-                    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
+                    while i > 0 {
+                        i -= 1;
+                        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());
+                        }
+                    }
+
+                    // 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;
+                    }
+                });
             });
 
-            // y axis
+            let mut neighbors = Vec::with_capacity(nxs.len());
-            s.spawn(|_| {
+            for (id, x) in nxs.into_iter() {
-                let ys = self.ys.read().unwrap();
+                if let Some(&y) = nys.get(&id) {
-                let p = PVal::new(0, point.y);
+                    neighbors.push(Point2d::new(x, y));
-                let ny = match ys.binary_search(&p) {
-                    Ok(i) | Err(i) => i,
-                };
-                let mut i = ny;
-                // backwards
-                while i > 0 {
-                    i -= 1;
-                    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());
-                    }
                 }
-
-                // 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
+        };
+        self.alloc.lock().unwrap().reset();
         neighbors
     }
 
@@ -269,13 +293,11 @@ impl Spindex2d {
             }
         }
 
-        points.sort_unstable_by(|a, b| {
+        points.into_iter().min_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> {