From 228c6ff976e7710dab46cea090479d2afe1a7d55 Mon Sep 17 00:00:00 2001 From: Joe Date: Tue, 20 Jan 2026 18:27:03 -0800 Subject: [PATCH] start moving rstar tree code into lib --- Cargo.lock | 13 +- Cargo.toml | 3 +- src/lib.rs | 5 + src/main.rs | 5 +- src/spindex.rs | 932 +++++++++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 940 insertions(+), 18 deletions(-) create mode 100644 src/lib.rs create mode 100644 src/spindex.rs diff --git a/Cargo.lock b/Cargo.lock index a6a3578..98bd6b2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -330,7 +330,7 @@ dependencies = [ ] [[package]] -name = "autobats" +name = "autobarts" version = "0.1.0" dependencies = [ "bevy", @@ -339,7 +339,6 @@ dependencies = [ "ordered-float", "rusqlite", "rusqlite_migration", - "spart", "steel-core", ] @@ -4795,16 +4794,6 @@ dependencies = [ "serde", ] -[[package]] -name = "spart" -version = "0.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a05d85e94c16e35b2f963dee2eaa6f6b2494ee5dad740b045b139c5614872d6" -dependencies = [ - "ordered-float", - "tracing", -] - [[package]] name = "spin" version = "0.10.0" diff --git a/Cargo.toml b/Cargo.toml index 34350fc..c6f0076 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,5 +1,5 @@ [package] -name = "autobats" +name = "autobarts" version = "0.1.0" edition = "2024" @@ -10,7 +10,6 @@ include_dir = "0.7.4" ordered-float = "5.1.0" rusqlite = { version = "0.37", default-features = false, features = ["bundled", "blob", "functions", "jiff"] } rusqlite_migration = { version = "2.3.0", features = ["from-directory"] } -spart = "0.5.0" steel-core = { git="https://github.com/mattwparas/steel.git", branch = "master" } # Enable a small amount of optimization in the dev profile. diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..e5e5eaf --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,5 @@ +pub mod db; +pub mod geom; +mod spindex; + +pub use spindex::*; diff --git a/src/main.rs b/src/main.rs index 672a0b5..dace6db 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,8 +1,5 @@ +use autobarts::db::init_db; use bevy::prelude::*; -use db::init_db; - -mod db; -mod geom; fn main() { App::new() diff --git a/src/spindex.rs b/src/spindex.rs new file mode 100644 index 0000000..b6e9db9 --- /dev/null +++ b/src/spindex.rs @@ -0,0 +1,932 @@ +use crate::geom::Point; +use bevy::math::bounding::Aabb2d; +use ordered_float::OrderedFloat; +use std::cmp::Ordering; +use std::collections::BinaryHeap; + +// Epsilon value for zero-sizes bounding boxes/cubes. +const EPSILON: f64 = 1e-10; + +/// An entry in the R*‑tree, which can be either a leaf or a node. +#[derive(Debug, Clone)] + +pub enum Entry { + Leaf { mbr: Aabb2d, object: Point }, + Node { mbr: Aabb2d, child: TreeNode }, +} + +impl Entry { + /// Returns a reference to the minimum bounding volume for this entry. + pub fn mbr(&self) -> &Aabb2d { + match self { + Entry::Leaf { mbr, .. } | Entry::Node { mbr, .. } => mbr, + } + } +} + +/// A node in the R*‑tree. +#[derive(Debug, Clone)] +pub struct TreeNode { + /// The entries stored in this node. + pub entries: Vec, + /// Indicates whether this node is a leaf. + pub is_leaf: bool, +} + +/// R*‑tree data structure for indexing 2D or 3D points. +/// +/// The tree is initialized with a maximum number of entries per node. If a node exceeds this +/// number, it will split. The tree supports insertion, deletion, and range searches. +#[derive(Debug, Clone)] +pub struct RStarTree { + root: TreeNode, + max_entries: usize, + min_entries: usize, +} + +// Common trait implementations for R*-tree to reuse shared algorithms. +impl Entry { + fn as_leaf_obj(&self) -> Option<&Point> { + match self { + Entry::Leaf { object, .. } => Some(object), + _ => None, + } + } + fn child(&self) -> Option<&Box> { + match self { + Entry::Node { child, .. } => Some(child), + _ => None, + } + } + fn child_mut(&mut self) -> Option<&mut ::Node> { + match self { + Entry::Node { child, .. } => Some(child), + _ => None, + } + } + fn set_mbr(&mut self, new_mbr: Self::BV) { + if let Entry::Node { mbr, .. } = self { + *mbr = new_mbr; + } + } + fn into_child(self) -> Option::Node>> + where + Self: Sized, + { + match self { + Entry::Node { child, .. } => Some(child), + _ => None, + } + } +} + +impl crate::rtree_common::NodeAccess for TreeNode { + type Entry = Entry; + fn is_leaf(&self) -> bool { + self.is_leaf + } + fn entries(&self) -> &Vec { + &self.entries + } + fn entries_mut(&mut self) -> &mut Vec { + &mut self.entries + } +} + +impl RStarTree { + /// Creates a new R*‑tree with the specified maximum number of entries per node. + /// + /// # Arguments + /// + /// * `max_entries` - The maximum number of entries allowed in a node. + /// + /// # Errors + /// + /// Returns `SpartError::InvalidCapacity` if `max_entries` is less than 2. + pub fn new(max_entries: usize) -> Result { + if max_entries < 2 { + return Err(SpartError::InvalidCapacity { + capacity: max_entries, + }); + } + info!("Creating new RStarTree with max_entries: {}", max_entries); + Ok(RStarTree { + root: TreeNode { + entries: Vec::new(), + is_leaf: true, + }, + max_entries, + min_entries: (max_entries as f64 * 0.4).ceil() as usize, + }) + } + + /// Inserts an object into the R*‑tree. + /// + /// # Arguments + /// + /// * `object` - The object to insert. + pub fn insert(&mut self, object: T) + where + T: Clone, + T::B: BSPBounds, + { + info!("Inserting object into RStarTree: {:?}", object); + let entry = Entry::Leaf { + mbr: object.mbr(), + object, + }; + self.insert_entry(entry, None); + } + + fn insert_entry(&mut self, entry: Entry, reinsert_from_level: Option) + where + T: Clone, + T::B: BSPBounds, + { + let mut to_insert = vec![(entry, 0)]; + let mut reinsert_level = reinsert_from_level; + + while let Some((item, level)) = to_insert.pop() { + let overflow = insert_recursive( + &mut self.root, + item, + self.max_entries, + level, + &mut reinsert_level, + &mut to_insert, + ); + + if let Some((overflowed_node, overflow_level)) = overflow { + if reinsert_level == Some(overflow_level) { + let old_entries = overflowed_node; + let (group1, group2) = split_entries(old_entries, self.max_entries); + let child1 = TreeNode { + entries: group1, + is_leaf: self.root.is_leaf, + }; + let child2 = TreeNode { + entries: group2, + is_leaf: self.root.is_leaf, + }; + let mbr1 = common_compute_group_mbr(&child1.entries) + .unwrap_or_else(|| unreachable!("non-empty group must have MBR")); + let mbr2 = common_compute_group_mbr(&child2.entries) + .unwrap_or_else(|| unreachable!("non-empty group must have MBR")); + self.root.is_leaf = false; + self.root.entries.clear(); + self.root.entries.push(Entry::Node { + mbr: mbr1, + child: Box::new(child1), + }); + self.root.entries.push(Entry::Node { + mbr: mbr2, + child: Box::new(child2), + }); + } else { + if reinsert_level.is_none() { + reinsert_level = Some(overflow_level); + } + let mut node = TreeNode { + entries: overflowed_node, + is_leaf: self.root.is_leaf, + }; + let reinserted_entries = forced_reinsert(&mut node, self.max_entries); + self.root.entries = node.entries; + for entry in reinserted_entries { + to_insert.push((entry, 0)); + } + } + } + } + } + + /// Performs a range search with a given query bounding volume. + /// + /// # Arguments + /// + /// * `query` - The bounding volume to search against. + /// + /// # Returns + /// + /// A vector of references to the objects whose minimum bounding volumes intersect the query. + pub fn range_search_bbox(&self, query: &T::B) -> Vec<&T> { + info!("Performing range search with query: {:?}", query); + let mut result = Vec::new(); + common_search_node(&self.root, query, &mut result); + result + } + + /// Inserts a bulk of objects into the R*-tree. + /// + /// # Arguments + /// + /// * `objects` - The objects to insert. + pub fn insert_bulk(&mut self, objects: Vec) + where + T: Clone, + T::B: BSPBounds, + { + if objects.is_empty() { + return; + } + + let mut entries: Vec> = objects + .into_iter() + .map(|obj| Entry::Leaf { + mbr: obj.mbr(), + object: obj, + }) + .collect(); + + while entries.len() > self.max_entries { + let mut new_level_entries = Vec::new(); + let chunks = entries.chunks(self.max_entries); + + for chunk in chunks { + let child_node = TreeNode { + entries: chunk.to_vec(), + is_leaf: self.root.is_leaf, + }; + if let Some(mbr) = common_compute_group_mbr(&child_node.entries) { + new_level_entries.push(Entry::Node { + mbr, + child: Box::new(child_node), + }); + } + } + entries = new_level_entries; + self.root.is_leaf = false; + } + + self.root.entries.extend(entries); + } + + #[doc(hidden)] + pub fn height(&self) -> usize { + let mut height = 1; + let mut current_node = &self.root; + while !current_node.is_leaf { + height += 1; + current_node = if let Some(Entry::Node { child, .. }) = current_node.entries.first() { + child + } else { + break; + }; + } + height + } +} + +fn choose_subtree(node: &TreeNode, entry: &Entry) -> usize { + let children_are_leaves = if let Some(Entry::Node { child, .. }) = node.entries.first() { + child.is_leaf + } else { + false + }; + + if children_are_leaves { + node.entries + .iter() + .enumerate() + .min_by(|&(_, a), &(_, b)| { + let mbr_a = a.mbr(); + let mbr_b = b.mbr(); + + let overlap_a = node + .entries + .iter() + .filter(|e| !std::ptr::eq(*e, a)) + .map(|e| e.mbr().union(entry.mbr()).overlap(e.mbr())) + .sum::(); + + let overlap_b = node + .entries + .iter() + .filter(|e| !std::ptr::eq(*e, b)) + .map(|e| e.mbr().union(entry.mbr()).overlap(e.mbr())) + .sum::(); + + let overlap_cmp = overlap_a.partial_cmp(&overlap_b).unwrap_or(Ordering::Equal); + if overlap_cmp != Ordering::Equal { + return overlap_cmp; + } + + let enlargement_a = mbr_a.enlargement(entry.mbr()); + let enlargement_b = mbr_b.enlargement(entry.mbr()); + let enlargement_cmp = enlargement_a + .partial_cmp(&enlargement_b) + .unwrap_or(Ordering::Equal); + if enlargement_cmp != Ordering::Equal { + return enlargement_cmp; + } + + mbr_a + .area() + .partial_cmp(&mbr_b.area()) + .unwrap_or(Ordering::Equal) + }) + .map(|(i, _)| i) + .unwrap_or(0) + } else { + node.entries + .iter() + .enumerate() + .min_by(|(_, a), (_, b)| { + let mbr_a = a.mbr(); + let mbr_b = b.mbr(); + + let enlargement_a = mbr_a.enlargement(entry.mbr()); + let enlargement_b = mbr_b.enlargement(entry.mbr()); + + let enlargement_cmp = enlargement_a + .partial_cmp(&enlargement_b) + .unwrap_or(Ordering::Equal); + if enlargement_cmp != Ordering::Equal { + return enlargement_cmp; + } + mbr_a + .area() + .partial_cmp(&mbr_b.area()) + .unwrap_or(Ordering::Equal) + }) + .map(|(i, _)| i) + .unwrap_or(0) + } +} + +fn insert_recursive( + node: &mut TreeNode, + entry: Entry, + max_entries: usize, + level: usize, + reinsert_level: &mut Option, + to_insert_queue: &mut Vec<(Entry, usize)>, +) -> Option<(Vec>, usize)> +where + T::B: BSPBounds, +{ + if node.is_leaf { + node.entries.push(entry); + } else { + let best_index = choose_subtree(node, &entry); + let child = if let Entry::Node { child, .. } = &mut node.entries[best_index] { + child + } else { + unreachable!() + }; + + if let Some((overflow, overflow_level)) = insert_recursive( + child, + entry, + max_entries, + level + 1, + reinsert_level, + to_insert_queue, + ) { + if reinsert_level.is_some() && *reinsert_level == Some(overflow_level) { + let (g1, g2) = split_entries(overflow, max_entries); + let child1 = TreeNode { + entries: g1, + is_leaf: child.is_leaf, + }; + let child2 = TreeNode { + entries: g2, + is_leaf: child.is_leaf, + }; + let mbr1 = common_compute_group_mbr(&child1.entries) + .unwrap_or_else(|| unreachable!("non-empty group must have MBR")); + let mbr2 = common_compute_group_mbr(&child2.entries) + .unwrap_or_else(|| unreachable!("non-empty group must have MBR")); + node.entries[best_index] = Entry::Node { + mbr: mbr1, + child: Box::new(child1), + }; + node.entries.push(Entry::Node { + mbr: mbr2, + child: Box::new(child2), + }); + } else { + if reinsert_level.is_none() { + *reinsert_level = Some(overflow_level); + } + let mut overflowed_node = TreeNode { + entries: overflow, + is_leaf: child.is_leaf, + }; + let reinserted = forced_reinsert(&mut overflowed_node, max_entries); + for item in reinserted { + to_insert_queue.push((item, 0)); + } + if let Entry::Node { child, .. } = &mut node.entries[best_index] { + child.entries = overflowed_node.entries; + } + } + } + if let Some(new_mbr) = common_compute_group_mbr( + if let Entry::Node { child, .. } = &node.entries[best_index] { + &child.entries + } else { + unreachable!() + }, + ) { + if let Entry::Node { mbr, .. } = &mut node.entries[best_index] { + *mbr = new_mbr; + } + } + } + + if node.entries.len() > max_entries { + return Some((std::mem::take(&mut node.entries), level)); + } + None +} + +fn forced_reinsert( + node: &mut TreeNode, + max_entries: usize, +) -> Vec> +where + T::B: BSPBounds, +{ + let node_mbr = if let Some(mbr) = common_compute_group_mbr(&node.entries) { + mbr + } else { + return Vec::new(); + }; + let reinsert_count = (max_entries as f64 * 0.3).ceil() as usize; + + node.entries.sort_by(|a, b| { + let center_a: Vec = (0..T::B::DIM) + .map(|d| { + a.mbr() + .center(d) + .unwrap_or_else(|_| unreachable!("dim valid")) + }) + .collect(); + let center_b: Vec = (0..T::B::DIM) + .map(|d| { + b.mbr() + .center(d) + .unwrap_or_else(|_| unreachable!("dim valid")) + }) + .collect(); + let node_center: Vec = (0..T::B::DIM) + .map(|d| { + node_mbr + .center(d) + .unwrap_or_else(|_| unreachable!("dim valid")) + }) + .collect(); + + let dist_a = center_a + .iter() + .zip(node_center.iter()) + .map(|(ca, cb)| (ca - cb).powi(2)) + .sum::(); + let dist_b = center_b + .iter() + .zip(node_center.iter()) + .map(|(ca, cb)| (ca - cb).powi(2)) + .sum::(); + + dist_b.partial_cmp(&dist_a).unwrap_or(Ordering::Equal) + }); + + node.entries.drain(0..reinsert_count).collect() +} + +fn split_entries( + mut entries: Vec>, + max_entries: usize, +) -> (Vec>, Vec>) +where + T::B: BSPBounds, +{ + let min_entries = (max_entries as f64 * 0.4).ceil() as usize; + let mut best_axis = 0; + let mut best_split_index = 0; + let mut min_margin = f64::INFINITY; + + for dim in 0..T::B::DIM { + entries.sort_by(|a, b| { + let ca = a + .mbr() + .center(dim) + .unwrap_or_else(|_| unreachable!("dim valid")); + let cb = b + .mbr() + .center(dim) + .unwrap_or_else(|_| unreachable!("dim valid")); + ca.partial_cmp(&cb).unwrap_or(Ordering::Equal) + }); + + for k in min_entries..=entries.len() - min_entries { + let group1 = &entries[..k]; + let group2 = &entries[k..]; + let mbr1 = common_compute_group_mbr(group1) + .unwrap_or_else(|| unreachable!("non-empty group must have MBR")); + let mbr2 = common_compute_group_mbr(group2) + .unwrap_or_else(|| unreachable!("non-empty group must have MBR")); + let margin = mbr1.margin() + mbr2.margin(); + if margin < min_margin { + min_margin = margin; + best_axis = dim; + best_split_index = k; + } + } + } + + entries.sort_by(|a, b| { + let ca = a + .mbr() + .center(best_axis) + .unwrap_or_else(|_| unreachable!("dim valid")); + let cb = b + .mbr() + .center(best_axis) + .unwrap_or_else(|_| unreachable!("dim valid")); + ca.partial_cmp(&cb).unwrap_or(Ordering::Equal) + }); + + let mut best_overlap = f64::INFINITY; + let mut best_area = f64::INFINITY; + + for k in min_entries..=entries.len() - min_entries { + let group1 = &entries[..k]; + let group2 = &entries[k..]; + let mbr1 = common_compute_group_mbr(group1) + .unwrap_or_else(|| unreachable!("non-empty group must have MBR")); + let mbr2 = common_compute_group_mbr(group2) + .unwrap_or_else(|| unreachable!("non-empty group must have MBR")); + let overlap = mbr1.overlap(&mbr2); + let area = mbr1.area() + mbr2.area(); + + if overlap < best_overlap { + best_overlap = overlap; + best_area = area; + best_split_index = k; + } else if (overlap - best_overlap).abs() < EPSILON && area < best_area { + best_area = area; + best_split_index = k; + } + } + + let (group1, group2) = entries.split_at(best_split_index); + (group1.to_vec(), group2.to_vec()) +} + +impl RStarTree +where + T: PartialEq + Clone, + T::B: BSPBounds, +{ + /// Deletes an object from the R*‑tree. + /// + /// # Arguments + /// + /// * `object` - The object to delete. + /// + /// # Returns + /// + /// `true` if at least one matching object was found and removed. + pub fn delete(&mut self, object: &T) -> bool { + info!("Attempting to delete object: {:?}", object); + let object_mbr = object.mbr(); + let mut reinsert_list = Vec::new(); + let deleted = common_delete_entry( + &mut self.root, + object, + &object_mbr, + self.min_entries, + &mut reinsert_list, + ); + + if deleted { + for entry in reinsert_list { + self.insert_entry(entry, None); + } + + if !self.root.is_leaf && self.root.entries.len() == 1 { + if let Some(Entry::Node { child, .. }) = self.root.entries.pop() { + self.root = *child; + } + } + } + deleted + } +} + +impl RStarTreeObject for Point2D { + type B = Rectangle; + fn mbr(&self) -> Self::B { + Rectangle { + x: self.x, + y: self.y, + width: EPSILON, + height: EPSILON, + } + } +} + +impl RStarTreeObject for Point3D { + type B = Cube; + fn mbr(&self) -> Self::B { + Cube { + x: self.x, + y: self.y, + z: self.z, + width: EPSILON, + height: EPSILON, + depth: EPSILON, + } + } +} + +impl RStarTree> { + /// Performs a k‑nearest neighbor search on an R*‑tree of 2D points. + /// + /// # Arguments + /// + /// * `query` - The 2D point to search near. + /// * `k` - The number of nearest neighbors to return. + /// + /// # Returns + /// + /// A vector of references to the k nearest 2D points. + /// + /// # Note + /// + /// The pruning logic for the search is based on Euclidean distance. Custom distance metrics + /// that are not compatible with Euclidean distance may lead to incorrect results or reduced + /// performance. + pub fn knn_search>>( + &self, + query: &Point2D, + k: usize, + ) -> Vec<&Point2D> { + if k == 0 { + return Vec::new(); + } + + let mut heap: BinaryHeap>>> = BinaryHeap::new(); + for entry in &self.root.entries { + let dist_sq = entry.mbr().min_distance(query).powi(2); + heap.push(KnnCandidate { + dist: dist_sq, + entry, + }); + } + + type OrdDist = OrderedFloat; + #[inline] + #[allow(non_snake_case)] + fn OrdDist(x: f64) -> OrderedFloat { + OrderedFloat(x) + } + + struct HeapItem<'a, P> { + key: OrdDist, + idx: usize, + obj: &'a P, + } + impl

PartialEq for HeapItem<'_, P> { + fn eq(&self, other: &Self) -> bool { + self.key == other.key && self.idx == other.idx + } + } + impl

Eq for HeapItem<'_, P> {} + impl

Ord for HeapItem<'_, P> { + fn cmp(&self, other: &Self) -> Ordering { + match self.key.cmp(&other.key) { + Ordering::Equal => self.idx.cmp(&other.idx), + ord => ord, + } + } + } + impl

PartialOrd for HeapItem<'_, P> { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } + } + + let mut results: BinaryHeap>> = BinaryHeap::new(); + let mut counter: usize = 0; + + while let Some(KnnCandidate { dist, entry }) = heap.pop() { + if results.len() >= k { + if let Some(worst_result) = results.peek() { + if dist > worst_result.key.0 { + break; + } + } + } + + match entry { + Entry::Leaf { object, .. } => { + let d_sq = M::distance_sq(query, object); + if results.len() < k { + counter += 1; + results.push(HeapItem { + key: OrdDist(d_sq), + idx: counter, + obj: object, + }); + } else if let Some(peek) = results.peek() { + if d_sq < peek.key.0 { + results.pop(); + counter += 1; + results.push(HeapItem { + key: OrdDist(d_sq), + idx: counter, + obj: object, + }); + } + } + } + Entry::Node { child, .. } => { + for child_entry in &child.entries { + let d_sq = child_entry.mbr().min_distance(query).powi(2); + if results.len() < k { + heap.push(KnnCandidate { + dist: d_sq, + entry: child_entry, + }); + } else if let Some(peek) = results.peek() { + if d_sq < peek.key.0 { + heap.push(KnnCandidate { + dist: d_sq, + entry: child_entry, + }); + } + } + } + } + } + } + + let mut sorted_results = results.into_vec(); + sorted_results.sort_by(|a, b| a.key.partial_cmp(&b.key).unwrap_or(Ordering::Equal)); + sorted_results.into_iter().map(|r| r.obj).collect() + } +} + +impl RStarTree> { + /// Performs a k‑nearest neighbor search on an R*‑tree of 3D points. + /// + /// # Arguments + /// + /// * `query` - The 3D point to search near. + /// * `k` - The number of nearest neighbors to return. + /// + /// # Returns + /// + /// A vector of references to the k nearest 3D points. + /// + /// # Note + /// + /// The pruning logic for the search is based on Euclidean distance. Custom distance metrics + /// that are not compatible with Euclidean distance may lead to incorrect results or reduced + /// performance. + pub fn knn_search>>( + &self, + query: &Point3D, + k: usize, + ) -> Vec<&Point3D> { + if k == 0 { + return Vec::new(); + } + + let mut heap: BinaryHeap>>> = BinaryHeap::new(); + for entry in &self.root.entries { + let dist_sq = entry.mbr().min_distance(query).powi(2); + heap.push(KnnCandidate { + dist: dist_sq, + entry, + }); + } + + type OrdDist = OrderedFloat; + #[inline] + #[allow(non_snake_case)] + fn OrdDist(x: f64) -> OrderedFloat { + OrderedFloat(x) + } + + struct HeapItem<'a, P> { + key: OrdDist, + idx: usize, + obj: &'a P, + } + impl

PartialEq for HeapItem<'_, P> { + fn eq(&self, other: &Self) -> bool { + self.key == other.key && self.idx == other.idx + } + } + impl

Eq for HeapItem<'_, P> {} + impl

Ord for HeapItem<'_, P> { + fn cmp(&self, other: &Self) -> Ordering { + match self.key.cmp(&other.key) { + Ordering::Equal => self.idx.cmp(&other.idx), + ord => ord, + } + } + } + impl

PartialOrd for HeapItem<'_, P> { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } + } + + let mut results: BinaryHeap>> = BinaryHeap::new(); + let mut counter: usize = 0; + + while let Some(KnnCandidate { dist, entry }) = heap.pop() { + if results.len() >= k { + if let Some(worst_result) = results.peek() { + if dist > worst_result.key.0 { + break; + } + } + } + + match entry { + Entry::Leaf { object, .. } => { + let d_sq = M::distance_sq(query, object); + if results.len() < k { + counter += 1; + results.push(HeapItem { + key: OrdDist(d_sq), + idx: counter, + obj: object, + }); + } else if let Some(peek) = results.peek() { + if d_sq < peek.key.0 { + results.pop(); + counter += 1; + results.push(HeapItem { + key: OrdDist(d_sq), + idx: counter, + obj: object, + }); + } + } + } + Entry::Node { child, .. } => { + for child_entry in &child.entries { + let d_sq = child_entry.mbr().min_distance(query).powi(2); + if results.len() < k { + heap.push(KnnCandidate { + dist: d_sq, + entry: child_entry, + }); + } else if let Some(peek) = results.peek() { + if d_sq < peek.key.0 { + heap.push(KnnCandidate { + dist: d_sq, + entry: child_entry, + }); + } + } + } + } + } + } + + let mut sorted_results = results.into_vec(); + sorted_results.sort_by(|a, b| a.key.partial_cmp(&b.key).unwrap_or(Ordering::Equal)); + sorted_results.into_iter().map(|r| r.obj).collect() + } +} + +impl RStarTree +where + T: RStarTreeObject + PartialEq + std::fmt::Debug, + T::B: BoundingVolumeFromPoint + HasMinDistance + Clone, +{ + /// Performs a range search on the R*‑tree using a query object and radius. + /// + /// The query object is wrapped into a bounding volume using `from_point_radius`. + /// + /// # Arguments + /// + /// * `query` - The query object. + /// * `radius` - The search radius. + /// + /// # Returns + /// + /// A vector of references to the objects within the given radius. + /// + /// # Note + /// + /// The pruning logic for the search is based on Euclidean distance. Custom distance metrics + /// that are not compatible with Euclidean distance may lead to incorrect results or reduced + /// performance. + pub fn range_search>(&self, query: &T, radius: f64) -> Vec<&T> { + let query_volume = T::B::from_point_radius(query, radius); + let candidates = self.range_search_bbox(&query_volume); + candidates + .into_iter() + .filter(|object| M::distance_sq(query, object) <= radius * radius) + .collect() + } +}