1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
use geo_traits::{CoordTrait, RectTrait};
use tinyvec::TinyVec;
use crate::indices::Indices;
use crate::kdtree::{KDTree, KDTreeMetadata, KDTreeRef, Node};
use crate::r#type::IndexableNum;
/// A trait for searching and accessing data out of a KDTree.
pub trait KDTreeIndex<N: IndexableNum>: Sized {
/// The underlying raw coordinate buffer of this tree
fn coords(&self) -> &[N];
/// The underlying raw indices buffer of this tree
fn indices(&self) -> Indices<'_>;
/// Access the metadata describing this KDTree
fn metadata(&self) -> &KDTreeMetadata<N>;
/// The number of items in this KDTree
fn num_items(&self) -> u32 {
self.metadata().num_items()
}
/// The node size of this KDTree
fn node_size(&self) -> u16 {
self.metadata().node_size()
}
/// Search the index for items within a given bounding box.
///
/// - min_x: bbox
/// - min_y: bbox
/// - max_x: bbox
/// - max_y: bbox
///
/// Returns indices of found items
fn range(&self, min_x: N, min_y: N, max_x: N, max_y: N) -> Vec<u32> {
let indices = self.indices();
let coords = self.coords();
let node_size = self.node_size();
// Use TinyVec to avoid heap allocations
let mut stack: TinyVec<[usize; 33]> = TinyVec::new();
stack.push(0);
stack.push(indices.len() - 1);
stack.push(0);
let mut result: Vec<u32> = vec![];
// recursively search for items in range in the kd-sorted arrays
while !stack.is_empty() {
let axis = stack.pop().unwrap_or(0);
let right = stack.pop().unwrap_or(0);
let left = stack.pop().unwrap_or(0);
// if we reached "tree node", search linearly
if right - left <= node_size as usize {
for i in left..=right {
let x = coords[2 * i];
let y = coords[2 * i + 1];
if x >= min_x && x <= max_x && y >= min_y && y <= max_y {
result.push(indices.get(i).try_into().unwrap());
}
}
continue;
}
// otherwise find the middle index
let m = (left + right) >> 1;
// include the middle item if it's in range
let x = coords[2 * m];
let y = coords[2 * m + 1];
if x >= min_x && x <= max_x && y >= min_y && y <= max_y {
result.push(indices.get(m).try_into().unwrap());
}
// queue search in halves that intersect the query
let lte = if axis == 0 { min_x <= x } else { min_y <= y };
if lte {
// Note: these are pushed in backwards order to what gets popped
stack.push(left);
stack.push(m - 1);
stack.push(1 - axis);
}
let gte = if axis == 0 { max_x >= x } else { max_y >= y };
if gte {
// Note: these are pushed in backwards order to what gets popped
stack.push(m + 1);
stack.push(right);
stack.push(1 - axis);
}
}
result
}
/// Search the index for items within a given bounding box.
///
/// Returns indices of found items
fn range_rect(&self, rect: &impl RectTrait<T = N>) -> Vec<u32> {
self.range(
rect.min().x(),
rect.min().y(),
rect.max().x(),
rect.max().y(),
)
}
/// Search the index for items within a given radius.
///
/// - qx: x value of query point
/// - qy: y value of query point
/// - r: radius
///
/// Returns indices of found items
fn within(&self, qx: N, qy: N, r: N) -> Vec<u32> {
let indices = self.indices();
let coords = self.coords();
let node_size = self.node_size();
// Use TinyVec to avoid heap allocations
let mut stack: TinyVec<[usize; 33]> = TinyVec::new();
stack.push(0);
stack.push(indices.len() - 1);
stack.push(0);
let mut result: Vec<u32> = vec![];
let r2 = r * r;
// recursively search for items within radius in the kd-sorted arrays
while !stack.is_empty() {
let axis = stack.pop().unwrap_or(0);
let right = stack.pop().unwrap_or(0);
let left = stack.pop().unwrap_or(0);
// if we reached "tree node", search linearly
if right - left <= node_size as usize {
for i in left..=right {
if sq_dist(coords[2 * i], coords[2 * i + 1], qx, qy) <= r2 {
result.push(indices.get(i).try_into().unwrap());
}
}
continue;
}
// otherwise find the middle index
let m = (left + right) >> 1;
// include the middle item if it's in range
let x = coords[2 * m];
let y = coords[2 * m + 1];
if sq_dist(x, y, qx, qy) <= r2 {
result.push(indices.get(m).try_into().unwrap());
}
// queue search in halves that intersect the query
let lte = if axis == 0 { qx - r <= x } else { qy - r <= y };
if lte {
stack.push(left);
stack.push(m - 1);
stack.push(1 - axis);
}
let gte = if axis == 0 { qx + r >= x } else { qy + r >= y };
if gte {
stack.push(m + 1);
stack.push(right);
stack.push(1 - axis);
}
}
result
}
/// Search the index for items within a given radius.
///
/// - coord: coordinate of query point
/// - r: radius
///
/// Returns indices of found items
fn within_coord(&self, coord: &impl CoordTrait<T = N>, r: N) -> Vec<u32> {
self.within(coord.x(), coord.y(), r)
}
/// Access the root node of the KDTree for manual traversal.
fn root(&self) -> Node<'_, N, Self> {
Node::from_root(self)
}
}
impl<N: IndexableNum> KDTreeIndex<N> for KDTree<N> {
fn coords(&self) -> &[N] {
self.metadata.coords_slice(&self.buffer)
}
fn indices(&self) -> Indices<'_> {
self.metadata.indices_slice(&self.buffer)
}
fn metadata(&self) -> &KDTreeMetadata<N> {
&self.metadata
}
}
impl<N: IndexableNum> KDTreeIndex<N> for KDTreeRef<'_, N> {
fn coords(&self) -> &[N] {
self.coords
}
fn indices(&self) -> Indices<'_> {
self.indices
}
fn metadata(&self) -> &KDTreeMetadata<N> {
&self.metadata
}
}
#[inline]
pub(crate) fn sq_dist<N: IndexableNum>(ax: N, ay: N, bx: N, by: N) -> N {
let dx = ax - bx;
let dy = ay - by;
dx * dx + dy * dy
}