Skip to main content

rstar/primitives/
geom_with_data.rs

1use crate::object::PointDistance;
2use crate::object::RTreeObject;
3use crate::{envelope::Envelope, object::Distance};
4
5/// An [RTreeObject] with a geometry and some associated data that can be inserted into an r-tree.
6///
7/// Often, adding metadata (like a database ID) to a geometry is required before adding it
8/// into an r-tree. This struct removes some of the boilerplate required to do so.
9///
10/// **Note:** while the container itself implements [RTreeObject], you will have to go through its
11/// [`geom`][Self::geom] method in order to access geometry-specific methods.
12///
13/// # Example
14/// ```
15/// use rstar::{RTree, PointDistance};
16/// use rstar::primitives::GeomWithData;
17///
18/// type RestaurantLocation = GeomWithData<[f64; 2], &'static str>;
19///
20/// let mut restaurants = RTree::new();
21/// restaurants.insert(RestaurantLocation::new([0.3, 0.2], "Pete's Pizza Place"));
22/// restaurants.insert(RestaurantLocation::new([-0.8, 0.0], "The Great Steak"));
23/// restaurants.insert(RestaurantLocation::new([0.2, -0.2], "Fishy Fortune"));
24///
25/// let my_location = [0.0, 0.0];
26///
27/// // Now find the closest restaurant!
28/// let place = restaurants.nearest_neighbor(my_location).unwrap();
29/// println!("Let's go to {}", place.data);
30/// println!("It's really close, only {} miles", place.distance_2(&my_location));
31/// ```
32#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
33#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
34pub struct GeomWithData<R: RTreeObject, T> {
35    geom: R,
36    /// Data to be associated with the geometry being stored in the [`RTree`](crate::RTree).
37    pub data: T,
38}
39
40impl<R: RTreeObject, T> RTreeObject for GeomWithData<R, T> {
41    type Envelope = R::Envelope;
42
43    fn envelope(&self) -> Self::Envelope {
44        self.geom.envelope()
45    }
46}
47
48impl<R: PointDistance, T> PointDistance for GeomWithData<R, T> {
49    fn distance_2(&self, point: &<Self::Envelope as Envelope>::Point) -> Distance<Self> {
50        self.geom.distance_2(point)
51    }
52
53    fn contains_point(&self, p: &<Self::Envelope as Envelope>::Point) -> bool {
54        self.geom.contains_point(p)
55    }
56
57    fn distance_2_if_less_or_equal(
58        &self,
59        point: &<Self::Envelope as Envelope>::Point,
60        max_distance_2: Distance<Self>,
61    ) -> Option<Distance<Self>> {
62        self.geom.distance_2_if_less_or_equal(point, max_distance_2)
63    }
64}
65
66impl<R: RTreeObject, T> GeomWithData<R, T> {
67    /// Create a new [GeomWithData] struct using the provided geometry and data.
68    pub fn new(geom: R, data: T) -> Self {
69        Self { geom, data }
70    }
71
72    /// Get a reference to the container's geometry.
73    pub fn geom(&self) -> &R {
74        &self.geom
75    }
76}
77
78#[cfg(test)]
79mod test {
80    use super::GeomWithData;
81    use crate::object::PointDistance;
82
83    use approx::*;
84
85    use crate::{primitives::Line, RTree};
86
87    #[test]
88    fn container_in_rtree() {
89        let line_1 = GeomWithData::new(Line::new([0.0, 0.0], [1.0, 1.0]), ());
90        let line_2 = GeomWithData::new(Line::new([0.0, 0.0], [-1.0, 1.0]), ());
91        let tree = RTree::bulk_load(vec![line_1, line_2]);
92
93        assert!(tree.contains(&line_1));
94    }
95
96    #[test]
97    fn container_edge_distance() {
98        let edge = GeomWithData::new(Line::new([0.5, 0.5], [0.5, 2.0]), 1usize);
99
100        assert_abs_diff_eq!(edge.distance_2(&[0.5, 0.5]), 0.0);
101        assert_abs_diff_eq!(edge.distance_2(&[0.0, 0.5]), 0.5 * 0.5);
102        assert_abs_diff_eq!(edge.distance_2(&[0.5, 1.0]), 0.0);
103        assert_abs_diff_eq!(edge.distance_2(&[0.0, 0.0]), 0.5);
104        assert_abs_diff_eq!(edge.distance_2(&[0.0, 1.0]), 0.5 * 0.5);
105        assert_abs_diff_eq!(edge.distance_2(&[1.0, 1.0]), 0.5 * 0.5);
106        assert_abs_diff_eq!(edge.distance_2(&[1.0, 3.0]), 0.5 * 0.5 + 1.0);
107    }
108
109    #[test]
110    fn container_length_2() {
111        let line = GeomWithData::new(Line::new([1, -1], [5, 5]), 1usize);
112
113        assert_eq!(line.geom().length_2(), 16 + 36);
114    }
115
116    #[test]
117    fn container_nearest_neighbour() {
118        let mut lines = RTree::new();
119        lines.insert(GeomWithData::new(
120            Line::new([0.0, 0.0], [1.0, 1.0]),
121            "Line A",
122        ));
123        lines.insert(GeomWithData::new(
124            Line::new([0.0, 0.0], [-1.0, 1.0]),
125            "Line B",
126        ));
127        let my_location = [0.0, 0.0];
128        // Now find the closest line
129        let place = lines.nearest_neighbor(my_location).unwrap();
130
131        assert_eq!(place.data, "Line A");
132    }
133}