1use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
11use core::f64;
12pub use error::Error;
13#[cfg(feature = "http")]
14use http_range_client::{AsyncBufferedHttpRangeClient, AsyncHttpRangeClient};
15use std::cmp::min;
16use std::collections::VecDeque;
17use std::io::{Cursor, Read, Seek, SeekFrom, Write};
18use std::mem::size_of;
19use std::ops::Range;
20
21mod error;
22
23#[derive(Clone, PartialEq, Debug)]
24#[repr(C)]
25pub struct NodeItem {
27 pub min_x: f64,
28 pub min_y: f64,
29 pub max_x: f64,
30 pub max_y: f64,
31 pub offset: u64,
33}
34
35impl NodeItem {
36 pub fn bounds(min_x: f64, min_y: f64, max_x: f64, max_y: f64) -> NodeItem {
37 NodeItem {
38 min_x,
39 min_y,
40 max_x,
41 max_y,
42 offset: 0,
43 }
44 }
45
46 pub fn create(offset: u64) -> NodeItem {
47 NodeItem {
48 min_x: f64::INFINITY,
49 min_y: f64::INFINITY,
50 max_x: f64::NEG_INFINITY,
51 max_y: f64::NEG_INFINITY,
52 offset,
53 }
54 }
55
56 pub fn from_reader(mut rdr: impl Read) -> Result<Self, Error> {
57 Ok(NodeItem {
58 min_x: rdr.read_f64::<LittleEndian>()?,
59 min_y: rdr.read_f64::<LittleEndian>()?,
60 max_x: rdr.read_f64::<LittleEndian>()?,
61 max_y: rdr.read_f64::<LittleEndian>()?,
62 offset: rdr.read_u64::<LittleEndian>()?,
63 })
64 }
65
66 fn from_bytes(raw: &[u8]) -> Result<Self, Error> {
67 Self::from_reader(&mut Cursor::new(raw))
68 }
69
70 pub fn write<W: Write>(&self, wtr: &mut W) -> std::io::Result<()> {
71 wtr.write_f64::<LittleEndian>(self.min_x)?;
72 wtr.write_f64::<LittleEndian>(self.min_y)?;
73 wtr.write_f64::<LittleEndian>(self.max_x)?;
74 wtr.write_f64::<LittleEndian>(self.max_y)?;
75 wtr.write_u64::<LittleEndian>(self.offset)?;
76 Ok(())
77 }
78
79 pub fn width(&self) -> f64 {
80 self.max_x - self.min_x
81 }
82
83 pub fn height(&self) -> f64 {
84 self.max_y - self.min_y
85 }
86
87 pub fn sum(mut a: NodeItem, b: &NodeItem) -> NodeItem {
88 a.expand(b);
89 a
90 }
91
92 pub fn expand(&mut self, r: &NodeItem) {
93 if r.min_x < self.min_x {
94 self.min_x = r.min_x;
95 }
96 if r.min_y < self.min_y {
97 self.min_y = r.min_y;
98 }
99 if r.max_x > self.max_x {
100 self.max_x = r.max_x;
101 }
102 if r.max_y > self.max_y {
103 self.max_y = r.max_y;
104 }
105 }
106
107 pub fn expand_xy(&mut self, x: f64, y: f64) {
108 if x < self.min_x {
109 self.min_x = x;
110 }
111 if y < self.min_y {
112 self.min_y = y;
113 }
114 if x > self.max_x {
115 self.max_x = x;
116 }
117 if y > self.max_y {
118 self.max_y = y;
119 }
120 }
121
122 pub fn intersects(&self, r: &NodeItem) -> bool {
123 if self.max_x < r.min_x {
124 return false;
125 }
126 if self.max_y < r.min_y {
127 return false;
128 }
129 if self.min_x > r.max_x {
130 return false;
131 }
132 if self.min_y > r.max_y {
133 return false;
134 }
135 true
136 }
137
138 pub fn contains_point(&self, x: f64, y: f64) -> bool {
140 x >= self.min_x && x <= self.max_x && y >= self.min_y && y <= self.max_y
141 }
142
143 pub fn centroid_distance_squared(&self, x: f64, y: f64) -> f64 {
145 let centroid_x = (self.min_x + self.max_x) / 2.0;
146 let centroid_y = (self.min_y + self.max_y) / 2.0;
147 let dx = x - centroid_x;
148 let dy = y - centroid_y;
149 dx * dx + dy * dy
150 }
151
152 pub fn min_distance_squared(&self, x: f64, y: f64) -> f64 {
155 if self.contains_point(x, y) {
156 return 0.0;
157 }
158
159 let closest_x = x.clamp(self.min_x, self.max_x);
161 let closest_y = y.clamp(self.min_y, self.max_y);
162
163 let dx = x - closest_x;
165 let dy = y - closest_y;
166 dx * dx + dy * dy
167 }
168}
169
170fn read_node_vec(node_items: &mut Vec<NodeItem>, mut data: impl Read) -> Result<(), Error> {
172 node_items.clear();
173 for _ in 0..node_items.capacity() {
174 node_items.push(NodeItem::from_reader(&mut data)?);
175 }
176 Ok(())
177}
178
179fn read_node_items<R: Read + Seek>(
181 data: &mut R,
182 base: u64,
183 node_index: usize,
184 length: usize,
185) -> Result<Vec<NodeItem>, Error> {
186 let mut node_items = Vec::with_capacity(length);
187 data.seek(SeekFrom::Start(
188 base + (node_index * size_of::<NodeItem>()) as u64,
189 ))?;
190 read_node_vec(&mut node_items, data)?;
191 Ok(node_items)
192}
193
194#[cfg(feature = "http")]
196async fn read_http_node_items<T: AsyncHttpRangeClient>(
197 client: &mut AsyncBufferedHttpRangeClient<T>,
198 base: usize,
199 node_ids: &Range<usize>,
200) -> Result<Vec<NodeItem>, Error> {
201 let begin = base + node_ids.start * size_of::<NodeItem>();
202 let length = node_ids.len() * size_of::<NodeItem>();
203 let bytes = client
204 .min_req_size(0)
206 .get_range(begin, length)
207 .await?;
208
209 let mut node_items = Vec::with_capacity(node_ids.len());
210 debug_assert_eq!(bytes.len(), length);
211 for node_item_bytes in bytes.chunks(size_of::<NodeItem>()) {
212 node_items.push(NodeItem::from_bytes(node_item_bytes)?);
213 }
214 Ok(node_items)
215}
216
217#[derive(Debug, Clone, Copy)]
218pub enum Query {
219 BBox(f64, f64, f64, f64),
220 PointIntersects(f64, f64),
221 PointNearest(f64, f64),
222}
223
224#[derive(Debug)]
225pub struct SearchResultItem {
227 pub offset: usize,
229 pub index: usize,
231}
232
233const HILBERT_MAX: u32 = (1 << 16) - 1;
234
235fn hilbert(x: u32, y: u32) -> u32 {
237 let mut a = x ^ y;
238 let mut b = 0xFFFF ^ a;
239 let mut c = 0xFFFF ^ (x | y);
240 let mut d = x & (y ^ 0xFFFF);
241
242 let mut aa = a | (b >> 1);
243 let mut bb = (a >> 1) ^ a;
244 let mut cc = ((c >> 1) ^ (b & (d >> 1))) ^ c;
245 let mut dd = ((a & (c >> 1)) ^ (d >> 1)) ^ d;
246
247 a = aa;
248 b = bb;
249 c = cc;
250 d = dd;
251 aa = (a & (a >> 2)) ^ (b & (b >> 2));
252 bb = (a & (b >> 2)) ^ (b & ((a ^ b) >> 2));
253 cc ^= (a & (c >> 2)) ^ (b & (d >> 2));
254 dd ^= (b & (c >> 2)) ^ ((a ^ b) & (d >> 2));
255
256 a = aa;
257 b = bb;
258 c = cc;
259 d = dd;
260 aa = (a & (a >> 4)) ^ (b & (b >> 4));
261 bb = (a & (b >> 4)) ^ (b & ((a ^ b) >> 4));
262 cc ^= (a & (c >> 4)) ^ (b & (d >> 4));
263 dd ^= (b & (c >> 4)) ^ ((a ^ b) & (d >> 4));
264
265 a = aa;
266 b = bb;
267 c = cc;
268 d = dd;
269 cc ^= (a & (c >> 8)) ^ (b & (d >> 8));
270 dd ^= (b & (c >> 8)) ^ ((a ^ b) & (d >> 8));
271
272 a = cc ^ (cc >> 1);
273 b = dd ^ (dd >> 1);
274
275 let mut i0 = x ^ y;
276 let mut i1 = b | (0xFFFF ^ (i0 | a));
277
278 i0 = (i0 | (i0 << 8)) & 0x00FF00FF;
279 i0 = (i0 | (i0 << 4)) & 0x0F0F0F0F;
280 i0 = (i0 | (i0 << 2)) & 0x33333333;
281 i0 = (i0 | (i0 << 1)) & 0x55555555;
282
283 i1 = (i1 | (i1 << 8)) & 0x00FF00FF;
284 i1 = (i1 | (i1 << 4)) & 0x0F0F0F0F;
285 i1 = (i1 | (i1 << 2)) & 0x33333333;
286 i1 = (i1 | (i1 << 1)) & 0x55555555;
287
288 (i1 << 1) | i0
289}
290
291fn hilbert_bbox(r: &NodeItem, hilbert_max: u32, extent: &NodeItem) -> u32 {
292 let x = (hilbert_max as f64 * ((r.min_x + r.max_x) / 2.0 - extent.min_x) / extent.width())
294 .floor() as u32;
295 let y = (hilbert_max as f64 * ((r.min_y + r.max_y) / 2.0 - extent.min_y) / extent.height())
296 .floor() as u32;
297 hilbert(x, y)
298}
299
300pub fn hilbert_sort(items: &mut [NodeItem], extent: &NodeItem) {
301 items.sort_by(|a, b| {
302 let ha = hilbert_bbox(a, HILBERT_MAX, extent);
303 let hb = hilbert_bbox(b, HILBERT_MAX, extent);
304 hb.partial_cmp(&ha).unwrap() });
306}
307
308pub fn calc_extent(nodes: &[NodeItem]) -> NodeItem {
309 nodes.iter().fold(NodeItem::create(0), |mut a, b| {
310 a.expand(b);
311 a
312 })
313}
314
315pub struct PackedRTree {
317 extent: NodeItem,
318 node_items: Vec<NodeItem>,
319 num_leaf_nodes: usize,
320 branching_factor: u16,
321 level_bounds: Vec<Range<usize>>,
322}
323
324impl PackedRTree {
325 pub const DEFAULT_NODE_SIZE: u16 = 16;
326
327 fn init(&mut self, node_size: u16) -> Result<(), Error> {
328 assert!(node_size >= 2, "Node size must be at least 2");
329 assert!(self.num_leaf_nodes > 0, "Cannot create empty tree");
330 self.branching_factor = node_size.clamp(2u16, 65535u16);
331 self.level_bounds =
332 PackedRTree::generate_level_bounds(self.num_leaf_nodes, self.branching_factor);
333 let num_nodes = self
334 .level_bounds
335 .first()
336 .expect("RTree has at least one level when node_size >= 2 and num_items > 0")
337 .end;
338 self.node_items = vec![NodeItem::create(0); num_nodes]; Ok(())
340 }
341
342 fn generate_level_bounds(num_items: usize, node_size: u16) -> Vec<Range<usize>> {
343 assert!(node_size >= 2, "Node size must be at least 2");
344 assert!(num_items > 0, "Cannot create empty tree");
345 assert!(
346 num_items <= usize::MAX - ((num_items / node_size as usize) * 2),
347 "Number of items too large"
348 );
349
350 let mut level_num_nodes: Vec<usize> = Vec::new();
352 let mut n = num_items;
353 let mut num_nodes = n;
354 level_num_nodes.push(n);
355 loop {
356 n = n.div_ceil(node_size as usize);
357 num_nodes += n;
358 level_num_nodes.push(n);
359 if n == 1 {
360 break;
361 }
362 }
363 let mut level_offsets: Vec<usize> = Vec::with_capacity(level_num_nodes.len());
365 n = num_nodes;
366 for size in &level_num_nodes {
367 level_offsets.push(n - size);
368 n -= size;
369 }
370 let mut level_bounds = Vec::with_capacity(level_num_nodes.len());
371 for i in 0..level_num_nodes.len() {
372 level_bounds.push(level_offsets[i]..level_offsets[i] + level_num_nodes[i]);
373 }
374 level_bounds
375 }
376
377 fn generate_nodes(&mut self) {
378 for level in 0..self.level_bounds.len() - 1 {
379 let children_level = &self.level_bounds[level];
380 let parent_level = &self.level_bounds[level + 1];
381
382 let mut parent_idx = parent_level.start;
383 let mut child_idx = children_level.start;
384 while child_idx < children_level.end {
385 let mut parent_node = NodeItem::create(child_idx as u64);
386 for _j in 0..self.branching_factor {
387 if child_idx >= children_level.end {
388 break;
389 }
390 parent_node.expand(&self.node_items[child_idx]);
391 child_idx += 1;
392 }
393 self.node_items[parent_idx] = parent_node;
394 parent_idx += 1;
395 }
396 }
397 }
398
399 fn read_data(&mut self, data: impl Read) -> Result<(), Error> {
400 read_node_vec(&mut self.node_items, data)?;
401 for node in &self.node_items {
402 self.extent.expand(node)
403 }
404 Ok(())
405 }
406
407 #[cfg(feature = "http")]
408 async fn read_http<T: AsyncHttpRangeClient>(
409 &mut self,
410 client: &mut AsyncBufferedHttpRangeClient<T>,
411 index_begin: usize,
412 ) -> Result<(), Error> {
413 let min_req_size = self.size(); let mut pos = index_begin;
415 for i in 0..self.num_nodes() {
416 let bytes = client
417 .min_req_size(min_req_size)
418 .get_range(pos, size_of::<NodeItem>())
419 .await?;
420 let n = NodeItem::from_bytes(bytes)?;
421 self.extent.expand(&n);
422 self.node_items[i] = n;
423 pos += size_of::<NodeItem>();
424 }
425 Ok(())
426 }
427
428 fn num_nodes(&self) -> usize {
429 self.node_items.len()
430 }
431
432 pub fn build(
433 nodes: &[NodeItem],
434 extent: &NodeItem,
435 node_size: u16,
436 ) -> Result<PackedRTree, Error> {
437 let mut tree = PackedRTree {
438 extent: extent.clone(),
439 node_items: Vec::new(),
440 num_leaf_nodes: nodes.len(),
441 branching_factor: 0,
442 level_bounds: Vec::new(),
443 };
444 tree.init(node_size)?;
445 let num_nodes = tree.num_nodes();
446 for (i, node) in nodes.iter().take(tree.num_leaf_nodes).cloned().enumerate() {
447 tree.node_items[num_nodes - tree.num_leaf_nodes + i] = node;
448 }
449 tree.generate_nodes();
450 Ok(tree)
451 }
452
453 pub fn from_buf(
454 data: impl Read,
455 num_items: usize,
456 node_size: u16,
457 ) -> Result<PackedRTree, Error> {
458 let node_size = node_size.clamp(2u16, 65535u16);
459 let level_bounds = PackedRTree::generate_level_bounds(num_items, node_size);
460 let num_nodes = level_bounds
461 .first()
462 .expect("RTree has at least one level when node_size >= 2 and num_items > 0")
463 .end;
464 let mut tree = PackedRTree {
465 extent: NodeItem::create(0),
466 node_items: Vec::with_capacity(num_nodes),
467 num_leaf_nodes: num_items,
468 branching_factor: node_size,
469 level_bounds,
470 };
471 tree.read_data(data)?;
472 Ok(tree)
473 }
474
475 #[cfg(feature = "http")]
476 pub async fn from_http<T: AsyncHttpRangeClient>(
477 client: &mut AsyncBufferedHttpRangeClient<T>,
478 index_begin: usize,
479 num_items: usize,
480 node_size: u16,
481 ) -> Result<PackedRTree, Error> {
482 let mut tree = PackedRTree {
483 extent: NodeItem::create(0),
484 node_items: Vec::new(),
485 num_leaf_nodes: num_items,
486 branching_factor: 0,
487 level_bounds: Vec::new(),
488 };
489 tree.init(node_size)?;
490 tree.read_http(client, index_begin).await?;
491 Ok(tree)
492 }
493
494 pub fn search(&self, query: Query) -> Result<Vec<SearchResultItem>, Error> {
496 let leaf_nodes_offset = self
497 .level_bounds
498 .first()
499 .expect("RTree has at least one level when node_size >= 2 and num_items > 0")
500 .start;
501
502 match query {
503 Query::BBox(min_x, min_y, max_x, max_y) => {
504 let bounds = NodeItem::bounds(min_x, min_y, max_x, max_y);
506 let mut results = Vec::new();
507 let mut queue = VecDeque::new();
508 queue.push_back((0, self.level_bounds.len() - 1));
509
510 while let Some(next) = queue.pop_front() {
511 let node_index = next.0;
512 let level = next.1;
513 let is_leaf_node = node_index >= self.num_nodes() - self.num_leaf_nodes;
514 let end = min(
516 node_index + self.branching_factor as usize,
517 self.level_bounds[level].end,
518 );
519 for pos in node_index..end {
521 let node_item = &self.node_items[pos];
522 if !bounds.intersects(node_item) {
523 continue;
524 }
525 if is_leaf_node {
526 results.push(SearchResultItem {
527 offset: node_item.offset as usize,
528 index: pos - leaf_nodes_offset,
529 });
530 } else {
531 queue.push_back((node_item.offset as usize, level - 1));
532 }
533 }
534 }
535 Ok(results)
536 }
537 Query::PointIntersects(x, y) => {
538 let mut results = Vec::new();
541 let mut queue = VecDeque::new();
542 queue.push_back((0, self.level_bounds.len() - 1));
543
544 while let Some(next) = queue.pop_front() {
545 let node_index = next.0;
546 let level = next.1;
547 let is_leaf_node = node_index >= self.num_nodes() - self.num_leaf_nodes;
548 let end = min(
550 node_index + self.branching_factor as usize,
551 self.level_bounds[level].end,
552 );
553 for pos in node_index..end {
555 let node_item = &self.node_items[pos];
556 if !node_item.contains_point(x, y) {
557 continue;
558 }
559 if is_leaf_node {
560 results.push(SearchResultItem {
561 offset: node_item.offset as usize,
562 index: pos - leaf_nodes_offset,
563 });
564 } else {
565 queue.push_back((node_item.offset as usize, level - 1));
566 }
567 }
568 }
569 Ok(results)
570 }
571 Query::PointNearest(x, y) => {
572 use std::cmp::Reverse;
575 use std::collections::BinaryHeap;
576
577 #[derive(PartialEq)]
578 struct QueueItem {
579 distance: f64,
580 node_index: usize,
581 level: usize,
582 }
583
584 impl Eq for QueueItem {}
585
586 impl PartialOrd for QueueItem {
587 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
588 self.distance.partial_cmp(&other.distance)
589 }
590 }
591
592 impl Ord for QueueItem {
593 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
594 self.partial_cmp(other).unwrap_or(std::cmp::Ordering::Equal)
595 }
596 }
597
598 let mut nearest: Option<(f64, SearchResultItem)> = None;
599 let mut queue = BinaryHeap::new();
600
601 queue.push(Reverse(QueueItem {
603 distance: 0.0,
604 node_index: 0,
605 level: self.level_bounds.len() - 1,
606 }));
607
608 while let Some(Reverse(next)) = queue.pop() {
609 if let Some((best_dist, _)) = nearest {
612 if next.distance > best_dist {
613 break;
614 }
615 }
616
617 let node_index = next.node_index;
618 let level = next.level;
619 let is_leaf_node = node_index >= self.num_nodes() - self.num_leaf_nodes;
620 let end = min(
622 node_index + self.branching_factor as usize,
623 self.level_bounds[level].end,
624 );
625
626 for pos in node_index..end {
628 let node_item = &self.node_items[pos];
629 let dist = node_item.min_distance_squared(x, y);
630
631 if let Some((best_dist, _)) = nearest {
633 if dist >= best_dist {
634 continue;
635 }
636 }
637
638 if is_leaf_node {
639 let result = SearchResultItem {
641 offset: node_item.offset as usize,
642 index: pos - leaf_nodes_offset,
643 };
644
645 let centroid_dist = node_item.centroid_distance_squared(x, y);
647
648 match nearest {
649 None => nearest = Some((centroid_dist, result)),
650 Some((best_dist, _)) if centroid_dist < best_dist => {
651 nearest = Some((centroid_dist, result))
652 }
653 _ => {}
654 }
655 } else {
656 queue.push(Reverse(QueueItem {
658 distance: dist,
659 node_index: node_item.offset as usize,
660 level: level - 1,
661 }));
662 }
663 }
664 }
665
666 Ok(nearest.map(|(_, item)| vec![item]).unwrap_or_default())
668 }
669 }
670 }
671
672 pub fn stream_search<R: Read + Seek>(
674 data: &mut R,
675 num_items: usize,
676 node_size: u16,
677 query: Query,
678 ) -> Result<Vec<SearchResultItem>, Error> {
679 let level_bounds = PackedRTree::generate_level_bounds(num_items, node_size);
680 let Range {
681 start: leaf_nodes_offset,
682 end: num_nodes,
683 } = level_bounds
684 .first()
685 .expect("RTree has at least one level when node_size >= 2 and num_items > 0");
686
687 let index_base = data.stream_position()?;
689
690 match query {
691 Query::BBox(min_x, min_y, max_x, max_y) => {
692 let bounds = NodeItem::bounds(min_x, min_y, max_x, max_y);
693
694 let mut queue = VecDeque::new();
696 queue.push_back((0, level_bounds.len() - 1));
697 let mut results = Vec::new();
698
699 while let Some(next) = queue.pop_front() {
700 let node_index = next.0;
701 let level = next.1;
702 let is_leaf_node = node_index >= num_nodes - num_items;
703 let end = min(node_index + node_size as usize, level_bounds[level].end);
705 let length = end - node_index;
706 let node_items = read_node_items(data, index_base, node_index, length)?;
707 for pos in node_index..end {
709 let node_pos = pos - node_index;
710 let node_item = &node_items[node_pos];
711 if !bounds.intersects(node_item) {
712 continue;
713 }
714 if is_leaf_node {
715 let index = pos - leaf_nodes_offset;
716 let offset = node_item.offset as usize;
717 results.push(SearchResultItem { offset, index });
718 } else {
719 let offset = node_item.offset as usize;
720 let prev_level = level - 1;
721 queue.push_back((offset, prev_level));
722 }
723 }
724 }
725
726 data.seek(SeekFrom::Start(
728 index_base + (num_nodes * size_of::<NodeItem>()) as u64,
729 ))?;
730 Ok(results)
731 }
732 Query::PointIntersects(x, y) => {
733 let mut queue = VecDeque::new();
735 queue.push_back((0, level_bounds.len() - 1));
736 let mut results = Vec::new();
737
738 while let Some(next) = queue.pop_front() {
739 let node_index = next.0;
740 let level = next.1;
741 let is_leaf_node = node_index >= num_nodes - num_items;
742 let end = min(node_index + node_size as usize, level_bounds[level].end);
744 let length = end - node_index;
745 let node_items = read_node_items(data, index_base, node_index, length)?;
746 for pos in node_index..end {
748 let node_pos = pos - node_index;
749 let node_item = &node_items[node_pos];
750 if !node_item.contains_point(x, y) {
751 continue;
752 }
753 if is_leaf_node {
754 let index = pos - leaf_nodes_offset;
755 let offset = node_item.offset as usize;
756 results.push(SearchResultItem { offset, index });
757 } else {
758 let offset = node_item.offset as usize;
759 let prev_level = level - 1;
760 queue.push_back((offset, prev_level));
761 }
762 }
763 }
764
765 data.seek(SeekFrom::Start(
767 index_base + (num_nodes * size_of::<NodeItem>()) as u64,
768 ))?;
769 Ok(results)
770 }
771 Query::PointNearest(x, y) => {
772 use std::cmp::Reverse;
773 use std::collections::BinaryHeap;
774
775 #[derive(PartialEq)]
776 struct QueueItem {
777 distance: f64,
778 node_index: usize,
779 level: usize,
780 }
781
782 impl Eq for QueueItem {}
783
784 impl PartialOrd for QueueItem {
785 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
786 self.distance.partial_cmp(&other.distance)
787 }
788 }
789
790 impl Ord for QueueItem {
791 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
792 self.partial_cmp(other).unwrap_or(std::cmp::Ordering::Equal)
793 }
794 }
795
796 let mut nearest: Option<(f64, SearchResultItem)> = None;
797 let mut queue = BinaryHeap::new();
798
799 queue.push(Reverse(QueueItem {
801 distance: 0.0,
802 node_index: 0,
803 level: level_bounds.len() - 1,
804 }));
805
806 while let Some(Reverse(next)) = queue.pop() {
807 if let Some((best_dist, _)) = nearest {
810 if next.distance > best_dist {
811 break;
812 }
813 }
814
815 let node_index = next.node_index;
816 let level = next.level;
817 let is_leaf_node = node_index >= num_nodes - num_items;
818
819 let end = min(node_index + node_size as usize, level_bounds[level].end);
821 let length = end - node_index;
822 let node_items = read_node_items(data, index_base, node_index, length)?;
823
824 for pos in node_index..end {
826 let node_pos = pos - node_index;
827 let node_item = &node_items[node_pos];
828 let dist = node_item.min_distance_squared(x, y);
829
830 if let Some((best_dist, _)) = nearest {
832 if dist >= best_dist {
833 continue;
834 }
835 }
836
837 if is_leaf_node {
838 let index = pos - leaf_nodes_offset;
840 let offset = node_item.offset as usize;
841 let result = SearchResultItem { offset, index };
842
843 let centroid_dist = node_item.centroid_distance_squared(x, y);
845
846 match nearest {
847 None => nearest = Some((centroid_dist, result)),
848 Some((best_dist, _)) if centroid_dist < best_dist => {
849 nearest = Some((centroid_dist, result))
850 }
851 _ => {}
852 }
853 } else {
854 queue.push(Reverse(QueueItem {
856 distance: dist,
857 node_index: node_item.offset as usize,
858 level: level - 1,
859 }));
860 }
861 }
862 }
863
864 data.seek(SeekFrom::Start(
866 index_base + (num_nodes * size_of::<NodeItem>()) as u64,
867 ))?;
868
869 Ok(nearest.map(|(_, item)| vec![item]).unwrap_or_default())
871 }
872 }
873 }
874
875 pub fn size(&self) -> usize {
876 self.num_nodes() * size_of::<NodeItem>()
877 }
878
879 pub fn index_size(num_items: usize, node_size: u16) -> usize {
880 assert!(node_size >= 2, "Node size must be at least 2");
881 assert!(num_items > 0, "Cannot create empty tree");
882 let node_size_min = node_size.clamp(2, 65535) as usize;
883 let mut n = num_items;
889 let mut num_nodes = n;
890 loop {
891 n = n.div_ceil(node_size_min);
892 num_nodes += n;
893 if n == 1 {
894 break;
895 }
896 }
897 num_nodes * size_of::<NodeItem>()
898 }
899
900 pub fn stream_write<W: Write>(&self, out: &mut W) -> std::io::Result<()> {
902 for item in &self.node_items {
903 item.write(out)?;
904 }
905 Ok(())
906 }
907
908 pub fn extent(&self) -> NodeItem {
909 self.extent.clone()
910 }
911
912 #[cfg(feature = "http")]
913 #[allow(clippy::too_many_arguments)]
914 pub async fn http_stream_search<T: AsyncHttpRangeClient>(
915 client: &mut AsyncBufferedHttpRangeClient<T>,
916 index_begin: usize,
917 attr_index_size: usize,
918 num_items: usize,
919 branching_factor: u16,
920 query: Query,
921 combine_request_threshold: usize,
922 ) -> Result<Vec<HttpSearchResultItem>, Error> {
923 use tracing::debug;
924
925 if num_items == 0 {
926 return Ok(vec![]);
927 }
928
929 let level_bounds = PackedRTree::generate_level_bounds(num_items, branching_factor);
930 let feature_begin =
931 index_begin + attr_index_size + PackedRTree::index_size(num_items, branching_factor);
932
933 match query {
934 Query::BBox(min_x, min_y, max_x, max_y) => {
935 debug!("http_stream_search - index_begin: {index_begin}, feature_begin: {feature_begin} num_items: {num_items}, branching_factor: {branching_factor}, level_bounds: {level_bounds:?}, GPS bounds:[({min_x}, {min_y}), ({max_x},{max_y})]");
936
937 let bounds = NodeItem::bounds(min_x, min_y, max_x, max_y);
938
939 #[derive(Debug, PartialEq, Eq)]
940 struct NodeRange {
941 level: usize,
942 nodes: Range<usize>,
943 has_sizing_tail: bool,
948 }
949
950 let mut queue = VecDeque::new();
951 queue.push_back(NodeRange {
952 nodes: 0..1,
953 level: level_bounds.len() - 1,
954 has_sizing_tail: false,
955 });
956 let mut results = Vec::new();
957
958 while let Some(node_range) = queue.pop_front() {
959 debug!("next: {node_range:?}. {} items left in queue", queue.len());
960 let node_items =
961 read_http_node_items(client, index_begin, &node_range.nodes).await?;
962 let emit_len = node_items.len()
966 - usize::from(node_range.level == 0 && node_range.has_sizing_tail);
967 for (node_pos, node_item) in node_items.iter().enumerate() {
968 if node_pos >= emit_len {
969 continue;
970 }
971 if !bounds.intersects(node_item) {
972 continue;
973 }
974
975 if node_range.level == 0 {
976 let start = feature_begin + node_item.offset as usize;
978 if let Some(next_node_item) = &node_items.get(node_pos + 1) {
979 let end = feature_begin + next_node_item.offset as usize;
980 results.push(HttpSearchResultItem {
981 range: HttpRange::Range(start..end),
982 });
983 } else {
984 debug_assert_eq!(
991 node_range.nodes.start + node_pos,
992 level_bounds[0].end - 1
993 );
994 results.push(HttpSearchResultItem {
995 range: HttpRange::RangeFrom(start..),
996 });
997 }
998 } else {
999 let children_level = node_range.level - 1;
1000 let mut children_nodes = node_item.offset as usize
1001 ..(node_item.offset + branching_factor as u64) as usize;
1002 let mut has_sizing_tail = false;
1003 if children_level == 0 {
1004 children_nodes.end += 1;
1011 has_sizing_tail = true;
1012 }
1013 let clamped_end =
1015 min(children_nodes.end, level_bounds[children_level].end);
1016 has_sizing_tail &= clamped_end == children_nodes.end;
1019 children_nodes.end = clamped_end;
1020
1021 let children_range = NodeRange {
1022 nodes: children_nodes,
1023 level: children_level,
1024 has_sizing_tail,
1025 };
1026
1027 let Some(tail) = queue.back_mut() else {
1028 debug!("Adding new request onto empty queue: {children_range:?}");
1029 queue.push_back(children_range);
1030 continue;
1031 };
1032
1033 if tail.level != children_level {
1034 debug!("Adding new request for new level: {children_range:?} (existing queue tail: {tail:?})");
1035 queue.push_back(children_range);
1036 continue;
1037 }
1038
1039 let wasted_bytes = {
1040 if children_range.nodes.start >= tail.nodes.end {
1041 (children_range.nodes.start - tail.nodes.end)
1042 * size_of::<NodeItem>()
1043 } else {
1044 debug_assert_eq!(
1050 children_range.nodes.start + 1,
1051 tail.nodes.end,
1052 "we only ever fetch one extra node"
1053 );
1054 debug_assert_eq!(
1055 children_level, 0,
1056 "extra node fetching only happens with leaf nodes"
1057 );
1058 0
1059 }
1060 };
1061 if wasted_bytes > combine_request_threshold {
1062 debug!("Adding new request for: {children_range:?} rather than merging with distant NodeRange: {tail:?} (would waste {wasted_bytes} bytes)");
1063 queue.push_back(children_range);
1064 continue;
1065 }
1066
1067 debug!("Extending existing request {tail:?} with nearby children: {:?} (wastes {wasted_bytes} bytes)", &children_range.nodes);
1069 tail.nodes.end = children_range.nodes.end;
1070 tail.has_sizing_tail = children_range.has_sizing_tail;
1073 }
1074 }
1075 }
1076 Ok(results)
1077 }
1078 Query::PointIntersects(x, y) => {
1079 debug!("http_stream_search point intersects - index_begin: {index_begin}, feature_begin: {feature_begin} num_items: {num_items}, branching_factor: {branching_factor}, level_bounds: {level_bounds:?}, point: ({x}, {y})");
1080
1081 #[derive(Debug, PartialEq, Eq)]
1082 struct NodeRange {
1083 level: usize,
1084 nodes: Range<usize>,
1085 has_sizing_tail: bool,
1090 }
1091
1092 let mut queue = VecDeque::new();
1093 queue.push_back(NodeRange {
1094 nodes: 0..1,
1095 level: level_bounds.len() - 1,
1096 has_sizing_tail: false,
1097 });
1098 let mut results = Vec::new();
1099
1100 while let Some(node_range) = queue.pop_front() {
1101 debug!("next: {node_range:?}. {} items left in queue", queue.len());
1102 let node_items =
1103 read_http_node_items(client, index_begin, &node_range.nodes).await?;
1104 let emit_len = node_items.len()
1108 - usize::from(node_range.level == 0 && node_range.has_sizing_tail);
1109 for (node_pos, node_item) in node_items.iter().enumerate() {
1110 if node_pos >= emit_len {
1111 continue;
1112 }
1113 if !node_item.contains_point(x, y) {
1114 continue;
1115 }
1116
1117 if node_range.level == 0 {
1118 let start = feature_begin + node_item.offset as usize;
1120 if let Some(next_node_item) = &node_items.get(node_pos + 1) {
1121 let end = feature_begin + next_node_item.offset as usize;
1122 results.push(HttpSearchResultItem {
1123 range: HttpRange::Range(start..end),
1124 });
1125 } else {
1126 debug_assert_eq!(
1130 node_range.nodes.start + node_pos,
1131 level_bounds[0].end - 1
1132 );
1133 results.push(HttpSearchResultItem {
1134 range: HttpRange::RangeFrom(start..),
1135 });
1136 }
1137 } else {
1138 let children_level = node_range.level - 1;
1139 let mut children_nodes = node_item.offset as usize
1140 ..(node_item.offset + branching_factor as u64) as usize;
1141 let mut has_sizing_tail = false;
1142 if children_level == 0 {
1143 children_nodes.end += 1;
1144 has_sizing_tail = true;
1145 }
1146 let clamped_end =
1147 min(children_nodes.end, level_bounds[children_level].end);
1148 has_sizing_tail &= clamped_end == children_nodes.end;
1151 children_nodes.end = clamped_end;
1152
1153 let children_range = NodeRange {
1154 nodes: children_nodes,
1155 level: children_level,
1156 has_sizing_tail,
1157 };
1158
1159 let Some(tail) = queue.back_mut() else {
1160 debug!("Adding new request onto empty queue: {children_range:?}");
1161 queue.push_back(children_range);
1162 continue;
1163 };
1164
1165 if tail.level != children_level {
1166 debug!("Adding new request for new level: {children_range:?} (existing queue tail: {tail:?})");
1167 queue.push_back(children_range);
1168 continue;
1169 }
1170
1171 let wasted_bytes = {
1172 if children_range.nodes.start >= tail.nodes.end {
1173 (children_range.nodes.start - tail.nodes.end)
1174 * size_of::<NodeItem>()
1175 } else {
1176 debug_assert_eq!(
1177 children_range.nodes.start + 1,
1178 tail.nodes.end,
1179 "we only ever fetch one extra node"
1180 );
1181 debug_assert_eq!(
1182 children_level, 0,
1183 "extra node fetching only happens with leaf nodes"
1184 );
1185 0
1186 }
1187 };
1188 if wasted_bytes > combine_request_threshold {
1189 debug!("Adding new request for: {children_range:?} rather than merging with distant NodeRange: {tail:?} (would waste {wasted_bytes} bytes)");
1190 queue.push_back(children_range);
1191 continue;
1192 }
1193
1194 tail.nodes.end = children_range.nodes.end;
1195 tail.has_sizing_tail = children_range.has_sizing_tail;
1198 }
1199 }
1200 }
1201 Ok(results)
1202 }
1203 Query::PointNearest(x, y) => {
1204 debug!("http_stream_search nearest neighbor - index_begin: {index_begin}, feature_begin: {feature_begin} num_items: {num_items}, branching_factor: {branching_factor}, level_bounds: {level_bounds:?}, point: ({x}, {y})");
1205
1206 use std::cmp::Reverse;
1207 use std::collections::BinaryHeap;
1208
1209 #[derive(PartialEq)]
1210 struct QueueItem {
1211 distance: f64,
1212 level: usize,
1213 nodes: Range<usize>,
1214 has_sizing_tail: bool,
1219 }
1220
1221 impl Eq for QueueItem {}
1222
1223 impl PartialOrd for QueueItem {
1224 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
1225 self.distance.partial_cmp(&other.distance)
1226 }
1227 }
1228
1229 impl Ord for QueueItem {
1230 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
1231 self.partial_cmp(other).unwrap_or(std::cmp::Ordering::Equal)
1232 }
1233 }
1234
1235 let mut nearest: Option<(f64, HttpSearchResultItem)> = None;
1236 let mut queue = BinaryHeap::new();
1237
1238 queue.push(Reverse(QueueItem {
1240 distance: 0.0,
1241 nodes: 0..1,
1242 level: level_bounds.len() - 1,
1243 has_sizing_tail: false,
1244 }));
1245
1246 while let Some(Reverse(next)) = queue.pop() {
1247 if let Some((best_dist, _)) = nearest {
1250 if next.distance > best_dist {
1251 break;
1252 }
1253 }
1254
1255 debug!(
1256 "next: node dist: {}, level: {}, nodes: {:?}, {} items left in queue",
1257 next.distance,
1258 next.level,
1259 next.nodes,
1260 queue.len()
1261 );
1262 let node_items = read_http_node_items(client, index_begin, &next.nodes).await?;
1263 let emit_len =
1267 node_items.len() - usize::from(next.level == 0 && next.has_sizing_tail);
1268 for (node_pos, node_item) in node_items.iter().enumerate() {
1269 if node_pos >= emit_len {
1270 continue;
1271 }
1272 let dist = node_item.min_distance_squared(x, y);
1273
1274 if let Some((best_dist, _)) = nearest {
1276 if dist >= best_dist {
1277 continue;
1278 }
1279 }
1280
1281 if next.level == 0 {
1282 let centroid_dist = node_item.centroid_distance_squared(x, y);
1284
1285 let start = feature_begin + node_item.offset as usize;
1287 let result = if let Some(next_node_item) = &node_items.get(node_pos + 1)
1288 {
1289 let end = feature_begin + next_node_item.offset as usize;
1290 HttpSearchResultItem {
1291 range: HttpRange::Range(start..end),
1292 }
1293 } else {
1294 debug_assert_eq!(
1298 next.nodes.start + node_pos,
1299 level_bounds[0].end - 1
1300 );
1301 HttpSearchResultItem {
1302 range: HttpRange::RangeFrom(start..),
1303 }
1304 };
1305
1306 match nearest {
1307 None => nearest = Some((centroid_dist, result)),
1308 Some((best_dist, _)) if centroid_dist < best_dist => {
1309 nearest = Some((centroid_dist, result))
1310 }
1311 _ => {}
1312 }
1313 } else {
1314 let children_level = next.level - 1;
1316 let mut children_nodes = node_item.offset as usize
1317 ..(node_item.offset + branching_factor as u64) as usize;
1318
1319 let mut has_sizing_tail = false;
1320 if children_level == 0 {
1321 children_nodes.end += 1;
1322 has_sizing_tail = true;
1323 }
1324
1325 let clamped_end =
1327 min(children_nodes.end, level_bounds[children_level].end);
1328 has_sizing_tail &= clamped_end == children_nodes.end;
1331 children_nodes.end = clamped_end;
1332
1333 queue.push(Reverse(QueueItem {
1334 distance: dist,
1335 nodes: children_nodes,
1336 level: children_level,
1337 has_sizing_tail,
1338 }));
1339 }
1340 }
1341 }
1342
1343 Ok(nearest.map(|(_, item)| vec![item]).unwrap_or_default())
1345 }
1346 }
1347 }
1348}
1349
1350#[cfg(feature = "http")]
1351pub mod http {
1352 use std::ops::{Range, RangeFrom};
1353
1354 #[derive(Debug, Clone)]
1356 pub enum HttpRange {
1357 Range(Range<usize>),
1358 RangeFrom(RangeFrom<usize>),
1359 }
1360
1361 impl HttpRange {
1362 pub fn start(&self) -> usize {
1363 match self {
1364 Self::Range(range) => range.start,
1365 Self::RangeFrom(range) => range.start,
1366 }
1367 }
1368
1369 pub fn end(&self) -> Option<usize> {
1370 match self {
1371 Self::Range(range) => Some(range.end),
1372 Self::RangeFrom(_) => None,
1373 }
1374 }
1375
1376 pub fn with_end(self, end: Option<usize>) -> Self {
1377 match end {
1378 Some(end) => Self::Range(self.start()..end),
1379 None => Self::RangeFrom(self.start()..),
1380 }
1381 }
1382
1383 pub fn length(&self) -> Option<usize> {
1384 match self {
1385 Self::Range(range) => Some(range.end - range.start),
1386 Self::RangeFrom(_) => None,
1387 }
1388 }
1389 }
1390
1391 #[derive(Debug)]
1392 pub struct HttpSearchResultItem {
1394 pub range: HttpRange,
1396 }
1397}
1398#[cfg(feature = "http")]
1399pub(crate) use http::*;
1400
1401#[cfg(test)]
1402mod tests {
1403 use super::*;
1404 type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;
1405
1406 #[test]
1407 fn tree_2items() -> Result<()> {
1408 let mut nodes = Vec::new();
1409 nodes.push(NodeItem::bounds(0.0, 0.0, 1.0, 1.0));
1410 nodes.push(NodeItem::bounds(2.0, 2.0, 3.0, 3.0));
1411 let extent = calc_extent(&nodes);
1412 assert_eq!(extent, NodeItem::bounds(0.0, 0.0, 3.0, 3.0));
1413 assert!(nodes[0].intersects(&NodeItem::bounds(0.0, 0.0, 1.0, 1.0)));
1414 assert!(nodes[1].intersects(&NodeItem::bounds(2.0, 2.0, 3.0, 3.0)));
1415 hilbert_sort(&mut nodes, &extent);
1416 let mut offset = 0;
1417 for node in &mut nodes {
1418 node.offset = offset;
1419 offset += size_of::<NodeItem>() as u64;
1420 }
1421 assert!(nodes[1].intersects(&NodeItem::bounds(0.0, 0.0, 1.0, 1.0)));
1422 assert!(nodes[0].intersects(&NodeItem::bounds(2.0, 2.0, 3.0, 3.0)));
1423 let tree = PackedRTree::build(&nodes, &extent, PackedRTree::DEFAULT_NODE_SIZE)?;
1424 let list = tree.search(Query::BBox(0.0, 0.0, 1.0, 1.0))?;
1425 assert_eq!(list.len(), 1);
1426 assert!(nodes[list[0].index].intersects(&NodeItem::bounds(0.0, 0.0, 1.0, 1.0)));
1427 Ok(())
1428 }
1429
1430 #[test]
1431 fn test_point_intersects_query() -> Result<()> {
1432 let mut nodes = Vec::new();
1434 nodes.push(NodeItem::bounds(0.0, 0.0, 1.0, 1.0)); nodes.push(NodeItem::bounds(2.0, 2.0, 3.0, 3.0)); nodes.push(NodeItem::bounds(5.0, 5.0, 10.0, 10.0)); nodes.push(NodeItem::bounds(-10.0, -10.0, 10.0, 10.0)); let extent = calc_extent(&nodes);
1440 hilbert_sort(&mut nodes, &extent);
1441
1442 let mut offset = 0;
1444 for node in &mut nodes {
1445 node.offset = offset;
1446 offset += size_of::<NodeItem>() as u64;
1447 }
1448
1449 let tree = PackedRTree::build(&nodes, &extent, PackedRTree::DEFAULT_NODE_SIZE)?;
1450
1451 let results = tree.search(Query::PointIntersects(0.5, 0.5))?;
1453 assert_eq!(results.len(), 2, "Point (0.5, 0.5) should be in 2 bboxes");
1454
1455 let results = tree.search(Query::PointIntersects(1.0, 1.0))?;
1457 assert_eq!(results.len(), 2, "Point (1.0, 1.0) should be in 2 bboxes");
1458
1459 let results = tree.search(Query::PointIntersects(7.5, 7.5))?;
1461 assert_eq!(results.len(), 2, "Point (7.5, 7.5) should be in 2 bboxes");
1462
1463 let results = tree.search(Query::PointIntersects(20.0, 20.0))?;
1465 assert_eq!(results.len(), 0, "Point (20.0, 20.0) should be in 0 bboxes");
1466
1467 let mut tree_data: Vec<u8> = Vec::new();
1469 tree.stream_write(&mut tree_data)?;
1470
1471 let mut reader = Cursor::new(&tree_data);
1472 let results = PackedRTree::stream_search(
1473 &mut reader,
1474 nodes.len(),
1475 PackedRTree::DEFAULT_NODE_SIZE,
1476 Query::PointIntersects(0.5, 0.5),
1477 )?;
1478 assert_eq!(
1479 results.len(),
1480 2,
1481 "Stream query: Point (0.5, 0.5) should be in 2 bboxes"
1482 );
1483
1484 Ok(())
1485 }
1486
1487 #[test]
1488 fn test_nearest_neighbor_query() -> Result<()> {
1489 let mut nodes = Vec::new();
1491 nodes.push(NodeItem::bounds(0.0, 0.0, 1.0, 1.0)); nodes.push(NodeItem::bounds(2.0, 2.0, 3.0, 3.0)); nodes.push(NodeItem::bounds(5.0, 5.0, 10.0, 10.0)); nodes.push(NodeItem::bounds(-10.0, -10.0, -5.0, -5.0)); let extent = calc_extent(&nodes);
1497 hilbert_sort(&mut nodes, &extent);
1498
1499 let mut offset = 0;
1501 for node in &mut nodes {
1502 node.offset = offset;
1503 offset += size_of::<NodeItem>() as u64;
1504 }
1505
1506 let tree = PackedRTree::build(&nodes, &extent, PackedRTree::DEFAULT_NODE_SIZE)?;
1507
1508 let results = tree.search(Query::PointNearest(0.0, 0.0))?;
1510 assert_eq!(results.len(), 1, "Should find exactly one nearest node");
1511
1512 let node0_centroid = (0.5, 0.5); let node3_centroid = (-7.5, -7.5); let dist_to_node0 = (node0_centroid.0 - 0.0) * (node0_centroid.0 - 0.0)
1517 + (node0_centroid.1 - 0.0) * (node0_centroid.1 - 0.0);
1518 let dist_to_node3 = (node3_centroid.0 - 0.0) * (node3_centroid.0 - 0.0)
1519 + (node3_centroid.1 - 0.0) * (node3_centroid.1 - 0.0);
1520
1521 assert!(
1522 dist_to_node0 < dist_to_node3,
1523 "Node 0 should be closer than Node 3"
1524 );
1525
1526 let node1_centroid = (2.5, 2.5); let node2_centroid = (7.5, 7.5); let dist_to_node1: f64 = (node1_centroid.0 - 4.0) * (node1_centroid.0 - 4.0)
1532 + (node1_centroid.1 - 4.0) * (node1_centroid.1 - 4.0);
1533 let dist_to_node2: f64 = (node2_centroid.0 - 4.0) * (node2_centroid.0 - 4.0)
1534 + (node2_centroid.1 - 4.0) * (node2_centroid.1 - 4.0);
1535
1536 let _expected_closest_distance = dist_to_node1.min(dist_to_node2);
1537
1538 let results = tree.search(Query::PointNearest(4.0, 4.0))?;
1540 assert_eq!(results.len(), 1, "Should find exactly one nearest node");
1541
1542 let mut tree_data: Vec<u8> = Vec::new();
1544 tree.stream_write(&mut tree_data)?;
1545
1546 let mut reader = Cursor::new(&tree_data);
1547 let results = PackedRTree::stream_search(
1548 &mut reader,
1549 nodes.len(),
1550 PackedRTree::DEFAULT_NODE_SIZE,
1551 Query::PointNearest(0.0, 0.0),
1552 )?;
1553 assert_eq!(
1554 results.len(),
1555 1,
1556 "Stream query: Should find exactly one nearest node"
1557 );
1558
1559 Ok(())
1560 }
1561
1562 #[test]
1563 fn test_node_item_helper_methods() -> Result<()> {
1564 let node = NodeItem::bounds(0.0, 0.0, 5.0, 5.0);
1566
1567 assert!(
1568 node.contains_point(0.0, 0.0),
1569 "Origin point should be contained (boundary)"
1570 );
1571 assert!(
1572 node.contains_point(5.0, 5.0),
1573 "Corner point should be contained (boundary)"
1574 );
1575 assert!(
1576 node.contains_point(2.5, 2.5),
1577 "Center point should be contained"
1578 );
1579 assert!(
1580 !node.contains_point(-1.0, 2.5),
1581 "Point outside should not be contained"
1582 );
1583 assert!(
1584 !node.contains_point(6.0, 2.5),
1585 "Point outside should not be contained"
1586 );
1587
1588 assert_eq!(
1590 node.min_distance_squared(2.5, 2.5),
1591 0.0,
1592 "Point inside should have zero distance"
1593 );
1594 assert_eq!(
1595 node.min_distance_squared(0.0, 0.0),
1596 0.0,
1597 "Point on boundary should have zero distance"
1598 );
1599
1600 let dist_to_outside_point = node.min_distance_squared(7.0, 8.0);
1601 let expected_dist = (7.0 - 5.0) * (7.0 - 5.0) + (8.0 - 5.0) * (8.0 - 5.0); assert_eq!(
1603 dist_to_outside_point, expected_dist,
1604 "Distance to outside point"
1605 );
1606
1607 let centroid_x = 2.5; let centroid_y = 2.5; let dist_to_point = node.centroid_distance_squared(0.0, 0.0);
1612 let expected_dist =
1613 (0.0 - centroid_x) * (0.0 - centroid_x) + (0.0 - centroid_y) * (0.0 - centroid_y);
1614 assert_eq!(
1615 dist_to_point, expected_dist,
1616 "Centroid distance calculation"
1617 );
1618
1619 Ok(())
1620 }
1621
1622 #[test]
1623 fn tree_19items_roundtrip_stream_search() -> Result<()> {
1624 let mut nodes = vec![
1625 NodeItem::bounds(0.0, 0.0, 1.0, 1.0),
1626 NodeItem::bounds(2.0, 2.0, 3.0, 3.0),
1627 NodeItem::bounds(100.0, 100.0, 110.0, 110.0),
1628 NodeItem::bounds(101.0, 101.0, 111.0, 111.0),
1629 NodeItem::bounds(102.0, 102.0, 112.0, 112.0),
1630 NodeItem::bounds(103.0, 103.0, 113.0, 113.0),
1631 NodeItem::bounds(104.0, 104.0, 114.0, 114.0),
1632 NodeItem::bounds(10010.0, 10010.0, 10110.0, 10110.0),
1633 NodeItem::bounds(10010.0, 10010.0, 10110.0, 10110.0),
1634 NodeItem::bounds(10010.0, 10010.0, 10110.0, 10110.0),
1635 NodeItem::bounds(10010.0, 10010.0, 10110.0, 10110.0),
1636 NodeItem::bounds(10010.0, 10010.0, 10110.0, 10110.0),
1637 NodeItem::bounds(10010.0, 10010.0, 10110.0, 10110.0),
1638 NodeItem::bounds(10010.0, 10010.0, 10110.0, 10110.0),
1639 NodeItem::bounds(10010.0, 10010.0, 10110.0, 10110.0),
1640 NodeItem::bounds(10010.0, 10010.0, 10110.0, 10110.0),
1641 NodeItem::bounds(10010.0, 10010.0, 10110.0, 10110.0),
1642 NodeItem::bounds(10010.0, 10010.0, 10110.0, 10110.0),
1643 NodeItem::bounds(10010.0, 10010.0, 10110.0, 10110.0),
1644 ];
1645
1646 let extent = calc_extent(&nodes);
1647 hilbert_sort(&mut nodes, &extent);
1648 let mut offset = 0;
1649 for node in &mut nodes {
1650 node.offset = offset;
1651 offset += size_of::<NodeItem>() as u64;
1652 }
1653 let tree = PackedRTree::build(&nodes, &extent, PackedRTree::DEFAULT_NODE_SIZE)?;
1654 let list = tree.search(Query::BBox(102.0, 102.0, 103.0, 103.0))?;
1655 assert_eq!(list.len(), 4);
1656
1657 let indexes: Vec<usize> = list.iter().map(|item| item.index).collect();
1658 let expected: Vec<usize> = vec![13, 14, 15, 16];
1659 assert_eq!(indexes, expected);
1660
1661 let mut tree_data: Vec<u8> = Vec::new();
1662 let res = tree.stream_write(&mut tree_data);
1663 assert!(res.is_ok());
1664 assert_eq!(tree_data.len(), (nodes.len() + 3) * size_of::<NodeItem>());
1665 assert_eq!(size_of::<NodeItem>(), 40);
1666
1667 let tree2 = PackedRTree::from_buf(
1668 &mut &tree_data[..],
1669 nodes.len(),
1670 PackedRTree::DEFAULT_NODE_SIZE,
1671 )?;
1672 let list = tree2.search(Query::BBox(102.0, 102.0, 103.0, 103.0))?;
1673 assert_eq!(list.len(), 4);
1674
1675 let indexes: Vec<usize> = list.iter().map(|item| item.index).collect();
1676 let expected: Vec<usize> = vec![13, 14, 15, 16];
1677 assert_eq!(indexes, expected);
1678
1679 let mut reader = Cursor::new(&tree_data);
1680 let list = PackedRTree::stream_search(
1681 &mut reader,
1682 nodes.len(),
1683 PackedRTree::DEFAULT_NODE_SIZE,
1684 Query::BBox(102.0, 102.0, 103.0, 103.0),
1685 )?;
1686 assert_eq!(list.len(), 4);
1687
1688 let indexes: Vec<usize> = list.iter().map(|item| item.index).collect();
1689 let expected: Vec<usize> = vec![13, 14, 15, 16];
1690 assert_eq!(indexes, expected);
1691
1692 Ok(())
1693 }
1694
1695 #[test]
1696 fn tree_100_000_items_in_denmark() -> Result<()> {
1697 use rand::distributions::{Distribution, Uniform};
1698
1699 let unifx = Uniform::from(466379..708929);
1700 let unify = Uniform::from(6096801..6322352);
1701 let mut rng = rand::thread_rng();
1702
1703 let mut nodes = Vec::new();
1704 for _ in 0..100000 {
1705 let x = unifx.sample(&mut rng) as f64;
1706 let y = unify.sample(&mut rng) as f64;
1707 nodes.push(NodeItem::bounds(x, y, x, y));
1708 }
1709
1710 let extent = calc_extent(&nodes);
1711 hilbert_sort(&mut nodes, &extent);
1712 let tree = PackedRTree::build(&nodes, &extent, PackedRTree::DEFAULT_NODE_SIZE)?;
1713 let list = tree.search(Query::BBox(690407.0, 6063692.0, 811682.0, 6176467.0))?;
1714
1715 for i in 0..list.len() {
1716 assert!(nodes[list[i].index]
1717 .intersects(&NodeItem::bounds(690407.0, 6063692.0, 811682.0, 6176467.0)));
1718 }
1719
1720 let mut tree_data: Vec<u8> = Vec::new();
1721 let res = tree.stream_write(&mut tree_data);
1722 assert!(res.is_ok());
1723
1724 let mut reader = Cursor::new(&tree_data);
1725 let list2 = PackedRTree::stream_search(
1726 &mut reader,
1727 nodes.len(),
1728 PackedRTree::DEFAULT_NODE_SIZE,
1729 Query::BBox(690407.0, 6063692.0, 811682.0, 6176467.0),
1730 )?;
1731 assert_eq!(list2.len(), list.len());
1732 for i in 0..list2.len() {
1733 assert!(nodes[list2[i].index]
1734 .intersects(&NodeItem::bounds(690407.0, 6063692.0, 811682.0, 6176467.0)));
1735 }
1736 Ok(())
1737 }
1738}
1739
1740#[cfg(all(test, feature = "http"))]
1741mod http_stream_search_tests {
1742 use super::*;
1754 use crate::static_btree::mocked_http_range_client::MockHttpRangeClient;
1755
1756 const FEATURE_SIZE: usize = 64;
1759
1760 type NormalizedRange = (usize, Option<usize>);
1763
1764 fn hilbert_sorted_leaves(num_items: usize) -> Vec<NodeItem> {
1767 let side = (num_items as f64).sqrt().ceil().max(1.0) as usize;
1768 let mut nodes: Vec<NodeItem> = (0..num_items)
1769 .map(|i| {
1770 let x = (i % side) as f64 * 10.0;
1771 let y = (i / side) as f64 * 10.0;
1772 NodeItem::bounds(x, y, x + 6.0, y + 6.0)
1773 })
1774 .collect();
1775 let extent = calc_extent(&nodes);
1776 hilbert_sort(&mut nodes, &extent);
1777 for (i, node) in nodes.iter_mut().enumerate() {
1778 node.offset = (i * FEATURE_SIZE) as u64;
1779 }
1780 nodes
1781 }
1782
1783 fn serialized_index(nodes: &[NodeItem], branching_factor: u16) -> Vec<u8> {
1784 let extent = calc_extent(nodes);
1785 let tree = PackedRTree::build(nodes, &extent, branching_factor).unwrap();
1786 let mut buf = Vec::new();
1787 tree.stream_write(&mut buf).unwrap();
1788 buf
1789 }
1790
1791 fn brute_force(
1793 nodes: &[NodeItem],
1794 feature_begin: usize,
1795 hits: impl Fn(&NodeItem) -> bool,
1796 ) -> Vec<NormalizedRange> {
1797 nodes
1798 .iter()
1799 .enumerate()
1800 .filter(|(_, node)| hits(node))
1801 .map(|(i, node)| {
1802 let start = feature_begin + node.offset as usize;
1803 let end = nodes
1804 .get(i + 1)
1805 .map(|next| feature_begin + next.offset as usize);
1806 (start, end)
1807 })
1808 .collect()
1809 }
1810
1811 fn normalize(results: &[HttpSearchResultItem]) -> Vec<NormalizedRange> {
1812 results
1813 .iter()
1814 .map(|item| (item.range.start(), item.range.end()))
1815 .collect()
1816 }
1817
1818 fn assert_strictly_increasing(actual: &[NormalizedRange], context: &str) {
1819 for pair in actual.windows(2) {
1820 let (prev, next) = (pair[0], pair[1]);
1821 assert!(
1822 prev.0 < next.0,
1823 "{context}: results must be strictly increasing, got {actual:?}"
1824 );
1825 if let Some(prev_end) = prev.1 {
1826 assert!(
1827 prev_end <= next.0,
1828 "{context}: results must not overlap, got {actual:?}"
1829 );
1830 }
1831 }
1832 }
1833
1834 async fn run_query(
1835 index: &[u8],
1836 num_items: usize,
1837 branching_factor: u16,
1838 query: Query,
1839 combine_request_threshold: usize,
1840 ) -> Vec<NormalizedRange> {
1841 let mut client = MockHttpRangeClient::new_mock_http_range_client(index);
1842 let results = PackedRTree::http_stream_search(
1843 &mut client,
1844 0,
1845 0,
1846 num_items,
1847 branching_factor,
1848 query,
1849 combine_request_threshold,
1850 )
1851 .await
1852 .expect("http_stream_search should succeed");
1853 normalize(&results)
1854 }
1855
1856 fn boundary_bboxes(nodes: &[NodeItem]) -> Vec<Query> {
1861 let mut queries = Vec::new();
1862 for (i, node) in nodes.iter().enumerate() {
1863 queries.push(Query::BBox(node.min_x, node.min_y, node.max_x, node.max_y));
1864 if let Some(next) = nodes.get(i + 1) {
1865 queries.push(Query::BBox(
1866 node.min_x.min(next.min_x),
1867 node.min_y.min(next.min_y),
1868 node.max_x.max(next.max_x),
1869 node.max_y.max(next.max_y),
1870 ));
1871 }
1872 }
1873 queries
1874 }
1875
1876 fn boundary_points(nodes: &[NodeItem]) -> Vec<Query> {
1879 nodes
1880 .iter()
1881 .flat_map(|node| {
1882 [
1883 Query::PointIntersects(
1884 (node.min_x + node.max_x) / 2.0,
1885 (node.min_y + node.max_y) / 2.0,
1886 ),
1887 Query::PointIntersects(node.min_x, node.min_y),
1888 Query::PointIntersects(node.max_x, node.max_y),
1889 ]
1890 })
1891 .collect()
1892 }
1893
1894 const THRESHOLDS: [usize; 2] = [0, 4096];
1898
1899 #[tokio::test]
1900 async fn bbox_search_matches_brute_force_at_leaf_boundaries() {
1901 for (num_items, branching_factor) in [(5, 2u16), (17, 4), (100, 4), (100, 16), (129, 8)] {
1902 let nodes = hilbert_sorted_leaves(num_items);
1903 let index = serialized_index(&nodes, branching_factor);
1904 let feature_begin = PackedRTree::index_size(num_items, branching_factor);
1905
1906 for query in boundary_bboxes(&nodes) {
1907 let Query::BBox(min_x, min_y, max_x, max_y) = query else {
1908 unreachable!("boundary_bboxes only yields bbox queries")
1909 };
1910 let bounds = NodeItem::bounds(min_x, min_y, max_x, max_y);
1911 let expected = brute_force(&nodes, feature_begin, |node| bounds.intersects(node));
1912
1913 for threshold in THRESHOLDS {
1914 let actual =
1915 run_query(&index, num_items, branching_factor, query, threshold).await;
1916 let context = format!(
1917 "num_items={num_items}, branching_factor={branching_factor}, \
1918 threshold={threshold}, query={query:?}"
1919 );
1920 assert_eq!(actual, expected, "{context}: hit set mismatch");
1921 assert_strictly_increasing(&actual, &context);
1922 }
1923 }
1924 }
1925 }
1926
1927 struct Lcg(u64);
1930
1931 impl Lcg {
1932 fn next_f64(&mut self, max: f64) -> f64 {
1933 self.0 = self
1934 .0
1935 .wrapping_mul(6364136223846793005)
1936 .wrapping_add(1442695040888963407);
1937 (self.0 >> 11) as f64 / (1u64 << 53) as f64 * max
1938 }
1939 }
1940
1941 #[tokio::test]
1942 async fn bbox_search_matches_brute_force_over_random_boxes() {
1943 let (num_items, branching_factor) = (200usize, 4u16);
1944 let nodes = hilbert_sorted_leaves(num_items);
1945 let index = serialized_index(&nodes, branching_factor);
1946 let feature_begin = PackedRTree::index_size(num_items, branching_factor);
1947 let extent = calc_extent(&nodes);
1948 let mut rng = Lcg(0x5EED);
1949
1950 for _ in 0..250 {
1951 let min_x = extent.min_x + rng.next_f64(extent.width());
1952 let min_y = extent.min_y + rng.next_f64(extent.height());
1953 let max_x = min_x + rng.next_f64(extent.width() / 4.0);
1954 let max_y = min_y + rng.next_f64(extent.height() / 4.0);
1955 let query = Query::BBox(min_x, min_y, max_x, max_y);
1956 let bounds = NodeItem::bounds(min_x, min_y, max_x, max_y);
1957 let expected = brute_force(&nodes, feature_begin, |node| bounds.intersects(node));
1958
1959 for threshold in THRESHOLDS {
1960 let actual = run_query(&index, num_items, branching_factor, query, threshold).await;
1961 let context = format!("threshold={threshold}, query={query:?}");
1962 assert_eq!(actual, expected, "{context}: hit set mismatch");
1963 assert_strictly_increasing(&actual, &context);
1964 }
1965 }
1966 }
1967
1968 fn brute_force_nearest(
1975 nodes: &[NodeItem],
1976 feature_begin: usize,
1977 x: f64,
1978 y: f64,
1979 ) -> Vec<NormalizedRange> {
1980 let best = nodes
1981 .iter()
1982 .map(|node| node.centroid_distance_squared(x, y))
1983 .fold(f64::INFINITY, f64::min);
1984 nodes
1985 .iter()
1986 .enumerate()
1987 .filter(|(_, node)| node.centroid_distance_squared(x, y) == best)
1988 .map(|(i, node)| {
1989 let start = feature_begin + node.offset as usize;
1990 let end = nodes
1991 .get(i + 1)
1992 .map(|next| feature_begin + next.offset as usize);
1993 (start, end)
1994 })
1995 .collect()
1996 }
1997
1998 #[tokio::test]
2007 async fn nearest_search_reaches_the_final_leaf_through_a_non_initial_range() {
2008 let (num_items, branching_factor) = (17usize, 4u16);
2009 let nodes = hilbert_sorted_leaves(num_items);
2010 let index = serialized_index(&nodes, branching_factor);
2011 let feature_begin = PackedRTree::index_size(num_items, branching_factor);
2012
2013 let last = nodes.last().expect("17 leaves");
2014 let (x, y) = (
2015 (last.min_x + last.max_x) / 2.0,
2016 (last.min_y + last.max_y) / 2.0,
2017 );
2018 let expected = brute_force_nearest(&nodes, feature_begin, x, y);
2019 assert_eq!(
2020 expected,
2021 vec![(feature_begin + last.offset as usize, None)],
2022 "the query point must single out the final leaf, whose length is \
2023 unknown (a RangeFrom)"
2024 );
2025
2026 for threshold in THRESHOLDS {
2027 let actual = run_query(
2028 &index,
2029 num_items,
2030 branching_factor,
2031 Query::PointNearest(x, y),
2032 threshold,
2033 )
2034 .await;
2035 assert_eq!(actual, expected, "threshold={threshold}");
2036 }
2037 }
2038
2039 #[tokio::test]
2049 async fn nearest_search_matches_brute_force_at_leaf_boundaries() {
2050 for (num_items, branching_factor) in [(5, 2u16), (17, 4), (100, 4), (129, 8)] {
2051 let nodes = hilbert_sorted_leaves(num_items);
2052 let index = serialized_index(&nodes, branching_factor);
2053 let feature_begin = PackedRTree::index_size(num_items, branching_factor);
2054
2055 for query in boundary_points(&nodes) {
2056 let Query::PointIntersects(x, y) = query else {
2057 unreachable!("boundary_points only yields point queries")
2058 };
2059 let query = Query::PointNearest(x, y);
2060 let expected = brute_force_nearest(&nodes, feature_begin, x, y);
2061
2062 for threshold in THRESHOLDS {
2063 let actual =
2064 run_query(&index, num_items, branching_factor, query, threshold).await;
2065 let context = format!(
2066 "num_items={num_items}, branching_factor={branching_factor}, \
2067 threshold={threshold}, query={query:?}"
2068 );
2069 assert_eq!(actual.len(), 1, "{context}: exactly one nearest expected");
2070 assert!(
2071 expected.contains(&actual[0]),
2072 "{context}: got {actual:?}, expected one of {expected:?}"
2073 );
2074 if actual[0].1.is_none() {
2078 let last = nodes.last().expect("non-empty");
2079 assert_eq!(
2080 actual[0].0,
2081 feature_begin + last.offset as usize,
2082 "{context}: an unbounded RangeFrom is a read to EOF \
2083 and is only correct for the level's last leaf"
2084 );
2085 }
2086 }
2087 }
2088 }
2089 }
2090
2091 #[tokio::test]
2092 async fn point_search_matches_brute_force_at_leaf_boundaries() {
2093 for (num_items, branching_factor) in [(5, 2u16), (17, 4), (100, 4), (129, 8)] {
2094 let nodes = hilbert_sorted_leaves(num_items);
2095 let index = serialized_index(&nodes, branching_factor);
2096 let feature_begin = PackedRTree::index_size(num_items, branching_factor);
2097
2098 for query in boundary_points(&nodes) {
2099 let Query::PointIntersects(x, y) = query else {
2100 unreachable!("boundary_points only yields point queries")
2101 };
2102 let expected = brute_force(&nodes, feature_begin, |node| node.contains_point(x, y));
2103
2104 for threshold in THRESHOLDS {
2105 let actual =
2106 run_query(&index, num_items, branching_factor, query, threshold).await;
2107 let context = format!(
2108 "num_items={num_items}, branching_factor={branching_factor}, \
2109 threshold={threshold}, query={query:?}"
2110 );
2111 assert_eq!(actual, expected, "{context}: hit set mismatch");
2112 assert_strictly_increasing(&actual, &context);
2113 }
2114 }
2115 }
2116 }
2117}