kdtree3d/
lib.rs

1use std::cmp::{min, max};
2
3pub trait Coord: Copy + Ord + std::ops::Sub<Output = Self> + std::ops::Add<Output = Self> + Default {
4    fn zero() -> Self { Self::default() }
5    fn from_i32(val: i32) -> Self;
6    fn min_value() -> Self;
7    fn max_value() -> Self;
8    fn to_f64(self) -> f64;
9}
10
11impl Coord for i32 {
12    #[inline]
13    fn from_i32(val: i32) -> Self { val }
14    #[inline]
15    fn min_value() -> Self { i32::MIN }
16    #[inline]
17    fn max_value() -> Self { i32::MAX }
18    #[inline]
19    fn to_f64(self) -> f64 { self as f64 }
20}
21
22impl Coord for i64 {
23    #[inline]
24    fn from_i32(val: i32) -> Self { val as i64 }
25    #[inline]
26    fn min_value() -> Self { i64::MIN }
27    #[inline]
28    fn max_value() -> Self { i64::MAX }
29    #[inline]
30    fn to_f64(self) -> f64 { self as f64 }
31}
32
33impl Coord for i128 {
34    #[inline]
35    fn from_i32(val: i32) -> Self { val as i128 }
36    #[inline]
37    fn min_value() -> Self { i128::MIN }
38    #[inline]
39    fn max_value() -> Self { i128::MAX }
40    #[inline]
41    fn to_f64(self) -> f64 { self as f64 }
42}
43
44// Box defines a 3D bounding box [left, bottom, floor, right, top, ceil]
45pub type KdBox<C = i32> = [C; 6];
46
47pub const LEFT: usize = 0;
48pub const BOTTOM: usize = 1;
49pub const FLOOR: usize = 2;
50pub const RIGHT: usize = 3;
51pub const TOP: usize = 4;
52pub const CEIL: usize = 5;
53
54#[derive(Debug, PartialEq, Clone, Copy)]
55pub enum Status {
56    Ok = 1,
57    NoMore = 2,
58    NotImpl = -3,
59    NotFound = -4,
60}
61
62pub struct Node<T, C = i32> {
63    pub item: Option<T>,
64    pub size: KdBox<C>,
65    pub lo_min_bound: C,
66    pub hi_max_bound: C,
67    pub other_bound: C,
68    pub sons: [Option<usize>; 2],
69}
70
71/// A 3D KD-Tree implementation using arena allocation.
72///
73/// The tree stores items of type `T` associated with a 6-element bounding box
74/// `[left, bottom, floor, right, top, ceil]`.
75/// It uses a `Vec`-backed arena for node storage to improve performance and cache locality.
76pub struct Tree<T, C = i32> {
77    pub arena: Vec<Node<T, C>>,
78    pub free_list: Vec<usize>,
79    pub root: Option<usize>,
80    pub item_count: i32,
81    pub dead_count: i32,
82    pub extent: KdBox<C>,
83    pub delete_flip: bool,
84}
85
86#[derive(Clone, Copy, PartialEq)]
87enum State { ThisOne, LoSon, HiSon, Done }
88
89struct Save<C> {
90    node_idx: usize,
91    disc: usize,
92    state: State,
93    bn: KdBox<C>,
94    bp: KdBox<C>,
95}
96
97impl<T: PartialEq + Clone, C: Coord> Tree<T, C> {
98    /// Creates a new, empty 3D KD-Tree.
99    pub fn new() -> Self {
100        Self {
101            arena: Vec::new(),
102            free_list: Vec::new(),
103            root: None,
104            item_count: 0,
105            dead_count: 0,
106            extent: [C::zero(); 6],
107            delete_flip: false,
108        }
109    }
110
111    /// Creates a new 3D KD-Tree with a pre-allocated capacity for the node arena.
112    pub fn with_capacity(capacity: usize) -> Self {
113        Self {
114            arena: Vec::with_capacity(capacity),
115            free_list: Vec::new(),
116            root: None,
117            item_count: 0,
118            dead_count: 0,
119            extent: [C::zero(); 6],
120            delete_flip: false,
121        }
122    }
123
124    fn allocate_node(&mut self, node: Node<T, C>) -> usize {
125        if let Some(index) = self.free_list.pop() {
126            self.arena[index] = node;
127            index
128        } else {
129            self.arena.push(node);
130            self.arena.len() - 1
131        }
132    }
133
134    fn free_node(&mut self, index: usize) {
135        self.arena[index].item = None;
136        self.free_list.push(index);
137    }
138
139    /// Inserts an item into the tree with the given 3D bounding box.
140    pub fn insert(&mut self, item: T, size: KdBox<C>) {
141        if self.root.is_none() {
142            let node = Node {
143                item: Some(item),
144                size,
145                lo_min_bound: size[0],
146                hi_max_bound: size[3],
147                other_bound: size[0],
148                sons: [None, None],
149            };
150            self.root = Some(self.allocate_node(node));
151            self.extent = size;
152            self.item_count = 1;
153            return;
154        }
155
156        let root_idx = self.root.unwrap();
157        if self.insert_recursive(root_idx, 0, item, &size) {
158            self.item_count += 1;
159            for i in 0..3 {
160                self.extent[i] = min(self.extent[i], size[i]);
161                self.extent[i + 3] = max(self.extent[i + 3], size[i + 3]);
162            }
163        }
164    }
165
166    fn insert_recursive(&mut self, node_idx: usize, disc: usize, item: T, size: &KdBox<C>) -> bool {
167        if let Some(ref node_item) = self.arena[node_idx].item {
168            if item == *node_item {
169                return false;
170            }
171        }
172
173        let mut val = size[disc] - self.arena[node_idx].size[disc];
174        if val == C::zero() {
175            let mut ndisc = next_disc(disc);
176            while ndisc != disc {
177                val = size[ndisc] - self.arena[node_idx].size[ndisc];
178                if val != C::zero() {
179                    break;
180                }
181                ndisc = next_disc(ndisc);
182            }
183            if val == C::zero() {
184                val = C::from_i32(1);
185            }
186        }
187
188        let child_idx = if val >= C::zero() { 1 } else { 0 };
189
190        if let Some(child_child_idx) = self.arena[node_idx].sons[child_idx] {
191            let inserted = self.insert_recursive(child_child_idx, next_disc(disc), item, size);
192            if inserted {
193                self.bounds_update(node_idx, disc, size);
194            }
195            inserted
196        } else {
197            let vert = next_disc(disc) % 3;
198            let mut new_node = Node {
199                item: Some(item),
200                size: *size,
201                lo_min_bound: size[vert],
202                hi_max_bound: size[vert + 3],
203                other_bound: C::zero(),
204                sons: [None, None],
205            };
206
207            if next_disc(disc) >= 3 {
208                new_node.other_bound = size[vert];
209            } else {
210                new_node.other_bound = size[vert + 3];
211            }
212
213            let new_idx = self.allocate_node(new_node);
214            self.arena[node_idx].sons[child_idx] = Some(new_idx);
215            self.bounds_update(node_idx, disc, size);
216            true
217        }
218    }
219
220    fn bounds_update(&mut self, node_idx: usize, disc: usize, size: &KdBox<C>) {
221        let vert = disc % 3;
222        self.arena[node_idx].lo_min_bound = min(self.arena[node_idx].lo_min_bound, size[vert]);
223        self.arena[node_idx].hi_max_bound = max(self.arena[node_idx].hi_max_bound, size[vert + 3]);
224        if disc >= 3 {
225            self.arena[node_idx].other_bound = min(self.arena[node_idx].other_bound, size[vert]);
226        } else {
227            self.arena[node_idx].other_bound = max(self.arena[node_idx].other_bound, size[vert + 3]);
228        }
229    }
230
231    fn find_recursive(&self, node_idx: Option<usize>, disc: usize, item: &T, size: &KdBox<C>) -> Option<usize> {
232        let idx = node_idx?;
233        let node = &self.arena[idx];
234
235        if let Some(ref node_item) = node.item {
236            if *item == *node_item {
237                return Some(idx);
238            }
239        }
240
241        let mut val = size[disc] - node.size[disc];
242        if val == C::zero() {
243            let mut ndisc = next_disc(disc);
244            while ndisc != disc {
245                val = size[ndisc] - node.size[ndisc];
246                if val != C::zero() {
247                    break;
248                }
249                ndisc = next_disc(ndisc);
250            }
251            if val == C::zero() {
252                val = C::from_i32(1);
253            }
254        }
255
256        let child_idx = if val >= C::zero() { 1 } else { 0 };
257        self.find_recursive(node.sons[child_idx], next_disc(disc), item, size)
258    }
259
260    /// Checks if the given item is stored in the tree with the specified bounding box.
261    pub fn is_member(&self, item: &T, size: &KdBox<C>) -> bool {
262        self.find_recursive(self.root, 0, item, size).is_some()
263    }
264
265    /// Marks an item as deleted (soft delete).
266    pub fn delete(&mut self, item: &T, size: &KdBox<C>) -> bool {
267        if let Some(node_idx) = self.find_recursive(self.root, 0, item, size) {
268            if self.arena[node_idx].item.is_some() {
269                self.arena[node_idx].item = None;
270                self.dead_count += 1;
271                return true;
272            }
273        }
274        false
275    }
276
277    /// Physically deletes an item from the tree and restructures the nodes (hard delete).
278    pub fn hard_delete(&mut self, item: &T, size: &KdBox<C>) -> bool {
279        let initial_count = self.item_count;
280        if self.root.is_none() {
281            return false;
282        }
283
284        let root_idx = self.root.unwrap();
285        self.root = self.hard_delete_recursive(Some(root_idx), 0, item, size);
286        self.item_count < initial_count
287    }
288
289    fn hard_delete_recursive(&mut self, node_idx: Option<usize>, disc: usize, item: &T, size: &KdBox<C>) -> Option<usize> {
290        let elem_idx = node_idx?;
291        
292        let is_match = match self.arena[elem_idx].item {
293            Some(ref node_item) => node_item == item,
294            None => false,
295        };
296
297        if is_match {
298            let mut stats = (0, 0); // (tries, del)
299            let q_opt = self.kd_do_delete(elem_idx, disc, &mut stats);
300            self.free_node(elem_idx);
301            self.item_count -= 1;
302            return q_opt;
303        }
304
305        let mut val = size[disc] - self.arena[elem_idx].size[disc];
306        if val == C::zero() {
307            let mut ndisc = next_disc(disc);
308            while ndisc != disc {
309                val = size[ndisc] - self.arena[elem_idx].size[ndisc];
310                if val != C::zero() {
311                    break;
312                }
313                ndisc = next_disc(ndisc);
314            }
315            if val == C::zero() {
316                val = C::from_i32(1);
317            }
318        }
319
320        let child_idx = if val >= C::zero() { 1 } else { 0 };
321        self.arena[elem_idx].sons[child_idx] = self.hard_delete_recursive(self.arena[elem_idx].sons[child_idx], next_disc(disc), item, size);
322        Some(elem_idx)
323    }
324
325    fn kd_do_delete(&mut self, elem_idx: usize, disc: usize, stats: &mut (i32, i32)) -> Option<usize> {
326        self.delete_flip = !self.delete_flip;
327
328        if self.arena[elem_idx].sons[0].is_none() && self.arena[elem_idx].sons[1].is_none() {
329            return None;
330        }
331
332        let mut q_idx: usize;
333        let mut q_dad_idx: usize = elem_idx;
334        let mut q_son: usize;
335        let mut newj: usize;
336        
337        if self.arena[elem_idx].sons[1].is_none() {
338            self.delete_flip = false;
339        } else if self.arena[elem_idx].sons[0].is_none() {
340            self.delete_flip = true;
341        }
342
343        if !self.delete_flip {
344            q_idx = self.arena[elem_idx].sons[0].unwrap();
345            q_son = 0;
346            newj = next_disc(disc);
347            stats.0 += self.find_min_max_node(disc, &mut q_idx, &mut q_dad_idx, &mut q_son, &mut newj, false);
348        } else {
349            q_idx = self.arena[elem_idx].sons[1].unwrap();
350            q_son = 1;
351            newj = next_disc(disc);
352            stats.num_tries_add(t_find_min_max_node_idx_tries(self, disc, &mut q_idx, &mut q_dad_idx, &mut q_son, &mut newj, true));
353        }
354
355        let q_replacement = self.kd_do_delete(q_idx, newj, stats);
356        self.arena[q_dad_idx].sons[q_son] = q_replacement;
357        stats.1 += 1;
358        
359        // Transfer elem data to q_idx
360        self.arena[q_idx].sons[0] = self.arena[elem_idx].sons[0];
361        self.arena[q_idx].sons[1] = self.arena[elem_idx].sons[1];
362        self.arena[q_idx].lo_min_bound = self.arena[elem_idx].lo_min_bound;
363        self.arena[q_idx].other_bound = self.arena[elem_idx].other_bound;
364        self.arena[q_idx].hi_max_bound = self.arena[elem_idx].hi_max_bound;
365        
366        Some(q_idx)
367    }
368
369    fn find_min_max_node(&self, j: usize, kd_minval_node_idx: &mut usize, kd_minval_nodesdad_idx: &mut usize, dir: &mut usize, newj: &mut usize, find_min: bool) -> i32 {
370        let mut kd_data_tries = 0;
371        let mut stack = vec![(*kd_minval_node_idx, next_disc(j), -1, 0)]; // (node_idx, m, state, dad_idx)
372
373        while let Some((node_idx, m, state, dad_idx)) = stack.pop() {
374            let node = &self.arena[node_idx];
375            match state {
376                -1 => {
377                    kd_data_tries += 1;
378                    let is_better = if find_min {
379                        node.size[j] < self.arena[*kd_minval_node_idx].size[j]
380                    } else {
381                        node.size[j] > self.arena[*kd_minval_node_idx].size[j]
382                    };
383
384                    if is_better && node_idx != *kd_minval_node_idx {
385                        *kd_minval_node_idx = node_idx;
386                        *kd_minval_nodesdad_idx = dad_idx;
387                        let dad = &self.arena[*kd_minval_nodesdad_idx];
388                        *dir = if dad.sons[0] == Some(node_idx) { 0 } else { 1 };
389                        *newj = m;
390                    }
391                    stack.push((node_idx, m, 0, dad_idx));
392                }
393                0 => {
394                    if let Some(lo) = node.sons[0] {
395                        stack.push((node_idx, m, 1, dad_idx));
396                        stack.push((lo, next_disc(m), -1, node_idx));
397                    } else {
398                        stack.push((node_idx, m, 1, dad_idx));
399                    }
400                }
401                1 => {
402                    let prune = if find_min {
403                        j == m && node.size[m] > self.arena[*kd_minval_node_idx].size[m]
404                    } else {
405                        j == m && node.size[m] < self.arena[*kd_minval_node_idx].size[m]
406                    };
407
408                    if !prune {
409                        if let Some(hi) = node.sons[1] {
410                            stack.push((hi, next_disc(m), -1, node_idx));
411                        }
412                    }
413                }
414                _ => {}
415            }
416        }
417        kd_data_tries
418    }
419
420    pub fn count(&self) -> i32 {
421        self.item_count - self.dead_count
422    }
423
424    pub fn badness(&self) {
425        let mut factor3 = 0;
426        let mut max_levels = 0;
427        
428        fn stats<T, C>(tree: &Tree<T, C>, node_idx: Option<usize>, level: i32, factor3: &mut i32, max_levels: &mut i32) {
429            if let Some(idx) = node_idx {
430                let node = &tree.arena[idx];
431                let has_lo = node.sons[0].is_some();
432                let has_hi = node.sons[1].is_some();
433                if (has_lo || has_hi) && !(has_lo && has_hi) {
434                    *factor3 += 1;
435                }
436                if level > *max_levels {
437                    *max_levels = level;
438                }
439                stats(tree, node.sons[0], level + 1, factor3, max_levels);
440                stats(tree, node.sons[1], level + 1, factor3, max_levels);
441            }
442        }
443
444        stats(self, self.root, 1, &mut factor3, &mut max_levels);
445
446        let mut targdepth = 0.0;
447        if self.item_count > 0 {
448            targdepth = (self.item_count as f64).log2().floor() + 1.0;
449        }
450
451        let ratio = if targdepth > 0.0 { (max_levels as f64) / targdepth } else { 0.0 };
452        let dead_pct = if self.item_count > 0 { (self.dead_count as f64 / self.item_count as f64) * 100.0 } else { 0.0 };
453        let factor3_pct = if self.item_count > 0 { (factor3 as f64 / self.item_count as f64) * 100.0 } else { 0.0 };
454
455        println!("balance ratio={:.1} (the closer to 1.0, the better), #of nodes with only one branch={} ({:.4}), max depth={}, dead={} ({:.4})",
456                 ratio, factor3, factor3_pct, max_levels, self.dead_count, dead_pct);
457    }
458}
459
460trait StatsHelper {
461    fn num_tries_add(&mut self, val: i32);
462}
463
464impl StatsHelper for (i32, i32) {
465    #[inline]
466    fn num_tries_add(&mut self, val: i32) {
467        self.0 += val;
468    }
469}
470
471#[inline]
472fn t_find_min_max_node_idx_tries<T: PartialEq + Clone, C: Coord>(tree: &Tree<T, C>, j: usize, kd_minval_node_idx: &mut usize, kd_minval_nodesdad_idx: &mut usize, dir: &mut usize, newj: &mut usize, find_min: bool) -> i32 {
473    tree.find_min_max_node(j, kd_minval_node_idx, kd_minval_nodesdad_idx, dir, newj, find_min)
474}
475
476pub fn next_disc(disc: usize) -> usize {
477    (disc + 1) % 6
478}
479
480#[derive(Clone, Copy, PartialEq)]
481enum GenState { ThisOne, LoSon, HiSon, Done }
482
483struct SimpleSave {
484    node_idx: usize,
485    disc: usize,
486    state: GenState,
487}
488
489pub struct Generator<'a, T, C = i32> {
490    arena: &'a Vec<Node<T, C>>,
491    extent: KdBox<C>,
492    stack: Vec<SimpleSave>,
493}
494
495impl<'a, T: 'a, C: Coord> Iterator for Generator<'a, T, C> {
496    type Item = (&'a T, KdBox<C>);
497
498    fn next(&mut self) -> Option<Self::Item> {
499        while let Some(top) = self.stack.last_mut() {
500            let node = &self.arena[top.node_idx];
501            let m = top.disc;
502            let hort = m % 3;
503
504            match top.state {
505                GenState::ThisOne => {
506                    top.state = GenState::LoSon;
507                    if let Some(ref item) = node.item {
508                        if intersect(&self.extent, &node.size) {
509                            return Some((item, node.size));
510                        }
511                    }
512                }
513                GenState::LoSon => {
514                    top.state = GenState::HiSon;
515                    if let Some(child_idx) = node.sons[0] {
516                        let mut should_push = false;
517                        if m >= 3 {
518                            if self.extent[hort] <= node.size[m] && self.extent[hort + 3] >= node.lo_min_bound {
519                                should_push = true;
520                            }
521                        } else {
522                            if self.extent[hort] <= node.other_bound && self.extent[hort + 3] >= node.lo_min_bound {
523                                should_push = true;
524                            }
525                        }
526                        if should_push {
527                            self.stack.push(SimpleSave {
528                                node_idx: child_idx,
529                                disc: next_disc(m),
530                                state: GenState::ThisOne,
531                            });
532                            continue;
533                        }
534                    }
535                }
536                GenState::HiSon => {
537                    top.state = GenState::Done;
538                    if let Some(child_idx) = node.sons[1] {
539                        let mut should_push = false;
540                        if m >= 3 {
541                            if self.extent[hort] <= node.hi_max_bound && self.extent[hort + 3] >= node.other_bound {
542                                should_push = true;
543                            }
544                        } else {
545                            if self.extent[hort] <= node.hi_max_bound && self.extent[hort + 3] >= node.size[m] {
546                                should_push = true;
547                            }
548                        }
549                        if should_push {
550                            self.stack.push(SimpleSave {
551                                node_idx: child_idx,
552                                disc: next_disc(m),
553                                state: GenState::ThisOne,
554                            });
555                            continue;
556                        }
557                    }
558                }
559                GenState::Done => {
560                    self.stack.pop();
561                }
562            }
563        }
564        None
565    }
566}
567
568impl<T: PartialEq + Clone, C: Coord> Tree<T, C> {
569    /// Returns an iterator over all items that intersect with the given 3D bounding box.
570    pub fn start(&self, area: KdBox<C>) -> Generator<'_, T, C> {
571        let mut stack = Vec::new();
572        if let Some(root_idx) = self.root {
573            stack.push(SimpleSave {
574                node_idx: root_idx,
575                disc: 0,
576                state: GenState::ThisOne,
577            });
578        }
579        Generator {
580            arena: &self.arena,
581            extent: area,
582            stack,
583        }
584    }
585
586    /// Finds the `m` nearest neighbors to the point `(x, y, z)`.
587    pub fn nearest(&self, x: C, y: C, z: C, m: usize) -> Vec<Priority<T>> {
588        if self.root.is_none() || m == 0 {
589            return Vec::new();
590        }
591
592        let mut list = vec![Priority { dist: f64::MAX, item: None }; m];
593        let xq = [x, y, z, x, y, z];
594        let bp = [C::max_value(); 6];
595        let bn = [C::min_value(); 6];
596
597        self.kd_neighbor(self.root.unwrap(), &xq, m, &mut list, bp, bn);
598
599        for p in &mut list {
600            if p.dist != f64::MAX {
601                p.dist = p.dist.sqrt();
602            }
603        }
604        list
605    }
606
607    fn kd_neighbor(&self, node_idx: usize, xq: &KdBox<C>, m: usize, list: &mut [Priority<T>], bp: KdBox<C>, bn: KdBox<C>) {
608        let mut stack = Vec::new();
609        stack.push(Save { node_idx, disc: 0, state: State::ThisOne, bn, bp });
610
611        while let Some(top) = stack.last_mut() {
612            let node = &self.arena[top.node_idx];
613            let d = top.disc;
614            let p = node.size[d];
615            let hort = d % 3;
616            let vert = d >= 3;
617
618            match top.state {
619                State::ThisOne => {
620                    top.state = State::LoSon;
621                    if let Some(ref item) = node.item {
622                        self.add_priority(m, list, xq, item, &node.size);
623                    }
624                }
625                State::LoSon => {
626                    top.state = State::HiSon;
627                    if xq[d] <= p {
628                        if let Some(child_idx) = node.sons[0] {
629                            let old_bn = top.bn[hort];
630                            let old_bp = top.bp[hort];
631                            if vert {
632                                top.bp[hort] = node.size[d];
633                                top.bn[hort] = node.lo_min_bound;
634                            } else {
635                                top.bp[hort] = node.other_bound;
636                                top.bn[hort] = node.lo_min_bound;
637                            }
638                            if self.bounds_overlap_ball(xq, &top.bp, &top.bn, m, list) {
639                                let (bn, bp) = (top.bn, top.bp);
640                                stack.push(Save { node_idx: child_idx, disc: next_disc(d), state: State::ThisOne, bn, bp });
641                                let last = stack.len() - 2;
642                                stack[last].bn[hort] = old_bn;
643                                stack[last].bp[hort] = old_bp;
644                                continue;
645                            }
646                            top.bn[hort] = old_bn;
647                            top.bp[hort] = old_bp;
648                        }
649                    } else {
650                        if let Some(child_idx) = node.sons[1] {
651                            let old_bn = top.bn[hort];
652                            let old_bp = top.bp[hort];
653                            if vert {
654                                top.bp[hort] = node.hi_max_bound;
655                                top.bn[hort] = node.other_bound;
656                            } else {
657                                top.bp[hort] = node.hi_max_bound;
658                                top.bn[hort] = node.size[d];
659                            }
660                            if self.bounds_overlap_ball(xq, &top.bp, &top.bn, m, list) {
661                                let (bn, bp) = (top.bn, top.bp);
662                                stack.push(Save { node_idx: child_idx, disc: next_disc(d), state: State::ThisOne, bn, bp });
663                                let last = stack.len() - 2;
664                                stack[last].bn[hort] = old_bn;
665                                stack[last].bp[hort] = old_bp;
666                                continue;
667                            }
668                            top.bn[hort] = old_bn;
669                            top.bp[hort] = old_bp;
670                        }
671                    }
672                }
673                State::HiSon => {
674                    top.state = State::Done;
675                    if xq[d] <= p {
676                        if let Some(child_idx) = node.sons[1] {
677                            let old_bn = top.bn[hort];
678                            let old_bp = top.bp[hort];
679                            if vert {
680                                top.bp[hort] = node.hi_max_bound;
681                                top.bn[hort] = node.other_bound;
682                            } else {
683                                top.bp[hort] = node.hi_max_bound;
684                                top.bn[hort] = node.size[d];
685                            }
686                            if self.bounds_overlap_ball(xq, &top.bp, &top.bn, m, list) {
687                                let (bn, bp) = (top.bn, top.bp);
688                                stack.push(Save { node_idx: child_idx, disc: next_disc(d), state: State::ThisOne, bn, bp });
689                                let last = stack.len() - 2;
690                                stack[last].bn[hort] = old_bn;
691                                stack[last].bp[hort] = old_bp;
692                                continue;
693                            }
694                            top.bn[hort] = old_bn;
695                            top.bp[hort] = old_bp;
696                        }
697                    } else {
698                        if let Some(child_idx) = node.sons[0] {
699                            let old_bn = top.bn[hort];
700                            let old_bp = top.bp[hort];
701                            if vert {
702                                top.bp[hort] = node.size[d];
703                                top.bn[hort] = node.lo_min_bound;
704                            } else {
705                                top.bp[hort] = node.other_bound;
706                                top.bn[hort] = node.lo_min_bound;
707                            }
708                            if self.bounds_overlap_ball(xq, &top.bp, &top.bn, m, list) {
709                                let (bn, bp) = (top.bn, top.bp);
710                                stack.push(Save { node_idx: child_idx, disc: next_disc(d), state: State::ThisOne, bn, bp });
711                                let last = stack.len() - 2;
712                                stack[last].bn[hort] = old_bn;
713                                stack[last].bp[hort] = old_bp;
714                                continue;
715                            }
716                            top.bn[hort] = old_bn;
717                            top.bp[hort] = old_bp;
718                        }
719                    }
720                }
721                State::Done => {
722                    stack.pop();
723                }
724            }
725        }
726    }
727
728    fn add_priority(&self, m: usize, list: &mut [Priority<T>], xq: &KdBox<C>, item: &T, size: &KdBox<C>) {
729        let d = kd_dist_sq(xq, size);
730        for x in (0..m).rev() {
731            if d < list[x].dist {
732                if x != m - 1 {
733                    list[x + 1] = list[x].clone();
734                }
735                list[x].dist = d;
736                list[x].item = Some(item.clone());
737            } else {
738                break;
739            }
740        }
741    }
742
743    fn bounds_overlap_ball(&self, xq: &KdBox<C>, bp: &KdBox<C>, bn: &KdBox<C>, m: usize, list: &[Priority<T>]) -> bool {
744        let mut sum = 0.0;
745        let max_dist = list[m - 1].dist;
746        for i in 0..3 {
747            if xq[i] < bn[i] {
748                let d = (xq[i] - bn[i]).to_f64();
749                sum += d * d;
750                if sum > max_dist {
751                    return false;
752                }
753            } else if xq[i] > bp[i] {
754                let d = (xq[i] - bp[i]).to_f64();
755                sum += d * d;
756                if sum > max_dist {
757                    return false;
758                }
759            }
760        }
761        true
762    }
763
764    /// Deletes structurally an item from the tree (really delete).
765    pub fn really_delete(&mut self, data: &T, old_size: &KdBox<C>) -> (Status, i32, i32) {
766        let mut path = Vec::new();
767        let elem_opt = self.find_item_with_path(self.root, 0, data, old_size, &mut path);
768        if elem_opt.is_none() {
769            return (Status::NotFound, 0, 0);
770        }
771
772        let elem_idx = elem_opt.unwrap();
773        let mut stats = (0, 1);
774
775        if Some(elem_idx) == self.root {
776            self.root = self.kd_do_delete(elem_idx, 0, &mut stats);
777        } else {
778            let parent_idx = path[path.len() - 1];
779            let j = path.len() % 6;
780            let new_elem = self.kd_do_delete(elem_idx, j, &mut stats);
781            let parent = &mut self.arena[parent_idx];
782            if parent.sons[1] == Some(elem_idx) {
783                parent.sons[1] = new_elem;
784            } else {
785                parent.sons[0] = new_elem;
786            }
787        }
788
789        self.item_count -= 1;
790        (Status::Ok, stats.0, stats.1)
791    }
792
793    fn find_item_with_path(&self, node_idx: Option<usize>, disc: usize, item: &T, size: &KdBox<C>, path: &mut Vec<usize>) -> Option<usize> {
794        let idx = node_idx?;
795        let node = &self.arena[idx];
796
797        if let Some(ref node_item) = node.item {
798            if *item == *node_item {
799                return Some(idx);
800            }
801        }
802
803        let mut val = size[disc] - node.size[disc];
804        if val == C::zero() {
805            let mut ndisc = next_disc(disc);
806            while ndisc != disc {
807                val = size[ndisc] - node.size[ndisc];
808                if val != C::zero() {
809                    break;
810                }
811                ndisc = next_disc(ndisc);
812            }
813            if val == C::zero() {
814                val = C::from_i32(1);
815            }
816        }
817
818        let child_idx = if val >= C::zero() { 1 } else { 0 };
819
820        if let Some(child_node_idx) = node.sons[child_idx] {
821            path.push(idx);
822            return self.find_item_with_path(Some(child_node_idx), next_disc(disc), item, size, path);
823        }
824
825        None
826    }
827}
828
829pub fn intersect<C: Coord>(b1: &KdBox<C>, b2: &KdBox<C>) -> bool {
830    b1[RIGHT] >= b2[LEFT] &&
831    b2[RIGHT] >= b1[LEFT] &&
832    b1[TOP] >= b2[BOTTOM] &&
833    b2[TOP] >= b1[BOTTOM] &&
834    b1[CEIL] >= b2[FLOOR] &&
835    b2[CEIL] >= b1[FLOOR]
836}
837
838pub fn kd_dist_sq<C: Coord>(xq: &KdBox<C>, box_val: &KdBox<C>) -> f64 {
839    let mut dx = 0.0;
840    let mut dy = 0.0;
841    let mut dz = 0.0;
842
843    if xq[LEFT] > box_val[RIGHT] {
844        dx = (xq[LEFT] - box_val[RIGHT]).to_f64();
845    } else if xq[RIGHT] < box_val[LEFT] {
846        dx = (box_val[LEFT] - xq[RIGHT]).to_f64();
847    }
848
849    if xq[BOTTOM] > box_val[TOP] {
850        dy = (xq[BOTTOM] - box_val[TOP]).to_f64();
851    } else if xq[TOP] < box_val[BOTTOM] {
852        dy = (box_val[BOTTOM] - xq[TOP]).to_f64();
853    }
854
855    if xq[FLOOR] > box_val[CEIL] {
856        dz = (xq[FLOOR] - box_val[CEIL]).to_f64();
857    } else if xq[CEIL] < box_val[FLOOR] {
858        dz = (box_val[FLOOR] - xq[CEIL]).to_f64();
859    }
860
861    dx * dx + dy * dy + dz * dz
862}
863
864#[derive(Clone)]
865pub struct Priority<T> {
866    pub dist: f64,
867    pub item: Option<T>,
868}
869
870#[cfg(test)]
871struct Lcg {
872    state: u32,
873}
874
875#[cfg(test)]
876impl Lcg {
877    fn next(&mut self) -> i32 {
878        self.state = self.state.wrapping_mul(1664525).wrapping_add(1013904223);
879        (self.state >> 16) as i32
880    }
881
882    fn next_range(&mut self, max: i32) -> i32 {
883        self.next().rem_euclid(max)
884    }
885}
886
887#[cfg(test)]
888mod tests {
889    use super::*;
890    use super::Lcg;
891
892    macro_rules! generate_tests {
893        ($t:ty, $mod_name:ident) => {
894            mod $mod_name {
895                use super::*;
896
897                const KD_BOXES: usize = 10000;
898                const MIN_RANGE: i32 = -100000;
899                const MAX_RANGE: i32 = 100000;
900                const RANGE_SPAN: i32 = MAX_RANGE - MIN_RANGE + 1;
901                const BOX_RANGE: i32 = 1000;
902
903                fn rand_box(rng: &mut Lcg) -> KdBox<$t> {
904                    let left = rng.next_range(RANGE_SPAN) + MIN_RANGE;
905                    let bottom = rng.next_range(RANGE_SPAN) + MIN_RANGE;
906                    let floor = rng.next_range(RANGE_SPAN) + MIN_RANGE;
907                    let right = left + rng.next_range(BOX_RANGE);
908                    let top = bottom + rng.next_range(BOX_RANGE);
909                    let ceil = floor + rng.next_range(BOX_RANGE);
910                    [
911                        <$t>::from_i32(left),
912                        <$t>::from_i32(bottom),
913                        <$t>::from_i32(floor),
914                        <$t>::from_i32(right),
915                        <$t>::from_i32(top),
916                        <$t>::from_i32(ceil),
917                    ]
918                }
919
920                #[test]
921                fn test_kd_3d() {
922                    let mut tree = Tree::<&str, $t>::new();
923                    let box1: KdBox<$t> = [<$t>::from_i32(0), <$t>::from_i32(0), <$t>::from_i32(0), <$t>::from_i32(10), <$t>::from_i32(10), <$t>::from_i32(10)];
924                    let box2: KdBox<$t> = [<$t>::from_i32(20), <$t>::from_i32(20), <$t>::from_i32(20), <$t>::from_i32(30), <$t>::from_i32(30), <$t>::from_i32(30)];
925                    let box3: KdBox<$t> = [<$t>::from_i32(5), <$t>::from_i32(5), <$t>::from_i32(5), <$t>::from_i32(15), <$t>::from_i32(15), <$t>::from_i32(15)];
926
927                    tree.insert("item1", box1);
928                    tree.insert("item2", box2);
929                    tree.insert("item3", box3);
930
931                    assert_eq!(tree.count(), 3);
932                    assert!(tree.is_member(&"item2", &box2));
933                }
934
935                #[test]
936                fn test_nearest() {
937                    let mut rng = Lcg { state: 42 };
938                    let mut boxes = Vec::new();
939                    let mut tree = Tree::<usize, $t>::new();
940
941                    for i in 0..KD_BOXES {
942                        let b = rand_box(&mut rng);
943                        boxes.push(b);
944                        tree.insert(i, b);
945                    }
946
947                    for m in [1, 2, 4, 8, 16] {
948                        for _ in 0..50 {
949                            let qx = <$t>::from_i32(rng.next_range(RANGE_SPAN) + MIN_RANGE);
950                            let qy = <$t>::from_i32(rng.next_range(RANGE_SPAN) + MIN_RANGE);
951                            let qz = <$t>::from_i32(rng.next_range(RANGE_SPAN) + MIN_RANGE);
952
953                            let list = tree.nearest(qx, qy, qz, m);
954                            assert_eq!(list.len(), m);
955
956                            for i in 1..m {
957                                assert!(list[i].dist >= list[i - 1].dist - 1e-9);
958                            }
959
960                            let mut brute: Vec<f64> = boxes.iter()
961                                .map(|b| kd_dist_sq(&[qx, qy, qz, qx, qy, qz], b).sqrt())
962                                .collect();
963                            brute.sort_by(|a, b| a.partial_cmp(b).unwrap());
964
965                            assert!(list[m - 1].dist <= brute[m - 1] + 1e-6);
966                        }
967                    }
968                }
969
970                #[test]
971                fn test_kd_tree_hard_delete() {
972                    let mut tree = Tree::<&str, $t>::new();
973                    let box1: KdBox<$t> = [<$t>::from_i32(0), <$t>::from_i32(0), <$t>::from_i32(0), <$t>::from_i32(10), <$t>::from_i32(10), <$t>::from_i32(10)];
974                    let box2: KdBox<$t> = [<$t>::from_i32(20), <$t>::from_i32(20), <$t>::from_i32(20), <$t>::from_i32(30), <$t>::from_i32(30), <$t>::from_i32(30)];
975                    let box3: KdBox<$t> = [<$t>::from_i32(5), <$t>::from_i32(5), <$t>::from_i32(5), <$t>::from_i32(15), <$t>::from_i32(15), <$t>::from_i32(15)];
976
977                    tree.insert("item1", box1);
978                    tree.insert("item2", box2);
979                    tree.insert("item3", box3);
980
981                    assert!(tree.hard_delete(&"item1", &box1));
982                    assert_eq!(tree.count(), 2);
983                    assert!(tree.is_member(&"item2", &box2));
984                    assert!(tree.is_member(&"item3", &box3));
985                }
986
987                #[test]
988                fn test_million_boxes() {
989                    let mut tree = Tree::<String, $t>::new();
990                    let mut rng = Lcg { state: 42 };
991                    let mut boxes_to_delete = Vec::new();
992
993                    for i in 0..100_000 { // Reduced to 100k to keep test suite fast
994                        let x1 = rng.next_range(100000);
995                        let y1 = rng.next_range(100000);
996                        let z1 = rng.next_range(100000);
997                        let x2 = x1 + rng.next_range(100) + 1;
998                        let y2 = y1 + rng.next_range(100) + 1;
999                        let z2 = z1 + rng.next_range(100) + 1;
1000                        let b: KdBox<$t> = [<$t>::from_i32(x1), <$t>::from_i32(y1), <$t>::from_i32(z1), <$t>::from_i32(x2), <$t>::from_i32(y2), <$t>::from_i32(z2)];
1001                        
1002                        if i < 1000 {
1003                            boxes_to_delete.push(b);
1004                        }
1005                        tree.insert(format!("box{}", i), b);
1006                    }
1007
1008                    assert_eq!(tree.count(), 100_000);
1009
1010                    let search_area: KdBox<$t> = [<$t>::from_i32(0), <$t>::from_i32(0), <$t>::from_i32(0), <$t>::from_i32(50000), <$t>::from_i32(50000), <$t>::from_i32(50000)];
1011                    let mut found_count = 0;
1012                    for _ in tree.start(search_area) {
1013                        found_count += 1;
1014                    }
1015                    assert!(found_count > 100);
1016
1017                    for i in 0..1000 {
1018                        let item_name = format!("box{}", i);
1019                        let deleted = tree.hard_delete(&item_name, &boxes_to_delete[i]);
1020                        assert!(deleted, "Failed to hard delete box{}", i);
1021                    }
1022
1023                    assert_eq!(tree.count(), 99_000);
1024                }
1025            }
1026        };
1027    }
1028
1029    generate_tests!(i32, tests_i32);
1030    generate_tests!(i64, tests_i64);
1031    generate_tests!(i128, tests_i128);
1032}
kdtree3d/lib.rs

kdtree3d/
lib.rs
