use dashmap::DashMap;
use g_math::fixed_point::FixedPoint;
use crate::hyperbolic_geometry::HyperbolicPoint;
use crate::metric_tree::{CachedNormPoint, HyperbolicMetric, Metric};
const MAX_BANDS: usize = 64;
const MAX_SECTORS: i64 = 1 << 28;
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug, PartialOrd, Ord)]
pub struct CellId {
pub band: i32,
pub sector: i64,
}
#[derive(Clone)]
struct Entry {
unique_id: String,
point: CachedNormPoint,
}
#[derive(Clone)]
struct Band {
cosh_lo: FixedPoint,
sinh_lo: FixedPoint,
tanh_lo: FixedPoint,
cosh_hi: FixedPoint,
sinh_hi: FixedPoint,
tanh_hi: FixedPoint,
sectors: i64,
}
pub struct CellIndex {
thresholds: Vec<FixedPoint>,
bands: Vec<Band>,
cells: DashMap<CellId, Vec<Entry>>,
node_cell: DashMap<String, CellId>,
band_load: DashMap<i32, usize>,
band_sectors: DashMap<i32, std::collections::BTreeSet<i64>>,
band_width: f64,
arc: f64,
}
impl Default for CellIndex {
fn default() -> Self {
Self::new(0.5, 0.5)
}
}
impl CellIndex {
pub fn new(w: f64, arc: f64) -> Self {
let thresholds = (0..=MAX_BANDS)
.map(|b| {
let edge = FixedPoint::from_f64((b as f64) * w / 2.0).tanh();
edge * edge
})
.collect();
let bands = (0..MAX_BANDS)
.map(|b| {
let lo = FixedPoint::from_f64((b as f64) * w);
let hi = FixedPoint::from_f64(((b + 1) as f64) * w);
let mid = ((b as f64) + 0.5) * w;
let ideal = 2.0 * std::f64::consts::PI * mid.sinh() / arc;
let sectors = if ideal >= MAX_SECTORS as f64 {
MAX_SECTORS
} else {
(ideal.ceil() as i64).clamp(1, MAX_SECTORS)
};
Band {
cosh_lo: lo.cosh(),
sinh_lo: lo.sinh(),
tanh_lo: lo.tanh(),
cosh_hi: hi.cosh(),
sinh_hi: hi.sinh(),
tanh_hi: hi.tanh(),
sectors,
}
})
.collect();
Self {
thresholds,
bands,
cells: DashMap::new(),
node_cell: DashMap::new(),
band_load: DashMap::new(),
band_sectors: DashMap::new(),
band_width: w,
arc,
}
}
pub fn parameters(&self) -> (f64, f64) {
(self.band_width, self.arc)
}
pub fn len(&self) -> usize {
self.node_cell.len()
}
pub fn is_empty(&self) -> bool {
self.node_cell.is_empty()
}
pub fn cell_count(&self) -> usize {
self.cells.len()
}
pub fn cell_of(&self, point: &HyperbolicPoint) -> CellId {
let norm_sq = planar_norm_sq(point);
let band = self.band_of(norm_sq);
let sectors = self.bands[band as usize].sectors;
CellId { band, sector: self.sector_of(point, sectors) }
}
fn band_of(&self, norm_sq: FixedPoint) -> i32 {
let idx = self.thresholds.partition_point(|t| *t <= norm_sq);
(idx.max(1) - 1).min(MAX_BANDS - 1) as i32
}
fn sector_of(&self, point: &HyperbolicPoint, sectors: i64) -> i64 {
let pseudo = pseudo_angle(point.coords()[0], point.coords()[1]);
sector_from_pseudo(pseudo.to_f64(), sectors)
}
pub fn insert(&self, unique_id: &str, point: &HyperbolicPoint) {
let cell = self.cell_of(point);
if let Some(previous) = self.node_cell.get(unique_id).map(|r| *r.value()) {
if previous == cell {
return;
}
self.detach(unique_id, previous);
}
self.cells.entry(cell).or_default().push(Entry {
unique_id: unique_id.to_string(),
point: CachedNormPoint::new(point.clone()),
});
self.node_cell.insert(unique_id.to_string(), cell);
*self.band_load.entry(cell.band).or_insert(0) += 1;
self.band_sectors.entry(cell.band).or_default().insert(cell.sector);
}
pub fn remove(&self, unique_id: &str) {
if let Some((_, cell)) = self.node_cell.remove(unique_id) {
self.detach(unique_id, cell);
}
}
fn detach(&self, unique_id: &str, cell: CellId) {
let emptied = match self.cells.get_mut(&cell) {
Some(mut members) => {
members.retain(|e| e.unique_id != unique_id);
members.is_empty()
}
None => false,
};
if emptied {
self.cells.remove(&cell);
if let Some(mut sectors) = self.band_sectors.get_mut(&cell.band) {
sectors.remove(&cell.sector);
}
}
if let Some(mut load) = self.band_load.get_mut(&cell.band) {
*load = load.saturating_sub(1);
}
}
pub fn knn(&self, query: &HyperbolicPoint, k: usize) -> Vec<(String, FixedPoint)> {
if k == 0 {
return Vec::new();
}
let probe = CachedNormPoint::new(query.clone());
let mut best: Vec<(FixedPoint, String)> = Vec::with_capacity(k + 1);
let ceiling = std::cell::Cell::new(None::<FixedPoint>);
self.expand(query, |cell| {
let Some(members) = self.cells.get(&cell) else { return };
for entry in members.iter() {
let score = HyperbolicMetric.proxy(&probe, &entry.point);
best.push((score, entry.unique_id.clone()));
}
best.sort_by(|a, b| {
a.0.partial_cmp(&b.0)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| a.1.cmp(&b.1))
});
best.dedup_by(|a, b| a.1 == b.1);
best.truncate(k);
if best.len() == k {
ceiling.set(Some(cosh_from_proxy(best[k - 1].0)));
}
}, || ceiling.get());
best.into_iter()
.map(|(score, id)| (id, ratio_sq_to_distance(score)))
.collect()
}
pub fn within_radius(&self, centre: &HyperbolicPoint, radius: FixedPoint) -> Vec<(String, FixedPoint)> {
let probe = CachedNormPoint::new(centre.clone());
let ceiling = radius.try_cosh().ok();
let mut found: Vec<(String, FixedPoint)> = Vec::new();
let limit = move || ceiling;
self.expand(centre, |cell| {
let Some(members) = self.cells.get(&cell) else { return };
for entry in members.iter() {
let distance = HyperbolicMetric.distance(&probe, &entry.point);
if distance <= radius {
found.push((entry.unique_id.clone(), distance));
}
}
}, limit);
found.sort_by(|a, b| {
a.1.partial_cmp(&b.1)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| a.0.cmp(&b.0))
});
found
}
fn expand<F, C>(&self, query: &HyperbolicPoint, mut visit: F, ceiling: C)
where
F: FnMut(CellId),
C: Fn() -> Option<FixedPoint>,
{
let norm_sq = planar_norm_sq(query);
let pseudo_q = pseudo_angle(query.coords()[0], query.coords()[1]).to_f64();
let one = FixedPoint::from_int(1);
let outside = one - norm_sq;
if outside <= FixedPoint::from_int(0) {
return;
}
let cosh_q = (one + norm_sq) / outside;
let sinh_q = FixedPoint::from_int(2) * norm_sq.sqrt() / outside;
let sinh_q_sq = sinh_q * sinh_q;
let mut order: Vec<(FixedPoint, i32)> = self
.band_load
.iter()
.filter(|r| *r.value() > 0)
.map(|r| (self.radial_bound(*r.key(), cosh_q, sinh_q), *r.key()))
.collect();
order.sort_by(|a, b| {
a.0.partial_cmp(&b.0)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| a.1.cmp(&b.1))
});
for (radial, band) in order {
if let Some(limit) = ceiling() {
if radial > limit {
break;
}
}
let data = &self.bands[band as usize];
let sectors = data.sectors;
let centre = sector_from_pseudo(pseudo_q, sectors);
let Some(occupied) = self.band_sectors.get(&band) else { continue };
let mut candidates: Vec<(i64, i64)> = occupied
.iter()
.map(|s| {
let raw = (s - centre).rem_euclid(sectors);
(raw.min(sectors - raw), *s)
})
.collect();
drop(occupied);
candidates.sort_unstable();
let sector_width = 4.0 / (sectors as f64);
for (offset, sector) in candidates {
if let Some(limit) = ceiling() {
let envelope_gap = ((offset - 1).max(0) as f64) * sector_width;
let envelope =
self.bound_for_gap(data, envelope_gap, cosh_q, sinh_q, sinh_q_sq);
if envelope.exceeds(limit) {
break;
}
let bound =
self.cell_bound(data, sector, pseudo_q, cosh_q, sinh_q, sinh_q_sq);
if bound.exceeds(limit) {
continue;
}
}
visit(CellId { band, sector });
}
}
}
fn radial_bound(&self, band: i32, cosh_q: FixedPoint, sinh_q: FixedPoint) -> FixedPoint {
let data = &self.bands[band as usize];
let one = FixedPoint::from_int(1);
let below = cosh_q * data.cosh_lo - sinh_q * data.sinh_lo;
let above = cosh_q * data.cosh_hi - sinh_q * data.sinh_hi;
if below < one && above < one {
one
} else if below >= one && above >= one {
if below < above { below } else { above }
} else {
one
}
}
fn cell_bound(
&self,
data: &Band,
sector: i64,
pseudo_q: f64,
cosh_q: FixedPoint,
sinh_q: FixedPoint,
sinh_q_sq: FixedPoint,
) -> Bound {
let sectors = data.sectors as f64;
let lo = (sector as f64) * 4.0 / sectors;
let hi = ((sector + 1) as f64) * 4.0 / sectors;
let circular = |a: f64, b: f64| {
let d = (a - b).abs();
d.min(4.0 - d)
};
let gap = if pseudo_q >= lo && pseudo_q < hi {
0.0
} else {
circular(pseudo_q, lo).min(circular(pseudo_q, hi))
};
self.bound_for_gap(data, gap, cosh_q, sinh_q, sinh_q_sq)
}
fn bound_for_gap(
&self,
data: &Band,
gap: f64,
cosh_q: FixedPoint,
sinh_q: FixedPoint,
sinh_q_sq: FixedPoint,
) -> Bound {
let delta_theta = gap.min(std::f64::consts::PI);
let (sin_dt, cos_dt) = FixedPoint::from_f64(delta_theta).sincos();
let zero = FixedPoint::from_int(0);
if cos_dt <= zero {
return Bound::Plain(cosh_q * data.cosh_lo - sinh_q * cos_dt * data.sinh_lo);
}
let b_term = sinh_q * cos_dt;
if b_term >= cosh_q * data.tanh_lo && b_term <= cosh_q * data.tanh_hi {
Bound::Squared(FixedPoint::from_int(1) + sinh_q_sq * sin_dt * sin_dt)
} else if b_term < cosh_q * data.tanh_lo {
Bound::Plain(cosh_q * data.cosh_lo - b_term * data.sinh_lo)
} else {
Bound::Plain(cosh_q * data.cosh_hi - b_term * data.sinh_hi)
}
}
}
enum Bound {
Plain(FixedPoint),
Squared(FixedPoint),
}
impl Bound {
fn exceeds(&self, limit: FixedPoint) -> bool {
match self {
Bound::Plain(v) => *v > limit,
Bound::Squared(v) => *v / limit > limit,
}
}
}
fn sector_from_pseudo(pseudo: f64, sectors: i64) -> i64 {
let scaled = (pseudo / 4.0) * sectors as f64;
(scaled.floor() as i64).rem_euclid(sectors)
}
fn planar_norm_sq(point: &HyperbolicPoint) -> FixedPoint {
let x = point.coords()[0];
let y = point.coords()[1];
x * x + y * y
}
fn pseudo_angle(x: FixedPoint, y: FixedPoint) -> FixedPoint {
let zero = FixedPoint::from_int(0);
if x == zero && y == zero {
return zero;
}
let one = FixedPoint::from_int(1);
if y >= zero {
if x >= zero {
y / (x + y)
} else {
one - x / (y - x)
}
} else if x < zero {
FixedPoint::from_int(2) - y / (-x - y)
} else {
FixedPoint::from_int(3) + x / (x - y)
}
}
fn cosh_from_proxy(s: FixedPoint) -> FixedPoint {
let one = FixedPoint::from_int(1);
let denominator = one - s;
if denominator <= FixedPoint::from_int(0) {
return crate::constants::near_boundary().cosh();
}
(one + s) / denominator
}
fn ratio_sq_to_distance(s: FixedPoint) -> FixedPoint {
crate::hyperbolic_geometry::ratio_to_distance(s.sqrt())
}
#[cfg(test)]
mod tests {
use super::*;
fn point(x: f64, y: f64) -> HyperbolicPoint {
HyperbolicPoint::from_slice(&[FixedPoint::from_f64(x), FixedPoint::from_f64(y)])
}
fn populated() -> (CellIndex, Vec<(String, HyperbolicPoint)>) {
let index = CellIndex::default();
let mut nodes = Vec::new();
for i in 0..400 {
let radius = 0.05 + 0.9 * ((i % 20) as f64) / 20.0;
let angle = 0.37 * i as f64;
let p = point(radius * angle.cos(), radius * angle.sin());
let id = format!("n{i}");
index.insert(&id, &p);
nodes.push((id, p));
}
(index, nodes)
}
fn brute_force(nodes: &[(String, HyperbolicPoint)], q: &HyperbolicPoint, k: usize) -> Vec<String> {
let mut all: Vec<(FixedPoint, String)> = nodes
.iter()
.map(|(id, p)| (q.hyperbolic_distance(p), id.clone()))
.collect();
all.sort_by(|a, b| {
a.0.partial_cmp(&b.0)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| a.1.cmp(&b.1))
});
all.into_iter().take(k).map(|(_, id)| id).collect()
}
#[test]
fn cell_assignment_is_stable_and_reproducible() {
let index = CellIndex::default();
let p = point(0.3, -0.4);
assert_eq!(index.cell_of(&p), index.cell_of(&p));
assert_eq!(index.parameters(), (0.5, 0.5));
}
#[test]
fn a_node_is_its_own_nearest_neighbour() {
let (index, nodes) = populated();
for (id, p) in &nodes {
let got = index.knn(p, 1);
assert_eq!(&got[0].0, id, "{id} did not find itself");
}
}
#[test]
fn matches_brute_force_for_k_greater_than_one() {
let (index, nodes) = populated();
let mut state: u64 = 0x9E37_79B9_7F4A_7C15;
let mut rand = || {
state ^= state << 13;
state ^= state >> 7;
state ^= state << 17;
(state >> 11) as f64 / (1u64 << 53) as f64
};
for k in [1usize, 5, 20] {
for _ in 0..40 {
let r = rand().sqrt() * 0.95;
let a = rand() * std::f64::consts::TAU;
let q = point(r * a.cos(), r * a.sin());
let got: Vec<String> = index.knn(&q, k).into_iter().map(|(id, _)| id).collect();
let want = brute_force(&nodes, &q, k);
assert_eq!(got, want, "k={k} disagreed with brute force");
}
}
}
#[test]
fn within_radius_matches_brute_force() {
let (index, nodes) = populated();
let q = point(0.2, 0.1);
for r in [0.5f64, 1.3, 2.7] {
let radius = FixedPoint::from_f64(r);
let mut got: Vec<String> =
index.within_radius(&q, radius).into_iter().map(|(id, _)| id).collect();
let mut want: Vec<String> = nodes
.iter()
.filter(|(_, p)| q.hyperbolic_distance(p) <= radius)
.map(|(id, _)| id.clone())
.collect();
got.sort();
want.sort();
assert_eq!(got, want, "radius {r} disagreed with brute force");
}
}
#[test]
fn concurrent_insert_query_and_remove_stay_coherent() {
use std::sync::Arc;
use std::thread;
const THREADS: usize = 8;
const PER_THREAD: usize = 200;
fn site(t: usize, i: usize) -> HyperbolicPoint {
let radius = 0.05 + 0.9 * (((t * 7 + i) % 20) as f64) / 20.0;
let angle = 0.37 * (i as f64) + 0.11 * (t as f64);
point(radius * angle.cos(), radius * angle.sin())
}
let index = Arc::new(CellIndex::default());
let mut handles = Vec::new();
for t in 0..THREADS {
let idx = Arc::clone(&index);
handles.push(thread::spawn(move || {
for i in 0..PER_THREAD {
idx.insert(&format!("t{t}n{i}"), &site(t, i));
}
}));
}
for h in handles {
h.join().expect("insert thread panicked");
}
assert_eq!(
index.len(),
THREADS * PER_THREAD,
"concurrent inserts lost or duplicated nodes",
);
let mut handles = Vec::new();
for t in 0..THREADS {
let idx = Arc::clone(&index);
handles.push(thread::spawn(move || -> usize {
let mut seen = 0usize;
if t % 2 == 0 {
for i in (1..PER_THREAD).step_by(2) {
idx.remove(&format!("t{t}n{i}"));
}
for i in (1..PER_THREAD).step_by(2) {
idx.insert(&format!("t{t}n{i}"), &site(t, i));
}
} else {
for i in 0..PER_THREAD {
let q = site(t, i);
for (id, _) in idx.knn(&q, 5) {
assert!(
id.starts_with('t') && id.contains('n'),
"query returned an id that was never inserted: {id}",
);
seen += 1;
}
let r = FixedPoint::from_f64(0.5);
for (id, _) in idx.within_radius(&q, r) {
assert!(
id.starts_with('t'),
"radius query returned a foreign id: {id}",
);
}
}
}
seen
}));
}
let seen: usize = handles
.into_iter()
.map(|h| h.join().expect("worker thread panicked"))
.sum();
assert!(seen > 0, "readers observed nothing at all");
assert_eq!(
index.len(),
THREADS * PER_THREAD,
"churn under contention changed the population",
);
for t in 0..THREADS {
for i in (0..PER_THREAD).step_by(37) {
let q = site(t, i);
let hit = index.knn(&q, 1);
assert!(!hit.is_empty(), "t{t}n{i} vanished from the index");
assert_eq!(
hit[0].1,
FixedPoint::from_int(0),
"querying at a node's own position must return distance 0",
);
}
}
}
#[test]
fn a_radius_too_large_for_cosh_returns_everything() {
let (index, nodes) = populated();
let q = point(0.2, 0.1);
for r in [30, 100, 1_000, 100_000] {
let got = index.within_radius(&q, FixedPoint::from_int(r));
assert_eq!(
got.len(),
nodes.len(),
"radius {r} should sweep the whole index",
);
}
}
#[test]
fn removal_takes_a_node_out_of_results() {
let (index, nodes) = populated();
let (victim, at) = nodes[17].clone();
assert_eq!(index.knn(&at, 1)[0].0, victim);
index.remove(&victim);
assert_eq!(index.len(), nodes.len() - 1);
let ids: Vec<String> = index.knn(&at, 5).into_iter().map(|(id, _)| id).collect();
assert!(!ids.contains(&victim), "removed node came back: {ids:?}");
}
#[test]
fn reinsert_moves_a_node_between_cells() {
let index = CellIndex::default();
let start = point(0.1, 0.1);
let end = point(-0.8, 0.2);
index.insert("drifter", &start);
let first = index.cell_of(&start);
index.insert("drifter", &end);
assert_ne!(first, index.cell_of(&end));
assert_eq!(index.len(), 1, "moving a node must not duplicate it");
assert_eq!(index.knn(&end, 1)[0].0, "drifter");
}
#[test]
fn off_plane_queries_stay_exact() {
let index = CellIndex::default();
let mut nodes = Vec::new();
for i in 0..400 {
let radius = 0.05 + 0.9 * ((i % 20) as f64) / 20.0;
let angle = 0.37 * i as f64;
let p = HyperbolicPoint::from_slice(&[
FixedPoint::from_f64(radius * angle.cos()),
FixedPoint::from_f64(radius * angle.sin()),
FixedPoint::from_int(0),
]);
let id = format!("n{i}");
index.insert(&id, &p);
nodes.push((id, p));
}
for (dx, dy, dz) in [(0.2, -0.1, 0.3), (-0.5, 0.25, 0.15), (0.0, 0.0, 0.6)] {
let q = HyperbolicPoint::from_slice(&[
FixedPoint::from_f64(dx),
FixedPoint::from_f64(dy),
FixedPoint::from_f64(dz),
]);
let got: Vec<String> = index.knn(&q, 5).into_iter().map(|(id, _)| id).collect();
let want = brute_force(&nodes, &q, 5);
assert_eq!(got, want, "off-plane query ({dx},{dy},{dz}) was not exact");
}
}
#[test]
fn every_cell_bound_is_a_true_lower_bound() {
let (index, nodes) = populated();
let mut state: u64 = 0x243F_6A88_85A3_08D3;
let mut rand = || {
state ^= state << 13;
state ^= state >> 7;
state ^= state << 17;
(state >> 11) as f64 / (1u64 << 53) as f64
};
const DEEP_NORMS: [f64; 3] = [1.0 - 2.5e-7, 1.0 - 1e-5, 1.0 - 1e-3];
for (i, depth_norm) in DEEP_NORMS.iter().enumerate() {
let angle = 0.9 * i as f64;
index.insert(
&format!("deep{i}"),
&point(depth_norm * angle.cos(), depth_norm * angle.sin()),
);
}
let nodes: Vec<(String, HyperbolicPoint)> = nodes
.into_iter()
.chain((0..3).map(|i| {
let angle = 0.9 * i as f64;
let n = DEEP_NORMS[i];
(format!("deep{i}"), point(n * angle.cos(), n * angle.sin()))
}))
.collect();
let one = FixedPoint::from_int(1);
let mut checked = 0usize;
let mut saturated = 0usize;
for round in 0..120 {
let r = if round % 4 == 0 {
1.0 - 2.5e-7 * (1.0 + rand())
} else {
rand().sqrt() * 0.96
};
let a = rand() * std::f64::consts::TAU;
let q = point(r * a.cos(), r * a.sin());
let norm_sq = planar_norm_sq(&q);
let outside = one - norm_sq;
let cosh_q = (one + norm_sq) / outside;
let sinh_q_sq = FixedPoint::from_int(4) * norm_sq / (outside * outside);
let sinh_q = sinh_q_sq.sqrt();
let pseudo_q = pseudo_angle(q.coords()[0], q.coords()[1]).to_f64();
for cell in index.cells.iter() {
let id = *cell.key();
let data = &index.bands[id.band as usize];
let bound = index.cell_bound(
data, id.sector, pseudo_q, cosh_q, sinh_q, sinh_q_sq,
);
for entry in cell.value().iter() {
let actual = &nodes
.iter()
.find(|(n, _)| *n == entry.unique_id)
.expect("indexed node must exist in the fixture")
.1;
let d = q.hyperbolic_distance(actual);
let cosh_d = d.cosh();
if d.to_f64() > 27.0 {
saturated += 1;
continue;
}
checked += 1;
let holds = match bound {
Bound::Plain(v) => cosh_d >= v,
Bound::Squared(v) => cosh_d >= v / cosh_d,
};
assert!(
holds,
"cell {id:?} claimed a bound that exceeds a member's true \
distance (cosh d = {}) — the ring expansion would prune \
a genuine neighbour",
cosh_d.to_f64()
);
}
}
}
assert!(checked > 10_000, "too few (query, point) pairs checked: {checked}");
assert!(
saturated * 4 < checked,
"{saturated} of {} pairs saturated the distance kernel — the fixture \
has drifted past the usable radius and is no longer testing anything",
saturated + checked
);
}
#[test]
fn occupancy_is_spread_not_concentrated() {
let (index, nodes) = populated();
assert!(
index.cell_count() > nodes.len() / 20,
"cells {} for {} nodes — occupancy collapsed",
index.cell_count(),
nodes.len()
);
}
}