packed_spatial_index/
index2d_soa.rs

1//! SoA index variant with SIMD searches (available with the `simd` feature).
2//!
3//! items are stored as four separate arrays (`min_x[]`, `min_y[]`, `max_x[]`,
4//! `max_y[]`). The tree is built exactly like the AoS version; only the layout
5//! and search implementation differ.
6
7use std::{collections::BinaryHeap, ops::ControlFlow};
8
9use wide::{CmpGe, CmpLe, f64x4};
10
11use crate::{
12    build::BuildError,
13    builder2d::BuildConfig,
14    config::{DEFAULT_NEIGHBOR_QUEUE_CAPACITY, DEFAULT_SEARCH_STACK_CAPACITY},
15    geometry::{Box2D, Point2D},
16    neighbors::{NeighborNodeState, NeighborState, NeighborWorkspace, max_distance_squared},
17    persistence::{
18        ByteWriter, LoadError, parse_index_bytes, read_f64_le_unchecked, read_u64_le_unchecked,
19        serialized_len,
20    },
21    sort2d::{SortKeyContext, encode_sort_by_key},
22    traversal::{SearchWorkspace, prefetch_read, upper_bound_level},
23    tree::{TreeLayout, try_compute_tree_layout},
24};
25
26type Num = f64;
27
28pub(crate) fn build_simd_index(
29    config: BuildConfig,
30    items: Vec<Box2D>,
31) -> Result<SimdIndex2D, BuildError> {
32    let node_size = config.node_size;
33    let num_items = config.num_items;
34    let TreeLayout {
35        level_bounds,
36        num_nodes,
37    } = try_compute_tree_layout(num_items, node_size)?;
38
39    if num_items == 0 {
40        return Ok(SimdIndex2D {
41            node_size,
42            num_items,
43            level_bounds,
44            min_xs: Vec::new(),
45            min_ys: Vec::new(),
46            max_xs: Vec::new(),
47            max_ys: Vec::new(),
48            indices: Vec::new(),
49        });
50    }
51
52    if num_items <= node_size {
53        return Ok(build_single_node_soa(
54            node_size,
55            num_items,
56            level_bounds,
57            items,
58        ));
59    }
60
61    let mut min_xs = vec![0.0f64; num_nodes];
62    let mut min_ys = vec![0.0f64; num_nodes];
63    let mut max_xs = vec![0.0f64; num_nodes];
64    let mut max_ys = vec![0.0f64; num_nodes];
65    let mut indices = vec![0usize; num_nodes];
66
67    let (mut e_min_x, mut e_min_y) = (f64::INFINITY, f64::INFINITY);
68    let (mut e_max_x, mut e_max_y) = (f64::NEG_INFINITY, f64::NEG_INFINITY);
69    for b in &items {
70        e_min_x = e_min_x.min(b.min_x);
71        e_min_y = e_min_y.min(b.min_y);
72        e_max_x = e_max_x.max(b.max_x);
73        e_max_y = e_max_y.max(b.max_y);
74    }
75    let scaled_width = u16::MAX as f64 / (e_max_x - e_min_x);
76    let scaled_height = u16::MAX as f64 / (e_max_y - e_min_y);
77
78    #[cfg(feature = "parallel")]
79    let use_parallel = config.parallel && num_items >= config.parallel_min_items;
80
81    let context = SortKeyContext {
82        scaled_width,
83        scaled_height,
84        min_x: e_min_x,
85        min_y: e_min_y,
86        radix: config.radix,
87        radix_bits: config.radix_bits,
88        #[cfg(feature = "parallel")]
89        use_parallel,
90    };
91    let order = encode_sort_by_key(&items, config.sort_key, context);
92
93    #[cfg(feature = "parallel")]
94    let scattered_in_parallel = if use_parallel {
95        reorder_parallel_soa_2d(
96            &mut min_xs[..num_items],
97            &mut min_ys[..num_items],
98            &mut max_xs[..num_items],
99            &mut max_ys[..num_items],
100            &mut indices[..num_items],
101            &order,
102            &items,
103        );
104        true
105    } else {
106        false
107    };
108    #[cfg(not(feature = "parallel"))]
109    let scattered_in_parallel = false;
110
111    if !scattered_in_parallel {
112        for (slot, &(_, orig)) in order.iter().enumerate() {
113            let b = items[orig as usize];
114            min_xs[slot] = b.min_x;
115            min_ys[slot] = b.min_y;
116            max_xs[slot] = b.max_x;
117            max_ys[slot] = b.max_y;
118            indices[slot] = orig as usize;
119        }
120    }
121
122    let mut read_pos = 0usize;
123    let mut write_pos = num_items;
124    for &level_end in &level_bounds[0..level_bounds.len() - 1] {
125        while read_pos < level_end {
126            let node_index = read_pos;
127            let (mut nmnx, mut nmny) = (f64::INFINITY, f64::INFINITY);
128            let (mut nmxx, mut nmxy) = (f64::NEG_INFINITY, f64::NEG_INFINITY);
129            let mut j = 0;
130            while j < node_size && read_pos < level_end {
131                nmnx = nmnx.min(min_xs[read_pos]);
132                nmny = nmny.min(min_ys[read_pos]);
133                nmxx = nmxx.max(max_xs[read_pos]);
134                nmxy = nmxy.max(max_ys[read_pos]);
135                read_pos += 1;
136                j += 1;
137            }
138            min_xs[write_pos] = nmnx;
139            min_ys[write_pos] = nmny;
140            max_xs[write_pos] = nmxx;
141            max_ys[write_pos] = nmxy;
142            indices[write_pos] = node_index;
143            write_pos += 1;
144        }
145    }
146
147    Ok(SimdIndex2D {
148        node_size,
149        num_items,
150        level_bounds,
151        min_xs,
152        min_ys,
153        max_xs,
154        max_ys,
155        indices,
156    })
157}
158
159fn build_single_node_soa(
160    node_size: usize,
161    num_items: usize,
162    level_bounds: Vec<usize>,
163    items: Vec<Box2D>,
164) -> SimdIndex2D {
165    let mut min_xs = Vec::with_capacity(num_items + 1);
166    let mut min_ys = Vec::with_capacity(num_items + 1);
167    let mut max_xs = Vec::with_capacity(num_items + 1);
168    let mut max_ys = Vec::with_capacity(num_items + 1);
169    let mut indices = Vec::with_capacity(num_items + 1);
170
171    let (mut root_min_x, mut root_min_y) = (f64::INFINITY, f64::INFINITY);
172    let (mut root_max_x, mut root_max_y) = (f64::NEG_INFINITY, f64::NEG_INFINITY);
173    for (idx, b) in items.into_iter().enumerate() {
174        min_xs.push(b.min_x);
175        min_ys.push(b.min_y);
176        max_xs.push(b.max_x);
177        max_ys.push(b.max_y);
178        indices.push(idx);
179
180        root_min_x = root_min_x.min(b.min_x);
181        root_min_y = root_min_y.min(b.min_y);
182        root_max_x = root_max_x.max(b.max_x);
183        root_max_y = root_max_y.max(b.max_y);
184    }
185
186    min_xs.push(root_min_x);
187    min_ys.push(root_min_y);
188    max_xs.push(root_max_x);
189    max_ys.push(root_max_y);
190    indices.push(0);
191
192    SimdIndex2D {
193        node_size,
194        num_items,
195        level_bounds,
196        min_xs,
197        min_ys,
198        max_xs,
199        max_ys,
200        indices,
201    }
202}
203
204/// Finished read-only SIMD index.
205///
206/// Created through [`Index2DBuilder::finish_simd`](crate::Index2DBuilder::finish_simd).
207/// It has the same public search and nearest-neighbor API as [`Index2D`](crate::Index2D),
208/// but stores box coordinates in structure-of-arrays form for SIMD traversal.
209///
210/// # Example
211///
212/// ```
213/// use packed_spatial_index::{Index2DBuilder, Box2D};
214///
215/// let mut builder = Index2DBuilder::new(1);
216/// builder.add(Box2D::new(0.0, 0.0, 1.0, 1.0));
217///
218/// let index = builder.finish_simd().unwrap();
219/// assert_eq!(index.search(Box2D::new(0.5, 0.5, 0.5, 0.5)), vec![0]);
220/// ```
221pub struct SimdIndex2D {
222    node_size: usize,
223    num_items: usize,
224    level_bounds: Vec<usize>,
225    min_xs: Vec<Num>,
226    min_ys: Vec<Num>,
227    max_xs: Vec<Num>,
228    max_ys: Vec<Num>,
229    indices: Vec<usize>,
230}
231
232impl SimdIndex2D {
233    /// Number of indexed items.
234    pub fn num_items(&self) -> usize {
235        self.num_items
236    }
237
238    /// Return the total extent of indexed items, or `None` for an empty index.
239    pub fn extent(&self) -> Option<Box2D> {
240        if self.num_items == 0 {
241            None
242        } else {
243            let last = self.min_xs.len() - 1;
244            Some(Box2D::new(
245                self.min_xs[last],
246                self.min_ys[last],
247                self.max_xs[last],
248                self.max_ys[last],
249            ))
250        }
251    }
252
253    /// Return the packed node size used by this index.
254    pub fn node_size(&self) -> usize {
255        self.node_size
256    }
257
258    /// Serialize into the stable little-endian `PSINDEX` format.
259    ///
260    /// The output is byte-identical to [`Index2D::to_bytes`](crate::Index2D::to_bytes)
261    /// for the same items, so a `SimdIndex2D` and an `Index2D` are interchangeable on
262    /// disk: either can load bytes produced by the other.
263    pub fn to_bytes(&self) -> Vec<u8> {
264        let mut out = Vec::new();
265        self.to_bytes_into(&mut out);
266        out
267    }
268
269    /// Serialize into a caller-provided buffer, reusing its allocation.
270    ///
271    /// Equivalent to [`to_bytes`](Self::to_bytes) but writes into `out` (cleared first).
272    pub fn to_bytes_into(&self, out: &mut Vec<u8>) {
273        let level_count = self.level_bounds.len();
274        let num_nodes = self.min_xs.len();
275        let len = serialized_len(level_count, num_nodes).expect("serialized index is too large");
276        let mut bytes = ByteWriter::new(out, len);
277        bytes.write_magic();
278        bytes.write_format_version();
279        bytes.write_header_len();
280        bytes.write_flags();
281        bytes.write_u64(self.node_size as u64);
282        bytes.write_u64(self.num_items as u64);
283        bytes.write_u64(num_nodes as u64);
284        bytes.write_u64(level_count as u64);
285        bytes.write_usize_slice_as_u64(&self.level_bounds);
286        bytes.write_soa_boxes_2d(&self.min_xs, &self.min_ys, &self.max_xs, &self.max_ys);
287        bytes.write_usize_slice_as_u64(&self.indices);
288        bytes.finish();
289    }
290
291    /// Load a SIMD index from bytes produced by [`to_bytes`](Self::to_bytes) or by
292    /// [`Index2D::to_bytes`](crate::Index2D::to_bytes); the AoS box records are
293    /// scattered into the structure-of-arrays columns.
294    pub fn from_bytes(bytes: &[u8]) -> Result<Self, LoadError> {
295        let parsed = parse_index_bytes(bytes)?;
296        let num_nodes = parsed.num_nodes;
297
298        let mut level_bounds = Vec::with_capacity(parsed.level_count);
299        for i in 0..parsed.level_count {
300            level_bounds.push(read_u64_le_unchecked(parsed.level_bounds, i * 8) as usize);
301        }
302
303        let mut min_xs = Vec::with_capacity(num_nodes);
304        let mut min_ys = Vec::with_capacity(num_nodes);
305        let mut max_xs = Vec::with_capacity(num_nodes);
306        let mut max_ys = Vec::with_capacity(num_nodes);
307        let mut indices = Vec::with_capacity(num_nodes);
308        for i in 0..num_nodes {
309            let off = i * 32; // four f64 per 2D box record
310            min_xs.push(read_f64_le_unchecked(parsed.entries, off));
311            min_ys.push(read_f64_le_unchecked(parsed.entries, off + 8));
312            max_xs.push(read_f64_le_unchecked(parsed.entries, off + 16));
313            max_ys.push(read_f64_le_unchecked(parsed.entries, off + 24));
314            indices.push(read_u64_le_unchecked(parsed.indices, i * 8) as usize);
315        }
316
317        Ok(SimdIndex2D {
318            node_size: parsed.node_size,
319            num_items: parsed.num_items,
320            level_bounds,
321            min_xs,
322            min_ys,
323            max_xs,
324            max_ys,
325            indices,
326        })
327    }
328
329    #[inline]
330    fn prefetch_node(&self, node_index: usize) {
331        if node_index < self.min_xs.len() {
332            prefetch_read(self.min_xs.as_ptr().wrapping_add(node_index));
333            prefetch_read(self.min_ys.as_ptr().wrapping_add(node_index));
334            prefetch_read(self.max_xs.as_ptr().wrapping_add(node_index));
335            prefetch_read(self.max_ys.as_ptr().wrapping_add(node_index));
336            prefetch_read(self.indices.as_ptr().wrapping_add(node_index));
337        }
338        let next_line = node_index.saturating_add((64 / std::mem::size_of::<Num>()).max(1));
339        if self.node_size > 1 && next_line < self.min_xs.len() {
340            prefetch_read(self.min_xs.as_ptr().wrapping_add(next_line));
341            prefetch_read(self.min_ys.as_ptr().wrapping_add(next_line));
342            prefetch_read(self.max_xs.as_ptr().wrapping_add(next_line));
343            prefetch_read(self.max_ys.as_ptr().wrapping_add(next_line));
344            prefetch_read(self.indices.as_ptr().wrapping_add(next_line));
345        }
346    }
347
348    /// Return the indices of all items whose boxes intersect `query`.
349    pub fn search(&self, query: Box2D) -> Vec<usize> {
350        let mut out = Vec::new();
351        self.search_into(query, &mut out);
352        out
353    }
354
355    /// Search with a reusable result buffer.
356    ///
357    /// This automatically chooses the widest available SIMD implementation: AVX-512
358    /// on supporting x86-64 CPUs, otherwise AVX2/SSE through `wide`.
359    pub fn search_into(&self, query: Box2D, out: &mut Vec<usize>) {
360        let mut stack = Vec::with_capacity(DEFAULT_SEARCH_STACK_CAPACITY);
361        self.search_avx512(query, out, &mut stack);
362    }
363
364    /// Search with reusable result and traversal buffers.
365    pub fn search_with<'a>(&self, query: Box2D, workspace: &'a mut SearchWorkspace) -> &'a [usize] {
366        self.search_avx512(query, &mut workspace.results, &mut workspace.stack);
367        &workspace.results
368    }
369
370    /// Return `true` if at least one item intersects `query`.
371    pub fn any(&self, query: Box2D) -> bool {
372        self.visit(query, |_| ControlFlow::Break(())).is_break()
373    }
374
375    /// Return one intersecting item, if any.
376    pub fn first(&self, query: Box2D) -> Option<usize> {
377        match self.visit(query, ControlFlow::Break) {
378            ControlFlow::Break(index) => Some(index),
379            ControlFlow::Continue(()) => None,
380        }
381    }
382
383    /// Return up to `max_results` item indices nearest to `point`.
384    pub fn neighbors(&self, point: Point2D, max_results: usize) -> Vec<usize> {
385        self.neighbors_within(point, max_results, f64::INFINITY)
386    }
387
388    /// Return up to `max_results` item indices within `max_distance` of `point`.
389    pub fn neighbors_within(
390        &self,
391        point: Point2D,
392        max_results: usize,
393        max_distance: f64,
394    ) -> Vec<usize> {
395        let mut results = Vec::new();
396        self.neighbors_into(point, max_results, max_distance, &mut results);
397        results
398    }
399
400    /// Nearest-neighbor search with a reusable result buffer.
401    pub fn neighbors_into(
402        &self,
403        point: Point2D,
404        max_results: usize,
405        max_distance: f64,
406        results: &mut Vec<usize>,
407    ) {
408        results.clear();
409        if max_results == 0 {
410            return;
411        }
412        if max_results == 1 {
413            let mut queue = BinaryHeap::with_capacity(DEFAULT_NEIGHBOR_QUEUE_CAPACITY);
414            if let Some(index) = self.nearest_one_with_queue(point, max_distance, &mut queue) {
415                results.push(index);
416            }
417            return;
418        }
419
420        let mut queue = BinaryHeap::with_capacity(DEFAULT_NEIGHBOR_QUEUE_CAPACITY);
421        self.collect_neighbors_with_queue(point, max_results, max_distance, results, &mut queue);
422    }
423
424    /// Nearest-neighbor search with reusable result and priority-queue buffers.
425    pub fn neighbors_with<'a>(
426        &self,
427        point: Point2D,
428        max_results: usize,
429        max_distance: f64,
430        workspace: &'a mut NeighborWorkspace,
431    ) -> &'a [usize] {
432        workspace.results.clear();
433        if max_results == 0 {
434            workspace.queue.clear();
435            workspace.node_queue.clear();
436            return &workspace.results;
437        }
438        if max_results == 1 {
439            workspace.queue.clear();
440            if let Some(index) =
441                self.nearest_one_with_queue(point, max_distance, &mut workspace.node_queue)
442            {
443                workspace.results.push(index);
444            }
445            return &workspace.results;
446        }
447
448        workspace.node_queue.clear();
449        self.collect_neighbors_with_queue(
450            point,
451            max_results,
452            max_distance,
453            &mut workspace.results,
454            &mut workspace.queue,
455        );
456        &workspace.results
457    }
458
459    /// Visit items in nondecreasing squared-distance order from `point`.
460    pub fn visit_neighbors<B, F>(
461        &self,
462        point: Point2D,
463        max_distance: f64,
464        mut visitor: F,
465    ) -> ControlFlow<B>
466    where
467        F: FnMut(usize, f64) -> ControlFlow<B>,
468    {
469        let mut queue = BinaryHeap::with_capacity(DEFAULT_NEIGHBOR_QUEUE_CAPACITY);
470        self.visit_neighbors_with_queue(point, max_distance, &mut queue, &mut visitor)
471    }
472
473    /// Visit intersecting items without collecting a result `Vec`.
474    pub fn visit<B, F>(&self, query: Box2D, visitor: F) -> ControlFlow<B>
475    where
476        F: FnMut(usize) -> ControlFlow<B>,
477    {
478        let mut stack = Vec::with_capacity(DEFAULT_SEARCH_STACK_CAPACITY);
479        self.visit_avx512(query, &mut stack, visitor)
480    }
481
482    fn collect_neighbors_with_queue(
483        &self,
484        point: Point2D,
485        max_results: usize,
486        max_distance: f64,
487        results: &mut Vec<usize>,
488        queue: &mut BinaryHeap<NeighborState>,
489    ) {
490        queue.clear();
491        let Some(max_dist_sq) = max_distance_squared(max_distance) else {
492            return;
493        };
494        if self.num_items == 0 {
495            return;
496        }
497
498        let mut node_index = self.min_xs.len() - 1;
499        loop {
500            let upper_bound_level = upper_bound_level(&self.level_bounds, node_index);
501            let end = (node_index + self.node_size).min(self.level_bounds[upper_bound_level]);
502            let is_leaf = node_index < self.num_items;
503
504            for pos in node_index..end {
505                let dist = self.distance_squared_to(pos, point);
506                if dist > max_dist_sq {
507                    continue;
508                }
509                queue.push(NeighborState::new(self.indices[pos], is_leaf, dist));
510            }
511
512            let mut continue_search = false;
513            while let Some(state) = queue.pop() {
514                if state.dist > max_dist_sq {
515                    queue.clear();
516                    return;
517                }
518                if state.is_leaf {
519                    results.push(state.index);
520                    if results.len() == max_results {
521                        return;
522                    }
523                } else {
524                    node_index = state.index;
525                    continue_search = true;
526                    break;
527                }
528            }
529
530            if !continue_search {
531                return;
532            }
533        }
534    }
535
536    fn nearest_one_with_queue(
537        &self,
538        point: Point2D,
539        max_distance: f64,
540        queue: &mut BinaryHeap<NeighborNodeState>,
541    ) -> Option<usize> {
542        queue.clear();
543        let mut best_dist = max_distance_squared(max_distance)?;
544        if self.num_items == 0 {
545            return None;
546        }
547
548        let mut best_index = None;
549        let mut node_index = self.min_xs.len() - 1;
550        loop {
551            let upper_bound_level = upper_bound_level(&self.level_bounds, node_index);
552            let end = (node_index + self.node_size).min(self.level_bounds[upper_bound_level]);
553            let is_leaf = node_index < self.num_items;
554
555            for pos in node_index..end {
556                let dist = self.distance_squared_to(pos, point);
557                if dist > best_dist {
558                    continue;
559                }
560                if is_leaf {
561                    if dist == 0.0 {
562                        return Some(self.indices[pos]);
563                    }
564                    best_dist = dist;
565                    best_index = Some(self.indices[pos]);
566                } else {
567                    queue.push(NeighborNodeState::new(self.indices[pos], dist));
568                }
569            }
570
571            match queue.pop() {
572                Some(state) if state.dist <= best_dist => node_index = state.index,
573                _ => return best_index,
574            }
575        }
576    }
577
578    fn visit_neighbors_with_queue<B, F>(
579        &self,
580        point: Point2D,
581        max_distance: f64,
582        queue: &mut BinaryHeap<NeighborState>,
583        visitor: &mut F,
584    ) -> ControlFlow<B>
585    where
586        F: FnMut(usize, f64) -> ControlFlow<B>,
587    {
588        queue.clear();
589        let Some(max_dist_sq) = max_distance_squared(max_distance) else {
590            return ControlFlow::Continue(());
591        };
592        if self.num_items == 0 {
593            return ControlFlow::Continue(());
594        }
595
596        let mut node_index = self.min_xs.len() - 1;
597        loop {
598            let upper_bound_level = upper_bound_level(&self.level_bounds, node_index);
599            let end = (node_index + self.node_size).min(self.level_bounds[upper_bound_level]);
600            let is_leaf = node_index < self.num_items;
601
602            for pos in node_index..end {
603                let dist = self.distance_squared_to(pos, point);
604                if dist > max_dist_sq {
605                    continue;
606                }
607                queue.push(NeighborState::new(self.indices[pos], is_leaf, dist));
608            }
609
610            let mut continue_search = false;
611            while let Some(state) = queue.pop() {
612                if state.dist > max_dist_sq {
613                    queue.clear();
614                    return ControlFlow::Continue(());
615                }
616                if state.is_leaf {
617                    visitor(state.index, state.dist)?;
618                } else {
619                    node_index = state.index;
620                    continue_search = true;
621                    break;
622                }
623            }
624
625            if !continue_search {
626                return ControlFlow::Continue(());
627            }
628        }
629    }
630
631    #[inline]
632    fn distance_squared_to(&self, pos: usize, point: Point2D) -> f64 {
633        let dx = axis_distance(point.x, self.min_xs[pos], self.max_xs[pos]);
634        let dy = axis_distance(point.y, self.min_ys[pos], self.max_ys[pos]);
635        dx * dx + dy * dy
636    }
637
638    /// Same as [`visit`](SimdIndex2D::visit), but the traversal stack is reused by the caller.
639    #[doc(hidden)]
640    pub fn visit_simd<B, F>(
641        &self,
642        query: Box2D,
643        stack: &mut Vec<usize>,
644        visitor: F,
645    ) -> ControlFlow<B>
646    where
647        F: FnMut(usize) -> ControlFlow<B>,
648    {
649        self.visit_simd_impl::<false, B, F>(query, stack, visitor)
650    }
651
652    /// Hidden prefetch variant of [`visit_simd`](SimdIndex2D::visit_simd).
653    #[doc(hidden)]
654    pub fn visit_simd_prefetch<B, F>(
655        &self,
656        query: Box2D,
657        stack: &mut Vec<usize>,
658        visitor: F,
659    ) -> ControlFlow<B>
660    where
661        F: FnMut(usize) -> ControlFlow<B>,
662    {
663        self.visit_simd_impl::<true, B, F>(query, stack, visitor)
664    }
665
666    /// AVX-512 visitor path, falling back to [`visit_simd`](SimdIndex2D::visit_simd).
667    #[doc(hidden)]
668    pub fn visit_avx512<B, F>(
669        &self,
670        query: Box2D,
671        stack: &mut Vec<usize>,
672        visitor: F,
673    ) -> ControlFlow<B>
674    where
675        F: FnMut(usize) -> ControlFlow<B>,
676    {
677        #[cfg(target_arch = "x86_64")]
678        {
679            if std::is_x86_feature_detected!("avx512f") {
680                // SAFETY: this branch is selected only after checking avx512f availability.
681                return unsafe { self.visit_avx512_impl::<B, F>(query, stack, visitor) };
682            }
683        }
684        self.visit_simd(query, stack, visitor)
685    }
686
687    /// Element-by-element traversal (SoA layout, branchless `overlaps`).
688    #[doc(hidden)]
689    pub fn search_scalar(&self, query: Box2D, out: &mut Vec<usize>, stack: &mut Vec<usize>) {
690        out.clear();
691        stack.clear();
692        if self.num_items == 0 {
693            return;
694        }
695        let mut node_index = self.min_xs.len() - 1;
696        let mut level = self.level_bounds.len() - 1;
697        loop {
698            let end = (node_index + self.node_size).min(self.level_bounds[level]);
699            let is_leaf = node_index < self.num_items;
700            for pos in node_index..end {
701                let hit = (self.min_xs[pos] <= query.max_x)
702                    & (self.max_xs[pos] >= query.min_x)
703                    & (self.min_ys[pos] <= query.max_y)
704                    & (self.max_ys[pos] >= query.min_y);
705                if !hit {
706                    continue;
707                }
708                let index = self.indices[pos];
709                if is_leaf {
710                    out.push(index);
711                } else {
712                    stack.push(index);
713                    stack.push(level - 1);
714                }
715            }
716            if stack.len() > 1 {
717                level = stack.pop().unwrap();
718                node_index = stack.pop().unwrap();
719            } else {
720                return;
721            }
722        }
723    }
724
725    /// AVX2/SSE path through `wide::f64x4`.
726    #[doc(hidden)]
727    pub fn search_simd(&self, query: Box2D, out: &mut Vec<usize>, stack: &mut Vec<usize>) {
728        self.search_simd_impl::<false>(query, out, stack);
729    }
730
731    /// AVX2/SSE path with prefetch for the next node from the stack.
732    #[doc(hidden)]
733    pub fn search_simd_prefetch(&self, query: Box2D, out: &mut Vec<usize>, stack: &mut Vec<usize>) {
734        self.search_simd_impl::<true>(query, out, stack);
735    }
736
737    fn search_simd_impl<const PREFETCH: bool>(
738        &self,
739        query: Box2D,
740        out: &mut Vec<usize>,
741        stack: &mut Vec<usize>,
742    ) {
743        out.clear();
744        stack.clear();
745        if self.num_items == 0 {
746            return;
747        }
748        let qmxx_v = f64x4::splat(query.max_x);
749        let qmnx_v = f64x4::splat(query.min_x);
750        let qmxy_v = f64x4::splat(query.max_y);
751        let qmny_v = f64x4::splat(query.min_y);
752
753        let mut node_index = self.min_xs.len() - 1;
754        let mut level = self.level_bounds.len() - 1;
755        loop {
756            let end = (node_index + self.node_size).min(self.level_bounds[level]);
757            let is_leaf = node_index < self.num_items;
758
759            let mut pos = node_index;
760            while pos + 4 <= end {
761                let mnx = load4(&self.min_xs, pos);
762                let mxx = load4(&self.max_xs, pos);
763                let mny = load4(&self.min_ys, pos);
764                let mxy = load4(&self.max_ys, pos);
765                let mask = mnx.simd_le(qmxx_v)
766                    & mxx.simd_ge(qmnx_v)
767                    & mny.simd_le(qmxy_v)
768                    & mxy.simd_ge(qmny_v);
769                let bits = mask.to_bitmask();
770                if bits != 0 {
771                    for k in 0..4 {
772                        if bits & (1 << k) != 0 {
773                            let p = pos + k;
774                            let index = self.indices[p];
775                            if is_leaf {
776                                out.push(index);
777                            } else {
778                                stack.push(index);
779                                stack.push(level - 1);
780                            }
781                        }
782                    }
783                }
784                pos += 4;
785            }
786
787            while pos < end {
788                let hit = (self.min_xs[pos] <= query.max_x)
789                    & (self.max_xs[pos] >= query.min_x)
790                    & (self.min_ys[pos] <= query.max_y)
791                    & (self.max_ys[pos] >= query.min_y);
792                if hit {
793                    let index = self.indices[pos];
794                    if is_leaf {
795                        out.push(index);
796                    } else {
797                        stack.push(index);
798                        stack.push(level - 1);
799                    }
800                }
801                pos += 1;
802            }
803
804            if stack.len() > 1 {
805                if PREFETCH {
806                    self.prefetch_node(stack[stack.len() - 2]);
807                }
808                level = stack.pop().unwrap();
809                node_index = stack.pop().unwrap();
810            } else {
811                return;
812            }
813        }
814    }
815
816    fn visit_simd_impl<const PREFETCH: bool, B, F>(
817        &self,
818        query: Box2D,
819        stack: &mut Vec<usize>,
820        mut visitor: F,
821    ) -> ControlFlow<B>
822    where
823        F: FnMut(usize) -> ControlFlow<B>,
824    {
825        stack.clear();
826        if self.num_items == 0 {
827            return ControlFlow::Continue(());
828        }
829        let qmxx_v = f64x4::splat(query.max_x);
830        let qmnx_v = f64x4::splat(query.min_x);
831        let qmxy_v = f64x4::splat(query.max_y);
832        let qmny_v = f64x4::splat(query.min_y);
833
834        let mut node_index = self.min_xs.len() - 1;
835        let mut level = self.level_bounds.len() - 1;
836        loop {
837            let end = (node_index + self.node_size).min(self.level_bounds[level]);
838            let is_leaf = node_index < self.num_items;
839
840            let mut pos = node_index;
841            while pos + 4 <= end {
842                let mnx = load4(&self.min_xs, pos);
843                let mxx = load4(&self.max_xs, pos);
844                let mny = load4(&self.min_ys, pos);
845                let mxy = load4(&self.max_ys, pos);
846                let mask = mnx.simd_le(qmxx_v)
847                    & mxx.simd_ge(qmnx_v)
848                    & mny.simd_le(qmxy_v)
849                    & mxy.simd_ge(qmny_v);
850                let bits = mask.to_bitmask();
851                if bits != 0 {
852                    for k in 0..4 {
853                        if bits & (1 << k) != 0 {
854                            let p = pos + k;
855                            let index = self.indices[p];
856                            if is_leaf {
857                                visitor(index)?;
858                            } else {
859                                stack.push(index);
860                                stack.push(level - 1);
861                            }
862                        }
863                    }
864                }
865                pos += 4;
866            }
867
868            while pos < end {
869                let hit = (self.min_xs[pos] <= query.max_x)
870                    & (self.max_xs[pos] >= query.min_x)
871                    & (self.min_ys[pos] <= query.max_y)
872                    & (self.max_ys[pos] >= query.min_y);
873                if hit {
874                    let index = self.indices[pos];
875                    if is_leaf {
876                        visitor(index)?;
877                    } else {
878                        stack.push(index);
879                        stack.push(level - 1);
880                    }
881                }
882                pos += 1;
883            }
884
885            if stack.len() > 1 {
886                if PREFETCH {
887                    self.prefetch_node(stack[stack.len() - 2]);
888                }
889                level = stack.pop().unwrap();
890                node_index = stack.pop().unwrap();
891            } else {
892                return ControlFlow::Continue(());
893            }
894        }
895    }
896
897    #[cfg(target_arch = "x86_64")]
898    #[target_feature(enable = "avx512f")]
899    unsafe fn visit_avx512_impl<B, F>(
900        &self,
901        query: Box2D,
902        stack: &mut Vec<usize>,
903        mut visitor: F,
904    ) -> ControlFlow<B>
905    where
906        F: FnMut(usize) -> ControlFlow<B>,
907    {
908        use std::arch::x86_64::*;
909
910        stack.clear();
911        if self.num_items == 0 {
912            return ControlFlow::Continue(());
913        }
914        let qmxx_v = _mm512_set1_pd(query.max_x);
915        let qmnx_v = _mm512_set1_pd(query.min_x);
916        let qmxy_v = _mm512_set1_pd(query.max_y);
917        let qmny_v = _mm512_set1_pd(query.min_y);
918
919        let mut node_index = self.min_xs.len() - 1;
920        let mut level = self.level_bounds.len() - 1;
921        loop {
922            let end = (node_index + self.node_size).min(self.level_bounds[level]);
923            let is_leaf = node_index < self.num_items;
924
925            let mut pos = node_index;
926            while pos + 8 <= end {
927                // SAFETY: `pos + 8 <= end`, and `end` is bounded by the array length.
928                let (mnx, mxx, mny, mxy) = unsafe {
929                    (
930                        _mm512_loadu_pd(self.min_xs.as_ptr().add(pos)),
931                        _mm512_loadu_pd(self.max_xs.as_ptr().add(pos)),
932                        _mm512_loadu_pd(self.min_ys.as_ptr().add(pos)),
933                        _mm512_loadu_pd(self.max_ys.as_ptr().add(pos)),
934                    )
935                };
936                let m1 = _mm512_cmp_pd_mask::<_CMP_LE_OQ>(mnx, qmxx_v);
937                let m2 = _mm512_cmp_pd_mask::<_CMP_GE_OQ>(mxx, qmnx_v);
938                let m3 = _mm512_cmp_pd_mask::<_CMP_LE_OQ>(mny, qmxy_v);
939                let m4 = _mm512_cmp_pd_mask::<_CMP_GE_OQ>(mxy, qmny_v);
940                let mut bits: u8 = m1 & m2 & m3 & m4;
941                while bits != 0 {
942                    let k = bits.trailing_zeros() as usize;
943                    let index = self.indices[pos + k];
944                    if is_leaf {
945                        visitor(index)?;
946                    } else {
947                        stack.push(index);
948                        stack.push(level - 1);
949                    }
950                    bits &= bits - 1;
951                }
952                pos += 8;
953            }
954
955            while pos < end {
956                let hit = (self.min_xs[pos] <= query.max_x)
957                    & (self.max_xs[pos] >= query.min_x)
958                    & (self.min_ys[pos] <= query.max_y)
959                    & (self.max_ys[pos] >= query.min_y);
960                if hit {
961                    let index = self.indices[pos];
962                    if is_leaf {
963                        visitor(index)?;
964                    } else {
965                        stack.push(index);
966                        stack.push(level - 1);
967                    }
968                }
969                pos += 1;
970            }
971
972            if stack.len() > 1 {
973                level = stack.pop().unwrap();
974                node_index = stack.pop().unwrap();
975            } else {
976                return ControlFlow::Continue(());
977            }
978        }
979    }
980
981    /// AVX-512 path, falling back to [`search_simd`](SimdIndex2D::search_simd).
982    #[doc(hidden)]
983    pub fn search_avx512(&self, query: Box2D, out: &mut Vec<usize>, stack: &mut Vec<usize>) {
984        #[cfg(target_arch = "x86_64")]
985        {
986            if std::is_x86_feature_detected!("avx512f") {
987                // SAFETY: this branch is selected only after checking avx512f availability.
988                unsafe { self.search_avx512_impl(query, out, stack) };
989                return;
990            }
991        }
992        self.search_simd(query, out, stack);
993    }
994
995    #[cfg(target_arch = "x86_64")]
996    #[target_feature(enable = "avx512f")]
997    unsafe fn search_avx512_impl(
998        &self,
999        query: Box2D,
1000        out: &mut Vec<usize>,
1001        stack: &mut Vec<usize>,
1002    ) {
1003        use std::arch::x86_64::*;
1004
1005        out.clear();
1006        stack.clear();
1007        if self.num_items == 0 {
1008            return;
1009        }
1010        let qmxx_v = _mm512_set1_pd(query.max_x);
1011        let qmnx_v = _mm512_set1_pd(query.min_x);
1012        let qmxy_v = _mm512_set1_pd(query.max_y);
1013        let qmny_v = _mm512_set1_pd(query.min_y);
1014
1015        let mut node_index = self.min_xs.len() - 1;
1016        let mut level = self.level_bounds.len() - 1;
1017        loop {
1018            let end = (node_index + self.node_size).min(self.level_bounds[level]);
1019            let is_leaf = node_index < self.num_items;
1020
1021            let mut pos = node_index;
1022            while pos + 8 <= end {
1023                // SAFETY: `pos + 8 <= end`, and `end` is bounded by the array length.
1024                let (mnx, mxx, mny, mxy) = unsafe {
1025                    (
1026                        _mm512_loadu_pd(self.min_xs.as_ptr().add(pos)),
1027                        _mm512_loadu_pd(self.max_xs.as_ptr().add(pos)),
1028                        _mm512_loadu_pd(self.min_ys.as_ptr().add(pos)),
1029                        _mm512_loadu_pd(self.max_ys.as_ptr().add(pos)),
1030                    )
1031                };
1032                let m1 = _mm512_cmp_pd_mask::<_CMP_LE_OQ>(mnx, qmxx_v);
1033                let m2 = _mm512_cmp_pd_mask::<_CMP_GE_OQ>(mxx, qmnx_v);
1034                let m3 = _mm512_cmp_pd_mask::<_CMP_LE_OQ>(mny, qmxy_v);
1035                let m4 = _mm512_cmp_pd_mask::<_CMP_GE_OQ>(mxy, qmny_v);
1036                let mut bits: u8 = m1 & m2 & m3 & m4;
1037                while bits != 0 {
1038                    let k = bits.trailing_zeros() as usize;
1039                    let index = self.indices[pos + k];
1040                    if is_leaf {
1041                        out.push(index);
1042                    } else {
1043                        stack.push(index);
1044                        stack.push(level - 1);
1045                    }
1046                    bits &= bits - 1;
1047                }
1048                pos += 8;
1049            }
1050
1051            while pos < end {
1052                let hit = (self.min_xs[pos] <= query.max_x)
1053                    & (self.max_xs[pos] >= query.min_x)
1054                    & (self.min_ys[pos] <= query.max_y)
1055                    & (self.max_ys[pos] >= query.min_y);
1056                if hit {
1057                    let index = self.indices[pos];
1058                    if is_leaf {
1059                        out.push(index);
1060                    } else {
1061                        stack.push(index);
1062                        stack.push(level - 1);
1063                    }
1064                }
1065                pos += 1;
1066            }
1067
1068            if stack.len() > 1 {
1069                level = stack.pop().unwrap();
1070                node_index = stack.pop().unwrap();
1071            } else {
1072                return;
1073            }
1074        }
1075    }
1076}
1077
1078#[inline]
1079fn axis_distance(point: f64, min: f64, max: f64) -> f64 {
1080    if point < min {
1081        min - point
1082    } else if point > max {
1083        point - max
1084    } else {
1085        0.0
1086    }
1087}
1088
1089#[inline]
1090fn load4(a: &[f64], p: usize) -> f64x4 {
1091    f64x4::from([a[p], a[p + 1], a[p + 2], a[p + 3]])
1092}
1093
1094/// Scatter the Hilbert-ordered items into the SoA leaf columns in parallel. Each
1095/// output slot is written exactly once, so the columns can be filled independently.
1096#[cfg(feature = "parallel")]
1097#[allow(clippy::too_many_arguments)]
1098fn reorder_parallel_soa_2d(
1099    min_xs: &mut [f64],
1100    min_ys: &mut [f64],
1101    max_xs: &mut [f64],
1102    max_ys: &mut [f64],
1103    indices: &mut [usize],
1104    order: &[(u32, u32)],
1105    items: &[Box2D],
1106) {
1107    use rayon::prelude::*;
1108
1109    min_xs
1110        .par_iter_mut()
1111        .zip(min_ys.par_iter_mut())
1112        .zip(max_xs.par_iter_mut())
1113        .zip(max_ys.par_iter_mut())
1114        .zip(indices.par_iter_mut())
1115        .zip(order.par_iter())
1116        .for_each(|(((((mnx, mny), mxx), mxy), idx), &(_, orig))| {
1117            let b = items[orig as usize];
1118            *mnx = b.min_x;
1119            *mny = b.min_y;
1120            *mxx = b.max_x;
1121            *mxy = b.max_y;
1122            *idx = orig as usize;
1123        });
1124}
1125
1126/// Byte size of one persisted 2D box record (`[min_x, min_y, max_x, max_y]`).
1127const RECORD_2D: usize = 32;
1128
1129/// Assemble one coordinate column for four consecutive 2D box records into a SIMD
1130/// vector. The four records are contiguous (128 bytes), so the strided reads stay
1131/// within the same cache lines.
1132#[inline]
1133fn lane4_2d(entries: &[u8], base: usize, field: usize) -> f64x4 {
1134    let o = base + field;
1135    f64x4::from([
1136        read_f64_le_unchecked(entries, o),
1137        read_f64_le_unchecked(entries, o + RECORD_2D),
1138        read_f64_le_unchecked(entries, o + 2 * RECORD_2D),
1139        read_f64_le_unchecked(entries, o + 3 * RECORD_2D),
1140    ])
1141}
1142
1143/// Zero-copy SIMD view over bytes produced by [`SimdIndex2D::to_bytes`] or
1144/// [`Index2D::to_bytes`](crate::Index2D::to_bytes).
1145///
1146/// Like [`Index2DView`](crate::Index2DView) it borrows the buffer without
1147/// allocating owned tree storage, but the traversal uses `wide::f64x4` overlap
1148/// tests, assembling lane vectors from four contiguous box records. Ideal for
1149/// querying memory-mapped indexes without allocating.
1150///
1151/// Nearest-neighbor results are returned in nondecreasing distance order. Ties
1152/// between equal-distance items are not stable across index layouts.
1153///
1154/// # Example
1155///
1156/// ```
1157/// use packed_spatial_index::{Index2DBuilder, SimdIndex2DView, Box2D};
1158///
1159/// let mut builder = Index2DBuilder::new(1);
1160/// builder.add(Box2D::new(0.0, 0.0, 1.0, 1.0));
1161/// let bytes = builder.finish_simd().unwrap().to_bytes();
1162///
1163/// let view = SimdIndex2DView::from_bytes(&bytes)?;
1164/// assert_eq!(view.search(Box2D::new(0.5, 0.5, 0.5, 0.5)), vec![0]);
1165/// # Ok::<(), Box<dyn std::error::Error>>(())
1166/// ```
1167pub struct SimdIndex2DView<'a> {
1168    node_size: usize,
1169    num_items: usize,
1170    num_nodes: usize,
1171    level_count: usize,
1172    level_bounds: &'a [u8],
1173    entries: &'a [u8],
1174    indices: &'a [u8],
1175}
1176
1177impl<'a> SimdIndex2DView<'a> {
1178    /// Borrow a zero-copy view over the canonical `PSINDEX` 2D bytes.
1179    pub fn from_bytes(bytes: &'a [u8]) -> Result<Self, LoadError> {
1180        let parsed = parse_index_bytes(bytes)?;
1181        Ok(Self {
1182            node_size: parsed.node_size,
1183            num_items: parsed.num_items,
1184            num_nodes: parsed.num_nodes,
1185            level_count: parsed.level_count,
1186            level_bounds: parsed.level_bounds,
1187            entries: parsed.entries,
1188            indices: parsed.indices,
1189        })
1190    }
1191
1192    /// Return the number of indexed items.
1193    pub fn num_items(&self) -> usize {
1194        self.num_items
1195    }
1196
1197    /// Return the packed node size.
1198    pub fn node_size(&self) -> usize {
1199        self.node_size
1200    }
1201
1202    /// Return the total extent of indexed items, or `None` for an empty view.
1203    pub fn extent(&self) -> Option<Box2D> {
1204        if self.num_items == 0 {
1205            None
1206        } else {
1207            Some(self.box_at(self.num_nodes - 1))
1208        }
1209    }
1210
1211    #[inline]
1212    fn index_at(&self, pos: usize) -> usize {
1213        read_u64_le_unchecked(self.indices, pos * 8) as usize
1214    }
1215
1216    #[inline]
1217    fn box_at(&self, pos: usize) -> Box2D {
1218        let b = pos * RECORD_2D;
1219        Box2D::new(
1220            read_f64_le_unchecked(self.entries, b),
1221            read_f64_le_unchecked(self.entries, b + 8),
1222            read_f64_le_unchecked(self.entries, b + 16),
1223            read_f64_le_unchecked(self.entries, b + 24),
1224        )
1225    }
1226
1227    #[inline]
1228    fn level_bound_unchecked(&self, index: usize) -> usize {
1229        read_u64_le_unchecked(self.level_bounds, index * 8) as usize
1230    }
1231
1232    #[inline]
1233    fn upper_bound_level(&self, node_index: usize) -> usize {
1234        let mut lo = 0usize;
1235        let mut hi = self.level_count - 1;
1236        while lo < hi {
1237            let mid = (lo + hi) / 2;
1238            if self.level_bound_unchecked(mid) > node_index {
1239                hi = mid;
1240            } else {
1241                lo = mid + 1;
1242            }
1243        }
1244        lo
1245    }
1246
1247    /// Return the indices of all items whose boxes intersect `query`.
1248    pub fn search(&self, query: Box2D) -> Vec<usize> {
1249        let mut out = Vec::new();
1250        self.search_into(query, &mut out);
1251        out
1252    }
1253
1254    /// Search with a reusable result buffer.
1255    pub fn search_into(&self, query: Box2D, out: &mut Vec<usize>) {
1256        let mut stack = Vec::with_capacity(DEFAULT_SEARCH_STACK_CAPACITY);
1257        out.clear();
1258        let _: ControlFlow<()> = self.try_visit(query, &mut stack, |index| {
1259            out.push(index);
1260            ControlFlow::Continue(())
1261        });
1262    }
1263
1264    /// Search with reusable result and traversal buffers.
1265    pub fn search_with<'b>(&self, query: Box2D, workspace: &'b mut SearchWorkspace) -> &'b [usize] {
1266        workspace.results.clear();
1267        let results = &mut workspace.results;
1268        let _: ControlFlow<()> = self.try_visit(query, &mut workspace.stack, |index| {
1269            results.push(index);
1270            ControlFlow::Continue(())
1271        });
1272        &workspace.results
1273    }
1274
1275    /// Return `true` if at least one item intersects `query`.
1276    pub fn any(&self, query: Box2D) -> bool {
1277        self.visit(query, |_| ControlFlow::Break(())).is_break()
1278    }
1279
1280    /// Return one intersecting item, if any.
1281    pub fn first(&self, query: Box2D) -> Option<usize> {
1282        match self.visit(query, ControlFlow::Break) {
1283            ControlFlow::Break(index) => Some(index),
1284            ControlFlow::Continue(()) => None,
1285        }
1286    }
1287
1288    /// Visit intersecting items without collecting a result `Vec`.
1289    pub fn visit<B, F>(&self, query: Box2D, visitor: F) -> ControlFlow<B>
1290    where
1291        F: FnMut(usize) -> ControlFlow<B>,
1292    {
1293        let mut stack = Vec::with_capacity(DEFAULT_SEARCH_STACK_CAPACITY);
1294        self.try_visit(query, &mut stack, visitor)
1295    }
1296
1297    fn try_visit<B, F>(
1298        &self,
1299        query: Box2D,
1300        stack: &mut Vec<usize>,
1301        mut visitor: F,
1302    ) -> ControlFlow<B>
1303    where
1304        F: FnMut(usize) -> ControlFlow<B>,
1305    {
1306        stack.clear();
1307        if self.num_items == 0 {
1308            return ControlFlow::Continue(());
1309        }
1310        let qmxx = f64x4::splat(query.max_x);
1311        let qmnx = f64x4::splat(query.min_x);
1312        let qmxy = f64x4::splat(query.max_y);
1313        let qmny = f64x4::splat(query.min_y);
1314
1315        let mut node_index = self.num_nodes - 1;
1316        let mut level = self.level_count - 1;
1317        loop {
1318            let end = (node_index + self.node_size).min(self.level_bound_unchecked(level));
1319            let is_leaf = node_index < self.num_items;
1320
1321            let mut pos = node_index;
1322            while pos + 4 <= end {
1323                let base = pos * RECORD_2D;
1324                let mask = lane4_2d(self.entries, base, 0).simd_le(qmxx)
1325                    & lane4_2d(self.entries, base, 16).simd_ge(qmnx)
1326                    & lane4_2d(self.entries, base, 8).simd_le(qmxy)
1327                    & lane4_2d(self.entries, base, 24).simd_ge(qmny);
1328                let bits = mask.to_bitmask();
1329                if bits != 0 {
1330                    for k in 0..4 {
1331                        if bits & (1 << k) != 0 {
1332                            let p = pos + k;
1333                            let index = self.index_at(p);
1334                            if is_leaf {
1335                                visitor(index)?;
1336                            } else {
1337                                stack.push(index);
1338                                stack.push(level - 1);
1339                            }
1340                        }
1341                    }
1342                }
1343                pos += 4;
1344            }
1345
1346            while pos < end {
1347                if self.box_at(pos).overlaps(query) {
1348                    let index = self.index_at(pos);
1349                    if is_leaf {
1350                        visitor(index)?;
1351                    } else {
1352                        stack.push(index);
1353                        stack.push(level - 1);
1354                    }
1355                }
1356                pos += 1;
1357            }
1358
1359            if stack.len() > 1 {
1360                level = stack.pop().unwrap();
1361                node_index = stack.pop().unwrap();
1362            } else {
1363                return ControlFlow::Continue(());
1364            }
1365        }
1366    }
1367
1368    #[inline]
1369    fn distance_squared_to(&self, pos: usize, point: Point2D) -> f64 {
1370        let b = self.box_at(pos);
1371        let dx = axis_distance(point.x, b.min_x, b.max_x);
1372        let dy = axis_distance(point.y, b.min_y, b.max_y);
1373        dx * dx + dy * dy
1374    }
1375
1376    /// Return up to `max_results` item indices nearest to `point`.
1377    pub fn neighbors(&self, point: Point2D, max_results: usize) -> Vec<usize> {
1378        self.neighbors_within(point, max_results, f64::INFINITY)
1379    }
1380
1381    /// Return up to `max_results` item indices within `max_distance` of `point`.
1382    pub fn neighbors_within(
1383        &self,
1384        point: Point2D,
1385        max_results: usize,
1386        max_distance: f64,
1387    ) -> Vec<usize> {
1388        let mut results = Vec::new();
1389        self.neighbors_into(point, max_results, max_distance, &mut results);
1390        results
1391    }
1392
1393    /// Nearest-neighbor search with a reusable result buffer.
1394    pub fn neighbors_into(
1395        &self,
1396        point: Point2D,
1397        max_results: usize,
1398        max_distance: f64,
1399        results: &mut Vec<usize>,
1400    ) {
1401        results.clear();
1402        if max_results == 0 {
1403            return;
1404        }
1405        if max_results == 1 {
1406            let mut queue = BinaryHeap::with_capacity(DEFAULT_NEIGHBOR_QUEUE_CAPACITY);
1407            if let Some(index) = self.nearest_one_with_queue(point, max_distance, &mut queue) {
1408                results.push(index);
1409            }
1410            return;
1411        }
1412        let mut queue = BinaryHeap::with_capacity(DEFAULT_NEIGHBOR_QUEUE_CAPACITY);
1413        self.collect_neighbors_with_queue(point, max_results, max_distance, results, &mut queue);
1414    }
1415
1416    /// Nearest-neighbor search with reusable result and priority-queue buffers.
1417    pub fn neighbors_with<'b>(
1418        &self,
1419        point: Point2D,
1420        max_results: usize,
1421        max_distance: f64,
1422        workspace: &'b mut NeighborWorkspace,
1423    ) -> &'b [usize] {
1424        workspace.results.clear();
1425        if max_results == 0 {
1426            workspace.queue.clear();
1427            workspace.node_queue.clear();
1428            return &workspace.results;
1429        }
1430        if max_results == 1 {
1431            workspace.queue.clear();
1432            if let Some(index) =
1433                self.nearest_one_with_queue(point, max_distance, &mut workspace.node_queue)
1434            {
1435                workspace.results.push(index);
1436            }
1437            return &workspace.results;
1438        }
1439        workspace.node_queue.clear();
1440        self.collect_neighbors_with_queue(
1441            point,
1442            max_results,
1443            max_distance,
1444            &mut workspace.results,
1445            &mut workspace.queue,
1446        );
1447        &workspace.results
1448    }
1449
1450    /// Visit items in nondecreasing squared-distance order from `point`.
1451    pub fn visit_neighbors<B, F>(
1452        &self,
1453        point: Point2D,
1454        max_distance: f64,
1455        mut visitor: F,
1456    ) -> ControlFlow<B>
1457    where
1458        F: FnMut(usize, f64) -> ControlFlow<B>,
1459    {
1460        let mut queue = BinaryHeap::with_capacity(DEFAULT_NEIGHBOR_QUEUE_CAPACITY);
1461        self.visit_neighbors_with_queue(point, max_distance, &mut queue, &mut visitor)
1462    }
1463
1464    fn collect_neighbors_with_queue(
1465        &self,
1466        point: Point2D,
1467        max_results: usize,
1468        max_distance: f64,
1469        results: &mut Vec<usize>,
1470        queue: &mut BinaryHeap<NeighborState>,
1471    ) {
1472        queue.clear();
1473        let Some(max_dist_sq) = max_distance_squared(max_distance) else {
1474            return;
1475        };
1476        if self.num_items == 0 {
1477            return;
1478        }
1479
1480        let mut node_index = self.num_nodes - 1;
1481        loop {
1482            let upper = self.upper_bound_level(node_index);
1483            let end = (node_index + self.node_size).min(self.level_bound_unchecked(upper));
1484            let is_leaf = node_index < self.num_items;
1485            for pos in node_index..end {
1486                let dist = self.distance_squared_to(pos, point);
1487                if dist > max_dist_sq {
1488                    continue;
1489                }
1490                queue.push(NeighborState::new(self.index_at(pos), is_leaf, dist));
1491            }
1492
1493            let mut continue_search = false;
1494            while let Some(state) = queue.pop() {
1495                if state.dist > max_dist_sq {
1496                    queue.clear();
1497                    return;
1498                }
1499                if state.is_leaf {
1500                    results.push(state.index);
1501                    if results.len() == max_results {
1502                        return;
1503                    }
1504                } else {
1505                    node_index = state.index;
1506                    continue_search = true;
1507                    break;
1508                }
1509            }
1510            if !continue_search {
1511                return;
1512            }
1513        }
1514    }
1515
1516    fn visit_neighbors_with_queue<B, F>(
1517        &self,
1518        point: Point2D,
1519        max_distance: f64,
1520        queue: &mut BinaryHeap<NeighborState>,
1521        visitor: &mut F,
1522    ) -> ControlFlow<B>
1523    where
1524        F: FnMut(usize, f64) -> ControlFlow<B>,
1525    {
1526        queue.clear();
1527        let Some(max_dist_sq) = max_distance_squared(max_distance) else {
1528            return ControlFlow::Continue(());
1529        };
1530        if self.num_items == 0 {
1531            return ControlFlow::Continue(());
1532        }
1533
1534        let mut node_index = self.num_nodes - 1;
1535        loop {
1536            let upper = self.upper_bound_level(node_index);
1537            let end = (node_index + self.node_size).min(self.level_bound_unchecked(upper));
1538            let is_leaf = node_index < self.num_items;
1539
1540            for pos in node_index..end {
1541                let dist = self.distance_squared_to(pos, point);
1542                if dist > max_dist_sq {
1543                    continue;
1544                }
1545                queue.push(NeighborState::new(self.index_at(pos), is_leaf, dist));
1546            }
1547
1548            let mut continue_search = false;
1549            while let Some(state) = queue.pop() {
1550                if state.dist > max_dist_sq {
1551                    queue.clear();
1552                    return ControlFlow::Continue(());
1553                }
1554                if state.is_leaf {
1555                    visitor(state.index, state.dist)?;
1556                } else {
1557                    node_index = state.index;
1558                    continue_search = true;
1559                    break;
1560                }
1561            }
1562            if !continue_search {
1563                return ControlFlow::Continue(());
1564            }
1565        }
1566    }
1567
1568    fn nearest_one_with_queue(
1569        &self,
1570        point: Point2D,
1571        max_distance: f64,
1572        queue: &mut BinaryHeap<NeighborNodeState>,
1573    ) -> Option<usize> {
1574        queue.clear();
1575        let mut best_dist = max_distance_squared(max_distance)?;
1576        if self.num_items == 0 {
1577            return None;
1578        }
1579        let mut best_index = None;
1580        let mut node_index = self.num_nodes - 1;
1581        loop {
1582            let upper = self.upper_bound_level(node_index);
1583            let end = (node_index + self.node_size).min(self.level_bound_unchecked(upper));
1584            let is_leaf = node_index < self.num_items;
1585            for pos in node_index..end {
1586                let dist = self.distance_squared_to(pos, point);
1587                if dist > best_dist {
1588                    continue;
1589                }
1590                if is_leaf {
1591                    if dist == 0.0 {
1592                        return Some(self.index_at(pos));
1593                    }
1594                    best_dist = dist;
1595                    best_index = Some(self.index_at(pos));
1596                } else {
1597                    queue.push(NeighborNodeState::new(self.index_at(pos), dist));
1598                }
1599            }
1600            match queue.pop() {
1601                Some(state) if state.dist <= best_dist => node_index = state.index,
1602                _ => return best_index,
1603            }
1604        }
1605    }
1606}
packed_spatial_index/index2d_soa.rs

packed_spatial_index/
index2d_soa.rs
