use std::collections::{HashMap, HashSet};
use std::fmt::{self, Debug, Formatter};
use std::sync::Mutex;
use dashmap::DashMap;
use sha3::{Sha3_512, Digest};
use g_math::fixed_point::{FixedPoint, FixedVector};
use super::hyperbolic_geometry::{PoincareDisk, HyperbolicPoint, distance_to_ratio};
use crate::metric_tree::{hyperbolic_ratio_sq, sq_ratio_separation_exceeds};
use crate::constants;
#[derive(Clone, PartialEq, Eq, Hash)]
pub struct GeometricSignature {
hash: String,
level: u32,
position_signature: Vec<i32>,
}
impl GeometricSignature {
pub fn new(hash: String, level: u32, position_signature: Vec<i32>) -> Self {
Self {
hash,
level,
position_signature,
}
}
pub fn hash(&self) -> &str {
&self.hash
}
pub fn level(&self) -> u32 {
self.level
}
pub fn position_signature(&self) -> &[i32] {
&self.position_signature
}
pub fn stub(unique_id: &str) -> Self {
Self {
hash: unique_id.to_string(),
level: 0,
position_signature: Vec::new(),
}
}
pub fn is_stub(&self) -> bool {
self.position_signature.is_empty()
}
pub fn unique_id(&self) -> String {
if self.position_signature.is_empty() {
return self.hash.clone();
}
use sha3::{Sha3_256, Digest as _};
let mut hasher = Sha3_256::new();
hasher.update(self.hash.as_bytes());
hasher.update(self.level.to_le_bytes());
for &v in &self.position_signature {
hasher.update(v.to_le_bytes());
}
hex::encode(&hasher.finalize()[..16])
}
}
impl Debug for GeometricSignature {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
write!(f, "GeometricSignature(hash={}, level={})",
&self.hash[0..8], self.level)
}
}
#[derive(Clone, Debug)]
pub struct HyperbolicRegion {
center: HyperbolicPoint,
radius: FixedPoint,
validation_mask: FixedVector,
center_norm_sq: FixedPoint,
radius_ratio_sq: FixedPoint,
}
impl HyperbolicRegion {
pub fn new(center: HyperbolicPoint, radius: FixedPoint) -> Self {
let dimension = center.dimension();
let validation_mask = {
let one = FixedPoint::from_int(1);
let mut mask = FixedVector::new(dimension);
for i in 0..dimension {
let x = center.coords()[i];
mask[i] = x * (one + x.tanh());
}
mask
};
let center_norm_sq = center.coords().length_squared();
let radius_ratio_sq = {
let r = distance_to_ratio(radius);
r * r
};
Self {
center,
radius,
validation_mask,
center_norm_sq,
radius_ratio_sq,
}
}
pub fn contains(&self, point: &HyperbolicPoint, _poincare_disk: &PoincareDisk) -> bool {
let point_norm_sq = point.coords().length_squared();
let s = hyperbolic_ratio_sq(&self.center, self.center_norm_sq, point, point_norm_sq);
s <= self.radius_ratio_sq
}
pub fn center(&self) -> &HyperbolicPoint {
&self.center
}
pub fn radius(&self) -> FixedPoint {
self.radius
}
pub fn validation_mask(&self) -> &FixedVector {
&self.validation_mask
}
pub fn quick_validate(&self, point: &HyperbolicPoint) -> bool {
let similarity = point.coords().dot(&self.validation_mask);
similarity > constants::epsilon()
}
}
#[derive(Clone, Debug)]
pub struct BucketEntry {
pub unique_id: String,
pub point: HyperbolicPoint,
pub level: u32,
pub norm_sq: FixedPoint,
}
impl BucketEntry {
pub fn new(unique_id: String, point: HyperbolicPoint, level: u32) -> Self {
let norm_sq = point.coords().length_squared();
Self { unique_id, point, level, norm_sq }
}
}
fn cmp_fp(a: FixedPoint, b: FixedPoint) -> std::cmp::Ordering {
if a < b { std::cmp::Ordering::Less }
else if a > b { std::cmp::Ordering::Greater }
else { std::cmp::Ordering::Equal }
}
fn euclidean_distance_sq(a: &HyperbolicPoint, b: &HyperbolicPoint) -> FixedPoint {
let mut sum = FixedPoint::from_int(0);
let d = a.dimension().min(b.dimension());
for i in 0..d {
let diff = a.coords()[i] - b.coords()[i];
sum = sum + diff * diff;
}
sum
}
const VP_BUFFER_THRESHOLD: usize = 32;
const VP_REBUILD_DIVISOR: usize = 16;
const VP_DELETE_THRESHOLD: usize = 32;
#[derive(Clone, Debug)]
struct VPNode {
entry: BucketEntry,
median: FixedPoint,
left: Option<Box<VPNode>>,
right: Option<Box<VPNode>>,
}
#[derive(Clone, Debug)]
pub struct VPTree {
root: Option<Box<VPNode>>,
buffer: Vec<BucketEntry>,
deleted: HashSet<String>,
tree_size: usize,
}
impl VPTree {
pub fn new() -> Self {
Self {
root: None,
buffer: Vec::new(),
deleted: HashSet::new(),
tree_size: 0,
}
}
pub fn insert(&mut self, entry: BucketEntry) {
if self.buffer.iter().any(|e| e.unique_id == entry.unique_id) {
return;
}
self.deleted.remove(&entry.unique_id);
self.buffer.push(entry);
let tree_len = self.tree_size;
let threshold = VP_BUFFER_THRESHOLD.max(tree_len / VP_REBUILD_DIVISOR);
if self.buffer.len() >= threshold {
self.rebuild();
}
}
pub fn remove(&mut self, unique_id: &str) {
let before = self.buffer.len();
self.buffer.retain(|e| e.unique_id != unique_id);
if self.buffer.len() < before {
return;
}
self.deleted.insert(unique_id.to_string());
if self.deleted.len() >= VP_DELETE_THRESHOLD {
self.rebuild();
}
}
pub fn live_count(&self) -> usize {
let tree_live = self.tree_size.saturating_sub(self.deleted.len());
tree_live + self.buffer.len()
}
pub fn is_empty(&self) -> bool {
self.live_count() == 0
}
pub fn find_in_radius(&self, center: &HyperbolicPoint, radius: FixedPoint) -> Vec<(String, FixedPoint)> {
let mut results = Vec::new();
let radius_sq = {
let r = distance_to_ratio(radius);
r * r
};
let center_norm_sq = center.coords().length_squared();
if let Some(ref root) = self.root {
Self::search_radius(root, center, center_norm_sq, radius_sq, &self.deleted, &mut results);
}
for entry in &self.buffer {
let s = hyperbolic_ratio_sq(center, center_norm_sq, &entry.point, entry.norm_sq);
if s <= radius_sq {
results.push((entry.unique_id.clone(), center.hyperbolic_distance(&entry.point)));
}
}
results
}
pub fn find_nearest(&self, point: &HyperbolicPoint, k: usize) -> Vec<(String, FixedPoint)> {
if k == 0 { return Vec::new(); }
let query_norm_sq = point.coords().length_squared();
let mut candidates: Vec<(FixedPoint, &BucketEntry)> = Vec::with_capacity(k + 1);
let mut tau = FixedPoint::from_int(1);
if let Some(ref root) = self.root {
Self::search_knn(root, point, query_norm_sq, k, &self.deleted, &mut candidates, &mut tau);
}
for entry in &self.buffer {
let s = hyperbolic_ratio_sq(point, query_norm_sq, &entry.point, entry.norm_sq);
if candidates.len() < k || s < tau {
candidates.push((s, entry));
candidates.sort_by(|a, b| cmp_fp(a.0, b.0).then_with(|| a.1.unique_id.cmp(&b.1.unique_id)));
if candidates.len() > k {
candidates.truncate(k);
}
if candidates.len() == k {
tau = candidates.last().unwrap().0;
}
}
}
candidates
.into_iter()
.map(|(_, e)| (e.unique_id.clone(), point.hyperbolic_distance(&e.point)))
.collect()
}
fn rebuild(&mut self) {
let mut entries = Vec::with_capacity(self.tree_size + self.buffer.len());
if let Some(root) = self.root.take() {
Self::collect_live(*root, &self.deleted, &mut entries);
}
entries.append(&mut self.buffer);
self.deleted.clear();
self.tree_size = entries.len();
self.root = Self::build_tree(entries);
}
fn collect_live(node: VPNode, deleted: &HashSet<String>, out: &mut Vec<BucketEntry>) {
if !deleted.contains(&node.entry.unique_id) {
out.push(node.entry);
}
if let Some(left) = node.left {
Self::collect_live(*left, deleted, out);
}
if let Some(right) = node.right {
Self::collect_live(*right, deleted, out);
}
}
pub fn farthest_from(&self, center: &HyperbolicPoint) -> Option<(String, FixedPoint)> {
let center_norm_sq = center.coords().length_squared();
let mut best: Option<(FixedPoint, String, HyperbolicPoint)> = None;
let mut consider = |entry: &BucketEntry| {
let s = hyperbolic_ratio_sq(center, center_norm_sq, &entry.point, entry.norm_sq);
if best.as_ref().is_none_or(|(m, _, _)| s > *m) {
best = Some((s, entry.unique_id.clone(), entry.point.clone()));
}
};
for entry in &self.buffer {
consider(entry);
}
if let Some(ref root) = self.root {
Self::visit_live(root, &self.deleted, &mut consider);
}
best.map(|(_, id, pt)| (id, center.hyperbolic_distance(&pt)))
}
fn visit_live<F: FnMut(&BucketEntry)>(node: &VPNode, deleted: &HashSet<String>, f: &mut F) {
if !deleted.contains(&node.entry.unique_id) {
f(&node.entry);
}
if let Some(ref left) = node.left {
Self::visit_live(left, deleted, f);
}
if let Some(ref right) = node.right {
Self::visit_live(right, deleted, f);
}
}
fn build_tree(mut entries: Vec<BucketEntry>) -> Option<Box<VPNode>> {
if entries.is_empty() {
return None;
}
if entries.len() == 1 {
return Some(Box::new(VPNode {
entry: entries.remove(0),
median: FixedPoint::from_int(0),
left: None,
right: None,
}));
}
let vp = entries.swap_remove(0);
let mut with_dists: Vec<(BucketEntry, FixedPoint)> = entries
.into_iter()
.map(|e| {
let s = hyperbolic_ratio_sq(&vp.point, vp.norm_sq, &e.point, e.norm_sq);
(e, s)
})
.collect();
with_dists.sort_by(|a, b| cmp_fp(a.1, b.1));
let median = with_dists[with_dists.len() / 2].1;
let (left_vec, right_vec): (Vec<_>, Vec<_>) = with_dists
.into_iter()
.partition(|(_, d)| *d < median);
let left = Self::build_tree(left_vec.into_iter().map(|(e, _)| e).collect());
let right = Self::build_tree(right_vec.into_iter().map(|(e, _)| e).collect());
Some(Box::new(VPNode {
entry: vp,
median,
left,
right,
}))
}
fn search_radius(
node: &VPNode,
center: &HyperbolicPoint,
center_norm_sq: FixedPoint,
radius_sq: FixedPoint,
deleted: &HashSet<String>,
results: &mut Vec<(String, FixedPoint)>,
) {
let s = hyperbolic_ratio_sq(center, center_norm_sq, &node.entry.point, node.entry.norm_sq);
if s <= radius_sq && !deleted.contains(&node.entry.unique_id) {
results.push((
node.entry.unique_id.clone(),
center.hyperbolic_distance(&node.entry.point),
));
}
if let Some(ref left) = node.left {
let prune = s > node.median && sq_ratio_separation_exceeds(s, node.median, radius_sq);
if !prune {
Self::search_radius(left, center, center_norm_sq, radius_sq, deleted, results);
}
}
if let Some(ref right) = node.right {
let prune = node.median > s && sq_ratio_separation_exceeds(node.median, s, radius_sq);
if !prune {
Self::search_radius(right, center, center_norm_sq, radius_sq, deleted, results);
}
}
}
#[allow(clippy::too_many_arguments)]
fn search_knn<'a>(
node: &'a VPNode,
center: &HyperbolicPoint,
center_norm_sq: FixedPoint,
k: usize,
deleted: &HashSet<String>,
candidates: &mut Vec<(FixedPoint, &'a BucketEntry)>,
tau: &mut FixedPoint,
) {
let s = hyperbolic_ratio_sq(center, center_norm_sq, &node.entry.point, node.entry.norm_sq);
if !deleted.contains(&node.entry.unique_id) {
if candidates.len() < k || s < *tau {
candidates.push((s, &node.entry));
candidates.sort_by(|a, b| cmp_fp(a.0, b.0).then_with(|| a.1.unique_id.cmp(&b.1.unique_id)));
if candidates.len() > k {
candidates.truncate(k);
}
if candidates.len() == k {
*tau = candidates.last().unwrap().0;
}
}
}
let search_left_first = s < node.median;
let prune_left = |s: FixedPoint, tau: FixedPoint| {
s > node.median && sq_ratio_separation_exceeds(s, node.median, tau)
};
let prune_right = |s: FixedPoint, tau: FixedPoint| {
node.median > s && sq_ratio_separation_exceeds(node.median, s, tau)
};
if search_left_first {
if let Some(ref left) = node.left {
if !prune_left(s, *tau) {
Self::search_knn(left, center, center_norm_sq, k, deleted, candidates, tau);
}
}
if let Some(ref right) = node.right {
if !prune_right(s, *tau) {
Self::search_knn(right, center, center_norm_sq, k, deleted, candidates, tau);
}
}
} else {
if let Some(ref right) = node.right {
if !prune_right(s, *tau) {
Self::search_knn(right, center, center_norm_sq, k, deleted, candidates, tau);
}
}
if let Some(ref left) = node.left {
if !prune_left(s, *tau) {
Self::search_knn(left, center, center_norm_sq, k, deleted, candidates, tau);
}
}
}
}
}
#[derive(Debug)]
pub struct HyperbolicHashBucket {
region: HyperbolicRegion,
position_signature: Vec<i32>,
_metrics: Vec<FixedPoint>,
vp_tree: Mutex<VPTree>,
eff: Mutex<EffRadius>,
}
#[derive(Clone, Debug)]
struct EffRadius {
nominal: FixedPoint,
current: FixedPoint,
max_uid: Option<String>,
}
impl EffRadius {
fn new(nominal: FixedPoint) -> Self {
Self { nominal, current: nominal, max_uid: None }
}
}
impl Clone for HyperbolicHashBucket {
fn clone(&self) -> Self {
let vp_tree = self.vp_tree.lock().unwrap_or_else(|e| e.into_inner()).clone();
let eff = self.eff.lock().unwrap_or_else(|e| e.into_inner()).clone();
Self {
region: self.region.clone(),
position_signature: self.position_signature.clone(),
_metrics: self._metrics.clone(),
vp_tree: Mutex::new(vp_tree),
eff: Mutex::new(eff),
}
}
}
impl HyperbolicHashBucket {
pub fn new(region: HyperbolicRegion, position_signature: Vec<i32>) -> Self {
let mut metrics = Vec::new();
let center = region.center();
metrics.push(center.euclidean_norm());
let sum_squares = center.coords().iter().enumerate().fold(
FixedPoint::from_int(0),
|acc, (_i, &x)| acc + x * x
);
metrics.push(sum_squares);
let nominal_radius = region.radius();
Self {
region,
position_signature,
_metrics: metrics,
vp_tree: Mutex::new(VPTree::new()),
eff: Mutex::new(EffRadius::new(nominal_radius)),
}
}
pub fn effective_radius(&self) -> FixedPoint {
self.eff.lock().unwrap_or_else(|e| e.into_inner()).current
}
fn note_node_distance(&self, unique_id: &str, center_dist: FixedPoint) {
let mut e = self.eff.lock().unwrap_or_else(|e| e.into_inner());
if center_dist > e.current {
e.current = center_dist;
e.max_uid = Some(unique_id.to_string());
}
}
fn forget_node(&self, unique_id: &str) {
let mut e = self.eff.lock().unwrap_or_else(|e| e.into_inner());
if e.max_uid.as_deref() != Some(unique_id) {
return;
}
let tree = self.vp_tree.lock().unwrap_or_else(|e| e.into_inner());
match tree.farthest_from(self.region.center()) {
Some((uid, dist)) if dist > e.nominal => {
e.current = dist;
e.max_uid = Some(uid);
}
_ => {
e.current = e.nominal;
e.max_uid = None;
}
}
}
pub fn contains(&self, point: &HyperbolicPoint, poincare_disk: &PoincareDisk) -> bool {
self.region.contains(point, poincare_disk)
}
pub fn region(&self) -> &HyperbolicRegion {
&self.region
}
pub fn position_signature(&self) -> &[i32] {
&self.position_signature
}
pub fn quick_validate(&self, point: &HyperbolicPoint) -> bool {
self.region.quick_validate(point)
}
}
#[derive(Clone)]
pub struct HyperbolicHashTable {
poincare_disk: PoincareDisk,
buckets: HashMap<String, HyperbolicHashBucket>,
signature_map: HashMap<Vec<i32>, String>,
node_to_bucket: DashMap<String, String>,
}
impl HyperbolicHashTable {
pub fn new(dimension: usize) -> Self {
let poincare_disk = PoincareDisk::new(dimension);
let mut table = Self {
poincare_disk,
buckets: HashMap::new(),
signature_map: HashMap::new(),
node_to_bucket: DashMap::new(),
};
table.initialize_buckets();
table
}
fn initialize_buckets(&mut self) {
let dimension = self.poincare_disk.dimension();
let distances = [
FixedPoint::from_int(0), constants::half(), FixedPoint::from_int(1), FixedPoint::from_int(3) / FixedPoint::from_int(2), FixedPoint::from_int(2), ];
let directions_per_distance = [
1, dimension * 2, dimension * 3, dimension * 4, dimension * 5, ];
for (dist_idx, &distance) in distances.iter().enumerate() {
let num_directions = directions_per_distance[dist_idx];
if dist_idx == 0 {
let origin = self.poincare_disk.origin();
let region = HyperbolicRegion::new(origin.clone(), constants::region_radius());
let position_signature = vec![0; dimension];
let bucket = HyperbolicHashBucket::new(region, position_signature.clone());
let signature = self.compute_geometric_signature(&origin);
let hash = self.compute_stable_hash(&signature);
self.buckets.insert(hash.clone(), bucket);
self.signature_map.insert(position_signature, hash);
continue;
}
for dir_idx in 0..num_directions {
let direction = self.generate_direction_vector(dir_idx, num_directions);
let center = self.poincare_disk.point_at_distance_from_origin(
&direction, distance
);
let one_fifth = FixedPoint::from_int(1) / FixedPoint::from_int(5);
let one_tenth = FixedPoint::from_int(1) / FixedPoint::from_int(10);
let radius = one_fifth + one_tenth * distance;
let region = HyperbolicRegion::new(center.clone(), radius);
let position_signature = self.generate_position_signature(¢er);
let bucket = HyperbolicHashBucket::new(region, position_signature.clone());
let signature = self.compute_geometric_signature(¢er);
let hash = self.compute_stable_hash(&signature);
self.buckets.insert(hash.clone(), bucket);
self.signature_map.insert(position_signature, hash);
}
}
}
fn generate_direction_vector(&self, index: usize, total: usize) -> FixedVector {
let dimension = self.poincare_disk.dimension();
let mut direction = FixedVector::new(dimension);
if dimension == 2 {
let angle = constants::two_pi()
* FixedPoint::from_int(index as i32)
/ FixedPoint::from_int(total as i32);
let (sin_a, cos_a) = angle.sincos();
direction[0] = cos_a;
direction[1] = sin_a;
return direction;
}
let phi = constants::golden_angle();
let idx = FixedPoint::from_int((index + 1) as i32);
for i in 0..dimension {
let phase = idx * phi * FixedPoint::from_int((i + 1) as i32);
direction[i] = phase.sin();
}
let norm_sq = direction.dot(&direction);
if norm_sq > constants::epsilon() {
direction.normalize();
} else {
direction[0] = FixedPoint::from_int(1);
}
direction
}
fn generate_position_signature(&self, point: &HyperbolicPoint) -> Vec<i32> {
let dimension = self.poincare_disk.dimension();
let mut signature = Vec::with_capacity(dimension);
for i in 0..dimension {
signature.push(constants::quantize_position(point.coords()[i]));
}
signature
}
fn compute_geometric_signature(&self, point: &HyperbolicPoint) -> Vec<i32> {
let dimension = self.poincare_disk.dimension();
let mut signature = Vec::with_capacity(dimension);
let one = FixedPoint::from_int(1);
for i in 0..dimension {
let x = point.coords()[i];
let transformed = x * (one + x.tanh());
signature.push(constants::quantize_1000(transformed));
}
signature
}
fn compute_stable_hash(&self, signature: &[i32]) -> String {
let mut hasher = Sha3_512::new();
for &value in signature {
hasher.update(value.to_le_bytes());
}
let hash = hasher.finalize();
hex::encode(&hash[..16])
}
pub fn find_bucket(&self, point: &HyperbolicPoint) -> Option<String> {
let position_signature = self.generate_position_signature(point);
if let Some(hash) = self.signature_map.get(&position_signature) {
return Some(hash.clone());
}
let mut candidates: Vec<(&String, FixedPoint)> = self.buckets.iter()
.map(|(hash, bucket)| {
(hash, euclidean_distance_sq(point, bucket.region().center()))
})
.collect();
candidates.sort_unstable_by(|a, b| cmp_fp(a.1, b.1).then_with(|| a.0.cmp(b.0)));
for (hash, _) in &candidates {
if let Some(bucket) = self.buckets.get(*hash) {
if bucket.contains(point, &self.poincare_disk) {
return Some((*hash).clone());
}
}
}
for (hash, _) in &candidates {
if let Some(bucket) = self.buckets.get(*hash) {
if bucket.quick_validate(point) {
return Some((*hash).clone());
}
}
}
None
}
pub fn create_signature(&self, point: &HyperbolicPoint, level: u32) -> Option<GeometricSignature> {
let position_signature = self.generate_position_signature(point);
let hash = if let Some(bucket_hash) = self.find_bucket(point) {
bucket_hash
} else {
let geo_sig = self.compute_geometric_signature(point);
self.compute_stable_hash(&geo_sig)
};
Some(GeometricSignature::new(hash, level, position_signature))
}
pub fn validate_point(&self, point: &HyperbolicPoint) -> bool {
let norm = point.euclidean_norm();
if norm >= FixedPoint::from_int(1) {
return false;
}
self.find_bucket(point).is_some()
}
pub fn poincare_disk(&self) -> &PoincareDisk {
&self.poincare_disk
}
pub fn bucket_count(&self) -> usize {
self.buckets.len()
}
pub fn register_node(&self, point: &HyperbolicPoint, unique_id: &str, level: u32) -> Option<String> {
self.register_node_with_hint(point, unique_id, level, None)
}
pub fn register_node_with_hint(&self, point: &HyperbolicPoint, unique_id: &str, level: u32, bucket_hint: Option<&str>) -> Option<String> {
if self.node_to_bucket.contains_key(unique_id) {
return self.node_to_bucket.get(unique_id).map(|r| r.value().clone());
}
let bucket_hash = match bucket_hint {
Some(hint) if self.buckets.contains_key(hint) => hint.to_string(),
_ => self.find_bucket(point)?,
};
if let Some(bucket) = self.buckets.get(&bucket_hash) {
let center_dist = self.poincare_disk.distance(point, bucket.region.center());
bucket.note_node_distance(unique_id, center_dist);
bucket.vp_tree.lock().unwrap_or_else(|e| e.into_inner()).insert(BucketEntry::new(unique_id.to_string(), point.clone(), level));
}
self.node_to_bucket.insert(unique_id.to_string(), bucket_hash.clone());
Some(bucket_hash)
}
pub fn unregister_node(&self, unique_id: &str) {
if let Some((_, bucket_hash)) = self.node_to_bucket.remove(unique_id) {
if let Some(bucket) = self.buckets.get(&bucket_hash) {
bucket.vp_tree.lock().unwrap_or_else(|e| e.into_inner()).remove(unique_id);
bucket.forget_node(unique_id);
}
}
}
pub fn find_nodes_in_radius(&self, center: &HyperbolicPoint, radius: FixedPoint) -> Vec<(String, FixedPoint)> {
let mut results = Vec::new();
for bucket in self.buckets.values() {
let bucket_center_dist = self.poincare_disk.distance(
center, bucket.region.center()
);
if bucket_center_dist > radius + bucket.effective_radius() {
continue;
}
let bucket_results = bucket.vp_tree.lock().unwrap_or_else(|e| e.into_inner()).find_in_radius(center, radius);
results.extend(bucket_results);
}
results
}
pub fn find_nearest_nodes(&self, point: &HyperbolicPoint, k: usize) -> Vec<(String, FixedPoint)> {
if k == 0 { return Vec::new(); }
let zero = FixedPoint::from_int(0);
let mut bucket_dists: Vec<(&String, FixedPoint)> = self.buckets.iter()
.map(|(hash, bucket)| {
let d = self.poincare_disk.distance(point, bucket.region.center());
let r = bucket.effective_radius();
let min_possible = if d > r { d - r } else { zero };
(hash, min_possible)
})
.collect();
bucket_dists.sort_by(|a, b| cmp_fp(a.1, b.1).then_with(|| a.0.cmp(b.0)));
let mut candidates: Vec<(String, FixedPoint)> = Vec::new();
for (hash, min_possible) in &bucket_dists {
if candidates.len() >= k {
let kth_dist = candidates.last().unwrap().1;
if *min_possible > kth_dist {
break;
}
}
if let Some(bucket) = self.buckets.get(*hash) {
let bucket_results = bucket.vp_tree.lock().unwrap_or_else(|e| e.into_inner()).find_nearest(point, k);
for result in bucket_results {
candidates.push(result);
}
candidates.sort_by(|a, b| cmp_fp(a.1, b.1));
candidates.truncate(k);
}
}
candidates
}
pub fn verify_integrity(&self) -> bool {
if self.buckets.is_empty() {
return false;
}
for (sig, hash) in &self.signature_map {
if !self.buckets.contains_key(hash) {
return false;
}
let bucket = &self.buckets[hash];
if bucket.position_signature() != sig.as_slice() {
return false;
}
}
true
}
}
impl Debug for HyperbolicHashTable {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
write!(f, "HyperbolicHashTable(dim={}, buckets={}, nodes={})",
self.poincare_disk.dimension(), self.buckets.len(),
self.node_to_bucket.len())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_hash_table_creation() {
let table = HyperbolicHashTable::new(2);
assert_eq!(table.poincare_disk().dimension(), 2);
assert!(table.bucket_count() > 0);
}
#[test]
fn test_geometric_signature() {
let table = HyperbolicHashTable::new(2);
let point = table.poincare_disk().point_from_f32_slice(&[0.5, 0.0]);
let signature = table.create_signature(&point, 0).unwrap();
assert_eq!(signature.level(), 0);
assert!(!signature.hash().is_empty());
assert!(!signature.position_signature().is_empty());
}
#[test]
fn test_bucket_finding() {
let table = HyperbolicHashTable::new(2);
let origin = table.poincare_disk().origin();
let bucket_hash = table.find_bucket(&origin);
assert!(bucket_hash.is_some());
}
#[test]
fn test_point_validation() {
let table = HyperbolicHashTable::new(2);
let valid_point = table.poincare_disk().point_from_f32_slice(&[0.5, 0.0]);
assert!(table.validate_point(&valid_point));
let projected_point = table.poincare_disk().point_from_f32_slice(&[1.5, 0.0]);
assert!(table.validate_point(&projected_point));
}
#[test]
fn test_hyperbolic_region() {
let disk = PoincareDisk::new(2);
let center = disk.point_from_f32_slice(&[0.5, 0.0]);
let radius = constants::half();
let region = HyperbolicRegion::new(center.clone(), radius);
assert!(region.contains(¢er, &disk));
assert!(!region.contains(&disk.origin(), &disk));
let far_point = disk.point_from_f32_slice(&[0.8, 0.0]);
assert!(!region.contains(&far_point, &disk));
}
#[test]
fn test_hash_bucket() {
let disk = PoincareDisk::new(2);
let center = disk.point_from_f32_slice(&[0.5, 0.0]);
let radius = constants::half();
let region = HyperbolicRegion::new(center.clone(), radius);
let position_signature = vec![500, 0];
let bucket = HyperbolicHashBucket::new(region, position_signature);
assert!(bucket.contains(¢er, &disk));
assert!(bucket.quick_validate(¢er));
}
#[test]
fn test_integrity_verification() {
let table = HyperbolicHashTable::new(2);
assert!(table.verify_integrity());
}
#[test]
fn test_vp_tree_empty() {
let vp = VPTree::new();
assert!(vp.is_empty());
assert_eq!(vp.live_count(), 0);
let origin = HyperbolicPoint::origin(2);
let results = vp.find_in_radius(&origin, FixedPoint::from_int(10));
assert!(results.is_empty());
let nearest = vp.find_nearest(&origin, 5);
assert!(nearest.is_empty());
}
#[test]
fn test_vp_tree_insert_and_find() {
let disk = PoincareDisk::new(2);
let mut vp = VPTree::new();
let points: Vec<(&str, [f32; 2])> = vec![
("a", [0.1, 0.0]),
("b", [0.2, 0.0]),
("c", [0.3, 0.0]),
("d", [0.0, 0.1]),
("e", [0.0, 0.2]),
];
for (id, coords) in &points {
vp.insert(BucketEntry::new(id.to_string(), disk.point_from_f32_slice(coords), 0));
}
assert_eq!(vp.live_count(), 5);
let origin = disk.origin();
let nearest = vp.find_nearest(&origin, 2);
assert_eq!(nearest.len(), 2);
assert!(nearest[0].1 <= nearest[1].1);
let all = vp.find_in_radius(&origin, FixedPoint::from_int(10));
assert_eq!(all.len(), 5);
let tiny = vp.find_in_radius(&origin, FixedPoint::from_int(1) / FixedPoint::from_int(10000));
assert!(tiny.len() <= 1);
}
#[test]
fn test_vp_tree_remove() {
let disk = PoincareDisk::new(2);
let mut vp = VPTree::new();
vp.insert(BucketEntry::new("x".to_string(), disk.point_from_f32_slice(&[0.1, 0.0]), 0));
vp.insert(BucketEntry::new("y".to_string(), disk.point_from_f32_slice(&[0.2, 0.0]), 0));
assert_eq!(vp.live_count(), 2);
vp.remove("x");
assert_eq!(vp.live_count(), 1);
let origin = disk.origin();
let results = vp.find_in_radius(&origin, FixedPoint::from_int(10));
assert_eq!(results.len(), 1);
assert_eq!(results[0].0, "y");
}
#[test]
fn test_vp_tree_rebuild_on_buffer_threshold() {
let disk = PoincareDisk::new(2);
let mut vp = VPTree::new();
for i in 0..(VP_BUFFER_THRESHOLD + 5) {
let angle = constants::two_pi()
* FixedPoint::from_int(i as i32)
/ FixedPoint::from_int((VP_BUFFER_THRESHOLD + 5) as i32);
let r = FixedPoint::from_int(3) / FixedPoint::from_int(10);
let mut coords = FixedVector::new(2);
let (sin_a, cos_a) = angle.sincos();
coords[0] = r * cos_a;
coords[1] = r * sin_a;
vp.insert(BucketEntry::new(format!("node_{}", i), HyperbolicPoint::new(coords), 0));
}
assert!(vp.root.is_some());
assert_eq!(vp.live_count(), VP_BUFFER_THRESHOLD + 5);
let origin = disk.origin();
let all = vp.find_in_radius(&origin, FixedPoint::from_int(10));
assert_eq!(all.len(), VP_BUFFER_THRESHOLD + 5);
}
#[test]
fn test_vp_tree_knn_ordering() {
let disk = PoincareDisk::new(2);
let mut vp = VPTree::new();
let distances = [0.05f32, 0.1, 0.2, 0.3, 0.5, 0.7];
for (i, &d) in distances.iter().enumerate() {
vp.insert(BucketEntry::new(format!("p{}", i), disk.point_from_f32_slice(&[d, 0.0]), 0));
}
let origin = disk.origin();
let nearest = vp.find_nearest(&origin, 3);
assert_eq!(nearest.len(), 3);
for i in 1..nearest.len() {
assert!(nearest[i].1 >= nearest[i - 1].1,
"Results not sorted: {:?} >= {:?}", nearest[i].1, nearest[i - 1].1);
}
let ids: Vec<&str> = nearest.iter().map(|(id, _)| id.as_str()).collect();
assert!(ids.contains(&"p0"));
assert!(ids.contains(&"p1"));
assert!(ids.contains(&"p2"));
}
#[test]
fn test_register_unregister_with_vp_tree() {
let table = HyperbolicHashTable::new(2);
let disk_clone = table.poincare_disk().clone();
let p1 = disk_clone.point_from_f32_slice(&[0.1, 0.0]);
let p2 = disk_clone.point_from_f32_slice(&[0.2, 0.0]);
let p3 = disk_clone.point_from_f32_slice(&[0.3, 0.0]);
table.register_node(&p1, "node1", 0);
table.register_node(&p2, "node2", 1);
table.register_node(&p3, "node3", 1);
let origin = disk_clone.origin();
let results = table.find_nodes_in_radius(&origin, FixedPoint::from_int(10));
assert!(results.len() >= 3, "Expected at least 3, got {}", results.len());
table.unregister_node("node2");
let results = table.find_nodes_in_radius(&origin, FixedPoint::from_int(10));
let ids: Vec<&str> = results.iter().map(|(id, _)| id.as_str()).collect();
assert!(!ids.contains(&"node2"), "node2 should be unregistered");
assert!(ids.contains(&"node1"));
assert!(ids.contains(&"node3"));
}
#[test]
fn test_find_nearest_with_early_termination() {
let table = HyperbolicHashTable::new(2);
let disk_clone = table.poincare_disk().clone();
let positions: Vec<(&str, [f32; 2])> = vec![
("close1", [0.05, 0.0]),
("close2", [0.0, 0.05]),
("mid1", [0.3, 0.0]),
("mid2", [0.0, 0.3]),
("far1", [0.7, 0.0]),
("far2", [0.0, 0.7]),
];
for (id, coords) in &positions {
let point = disk_clone.point_from_f32_slice(coords);
table.register_node(&point, id, 0);
}
let origin = disk_clone.origin();
let nearest = table.find_nearest_nodes(&origin, 2);
assert_eq!(nearest.len(), 2);
let ids: Vec<&str> = nearest.iter().map(|(id, _)| id.as_str()).collect();
assert!(ids.contains(&"close1"));
assert!(ids.contains(&"close2"));
assert!(nearest[0].1 <= nearest[1].1);
}
#[test]
fn test_duplicate_registration_prevented() {
let table = HyperbolicHashTable::new(2);
let point = table.poincare_disk().point_from_f32_slice(&[0.1, 0.0]);
let h1 = table.register_node(&point, "dup_node", 0);
let h2 = table.register_node(&point, "dup_node", 0);
assert_eq!(h1, h2);
let origin = table.poincare_disk().origin();
let results = table.find_nodes_in_radius(&origin, FixedPoint::from_int(10));
let count = results.iter().filter(|(id, _)| id == "dup_node").count();
assert_eq!(count, 1, "Duplicate registration should be prevented");
}
#[test]
fn effective_radius_returns_to_nominal_when_lone_outlier_removed() {
let table = HyperbolicHashTable::new(2);
let disk = table.poincare_disk().clone();
let deep = disk.point_from_f32_slice(&[0.95, 0.0]);
let bucket_hash = table.register_node(&deep, "deep", 5).unwrap();
let nominal = table.buckets.get(&bucket_hash).unwrap().region().radius();
let inflated = table.buckets.get(&bucket_hash).unwrap().effective_radius();
assert!(
inflated > nominal,
"deep node should widen the bucket past nominal (inflated={:?}, nominal={:?})",
inflated, nominal
);
table.unregister_node("deep");
let after = table.buckets.get(&bucket_hash).unwrap().effective_radius();
assert_eq!(
after, nominal,
"with the only out-of-region member gone, the bound must return to nominal"
);
}
#[test]
fn effective_radius_falls_to_second_farthest_not_nominal() {
let table = HyperbolicHashTable::new(2);
let disk = table.poincare_disk().clone();
let near_deep = disk.point_from_f32_slice(&[0.85, 0.0]);
let far_deep = disk.point_from_f32_slice(&[0.97, 0.0]);
let h_near = table.register_node(&near_deep, "near_deep", 4).unwrap();
let h_far = table.register_node(&far_deep, "far_deep", 6).unwrap();
if h_near != h_far {
return;
}
let bucket = || table.buckets.get(&h_near).unwrap();
let nominal = bucket().region().radius();
let with_both = bucket().effective_radius();
let center = bucket().region().center().clone();
let near_dist = center.hyperbolic_distance(&near_deep);
table.unregister_node("far_deep");
let after = bucket().effective_radius();
assert!(after < with_both, "removing the farther node must shrink the bound");
assert!(after > nominal, "the remaining out-of-region node must keep the bound above nominal");
assert_eq!(after, near_dist, "the bound must equal the remaining node's center distance");
}
}