geo_index/rtree/
trait.rs

1use std::cmp::Reverse;
2use std::collections::{BinaryHeap, VecDeque};
3use std::vec;
4
5#[cfg(feature = "use-geo_0_31")]
6use geo_0_31::algorithm::BoundingRect;
7#[cfg(feature = "use-geo_0_31")]
8use geo_0_31::Geometry;
9use geo_traits::{CoordTrait, RectTrait};
10
11use crate::error::Result;
12use crate::indices::Indices;
13use crate::r#type::IndexableNum;
14#[cfg(feature = "use-geo_0_31")]
15use crate::rtree::distance::DistanceMetric;
16use crate::rtree::index::{RTree, RTreeRef};
17use crate::rtree::traversal::{IntersectionIterator, Node};
18use crate::rtree::util::upper_bound;
19use crate::rtree::RTreeMetadata;
20use crate::GeoIndexError;
21
22/// A simple distance metric trait that doesn't depend on geo.
23///
24/// This trait is used for basic distance calculations without geometry support.
25pub trait SimpleDistanceMetric<N: IndexableNum> {
26    /// Calculate the distance between two points (x1, y1) and (x2, y2).
27    fn distance(&self, x1: N, y1: N, x2: N, y2: N) -> N;
28
29    /// Calculate the distance from a point to a bounding box.
30    fn distance_to_bbox(&self, x: N, y: N, min_x: N, min_y: N, max_x: N, max_y: N) -> N;
31
32    /// Return the maximum distance value for this metric.
33    fn max_distance(&self) -> N {
34        N::max_value()
35    }
36}
37
38/// A trait for accessing geometries by index.
39///
40/// This trait allows different storage strategies for geometries (direct storage,
41/// WKB decoding, caching, etc.) to be used with spatial index queries.
42#[cfg(feature = "use-geo_0_31")]
43pub trait GeometryAccessor {
44    /// Get the geometry at the given index.
45    ///
46    /// # Arguments
47    /// * `item_index` - Index of the item to retrieve
48    ///
49    /// # Returns
50    /// A reference to the geometry at the given index, or None if the index is out of bounds
51    fn get_geometry(&self, item_index: usize) -> Option<&Geometry<f64>>;
52}
53
54/// A trait for searching and accessing data out of an RTree.
55pub trait RTreeIndex<N: IndexableNum>: Sized {
56    /// A slice representing all the bounding boxes of all elements contained within this tree,
57    /// including the bounding boxes of each internal node.
58    fn boxes(&self) -> &[N];
59
60    /// A slice representing the indices within the `boxes` slice, including internal nodes.
61    fn indices(&self) -> Indices<'_>;
62
63    /// Access the metadata describing this RTree
64    fn metadata(&self) -> &RTreeMetadata<N>;
65
66    /// The total number of items contained in this RTree.
67    fn num_items(&self) -> u32 {
68        self.metadata().num_items()
69    }
70
71    /// The total number of nodes in this RTree, including both leaf and intermediate nodes.
72    fn num_nodes(&self) -> usize {
73        self.metadata().num_nodes()
74    }
75
76    /// The maximum number of elements in each node.
77    fn node_size(&self) -> u16 {
78        self.metadata().node_size()
79    }
80
81    /// The offsets into [RTreeIndex::boxes] where each level's boxes starts and ends. The tree is
82    /// laid out bottom-up, and there's an implicit initial 0. So the boxes of the lowest level of
83    /// the tree are located from `boxes[0..self.level_bounds()[0]]`.
84    fn level_bounds(&self) -> &[usize] {
85        self.metadata().level_bounds()
86    }
87
88    /// The number of levels (height) of the tree.
89    fn num_levels(&self) -> usize {
90        self.level_bounds().len()
91    }
92
93    /// The tree is laid out from bottom to top. Level 0 is the _base_ of the tree. Each integer
94    /// higher is one level higher of the tree.
95    fn boxes_at_level(&self, level: usize) -> Result<&[N]> {
96        let level_bounds = self.level_bounds();
97        if level >= level_bounds.len() {
98            return Err(GeoIndexError::General("Level out of bounds".to_string()));
99        }
100        let result = if level == 0 {
101            &self.boxes()[0..level_bounds[0]]
102        } else if level == level_bounds.len() {
103            &self.boxes()[level_bounds[level]..]
104        } else {
105            &self.boxes()[level_bounds[level - 1]..level_bounds[level]]
106        };
107        Ok(result)
108    }
109
110    /// Search an RTree given the provided bounding box.
111    ///
112    /// Results are the indexes of the inserted objects in insertion order.
113    fn search(&self, min_x: N, min_y: N, max_x: N, max_y: N) -> Vec<u32> {
114        let boxes = self.boxes();
115        let indices = self.indices();
116        if boxes.is_empty() {
117            return vec![];
118        }
119
120        let mut outer_node_index = boxes.len().checked_sub(4);
121
122        let mut queue = VecDeque::with_capacity(self.node_size() as usize);
123        let mut results = vec![];
124
125        while let Some(node_index) = outer_node_index {
126            // find the end index of the node
127            let end = (node_index + self.node_size() as usize * 4)
128                .min(upper_bound(node_index, self.level_bounds()));
129
130            // search through child nodes
131            for pos in (node_index..end).step_by(4) {
132                // Safety: pos was checked before to be within bounds
133                // Justification: avoiding bounds check improves performance by up to 30%
134                let (node_min_x, node_min_y, node_max_x, node_max_y) = unsafe {
135                    let node_min_x = *boxes.get_unchecked(pos);
136                    let node_min_y = *boxes.get_unchecked(pos + 1);
137                    let node_max_x = *boxes.get_unchecked(pos + 2);
138                    let node_max_y = *boxes.get_unchecked(pos + 3);
139                    (node_min_x, node_min_y, node_max_x, node_max_y)
140                };
141
142                // check if the query box disjoint with the node box
143                if max_x < node_min_x
144                    || max_y < node_min_y
145                    || min_x > node_max_x
146                    || min_y > node_max_y
147                {
148                    continue;
149                }
150
151                let index = indices.get(pos >> 2);
152
153                if node_index >= self.num_items() as usize * 4 {
154                    queue.push_back(index); // node; add it to the search queue
155                } else {
156                    // Since the max items of the index is u32, we can coerce to u32
157                    results.push(index.try_into().unwrap()); // leaf item
158                }
159            }
160
161            outer_node_index = queue.pop_front();
162        }
163
164        results
165    }
166
167    /// Search an RTree given the provided bounding box.
168    ///
169    /// Results are the indexes of the inserted objects in insertion order.
170    fn search_rect(&self, rect: &impl RectTrait<T = N>) -> Vec<u32> {
171        self.search(
172            rect.min().x(),
173            rect.min().y(),
174            rect.max().x(),
175            rect.max().y(),
176        )
177    }
178
179    /// Search items in order of distance from the given point.
180    ///
181    /// This method uses Euclidean distance by default. For other distance metrics,
182    /// use [`neighbors_with_distance`].
183    ///
184    /// ```
185    /// use geo_index::rtree::{RTreeBuilder, RTreeIndex, RTreeRef};
186    /// use geo_index::rtree::sort::HilbertSort;
187    ///
188    /// // Create an RTree
189    /// let mut builder = RTreeBuilder::<f64>::new(3);
190    /// builder.add(0., 0., 2., 2.);
191    /// builder.add(1., 1., 3., 3.);
192    /// builder.add(2., 2., 4., 4.);
193    /// let tree = builder.finish::<HilbertSort>();
194    ///
195    /// let results = tree.neighbors(5., 5., None, None);
196    /// assert_eq!(results, vec![2, 1, 0]);
197    /// ```
198    fn neighbors(
199        &self,
200        x: N,
201        y: N,
202        max_results: Option<usize>,
203        max_distance: Option<N>,
204    ) -> Vec<u32> {
205        // Use simple squared distance for backward compatibility
206        struct SimpleSquaredDistance;
207        impl<N: IndexableNum> SimpleDistanceMetric<N> for SimpleSquaredDistance {
208            fn distance(&self, x1: N, y1: N, x2: N, y2: N) -> N {
209                let dx = x2 - x1;
210                let dy = y2 - y1;
211                dx * dx + dy * dy
212            }
213            fn distance_to_bbox(&self, x: N, y: N, min_x: N, min_y: N, max_x: N, max_y: N) -> N {
214                let dx = axis_dist(x, min_x, max_x);
215                let dy = axis_dist(y, min_y, max_y);
216                dx * dx + dy * dy
217            }
218        }
219        let simple_distance = SimpleSquaredDistance;
220        self.neighbors_with_simple_distance(x, y, max_results, max_distance, &simple_distance)
221    }
222
223    /// Search items in order of distance from the given point using a simple distance metric.
224    ///
225    /// This is the base method for distance-based neighbor searches that works without the geo feature.
226    fn neighbors_with_simple_distance<M: SimpleDistanceMetric<N> + ?Sized>(
227        &self,
228        x: N,
229        y: N,
230        max_results: Option<usize>,
231        max_distance: Option<N>,
232        distance_metric: &M,
233    ) -> Vec<u32> {
234        let boxes = self.boxes();
235        let indices = self.indices();
236        let max_distance = max_distance.unwrap_or(distance_metric.max_distance());
237
238        let mut outer_node_index = Some(boxes.len() - 4);
239        let mut queue = BinaryHeap::new();
240        let mut results: Vec<u32> = vec![];
241
242        'outer: while let Some(node_index) = outer_node_index {
243            // find the end index of the node
244            let end = (node_index + self.node_size() as usize * 4)
245                .min(upper_bound(node_index, self.level_bounds()));
246
247            // add child nodes to the queue
248            for pos in (node_index..end).step_by(4) {
249                let index = indices.get(pos >> 2);
250
251                // Use the custom distance metric for bbox distance calculation
252                let dist = distance_metric.distance_to_bbox(
253                    x,
254                    y,
255                    boxes[pos],
256                    boxes[pos + 1],
257                    boxes[pos + 2],
258                    boxes[pos + 3],
259                );
260
261                if dist > max_distance {
262                    continue;
263                }
264
265                if node_index >= self.num_items() as usize * 4 {
266                    // node (use even id)
267                    queue.push(Reverse(NeighborNode {
268                        id: index << 1,
269                        dist,
270                    }));
271                } else {
272                    // leaf item (use odd id)
273                    // Use consistent distance calculation for both nodes and leaf items
274                    queue.push(Reverse(NeighborNode {
275                        id: (index << 1) + 1,
276                        dist,
277                    }));
278                }
279            }
280
281            // pop items from the queue
282            while !queue.is_empty() && queue.peek().is_some_and(|val| (val.0.id & 1) != 0) {
283                let dist = queue.peek().unwrap().0.dist;
284                if dist > max_distance {
285                    break 'outer;
286                }
287                let item = queue.pop().unwrap();
288                results.push((item.0.id >> 1).try_into().unwrap());
289                if max_results.is_some_and(|max_results| results.len() == max_results) {
290                    break 'outer;
291                }
292            }
293
294            if let Some(item) = queue.pop() {
295                outer_node_index = Some(item.0.id >> 1);
296            } else {
297                outer_node_index = None;
298            }
299        }
300
301        results
302    }
303
304    /// Search items in order of distance from the given point using a custom distance metric.
305    ///
306    /// This method allows you to specify a custom distance calculation method, such as
307    /// Euclidean, Haversine, or Spheroid distance.
308    ///
309    /// ```
310    /// use geo_index::rtree::{RTreeBuilder, RTreeIndex};
311    /// use geo_index::rtree::distance::{EuclideanDistance, HaversineDistance};
312    /// use geo_index::rtree::sort::HilbertSort;
313    ///
314    /// // Create an RTree with geographic coordinates (longitude, latitude)
315    /// let mut builder = RTreeBuilder::<f64>::new(3);
316    /// builder.add(-74.0, 40.7, -74.0, 40.7); // New York
317    /// builder.add(-0.1, 51.5, -0.1, 51.5);   // London
318    /// builder.add(139.7, 35.7, 139.7, 35.7); // Tokyo
319    /// let tree = builder.finish::<HilbertSort>();
320    ///
321    /// // Find nearest neighbors using Haversine distance (great-circle distance)
322    /// let haversine = HaversineDistance::default();
323    /// let results = tree.neighbors_with_distance(-74.0, 40.7, Some(2), None, &haversine);
324    /// ```
325    #[cfg(feature = "use-geo_0_31")]
326    fn neighbors_with_distance<M: DistanceMetric<N> + ?Sized>(
327        &self,
328        x: N,
329        y: N,
330        max_results: Option<usize>,
331        max_distance: Option<N>,
332        distance_metric: &M,
333    ) -> Vec<u32> {
334        self.neighbors_with_simple_distance(x, y, max_results, max_distance, distance_metric)
335    }
336
337    /// Search items in order of distance from the given coordinate.
338    fn neighbors_coord(
339        &self,
340        coord: &impl CoordTrait<T = N>,
341        max_results: Option<usize>,
342        max_distance: Option<N>,
343    ) -> Vec<u32> {
344        self.neighbors(coord.x(), coord.y(), max_results, max_distance)
345    }
346
347    /// Search items in order of distance from the given coordinate using a custom distance metric.
348    #[cfg(feature = "use-geo_0_31")]
349    fn neighbors_coord_with_distance<M: DistanceMetric<N> + ?Sized>(
350        &self,
351        coord: &impl CoordTrait<T = N>,
352        max_results: Option<usize>,
353        max_distance: Option<N>,
354        distance_metric: &M,
355    ) -> Vec<u32> {
356        self.neighbors_with_distance(
357            coord.x(),
358            coord.y(),
359            max_results,
360            max_distance,
361            distance_metric,
362        )
363    }
364
365    /// Search items in order of distance from a query geometry using a distance metric and geometry accessor.
366    ///
367    /// This method allows searching with geometry-to-geometry distance calculations.
368    /// The distance metric defines how distances are computed, and the geometry accessor
369    /// provides access to the actual geometries by index.
370    ///
371    /// ```
372    /// use geo_index::rtree::{RTreeBuilder, RTreeIndex};
373    /// use geo_index::rtree::distance::{EuclideanDistance, SliceGeometryAccessor};
374    /// use geo_index::rtree::sort::HilbertSort;
375    /// use geo_0_31::{Point, Geometry};
376    ///
377    /// // Create an RTree
378    /// let mut builder = RTreeBuilder::<f64>::new(3);
379    /// builder.add(0., 0., 2., 2.);
380    /// builder.add(5., 5., 7., 7.);
381    /// builder.add(10., 10., 12., 12.);
382    /// let tree = builder.finish::<HilbertSort>();
383    ///
384    /// // Example geometries
385    /// let geometries: Vec<Geometry<f64>> = vec![
386    ///     Geometry::Point(Point::new(1.0, 1.0)),
387    ///     Geometry::Point(Point::new(6.0, 6.0)),
388    ///     Geometry::Point(Point::new(11.0, 11.0)),
389    /// ];
390    ///
391    /// let metric = EuclideanDistance;
392    /// let accessor = SliceGeometryAccessor::new(&geometries);
393    /// let query_geom = Geometry::Point(Point::new(3.0, 3.0));
394    /// let results = tree.neighbors_geometry(&query_geom, None, None, &metric, &accessor);
395    /// ```
396    #[cfg(feature = "use-geo_0_31")]
397    fn neighbors_geometry<M: DistanceMetric<N> + ?Sized, A: GeometryAccessor + ?Sized>(
398        &self,
399        query_geometry: &Geometry<f64>,
400        max_results: Option<usize>,
401        max_distance: Option<N>,
402        distance_metric: &M,
403        accessor: &A,
404    ) -> Vec<u32> {
405        let boxes = self.boxes();
406        let indices = self.indices();
407        let max_distance = max_distance.unwrap_or(distance_metric.max_distance());
408
409        // Get the bounding box of the query geometry
410        let bounds = query_geometry.bounding_rect();
411        let (query_min_x, query_min_y, query_max_x, query_max_y) = if let Some(rect) = bounds {
412            let min = rect.min();
413            let max = rect.max();
414            (
415                N::from_f64(min.x).unwrap_or(N::zero()),
416                N::from_f64(min.y).unwrap_or(N::zero()),
417                N::from_f64(max.x).unwrap_or(N::zero()),
418                N::from_f64(max.y).unwrap_or(N::zero()),
419            )
420        } else {
421            // If no bounding box, use origin
422            (N::zero(), N::zero(), N::zero(), N::zero())
423        };
424
425        let mut outer_node_index = Some(boxes.len() - 4);
426        let mut queue = BinaryHeap::new();
427        let mut results: Vec<u32> = vec![];
428
429        'outer: while let Some(node_index) = outer_node_index {
430            // find the end index of the node
431            let end = (node_index + self.node_size() as usize * 4)
432                .min(upper_bound(node_index, self.level_bounds()));
433
434            // add child nodes to the queue
435            for pos in (node_index..end).step_by(4) {
436                let index = indices.get(pos >> 2);
437
438                let dist = if node_index >= self.num_items() as usize * 4 {
439                    // For internal nodes, use bbox-to-bbox distance as approximation
440                    let center_x = (query_min_x + query_max_x) / (N::one() + N::one());
441                    let center_y = (query_min_y + query_max_y) / (N::one() + N::one());
442
443                    distance_metric.distance_to_bbox(
444                        center_x,
445                        center_y,
446                        boxes[pos],
447                        boxes[pos + 1],
448                        boxes[pos + 2],
449                        boxes[pos + 3],
450                    )
451                } else {
452                    // For leaf items, use geometry-to-geometry distance
453                    if let Some(item_geom) = accessor.get_geometry(index) {
454                        distance_metric.distance_to_geometry(query_geometry, item_geom)
455                    } else {
456                        distance_metric.max_distance()
457                    }
458                };
459
460                if dist > max_distance {
461                    continue;
462                }
463
464                if node_index >= self.num_items() as usize * 4 {
465                    // node (use even id)
466                    queue.push(Reverse(NeighborNode {
467                        id: index << 1,
468                        dist,
469                    }));
470                } else {
471                    // leaf item (use odd id)
472                    queue.push(Reverse(NeighborNode {
473                        id: (index << 1) + 1,
474                        dist,
475                    }));
476                }
477            }
478
479            // pop items from the queue
480            while !queue.is_empty() && queue.peek().is_some_and(|val| (val.0.id & 1) != 0) {
481                let dist = queue.peek().unwrap().0.dist;
482                if dist > max_distance {
483                    break 'outer;
484                }
485                let item = queue.pop().unwrap();
486                results.push((item.0.id >> 1).try_into().unwrap());
487                if max_results.is_some_and(|max_results| results.len() == max_results) {
488                    break 'outer;
489                }
490            }
491
492            if let Some(item) = queue.pop() {
493                outer_node_index = Some(item.0.id >> 1);
494            } else {
495                outer_node_index = None;
496            }
497        }
498
499        results
500    }
501
502    /// Returns an iterator over the indexes of objects in this and another tree that intersect.
503    ///
504    /// Each returned object is of the form `(u32, u32)`, where the first is the positional
505    /// index of the "left" tree and the second is the index of the "right" tree.
506    fn intersection_candidates_with_other_tree<'a>(
507        &'a self,
508        other: &'a impl RTreeIndex<N>,
509    ) -> impl Iterator<Item = (u32, u32)> + 'a {
510        IntersectionIterator::from_trees(self, other)
511    }
512
513    /// Access the root node of the RTree for manual traversal.
514    fn root(&self) -> Node<'_, N, Self> {
515        Node::from_root(self)
516    }
517}
518
519/// A wrapper around a node and its distance for use in the priority queue.
520#[derive(Debug, Clone, Copy, PartialEq)]
521struct NeighborNode<N: IndexableNum> {
522    id: usize,
523    dist: N,
524}
525
526impl<N: IndexableNum> Eq for NeighborNode<N> {}
527
528impl<N: IndexableNum> Ord for NeighborNode<N> {
529    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
530        // We don't allow NaN. This should only panic on NaN
531        self.dist.partial_cmp(&other.dist).unwrap()
532    }
533}
534
535impl<N: IndexableNum> PartialOrd for NeighborNode<N> {
536    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
537        Some(self.cmp(other))
538    }
539}
540
541impl<N: IndexableNum> RTreeIndex<N> for RTree<N> {
542    fn boxes(&self) -> &[N] {
543        self.metadata.boxes_slice(&self.buffer)
544    }
545
546    fn indices(&self) -> Indices<'_> {
547        self.metadata.indices_slice(&self.buffer)
548    }
549
550    fn metadata(&self) -> &RTreeMetadata<N> {
551        &self.metadata
552    }
553}
554
555impl<N: IndexableNum> RTreeIndex<N> for RTreeRef<'_, N> {
556    fn boxes(&self) -> &[N] {
557        self.boxes
558    }
559
560    fn indices(&self) -> Indices<'_> {
561        self.indices
562    }
563
564    fn metadata(&self) -> &RTreeMetadata<N> {
565        &self.metadata
566    }
567}
568
569/// 1D distance from a value to a range.
570#[inline]
571pub(crate) fn axis_dist<N: IndexableNum>(k: N, min: N, max: N) -> N {
572    if k < min {
573        min - k
574    } else if k <= max {
575        N::zero()
576    } else {
577        k - max
578    }
579}
580
581#[cfg(test)]
582mod test {
583    // Replication of tests from flatbush js
584    mod js {
585        use crate::rtree::RTreeIndex;
586        use crate::test::{flatbush_js_test_data, flatbush_js_test_index};
587
588        #[test]
589        fn performs_bbox_search() {
590            let data = flatbush_js_test_data();
591            let index = flatbush_js_test_index();
592            let ids = index.search(40., 40., 60., 60.);
593
594            let mut results: Vec<usize> = vec![];
595            for id in ids {
596                results.push(data[4 * id as usize] as usize);
597                results.push(data[4 * id as usize + 1] as usize);
598                results.push(data[4 * id as usize + 2] as usize);
599                results.push(data[4 * id as usize + 3] as usize);
600            }
601
602            results.sort();
603
604            let mut expected = vec![
605                57, 59, 58, 59, 48, 53, 52, 56, 40, 42, 43, 43, 43, 41, 47, 43,
606            ];
607            expected.sort();
608
609            assert_eq!(results, expected);
610        }
611    }
612
613    #[cfg(feature = "use-geo_0_31")]
614    mod distance_metrics {
615        use crate::rtree::distance::{EuclideanDistance, HaversineDistance};
616        use crate::rtree::r#trait::SimpleDistanceMetric;
617        use crate::rtree::sort::HilbertSort;
618        use crate::rtree::{RTreeBuilder, RTreeIndex};
619
620        #[test]
621        fn test_euclidean_distance_neighbors() {
622            let mut builder = RTreeBuilder::<f64>::new(3);
623            builder.add(0., 0., 1., 1.);
624            builder.add(2., 2., 3., 3.);
625            builder.add(4., 4., 5., 5.);
626            let tree = builder.finish::<HilbertSort>();
627
628            let euclidean = EuclideanDistance;
629            let results = tree.neighbors_with_distance(0., 0., None, None, &euclidean);
630
631            // Should return items in order of distance from (0,0)
632            assert_eq!(results, vec![0, 1, 2]);
633        }
634
635        #[test]
636        fn test_haversine_distance_neighbors() {
637            let mut builder = RTreeBuilder::<f64>::new(3);
638            // Add some geographic points (longitude, latitude)
639            builder.add(-74.0, 40.7, -74.0, 40.7); // New York
640            builder.add(-0.1, 51.5, -0.1, 51.5); // London
641            builder.add(139.7, 35.7, 139.7, 35.7); // Tokyo
642            let tree = builder.finish::<HilbertSort>();
643
644            let haversine = HaversineDistance::default();
645            let results = tree.neighbors_with_distance(-74.0, 40.7, None, None, &haversine);
646
647            // From New York, should find New York first, then London, then Tokyo
648            assert_eq!(results, vec![0, 1, 2]);
649        }
650
651        #[test]
652        fn test_backward_compatibility() {
653            let mut builder = RTreeBuilder::<f64>::new(3);
654            builder.add(0., 0., 1., 1.);
655            builder.add(2., 2., 3., 3.);
656            builder.add(4., 4., 5., 5.);
657            let tree = builder.finish::<HilbertSort>();
658
659            // Test that original neighbors method still works
660            let results_original = tree.neighbors(0., 0., None, None);
661
662            // Test that new method with Euclidean distance gives same results
663            let euclidean = EuclideanDistance;
664            let results_new = tree.neighbors_with_distance(0., 0., None, None, &euclidean);
665
666            assert_eq!(results_original, results_new);
667        }
668
669        #[test]
670        fn test_max_distance_filtering() {
671            let mut builder = RTreeBuilder::<f64>::new(3);
672            builder.add(0., 0., 1., 1.);
673            builder.add(2., 2., 3., 3.);
674            builder.add(10., 10., 11., 11.);
675            let tree = builder.finish::<HilbertSort>();
676
677            let euclidean = EuclideanDistance;
678            // Only find neighbors within distance 5
679            let results = tree.neighbors_with_distance(0., 0., None, Some(5.0), &euclidean);
680
681            // Should only find first two items, not the distant third one
682            assert_eq!(results.len(), 2);
683            assert_eq!(results, vec![0, 1]);
684        }
685
686        #[test]
687        #[cfg(feature = "use-geo_0_31")]
688        fn test_geometry_neighbors_euclidean() {
689            use crate::r#type::IndexableNum;
690            use crate::rtree::distance::{DistanceMetric, SliceGeometryAccessor};
691            use geo_0_31::algorithm::{Distance, Euclidean};
692            use geo_0_31::{Geometry, Point};
693
694            let mut builder = RTreeBuilder::<f64>::new(3);
695            builder.add(0., 0., 2., 2.); // Item 0
696            builder.add(5., 5., 7., 7.); // Item 1
697            builder.add(10., 10., 12., 12.); // Item 2
698            let tree = builder.finish::<HilbertSort>();
699
700            // Geometries corresponding to the bboxes
701            let geometries: Vec<Geometry<f64>> = vec![
702                Geometry::Point(Point::new(1.0, 1.0)),   // Item 0
703                Geometry::Point(Point::new(6.0, 6.0)),   // Item 1
704                Geometry::Point(Point::new(11.0, 11.0)), // Item 2
705            ];
706
707            struct SimpleMetric;
708            impl<N: IndexableNum> SimpleDistanceMetric<N> for SimpleMetric {
709                fn distance(&self, x1: N, y1: N, x2: N, y2: N) -> N {
710                    let dx = x2 - x1;
711                    let dy = y2 - y1;
712                    (dx * dx + dy * dy).sqrt().unwrap_or(N::max_value())
713                }
714                fn distance_to_bbox(
715                    &self,
716                    x: N,
717                    y: N,
718                    min_x: N,
719                    min_y: N,
720                    max_x: N,
721                    max_y: N,
722                ) -> N {
723                    let dx = if x < min_x {
724                        min_x - x
725                    } else if x > max_x {
726                        x - max_x
727                    } else {
728                        N::zero()
729                    };
730                    let dy = if y < min_y {
731                        min_y - y
732                    } else if y > max_y {
733                        y - max_y
734                    } else {
735                        N::zero()
736                    };
737                    (dx * dx + dy * dy).sqrt().unwrap_or(N::max_value())
738                }
739            }
740            impl<N: IndexableNum> DistanceMetric<N> for SimpleMetric {
741                fn distance_to_geometry(&self, geom1: &Geometry<f64>, geom2: &Geometry<f64>) -> N {
742                    N::from_f64(Euclidean.distance(geom1, geom2)).unwrap_or(N::max_value())
743                }
744            }
745
746            let query_geom = Geometry::Point(Point::new(3.0, 3.0));
747            let metric = SimpleMetric;
748            let accessor = SliceGeometryAccessor::new(&geometries);
749            let results = tree.neighbors_geometry(&query_geom, None, None, &metric, &accessor);
750
751            // Item 0 should be closest to query point (3,3)
752            assert_eq!(results[0], 0);
753            assert_eq!(results[1], 1);
754            assert_eq!(results[2], 2);
755        }
756
757        #[test]
758        #[cfg(feature = "use-geo_0_31")]
759        fn test_geometry_neighbors_linestring() {
760            use crate::r#type::IndexableNum;
761            use crate::rtree::distance::{DistanceMetric, SliceGeometryAccessor};
762            use geo_0_31::algorithm::{Distance, Euclidean};
763            use geo_0_31::{coord, Geometry, LineString, Point};
764
765            let mut builder = RTreeBuilder::<f64>::new(3);
766            builder.add(0., 0., 10., 0.); // Item 0 - horizontal line
767            builder.add(5., 5., 15., 5.); // Item 1 - horizontal line higher up
768            builder.add(0., 10., 10., 10.); // Item 2 - horizontal line at top
769            let tree = builder.finish::<HilbertSort>();
770
771            // Geometries corresponding to the bboxes
772            let geometries: Vec<Geometry<f64>> = vec![
773                Geometry::LineString(LineString::new(vec![
774                    coord! { x: 0.0, y: 0.0 },
775                    coord! { x: 10.0, y: 0.0 },
776                ])),
777                Geometry::LineString(LineString::new(vec![
778                    coord! { x: 5.0, y: 5.0 },
779                    coord! { x: 15.0, y: 5.0 },
780                ])),
781                Geometry::LineString(LineString::new(vec![
782                    coord! { x: 0.0, y: 10.0 },
783                    coord! { x: 10.0, y: 10.0 },
784                ])),
785            ];
786
787            struct SimpleMetric;
788            impl<N: IndexableNum> SimpleDistanceMetric<N> for SimpleMetric {
789                fn distance(&self, x1: N, y1: N, x2: N, y2: N) -> N {
790                    let dx = x2 - x1;
791                    let dy = y2 - y1;
792                    (dx * dx + dy * dy).sqrt().unwrap_or(N::max_value())
793                }
794                fn distance_to_bbox(
795                    &self,
796                    x: N,
797                    y: N,
798                    min_x: N,
799                    min_y: N,
800                    max_x: N,
801                    max_y: N,
802                ) -> N {
803                    let dx = if x < min_x {
804                        min_x - x
805                    } else if x > max_x {
806                        x - max_x
807                    } else {
808                        N::zero()
809                    };
810                    let dy = if y < min_y {
811                        min_y - y
812                    } else if y > max_y {
813                        y - max_y
814                    } else {
815                        N::zero()
816                    };
817                    (dx * dx + dy * dy).sqrt().unwrap_or(N::max_value())
818                }
819            }
820            impl<N: IndexableNum> DistanceMetric<N> for SimpleMetric {
821                fn distance_to_geometry(&self, geom1: &Geometry<f64>, geom2: &Geometry<f64>) -> N {
822                    N::from_f64(Euclidean.distance(geom1, geom2)).unwrap_or(N::max_value())
823                }
824            }
825
826            let query_geom = Geometry::Point(Point::new(5.0, 2.0));
827            let metric = SimpleMetric;
828            let accessor = SliceGeometryAccessor::new(&geometries);
829            let results = tree.neighbors_geometry(&query_geom, None, None, &metric, &accessor);
830
831            // Item 0 (bottom line) should be closest to point (5, 2)
832            assert_eq!(results[0], 0);
833        }
834
835        #[test]
836        #[cfg(feature = "use-geo_0_31")]
837        fn test_geometry_neighbors_with_max_results() {
838            use crate::r#type::IndexableNum;
839            use crate::rtree::distance::{DistanceMetric, SliceGeometryAccessor};
840            use geo_0_31::algorithm::{Distance, Euclidean};
841            use geo_0_31::{Geometry, Point};
842
843            let mut builder = RTreeBuilder::<f64>::new(5);
844            for i in 0..5 {
845                let x = (i * 3) as f64;
846                builder.add(x, x, x + 1., x + 1.);
847            }
848            let tree = builder.finish::<HilbertSort>();
849
850            // Create geometries for each bbox
851            let geometries: Vec<Geometry<f64>> = (0..5)
852                .map(|i| {
853                    let x = (i * 3) as f64;
854                    Geometry::Point(Point::new(x + 0.5, x + 0.5))
855                })
856                .collect();
857
858            struct SimpleMetric;
859            impl<N: IndexableNum> SimpleDistanceMetric<N> for SimpleMetric {
860                fn distance(&self, x1: N, y1: N, x2: N, y2: N) -> N {
861                    let dx = x2 - x1;
862                    let dy = y2 - y1;
863                    (dx * dx + dy * dy).sqrt().unwrap_or(N::max_value())
864                }
865                fn distance_to_bbox(
866                    &self,
867                    x: N,
868                    y: N,
869                    min_x: N,
870                    min_y: N,
871                    max_x: N,
872                    max_y: N,
873                ) -> N {
874                    let dx = if x < min_x {
875                        min_x - x
876                    } else if x > max_x {
877                        x - max_x
878                    } else {
879                        N::zero()
880                    };
881                    let dy = if y < min_y {
882                        min_y - y
883                    } else if y > max_y {
884                        y - max_y
885                    } else {
886                        N::zero()
887                    };
888                    (dx * dx + dy * dy).sqrt().unwrap_or(N::max_value())
889                }
890            }
891            impl<N: IndexableNum> DistanceMetric<N> for SimpleMetric {
892                fn distance_to_geometry(&self, geom1: &Geometry<f64>, geom2: &Geometry<f64>) -> N {
893                    N::from_f64(Euclidean.distance(geom1, geom2)).unwrap_or(N::max_value())
894                }
895            }
896
897            let query_geom = Geometry::Point(Point::new(5.0, 5.0));
898            let metric = SimpleMetric;
899            let accessor = SliceGeometryAccessor::new(&geometries);
900            let results = tree.neighbors_geometry(&query_geom, Some(3), None, &metric, &accessor);
901
902            assert_eq!(results.len(), 3);
903            // Should get the 3 closest items
904        }
905
906        #[test]
907        #[cfg(feature = "use-geo_0_31")]
908        fn test_geometry_neighbors_haversine() {
909            use crate::r#type::IndexableNum;
910            use crate::rtree::distance::{DistanceMetric, SliceGeometryAccessor};
911            use geo_0_31::algorithm::{Centroid, Distance, Haversine};
912            use geo_0_31::{Geometry, Point};
913
914            let mut builder = RTreeBuilder::<f64>::new(3);
915            // Geographic bounding boxes (lon, lat)
916            builder.add(-74.1, 40.6, -74.0, 40.7); // New York area
917            builder.add(-0.2, 51.4, -0.1, 51.5); // London area
918            builder.add(139.6, 35.6, 139.7, 35.7); // Tokyo area
919            let tree = builder.finish::<HilbertSort>();
920
921            let geometries: Vec<Geometry<f64>> = vec![
922                Geometry::Point(Point::new(-74.0, 40.7)), // New York
923                Geometry::Point(Point::new(-0.1, 51.5)),  // London
924                Geometry::Point(Point::new(139.7, 35.7)), // Tokyo
925            ];
926
927            struct HaversineMetric;
928            impl<N: IndexableNum> SimpleDistanceMetric<N> for HaversineMetric {
929                fn distance(&self, lon1: N, lat1: N, lon2: N, lat2: N) -> N {
930                    let p1 = Point::new(lon1.to_f64().unwrap_or(0.0), lat1.to_f64().unwrap_or(0.0));
931                    let p2 = Point::new(lon2.to_f64().unwrap_or(0.0), lat2.to_f64().unwrap_or(0.0));
932                    N::from_f64(Haversine.distance(p1, p2)).unwrap_or(N::max_value())
933                }
934                fn distance_to_bbox(
935                    &self,
936                    lon: N,
937                    lat: N,
938                    min_lon: N,
939                    min_lat: N,
940                    max_lon: N,
941                    max_lat: N,
942                ) -> N {
943                    let lon_f = lon.to_f64().unwrap_or(0.0);
944                    let lat_f = lat.to_f64().unwrap_or(0.0);
945                    let min_lon_f = min_lon.to_f64().unwrap_or(0.0);
946                    let min_lat_f = min_lat.to_f64().unwrap_or(0.0);
947                    let max_lon_f = max_lon.to_f64().unwrap_or(0.0);
948                    let max_lat_f = max_lat.to_f64().unwrap_or(0.0);
949                    let closest_lon = lon_f.clamp(min_lon_f, max_lon_f);
950                    let closest_lat = lat_f.clamp(min_lat_f, max_lat_f);
951                    let point = Point::new(lon_f, lat_f);
952                    let closest_point = Point::new(closest_lon, closest_lat);
953                    N::from_f64(Haversine.distance(point, closest_point)).unwrap_or(N::max_value())
954                }
955            }
956            impl<N: IndexableNum> DistanceMetric<N> for HaversineMetric {
957                fn distance_to_geometry(&self, geom1: &Geometry<f64>, geom2: &Geometry<f64>) -> N {
958                    let c1 = geom1.centroid().unwrap_or(Point::new(0.0, 0.0));
959                    let c2 = geom2.centroid().unwrap_or(Point::new(0.0, 0.0));
960                    N::from_f64(Haversine.distance(c1, c2)).unwrap_or(N::max_value())
961                }
962            }
963
964            let query_geom = Geometry::Point(Point::new(-74.0, 40.7)); // New York
965            let metric = HaversineMetric;
966            let accessor = SliceGeometryAccessor::new(&geometries);
967            let results = tree.neighbors_geometry(&query_geom, None, None, &metric, &accessor);
968
969            // New York should be closest (distance 0)
970            assert_eq!(results[0], 0);
971        }
972    }
973}
geo_index/rtree/trait.rs

geo_index/rtree/
trait.rs
