use std::{collections::BinaryHeap, ops::ControlFlow};
use crate::config::{DEFAULT_NEIGHBOR_QUEUE_CAPACITY, DEFAULT_SEARCH_STACK_CAPACITY};
use crate::geometry::{Box2D, Point2D};
use crate::neighbors::{NeighborNodeState, NeighborState, NeighborWorkspace, max_distance_squared};
use crate::persistence::{
ByteWriter, LoadError, parse_index_bytes, read_f64_le_unchecked, read_u64_le_unchecked,
serialized_len,
};
use crate::traversal::{SearchWorkspace, prefetch_read, upper_bound_level};
#[inline]
fn prefetch_aos_node(entries: &[Box2D], indices: &[usize], node_index: usize, node_size: usize) {
if node_index < entries.len() {
prefetch_read(entries.as_ptr().wrapping_add(node_index));
prefetch_read(indices.as_ptr().wrapping_add(node_index));
}
let next_line = node_index.saturating_add((64 / std::mem::size_of::<Box2D>()).max(1));
if node_size > 1 && next_line < entries.len() {
prefetch_read(entries.as_ptr().wrapping_add(next_line));
prefetch_read(indices.as_ptr().wrapping_add(next_line));
}
}
pub struct Index2D {
pub(crate) node_size: usize,
pub(crate) num_items: usize,
pub(crate) level_bounds: Vec<usize>,
pub(crate) entries: Vec<Box2D>,
pub(crate) indices: Vec<usize>,
}
impl Index2D {
pub fn num_items(&self) -> usize {
self.num_items
}
pub fn extent(&self) -> Option<Box2D> {
self.entries.last().copied()
}
pub fn node_size(&self) -> usize {
self.node_size
}
pub fn to_bytes(&self) -> Vec<u8> {
let mut out = Vec::new();
self.to_bytes_into(&mut out);
out
}
pub fn to_bytes_into(&self, out: &mut Vec<u8>) {
let level_count = self.level_bounds.len();
let num_nodes = self.entries.len();
let len = serialized_len(level_count, num_nodes).expect("serialized index is too large");
let mut bytes = ByteWriter::new(out, len);
bytes.write_magic();
bytes.write_format_version();
bytes.write_header_len();
bytes.write_flags();
bytes.write_u64(self.node_size as u64);
bytes.write_u64(self.num_items as u64);
bytes.write_u64(num_nodes as u64);
bytes.write_u64(level_count as u64);
bytes.write_usize_slice_as_u64(&self.level_bounds);
bytes.write_box2d_slice(&self.entries);
bytes.write_usize_slice_as_u64(&self.indices);
bytes.finish();
}
pub fn from_bytes(bytes: &[u8]) -> Result<Self, LoadError> {
let view = Index2DView::from_bytes(bytes)?;
let mut level_bounds = Vec::with_capacity(view.level_count);
for i in 0..view.level_count {
level_bounds.push(view.level_bound_unchecked(i));
}
let mut entries = Vec::with_capacity(view.num_nodes);
for i in 0..view.num_nodes {
entries.push(view.entry_at_unchecked(i));
}
let mut indices = Vec::with_capacity(view.num_nodes);
for i in 0..view.num_nodes {
indices.push(view.index_at_unchecked(i));
}
Ok(Self {
node_size: view.node_size,
num_items: view.num_items,
level_bounds,
entries,
indices,
})
}
pub fn search(&self, query: Box2D) -> Vec<usize> {
let mut results = Vec::new();
self.search_into(query, &mut results);
results
}
pub fn search_into(&self, query: Box2D, results: &mut Vec<usize>) {
let mut stack: Vec<usize> = Vec::with_capacity(DEFAULT_SEARCH_STACK_CAPACITY);
self.search_into_stack(query, results, &mut stack);
}
pub fn search_with<'a>(&self, query: Box2D, workspace: &'a mut SearchWorkspace) -> &'a [usize] {
self.search_into_stack(query, &mut workspace.results, &mut workspace.stack);
&workspace.results
}
pub fn any(&self, query: Box2D) -> bool {
self.visit(query, |_| ControlFlow::Break(())).is_break()
}
pub fn first(&self, query: Box2D) -> Option<usize> {
match self.visit(query, ControlFlow::Break) {
ControlFlow::Break(index) => Some(index),
ControlFlow::Continue(()) => None,
}
}
pub fn neighbors(&self, point: Point2D, max_results: usize) -> Vec<usize> {
self.neighbors_within(point, max_results, f64::INFINITY)
}
pub fn neighbors_within(
&self,
point: Point2D,
max_results: usize,
max_distance: f64,
) -> Vec<usize> {
let mut results = Vec::new();
self.neighbors_into(point, max_results, max_distance, &mut results);
results
}
pub fn neighbors_into(
&self,
point: Point2D,
max_results: usize,
max_distance: f64,
results: &mut Vec<usize>,
) {
results.clear();
if max_results == 0 {
return;
}
if max_results == 1 {
let mut queue = BinaryHeap::with_capacity(DEFAULT_NEIGHBOR_QUEUE_CAPACITY);
if let Some(index) = self.nearest_one_with_queue(point, max_distance, &mut queue) {
results.push(index);
}
return;
}
let mut item_queue = BinaryHeap::with_capacity(DEFAULT_NEIGHBOR_QUEUE_CAPACITY);
let mut node_queue = BinaryHeap::with_capacity(DEFAULT_NEIGHBOR_QUEUE_CAPACITY);
self.collect_neighbors_with_queues(
point,
max_results,
max_distance,
results,
&mut item_queue,
&mut node_queue,
);
}
pub fn neighbors_with<'a>(
&self,
point: Point2D,
max_results: usize,
max_distance: f64,
workspace: &'a mut NeighborWorkspace,
) -> &'a [usize] {
workspace.results.clear();
if max_results == 0 {
workspace.queue.clear();
workspace.node_queue.clear();
return &workspace.results;
}
if max_results == 1 {
workspace.queue.clear();
if let Some(index) =
self.nearest_one_with_queue(point, max_distance, &mut workspace.node_queue)
{
workspace.results.push(index);
}
return &workspace.results;
}
self.collect_neighbors_with_queues(
point,
max_results,
max_distance,
&mut workspace.results,
&mut workspace.queue,
&mut workspace.node_queue,
);
&workspace.results
}
pub fn visit_neighbors<B, F>(
&self,
point: Point2D,
max_distance: f64,
mut visitor: F,
) -> ControlFlow<B>
where
F: FnMut(usize, f64) -> ControlFlow<B>,
{
let mut queue = BinaryHeap::with_capacity(DEFAULT_NEIGHBOR_QUEUE_CAPACITY);
self.visit_neighbors_with_queue(point, max_distance, &mut queue, &mut visitor)
}
pub fn visit<B, F>(&self, query: Box2D, visitor: F) -> ControlFlow<B>
where
F: FnMut(usize) -> ControlFlow<B>,
{
let mut stack: Vec<usize> = Vec::with_capacity(DEFAULT_SEARCH_STACK_CAPACITY);
self.visit_with_stack(query, &mut stack, visitor)
}
fn collect_neighbors_with_queues(
&self,
point: Point2D,
max_results: usize,
max_distance: f64,
results: &mut Vec<usize>,
item_queue: &mut BinaryHeap<NeighborState>,
node_queue: &mut BinaryHeap<NeighborNodeState>,
) {
item_queue.clear();
node_queue.clear();
let Some(max_dist_sq) = max_distance_squared(max_distance) else {
return;
};
if self.num_items == 0 {
return;
}
let root_index = self.entries.len() - 1;
let root_dist = self.entries[root_index].distance_squared_to(point);
if root_dist > max_dist_sq {
return;
}
node_queue.push(NeighborNodeState::new(root_index, root_dist));
while results.len() < max_results {
while let Some(&node) = node_queue.peek() {
if node.dist > max_dist_sq {
node_queue.clear();
break;
}
if item_queue.peek().is_some_and(|item| item.dist < node.dist) {
break;
}
let node = node_queue.pop().unwrap();
let upper_bound_level = upper_bound_level(&self.level_bounds, node.index);
let end = (node.index + self.node_size).min(self.level_bounds[upper_bound_level]);
let is_leaf = node.index < self.num_items;
if is_leaf {
for pos in node.index..end {
let b = self.entries[pos];
let dist = b.distance_squared_to(point);
if dist <= max_dist_sq {
item_queue.push(NeighborState::new(self.indices[pos], true, dist));
}
}
} else {
for pos in node.index..end {
let b = self.entries[pos];
let dist = b.distance_squared_to(point);
if dist <= max_dist_sq {
node_queue.push(NeighborNodeState::new(self.indices[pos], dist));
}
}
}
}
match item_queue.pop() {
Some(state) if state.dist <= max_dist_sq => results.push(state.index),
_ => return,
}
}
}
fn nearest_one_with_queue(
&self,
point: Point2D,
max_distance: f64,
queue: &mut BinaryHeap<NeighborNodeState>,
) -> Option<usize> {
queue.clear();
let mut best_dist = max_distance_squared(max_distance)?;
if self.num_items == 0 {
return None;
}
let mut best_index = None;
let mut node_index = self.entries.len() - 1;
loop {
let upper_bound_level = upper_bound_level(&self.level_bounds, node_index);
let end = (node_index + self.node_size).min(self.level_bounds[upper_bound_level]);
let is_leaf = node_index < self.num_items;
for pos in node_index..end {
let b = self.entries[pos];
let dist = b.distance_squared_to(point);
if dist > best_dist {
continue;
}
if is_leaf {
if dist == 0.0 {
return Some(self.indices[pos]);
}
best_dist = dist;
best_index = Some(self.indices[pos]);
} else {
queue.push(NeighborNodeState::new(self.indices[pos], dist));
}
}
match queue.pop() {
Some(state) if state.dist <= best_dist => node_index = state.index,
_ => return best_index,
}
}
}
fn visit_neighbors_with_queue<B, F>(
&self,
point: Point2D,
max_distance: f64,
queue: &mut BinaryHeap<NeighborState>,
visitor: &mut F,
) -> ControlFlow<B>
where
F: FnMut(usize, f64) -> ControlFlow<B>,
{
queue.clear();
let Some(max_dist_sq) = max_distance_squared(max_distance) else {
return ControlFlow::Continue(());
};
if self.num_items == 0 {
return ControlFlow::Continue(());
}
let mut node_index = self.entries.len() - 1;
loop {
let upper_bound_level = upper_bound_level(&self.level_bounds, node_index);
let end = (node_index + self.node_size).min(self.level_bounds[upper_bound_level]);
let is_leaf = node_index < self.num_items;
for pos in node_index..end {
let b = self.entries[pos];
let dist = b.distance_squared_to(point);
if dist > max_dist_sq {
continue;
}
queue.push(NeighborState::new(self.indices[pos], is_leaf, dist));
}
let mut continue_search = false;
while let Some(state) = queue.pop() {
if state.dist > max_dist_sq {
queue.clear();
return ControlFlow::Continue(());
}
if state.is_leaf {
visitor(state.index, state.dist)?;
} else {
node_index = state.index;
continue_search = true;
break;
}
}
if !continue_search {
return ControlFlow::Continue(());
}
}
}
#[doc(hidden)]
pub fn visit_with_stack<B, F>(
&self,
query: Box2D,
stack: &mut Vec<usize>,
visitor: F,
) -> ControlFlow<B>
where
F: FnMut(usize) -> ControlFlow<B>,
{
self.visit_with_stack_impl::<false, B, F>(query, stack, visitor)
}
#[doc(hidden)]
pub fn visit_with_stack_prefetch<B, F>(
&self,
query: Box2D,
stack: &mut Vec<usize>,
visitor: F,
) -> ControlFlow<B>
where
F: FnMut(usize) -> ControlFlow<B>,
{
self.visit_with_stack_impl::<true, B, F>(query, stack, visitor)
}
#[doc(hidden)]
pub fn search_into_stack(
&self,
query: Box2D,
results: &mut Vec<usize>,
stack: &mut Vec<usize>,
) {
self.search_into_stack_impl::<false>(query, results, stack);
}
#[doc(hidden)]
pub fn search_into_stack_prefetch(
&self,
query: Box2D,
results: &mut Vec<usize>,
stack: &mut Vec<usize>,
) {
self.search_into_stack_impl::<true>(query, results, stack);
}
fn search_into_stack_impl<const PREFETCH: bool>(
&self,
query: Box2D,
results: &mut Vec<usize>,
stack: &mut Vec<usize>,
) {
results.clear();
stack.clear();
if self.num_items == 0 {
return;
}
let mut node_index = self.entries.len() - 1;
let mut level = self.level_bounds.len() - 1;
loop {
let end = (node_index + self.node_size).min(self.level_bounds[level]);
let is_leaf = node_index < self.num_items;
let node_entries = &self.entries[node_index..end];
let node_indices = &self.indices[node_index..end];
if is_leaf {
for (b, &index) in node_entries.iter().zip(node_indices) {
if !b.overlaps(query) {
continue;
}
results.push(index);
}
} else {
let child_level = level - 1;
for (b, &index) in node_entries.iter().zip(node_indices).rev() {
if !b.overlaps(query) {
continue;
}
stack.push(index);
stack.push(child_level);
}
}
if stack.len() > 1 {
if PREFETCH {
prefetch_aos_node(
&self.entries,
&self.indices,
stack[stack.len() - 2],
self.node_size,
);
}
level = stack.pop().unwrap();
node_index = stack.pop().unwrap();
} else {
return;
}
}
}
fn visit_with_stack_impl<const PREFETCH: bool, B, F>(
&self,
query: Box2D,
stack: &mut Vec<usize>,
mut visitor: F,
) -> ControlFlow<B>
where
F: FnMut(usize) -> ControlFlow<B>,
{
stack.clear();
if self.num_items == 0 {
return ControlFlow::Continue(());
}
let mut node_index = self.entries.len() - 1;
let mut level = self.level_bounds.len() - 1;
loop {
let end = (node_index + self.node_size).min(self.level_bounds[level]);
let is_leaf = node_index < self.num_items;
let node_entries = &self.entries[node_index..end];
let node_indices = &self.indices[node_index..end];
if is_leaf {
for (b, &index) in node_entries.iter().zip(node_indices) {
if !b.overlaps(query) {
continue;
}
visitor(index)?;
}
} else {
let child_level = level - 1;
for (b, &index) in node_entries.iter().zip(node_indices).rev() {
if !b.overlaps(query) {
continue;
}
stack.push(index);
stack.push(child_level);
}
}
if stack.len() > 1 {
if PREFETCH {
prefetch_aos_node(
&self.entries,
&self.indices,
stack[stack.len() - 2],
self.node_size,
);
}
level = stack.pop().unwrap();
node_index = stack.pop().unwrap();
} else {
return ControlFlow::Continue(());
}
}
}
#[doc(hidden)]
pub fn search_visited(&self, query: Box2D) -> (usize, usize) {
let mut results = 0usize;
let mut visited = 0usize;
if self.num_items == 0 {
return (0, 0);
}
let mut node_index = self.entries.len() - 1;
let mut level = self.level_bounds.len() - 1;
let mut stack: Vec<usize> = Vec::with_capacity(DEFAULT_SEARCH_STACK_CAPACITY);
loop {
let end = (node_index + self.node_size).min(self.level_bounds[level]);
let is_leaf = node_index < self.num_items;
for pos in node_index..end {
visited += 1;
let b = &self.entries[pos];
if !b.overlaps(query) {
continue;
}
let index = self.indices[pos];
if is_leaf {
results += 1;
} else {
stack.push(index);
stack.push(level - 1);
}
}
if stack.len() > 1 {
level = stack.pop().unwrap();
node_index = stack.pop().unwrap();
} else {
return (results, visited);
}
}
}
}
pub struct Index2DView<'a> {
node_size: usize,
num_items: usize,
num_nodes: usize,
level_count: usize,
level_bounds: &'a [u8],
entries: &'a [u8],
indices: &'a [u8],
}
impl<'a> Index2DView<'a> {
pub fn from_bytes(bytes: &'a [u8]) -> Result<Self, LoadError> {
let parsed = parse_index_bytes(bytes)?;
Ok(Self {
node_size: parsed.node_size,
num_items: parsed.num_items,
num_nodes: parsed.num_nodes,
level_count: parsed.level_count,
level_bounds: parsed.level_bounds,
entries: parsed.entries,
indices: parsed.indices,
})
}
pub fn num_items(&self) -> usize {
self.num_items
}
pub fn extent(&self) -> Option<Box2D> {
if self.num_items == 0 {
None
} else {
Some(self.entry_at_unchecked(self.num_nodes - 1))
}
}
pub fn node_size(&self) -> usize {
self.node_size
}
pub fn search(&self, query: Box2D) -> Vec<usize> {
let mut results = Vec::new();
self.search_into(query, &mut results);
results
}
pub fn search_into(&self, query: Box2D, results: &mut Vec<usize>) {
let mut stack: Vec<usize> = Vec::with_capacity(DEFAULT_SEARCH_STACK_CAPACITY);
self.search_into_stack(query, results, &mut stack);
}
pub fn search_with<'b>(&self, query: Box2D, workspace: &'b mut SearchWorkspace) -> &'b [usize] {
self.search_into_stack(query, &mut workspace.results, &mut workspace.stack);
&workspace.results
}
pub fn any(&self, query: Box2D) -> bool {
self.visit(query, |_| ControlFlow::Break(())).is_break()
}
pub fn first(&self, query: Box2D) -> Option<usize> {
match self.visit(query, ControlFlow::Break) {
ControlFlow::Break(index) => Some(index),
ControlFlow::Continue(()) => None,
}
}
pub fn neighbors(&self, point: Point2D, max_results: usize) -> Vec<usize> {
self.neighbors_within(point, max_results, f64::INFINITY)
}
pub fn neighbors_within(
&self,
point: Point2D,
max_results: usize,
max_distance: f64,
) -> Vec<usize> {
let mut results = Vec::new();
self.neighbors_into(point, max_results, max_distance, &mut results);
results
}
pub fn neighbors_into(
&self,
point: Point2D,
max_results: usize,
max_distance: f64,
results: &mut Vec<usize>,
) {
results.clear();
if max_results == 0 {
return;
}
if max_results == 1 {
let mut queue = BinaryHeap::with_capacity(DEFAULT_NEIGHBOR_QUEUE_CAPACITY);
if let Some(index) = self.nearest_one_with_queue(point, max_distance, &mut queue) {
results.push(index);
}
return;
}
let mut item_queue = BinaryHeap::with_capacity(DEFAULT_NEIGHBOR_QUEUE_CAPACITY);
let mut node_queue = BinaryHeap::with_capacity(DEFAULT_NEIGHBOR_QUEUE_CAPACITY);
self.collect_neighbors_with_queues(
point,
max_results,
max_distance,
results,
&mut item_queue,
&mut node_queue,
);
}
pub fn neighbors_with<'b>(
&self,
point: Point2D,
max_results: usize,
max_distance: f64,
workspace: &'b mut NeighborWorkspace,
) -> &'b [usize] {
workspace.results.clear();
if max_results == 0 {
workspace.queue.clear();
workspace.node_queue.clear();
return &workspace.results;
}
if max_results == 1 {
workspace.queue.clear();
if let Some(index) =
self.nearest_one_with_queue(point, max_distance, &mut workspace.node_queue)
{
workspace.results.push(index);
}
return &workspace.results;
}
self.collect_neighbors_with_queues(
point,
max_results,
max_distance,
&mut workspace.results,
&mut workspace.queue,
&mut workspace.node_queue,
);
&workspace.results
}
pub fn visit_neighbors<B, F>(
&self,
point: Point2D,
max_distance: f64,
mut visitor: F,
) -> ControlFlow<B>
where
F: FnMut(usize, f64) -> ControlFlow<B>,
{
let mut queue = BinaryHeap::with_capacity(DEFAULT_NEIGHBOR_QUEUE_CAPACITY);
self.visit_neighbors_with_queue(point, max_distance, &mut queue, &mut visitor)
}
pub fn visit<B, F>(&self, query: Box2D, visitor: F) -> ControlFlow<B>
where
F: FnMut(usize) -> ControlFlow<B>,
{
let mut stack: Vec<usize> = Vec::with_capacity(DEFAULT_SEARCH_STACK_CAPACITY);
self.visit_with_stack(query, &mut stack, visitor)
}
fn collect_neighbors_with_queues(
&self,
point: Point2D,
max_results: usize,
max_distance: f64,
results: &mut Vec<usize>,
item_queue: &mut BinaryHeap<NeighborState>,
node_queue: &mut BinaryHeap<NeighborNodeState>,
) {
item_queue.clear();
node_queue.clear();
let Some(max_dist_sq) = max_distance_squared(max_distance) else {
return;
};
if self.num_items == 0 {
return;
}
let root_index = self.num_nodes - 1;
let root_dist = self
.entry_at_unchecked(root_index)
.distance_squared_to(point);
if root_dist > max_dist_sq {
return;
}
node_queue.push(NeighborNodeState::new(root_index, root_dist));
while results.len() < max_results {
while let Some(&node) = node_queue.peek() {
if node.dist > max_dist_sq {
node_queue.clear();
break;
}
if item_queue.peek().is_some_and(|item| item.dist < node.dist) {
break;
}
let node = node_queue.pop().unwrap();
let upper_bound_level = self.upper_bound_level(node.index);
let end = (node.index + self.node_size)
.min(self.level_bound_unchecked(upper_bound_level));
let is_leaf = node.index < self.num_items;
if is_leaf {
for pos in node.index..end {
let b = self.entry_at_unchecked(pos);
let dist = b.distance_squared_to(point);
if dist <= max_dist_sq {
item_queue.push(NeighborState::new(
self.index_at_unchecked(pos),
true,
dist,
));
}
}
} else {
for pos in node.index..end {
let b = self.entry_at_unchecked(pos);
let dist = b.distance_squared_to(point);
if dist <= max_dist_sq {
node_queue
.push(NeighborNodeState::new(self.index_at_unchecked(pos), dist));
}
}
}
}
match item_queue.pop() {
Some(state) if state.dist <= max_dist_sq => results.push(state.index),
_ => return,
}
}
}
fn nearest_one_with_queue(
&self,
point: Point2D,
max_distance: f64,
queue: &mut BinaryHeap<NeighborNodeState>,
) -> Option<usize> {
queue.clear();
let mut best_dist = max_distance_squared(max_distance)?;
if self.num_items == 0 {
return None;
}
let mut best_index = None;
let mut node_index = self.num_nodes - 1;
loop {
let upper_bound_level = self.upper_bound_level(node_index);
let end =
(node_index + self.node_size).min(self.level_bound_unchecked(upper_bound_level));
let is_leaf = node_index < self.num_items;
for pos in node_index..end {
let b = self.entry_at_unchecked(pos);
let dist = b.distance_squared_to(point);
if dist > best_dist {
continue;
}
if is_leaf {
if dist == 0.0 {
return Some(self.index_at_unchecked(pos));
}
best_dist = dist;
best_index = Some(self.index_at_unchecked(pos));
} else {
queue.push(NeighborNodeState::new(self.index_at_unchecked(pos), dist));
}
}
match queue.pop() {
Some(state) if state.dist <= best_dist => node_index = state.index,
_ => return best_index,
}
}
}
#[doc(hidden)]
pub fn search_into_stack(
&self,
query: Box2D,
results: &mut Vec<usize>,
stack: &mut Vec<usize>,
) {
results.clear();
stack.clear();
if self.num_items == 0 {
return;
}
let mut node_index = self.num_nodes - 1;
let mut level = self.level_count - 1;
loop {
let end = (node_index + self.node_size).min(self.level_bound_unchecked(level));
let is_leaf = node_index < self.num_items;
if is_leaf {
for pos in node_index..end {
let b = self.entry_at_unchecked(pos);
if !b.overlaps(query) {
continue;
}
let index = self.index_at_unchecked(pos);
results.push(index);
}
} else {
let child_level = level - 1;
for pos in (node_index..end).rev() {
let b = self.entry_at_unchecked(pos);
if !b.overlaps(query) {
continue;
}
let index = self.index_at_unchecked(pos);
stack.push(index);
stack.push(child_level);
}
}
if stack.len() > 1 {
level = stack.pop().unwrap();
node_index = stack.pop().unwrap();
} else {
return;
}
}
}
#[doc(hidden)]
pub fn visit_with_stack<B, F>(
&self,
query: Box2D,
stack: &mut Vec<usize>,
mut visitor: F,
) -> ControlFlow<B>
where
F: FnMut(usize) -> ControlFlow<B>,
{
stack.clear();
if self.num_items == 0 {
return ControlFlow::Continue(());
}
let mut node_index = self.num_nodes - 1;
let mut level = self.level_count - 1;
loop {
let end = (node_index + self.node_size).min(self.level_bound_unchecked(level));
let is_leaf = node_index < self.num_items;
if is_leaf {
for pos in node_index..end {
let b = self.entry_at_unchecked(pos);
if !b.overlaps(query) {
continue;
}
let index = self.index_at_unchecked(pos);
visitor(index)?;
}
} else {
let child_level = level - 1;
for pos in (node_index..end).rev() {
let b = self.entry_at_unchecked(pos);
if !b.overlaps(query) {
continue;
}
let index = self.index_at_unchecked(pos);
stack.push(index);
stack.push(child_level);
}
}
if stack.len() > 1 {
level = stack.pop().unwrap();
node_index = stack.pop().unwrap();
} else {
return ControlFlow::Continue(());
}
}
}
fn visit_neighbors_with_queue<B, F>(
&self,
point: Point2D,
max_distance: f64,
queue: &mut BinaryHeap<NeighborState>,
visitor: &mut F,
) -> ControlFlow<B>
where
F: FnMut(usize, f64) -> ControlFlow<B>,
{
queue.clear();
let Some(max_dist_sq) = max_distance_squared(max_distance) else {
return ControlFlow::Continue(());
};
if self.num_items == 0 {
return ControlFlow::Continue(());
}
let mut node_index = self.num_nodes - 1;
loop {
let upper_bound_level = self.upper_bound_level(node_index);
let end =
(node_index + self.node_size).min(self.level_bound_unchecked(upper_bound_level));
let is_leaf = node_index < self.num_items;
for pos in node_index..end {
let b = self.entry_at_unchecked(pos);
let dist = b.distance_squared_to(point);
if dist > max_dist_sq {
continue;
}
queue.push(NeighborState::new(
self.index_at_unchecked(pos),
is_leaf,
dist,
));
}
let mut continue_search = false;
while let Some(state) = queue.pop() {
if state.dist > max_dist_sq {
queue.clear();
return ControlFlow::Continue(());
}
if state.is_leaf {
visitor(state.index, state.dist)?;
} else {
node_index = state.index;
continue_search = true;
break;
}
}
if !continue_search {
return ControlFlow::Continue(());
}
}
}
fn upper_bound_level(&self, node_index: usize) -> usize {
let mut lo = 0usize;
let mut hi = self.level_count - 1;
while lo < hi {
let mid = (lo + hi) / 2;
if self.level_bound_unchecked(mid) > node_index {
hi = mid;
} else {
lo = mid + 1;
}
}
lo
}
#[inline]
fn level_bound_unchecked(&self, index: usize) -> usize {
read_u64_le_unchecked(self.level_bounds, index * 8) as usize
}
#[inline]
fn entry_at_unchecked(&self, index: usize) -> Box2D {
let offset = index * 32;
Box2D::new(
read_f64_le_unchecked(self.entries, offset),
read_f64_le_unchecked(self.entries, offset + 8),
read_f64_le_unchecked(self.entries, offset + 16),
read_f64_le_unchecked(self.entries, offset + 24),
)
}
#[inline]
fn index_at_unchecked(&self, index: usize) -> usize {
read_u64_le_unchecked(self.indices, index * 8) as usize
}
}