1use crate::types::{PolyVert, Rect};
18
19fn build_two_d_tree_impl(points: &mut [PolyVert], sort_x: bool) {
20 if sort_x {
21 points.sort_by(|a, b| a.pos.x.partial_cmp(&b.pos.x).unwrap_or(std::cmp::Ordering::Equal));
22 } else {
23 points.sort_by(|a, b| a.pos.y.partial_cmp(&b.pos.y).unwrap_or(std::cmp::Ordering::Equal));
24 }
25 if points.len() < 2 {
26 return;
27 }
28
29 let mid = points.len() / 2;
30 build_two_d_tree_impl(&mut points[..mid], !sort_x);
31 if mid + 1 < points.len() {
32 build_two_d_tree_impl(&mut points[mid + 1..], !sort_x);
33 }
34}
35
36pub fn build_two_d_tree(points: &mut [PolyVert]) {
37 if points.len() <= 8 {
38 return;
39 }
40 build_two_d_tree_impl(points, true);
41}
42
43pub fn query_two_d_tree<F: FnMut(PolyVert)>(points: &[PolyVert], rect: Rect, mut f: F) {
44 if points.len() <= 8 {
45 for &p in points {
46 if rect.contains_point(p.pos) {
47 f(p);
48 }
49 }
50 return;
51 }
52
53 let mut current = Rect::new();
54 current.min = crate::linalg::Vec2::splat(f64::NEG_INFINITY);
55 current.max = crate::linalg::Vec2::splat(f64::INFINITY);
56
57 let mut level = 0usize;
58 let mut current_range = (0usize, points.len());
59 let mut rect_stack: Vec<Rect> = Vec::new();
60 let mut range_stack: Vec<(usize, usize)> = Vec::new();
61 let mut level_stack: Vec<usize> = Vec::new();
62
63 loop {
64 let current_view = &points[current_range.0..current_range.1];
65 if current_view.len() <= 8 {
66 for &p in current_view {
67 if rect.contains_point(p.pos) {
68 f(p);
69 }
70 }
71 if let Some(next_level) = level_stack.pop() {
72 level = next_level;
73 current_range = range_stack.pop().unwrap();
74 current = rect_stack.pop().unwrap();
75 continue;
76 }
77 break;
78 }
79
80 let mut left = current;
81 let mut right = current;
82 let mid = current_view.len() / 2;
83 let middle = current_view[mid];
84 if level % 2 == 0 {
85 left.max.x = middle.pos.x;
86 right.min.x = middle.pos.x;
87 } else {
88 left.max.y = middle.pos.y;
89 right.min.y = middle.pos.y;
90 }
91
92 if rect.contains_point(middle.pos) {
93 f(middle);
94 }
95
96 let left_overlaps = left.does_overlap(&rect);
97 let right_overlaps = right.does_overlap(&rect);
98 if left_overlaps {
99 if right_overlaps {
100 rect_stack.push(right);
101 range_stack.push((current_range.0 + mid + 1, current_range.1));
102 level_stack.push(level + 1);
103 }
104 current = left;
105 current_range = (current_range.0, current_range.0 + mid);
106 level += 1;
107 } else if right_overlaps {
108 current = right;
109 current_range = (current_range.0 + mid + 1, current_range.1);
110 level += 1;
111 } else if let Some(next_level) = level_stack.pop() {
112 level = next_level;
113 current_range = range_stack.pop().unwrap();
114 current = rect_stack.pop().unwrap();
115 } else {
116 break;
117 }
118 }
119}
120
121#[cfg(test)]
122mod tests {
123 use super::*;
124 use crate::linalg::Vec2;
125
126 #[test]
127 fn test_build_two_d_tree_and_query() {
128 let mut points = vec![
129 PolyVert { pos: Vec2::new(0.0, 0.0), idx: 0 },
130 PolyVert { pos: Vec2::new(1.0, 1.0), idx: 1 },
131 PolyVert { pos: Vec2::new(2.0, 2.0), idx: 2 },
132 PolyVert { pos: Vec2::new(3.0, 3.0), idx: 3 },
133 PolyVert { pos: Vec2::new(4.0, 4.0), idx: 4 },
134 PolyVert { pos: Vec2::new(5.0, 5.0), idx: 5 },
135 PolyVert { pos: Vec2::new(6.0, 6.0), idx: 6 },
136 PolyVert { pos: Vec2::new(7.0, 7.0), idx: 7 },
137 PolyVert { pos: Vec2::new(8.0, 8.0), idx: 8 },
138 ];
139 build_two_d_tree(&mut points);
140 let rect = Rect::from_points(Vec2::new(1.5, 1.5), Vec2::new(6.5, 6.5));
141 let mut out = Vec::new();
142 query_two_d_tree(&points, rect, |p| out.push(p.idx));
143 out.sort_unstable();
144 assert_eq!(out, vec![2, 3, 4, 5, 6]);
145 }
146
147 #[test]
148 fn test_query_two_d_tree_small_input() {
149 let points = vec![
150 PolyVert { pos: Vec2::new(0.0, 0.0), idx: 0 },
151 PolyVert { pos: Vec2::new(2.0, 2.0), idx: 1 },
152 ];
153 let rect = Rect::from_points(Vec2::new(-1.0, -1.0), Vec2::new(1.0, 1.0));
154 let mut out = Vec::new();
155 query_two_d_tree(&points, rect, |p| out.push(p.idx));
156 assert_eq!(out, vec![0]);
157 }
158}