packed_spatial_index/
builder2d.rs

1#[cfg(feature = "parallel")]
2use crate::config::DEFAULT_PARALLEL_MIN_ITEMS;
3use crate::{
4    build::BuildError,
5    config::DEFAULT_NODE_SIZE,
6    geometry::{Box2D, Num},
7    index2d::Index2D,
8    sort2d::{
9        DEFAULT_RADIX_BITS, ExperimentalSortKey2D, SortKey2D, SortKeyContext, encode_sort_by_key,
10        normalize_radix_bits,
11    },
12    tree::{TreeLayout, compute_tree_layout, normalize_node_size},
13};
14
15/// Builder for [`Index2D`] and, with the `simd` feature, `SimdIndex2D`.
16///
17/// # Example
18///
19/// ```
20/// use packed_spatial_index::{Index2DBuilder, Box2D};
21///
22/// let mut builder = Index2DBuilder::new(2);
23/// builder.add(Box2D::new(0.0, 0.0, 1.0, 1.0));
24/// builder.add(Box2D::new(5.0, 5.0, 6.0, 6.0));
25///
26/// let index = builder.finish().unwrap();
27/// assert_eq!(index.search(Box2D::new(0.0, 0.0, 2.0, 2.0)), vec![0]);
28/// ```
29#[must_use = "Index2DBuilder methods return an updated builder; assign the result or chain the call"]
30pub struct Index2DBuilder {
31    node_size: usize,
32    num_items: usize,
33    sort_key: ExperimentalSortKey2D,
34    radix: bool,
35    radix_bits: u32,
36    #[cfg(feature = "parallel")]
37    parallel: bool,
38    #[cfg(feature = "parallel")]
39    parallel_min_items: usize,
40    items: Vec<Box2D>,
41}
42
43#[derive(Clone, Copy)]
44#[cfg(feature = "simd")]
45pub(crate) struct BuildConfig {
46    pub(crate) node_size: usize,
47    pub(crate) num_items: usize,
48    pub(crate) sort_key: ExperimentalSortKey2D,
49    pub(crate) radix: bool,
50    pub(crate) radix_bits: u32,
51    #[cfg(feature = "parallel")]
52    pub(crate) parallel: bool,
53    #[cfg(feature = "parallel")]
54    pub(crate) parallel_min_items: usize,
55}
56
57impl Index2DBuilder {
58    /// Create a builder for exactly `count` items with [`DEFAULT_NODE_SIZE`].
59    pub fn new(count: usize) -> Self {
60        Index2DBuilder {
61            node_size: DEFAULT_NODE_SIZE,
62            num_items: count,
63            sort_key: SortKey2D::Hilbert.into(),
64            radix: true,
65            radix_bits: DEFAULT_RADIX_BITS,
66            #[cfg(feature = "parallel")]
67            parallel: false,
68            #[cfg(feature = "parallel")]
69            parallel_min_items: DEFAULT_PARALLEL_MIN_ITEMS,
70            items: Vec::with_capacity(count.saturating_add(1)),
71        }
72    }
73
74    /// Set the maximum number of children per tree node (clamped to `[2, 65535]`).
75    pub fn node_size(mut self, node_size: usize) -> Self {
76        self.node_size = normalize_node_size(node_size);
77        self
78    }
79
80    /// Choose the sort key (default: [`SortKey2D::Hilbert`]).
81    pub fn sort_key(mut self, key: SortKey2D) -> Self {
82        self.sort_key = key.into();
83        self
84    }
85
86    /// Choose an experimental sort-key implementation.
87    #[doc(hidden)]
88    pub fn experimental_sort_key(mut self, key: ExperimentalSortKey2D) -> Self {
89        self.sort_key = key;
90        self
91    }
92
93    /// Use LSD radix sort on the u32 Hilbert key instead of comparison-based sorting.
94    #[doc(hidden)]
95    pub fn radix(mut self, radix: bool) -> Self {
96        self.radix = radix;
97        self
98    }
99
100    /// Set the LSD radix-sort digit width for benchmarks and tuning.
101    ///
102    /// Values are clamped to `1..=16`; the default is 8.
103    #[doc(hidden)]
104    pub fn experimental_radix_bits(mut self, bits: u32) -> Self {
105        self.radix_bits = normalize_radix_bits(bits);
106        self
107    }
108
109    /// Allow adaptive parallel builds through rayon.
110    ///
111    /// Default is `false`. When set to `true`, rayon is used only when the item
112    /// count is at least the current `parallel_min_items` threshold.
113    #[cfg(feature = "parallel")]
114    pub fn parallel(mut self, parallel: bool) -> Self {
115        self.parallel = parallel;
116        self
117    }
118
119    /// Set the minimum `count` at which [`parallel(true)`](Self::parallel)
120    /// actually enables rayon.
121    #[cfg(feature = "parallel")]
122    pub fn parallel_min_items(mut self, min_items: usize) -> Self {
123        self.parallel_min_items = min_items;
124        self
125    }
126
127    /// Add one indexed box.
128    #[inline]
129    pub fn add(&mut self, item: Box2D) {
130        self.items.push(item);
131    }
132
133    /// Pack the tree and return the finished index.
134    pub fn finish(self) -> Result<Index2D, BuildError> {
135        if self.items.len() != self.num_items {
136            return Err(BuildError::ItemCount {
137                added: self.items.len(),
138                expected: self.num_items,
139            });
140        }
141        Ok(self.build_unchecked())
142    }
143
144    /// Pack the tree into the SIMD-accelerated SoA index.
145    ///
146    /// # Example
147    ///
148    /// ```
149    /// use packed_spatial_index::{Index2DBuilder, Box2D};
150    ///
151    /// let mut builder = Index2DBuilder::new(1);
152    /// builder.add(Box2D::new(0.0, 0.0, 1.0, 1.0));
153    ///
154    /// let index = builder.finish_simd().unwrap();
155    /// assert_eq!(index.search(Box2D::new(0.5, 0.5, 0.5, 0.5)), vec![0]);
156    /// ```
157    #[cfg(feature = "simd")]
158    pub fn finish_simd(self) -> Result<crate::SimdIndex2D, BuildError> {
159        if self.items.len() != self.num_items {
160            return Err(BuildError::ItemCount {
161                added: self.items.len(),
162                expected: self.num_items,
163            });
164        }
165        Ok(crate::index2d_soa::build_simd_index(
166            self.config(),
167            self.items,
168        ))
169    }
170
171    #[cfg(feature = "simd")]
172    fn config(&self) -> BuildConfig {
173        BuildConfig {
174            node_size: self.node_size,
175            num_items: self.num_items,
176            sort_key: self.sort_key,
177            radix: self.radix,
178            radix_bits: self.radix_bits,
179            #[cfg(feature = "parallel")]
180            parallel: self.parallel,
181            #[cfg(feature = "parallel")]
182            parallel_min_items: self.parallel_min_items,
183        }
184    }
185
186    fn build_unchecked(self) -> Index2D {
187        let node_size = self.node_size;
188        let num_items = self.num_items;
189        let TreeLayout {
190            level_bounds,
191            num_nodes,
192        } = compute_tree_layout(num_items, node_size);
193
194        if num_items == 0 {
195            return Index2D {
196                node_size,
197                num_items,
198                level_bounds,
199                entries: Vec::new(),
200                indices: Vec::new(),
201            };
202        }
203
204        if num_items <= node_size {
205            return build_single_node_index(node_size, num_items, level_bounds, self.items);
206        }
207
208        let mut entries: Vec<Box2D> = vec![Box2D::new(0.0, 0.0, 0.0, 0.0); num_nodes];
209        let mut indices: Vec<usize> = vec![0usize; num_nodes];
210        let items = &self.items;
211
212        #[cfg(feature = "parallel")]
213        let use_parallel = self.parallel && num_items >= self.parallel_min_items;
214
215        let mut min_x = Num::INFINITY;
216        let mut min_y = Num::INFINITY;
217        let mut max_x = Num::NEG_INFINITY;
218        let mut max_y = Num::NEG_INFINITY;
219        for b in items {
220            min_x = min_x.min(b.min_x);
221            min_y = min_y.min(b.min_y);
222            max_x = max_x.max(b.max_x);
223            max_y = max_y.max(b.max_y);
224        }
225
226        let scaled_width = u16::MAX as f64 / (max_x - min_x);
227        let scaled_height = u16::MAX as f64 / (max_y - min_y);
228        let context = SortKeyContext {
229            scaled_width,
230            scaled_height,
231            min_x,
232            min_y,
233            radix: self.radix,
234            radix_bits: self.radix_bits,
235            #[cfg(feature = "parallel")]
236            use_parallel,
237        };
238
239        let order = encode_sort_by_key(items, self.sort_key, context);
240
241        #[cfg(feature = "parallel")]
242        if use_parallel {
243            reorder_parallel(&mut entries, &mut indices, &order, items, num_items);
244        } else {
245            reorder_serial(&mut entries, &mut indices, &order, items);
246        }
247        #[cfg(not(feature = "parallel"))]
248        reorder_serial(&mut entries, &mut indices, &order, items);
249
250        let mut read_pos = 0usize;
251        let mut write_pos = num_items;
252        for &level_end in &level_bounds[0..level_bounds.len() - 1] {
253            while read_pos < level_end {
254                let node_index = read_pos;
255                let mut node_bounds = Box2D::new(
256                    Num::INFINITY,
257                    Num::INFINITY,
258                    Num::NEG_INFINITY,
259                    Num::NEG_INFINITY,
260                );
261                let mut j = 0;
262                while j < node_size && read_pos < level_end {
263                    let b = entries[read_pos];
264                    read_pos += 1;
265                    node_bounds.min_x = node_bounds.min_x.min(b.min_x);
266                    node_bounds.min_y = node_bounds.min_y.min(b.min_y);
267                    node_bounds.max_x = node_bounds.max_x.max(b.max_x);
268                    node_bounds.max_y = node_bounds.max_y.max(b.max_y);
269                    j += 1;
270                }
271                entries[write_pos] = node_bounds;
272                indices[write_pos] = node_index;
273                write_pos += 1;
274            }
275        }
276
277        Index2D {
278            node_size,
279            num_items,
280            level_bounds,
281            entries,
282            indices,
283        }
284    }
285}
286
287fn reorder_serial(
288    entries: &mut [Box2D],
289    indices: &mut [usize],
290    order: &[(u32, u32)],
291    items: &[Box2D],
292) {
293    for (slot, &(_, orig)) in order.iter().enumerate() {
294        entries[slot] = items[orig as usize];
295        indices[slot] = orig as usize;
296    }
297}
298
299fn build_single_node_index(
300    node_size: usize,
301    num_items: usize,
302    level_bounds: Vec<usize>,
303    mut entries: Vec<Box2D>,
304) -> Index2D {
305    let mut root = Box2D::new(
306        Num::INFINITY,
307        Num::INFINITY,
308        Num::NEG_INFINITY,
309        Num::NEG_INFINITY,
310    );
311    for b in &entries {
312        root.min_x = root.min_x.min(b.min_x);
313        root.min_y = root.min_y.min(b.min_y);
314        root.max_x = root.max_x.max(b.max_x);
315        root.max_y = root.max_y.max(b.max_y);
316    }
317    entries.push(root);
318
319    let mut indices = Vec::with_capacity(num_items + 1);
320    indices.extend(0..num_items);
321    indices.push(0);
322
323    Index2D {
324        node_size,
325        num_items,
326        level_bounds,
327        entries,
328        indices,
329    }
330}
331
332#[cfg(feature = "parallel")]
333fn reorder_parallel(
334    entries: &mut [Box2D],
335    indices: &mut [usize],
336    order: &[(u32, u32)],
337    items: &[Box2D],
338    num_items: usize,
339) {
340    use rayon::prelude::*;
341
342    entries[..num_items]
343        .par_iter_mut()
344        .zip(indices[..num_items].par_iter_mut())
345        .zip(order.par_iter())
346        .for_each(|((slot_box, slot_idx), &(_, orig))| {
347            *slot_box = items[orig as usize];
348            *slot_idx = orig as usize;
349        });
350}
packed_spatial_index/builder2d.rs

packed_spatial_index/
builder2d.rs
