geo_index/rtree/
traversal.rs

1//! Utilities to traverse the RTree structure.
2
3use geo_traits::{
4    GeometryTrait, RectTrait, UnimplementedGeometryCollection, UnimplementedLine,
5    UnimplementedLineString, UnimplementedMultiLineString, UnimplementedMultiPoint,
6    UnimplementedMultiPolygon, UnimplementedPoint, UnimplementedPolygon, UnimplementedTriangle,
7};
8
9use crate::r#type::{Coord, IndexableNum};
10use crate::rtree::util::upper_bound;
11use crate::rtree::RTreeIndex;
12use core::mem::take;
13use std::marker::PhantomData;
14
15/// An internal node in the RTree.
16#[derive(Debug, Clone)]
17pub struct Node<'a, N: IndexableNum, T: RTreeIndex<N>> {
18    /// The tree that this node is a reference onto
19    tree: &'a T,
20
21    /// This points to the position in the full `boxes` slice of the **first** coordinate of
22    /// this node. So
23    /// ```notest
24    /// self.tree.boxes()[self.pos]
25    /// ```
26    /// accesses the `min_x` coordinate of this node.
27    ///
28    /// This also relates to the children and the insertion index. When this is `<
29    /// self.tree.num_items() * 4`, it means it's a leaf node at the bottom of the tree. In this
30    /// case, calling `>> 2` on this finds the original insertion index.
31    ///
32    /// When this is `>= self.tree.num_items() * 4`, it means it's _not_ a leaf node, and calling
33    /// `>> 2` retrieves the `pos` of the first of its children.
34    pos: usize,
35
36    phantom: PhantomData<N>,
37}
38
39impl<'a, N: IndexableNum, T: RTreeIndex<N>> Node<'a, N, T> {
40    fn new(tree: &'a T, pos: usize) -> Self {
41        Self {
42            tree,
43            pos,
44            phantom: PhantomData,
45        }
46    }
47
48    pub(crate) fn from_root(tree: &'a T) -> Self {
49        let root_index = tree.boxes().len() - 4;
50        Self {
51            tree,
52            pos: root_index,
53            phantom: PhantomData,
54        }
55    }
56
57    /// Get the minimum `x` value of this node.
58    #[inline]
59    pub fn min_x(&self) -> N {
60        self.tree.boxes()[self.pos]
61    }
62
63    /// Get the minimum `y` value of this node.
64    #[inline]
65    pub fn min_y(&self) -> N {
66        self.tree.boxes()[self.pos + 1]
67    }
68
69    /// Get the maximum `x` value of this node.
70    #[inline]
71    pub fn max_x(&self) -> N {
72        self.tree.boxes()[self.pos + 2]
73    }
74
75    /// Get the maximum `y` value of this node.
76    #[inline]
77    pub fn max_y(&self) -> N {
78        self.tree.boxes()[self.pos + 3]
79    }
80
81    /// Returns `true` if this is a leaf node without children.
82    #[inline]
83    pub fn is_leaf(&self) -> bool {
84        self.pos < self.tree.num_items() as usize * 4
85    }
86
87    /// Returns `true` if this is an intermediate node with children.
88    #[inline]
89    pub fn is_parent(&self) -> bool {
90        !self.is_leaf()
91    }
92
93    /// Returns `true` if this node intersects another node.
94    #[inline]
95    pub fn intersects<T2: RTreeIndex<N>>(&self, other: &Node<N, T2>) -> bool {
96        if self.max_x() < other.min_x() {
97            return false;
98        }
99
100        if self.max_y() < other.min_y() {
101            return false;
102        }
103
104        if self.min_x() > other.max_x() {
105            return false;
106        }
107
108        if self.min_y() > other.max_y() {
109            return false;
110        }
111
112        true
113    }
114
115    /// Returns an iterator over the child nodes of this node.
116    ///
117    /// Returns `None` if [`Self::is_parent`] is `false`.
118    pub fn children(&self) -> Option<impl Iterator<Item = Node<'_, N, T>>> {
119        if self.is_parent() {
120            Some(self.children_unchecked())
121        } else {
122            None
123        }
124    }
125
126    /// Returns an iterator over the child nodes of this node. This is only valid when
127    /// [`Self::is_parent`] is `true`.
128    pub fn children_unchecked(&self) -> impl Iterator<Item = Node<'_, N, T>> {
129        debug_assert!(self.is_parent());
130
131        // find the start and end indexes of the children of this node
132        let start_child_pos = self.tree.indices().get(self.pos >> 2);
133        let end_children_pos = (start_child_pos + self.tree.node_size() as usize * 4)
134            .min(upper_bound(start_child_pos, self.tree.level_bounds()));
135
136        (start_child_pos..end_children_pos)
137            .step_by(4)
138            .map(|pos| Node::new(self.tree, pos))
139    }
140
141    /// The original insertion index. This is only valid when this is a leaf node, which you can
142    /// check with `Self::is_leaf`.
143    ///
144    /// Returns `None` if [`Self::is_leaf`] is `false`.
145    #[inline]
146    pub fn insertion_index(&self) -> Option<u32> {
147        if self.is_leaf() {
148            Some(self.insertion_index_unchecked())
149        } else {
150            None
151        }
152    }
153
154    /// The original insertion index. This is only valid when this is a leaf node, which you can
155    /// check with `Self::is_leaf`.
156    #[inline]
157    pub fn insertion_index_unchecked(&self) -> u32 {
158        debug_assert!(self.is_leaf());
159        self.tree.indices().get(self.pos >> 2) as u32
160    }
161}
162
163impl<N: IndexableNum, T: RTreeIndex<N>> RectTrait for Node<'_, N, T> {
164    type CoordType<'a>
165        = Coord<N>
166    where
167        Self: 'a;
168
169    fn min(&self) -> Self::CoordType<'_> {
170        Coord {
171            x: self.min_x(),
172            y: self.min_y(),
173        }
174    }
175
176    fn max(&self) -> Self::CoordType<'_> {
177        Coord {
178            x: self.max_x(),
179            y: self.max_y(),
180        }
181    }
182}
183
184impl<N: IndexableNum, T: RTreeIndex<N>> GeometryTrait for Node<'_, N, T> {
185    type T = N;
186
187    type PointType<'a>
188        = UnimplementedPoint<N>
189    where
190        Self: 'a;
191
192    type LineStringType<'a>
193        = UnimplementedLineString<N>
194    where
195        Self: 'a;
196
197    type PolygonType<'a>
198        = UnimplementedPolygon<N>
199    where
200        Self: 'a;
201
202    type MultiPointType<'a>
203        = UnimplementedMultiPoint<N>
204    where
205        Self: 'a;
206
207    type MultiLineStringType<'a>
208        = UnimplementedMultiLineString<N>
209    where
210        Self: 'a;
211
212    type MultiPolygonType<'a>
213        = UnimplementedMultiPolygon<N>
214    where
215        Self: 'a;
216
217    type GeometryCollectionType<'a>
218        = UnimplementedGeometryCollection<N>
219    where
220        Self: 'a;
221
222    type RectType<'a>
223        = Node<'a, N, T>
224    where
225        Self: 'a;
226
227    type TriangleType<'a>
228        = UnimplementedTriangle<N>
229    where
230        Self: 'a;
231
232    type LineType<'a>
233        = UnimplementedLine<N>
234    where
235        Self: 'a;
236
237    fn dim(&self) -> geo_traits::Dimensions {
238        geo_traits::Dimensions::Xy
239    }
240
241    fn as_type(
242        &self,
243    ) -> geo_traits::GeometryType<
244        '_,
245        Self::PointType<'_>,
246        Self::LineStringType<'_>,
247        Self::PolygonType<'_>,
248        Self::MultiPointType<'_>,
249        Self::MultiLineStringType<'_>,
250        Self::MultiPolygonType<'_>,
251        Self::GeometryCollectionType<'_>,
252        Self::RectType<'_>,
253        Self::TriangleType<'_>,
254        Self::LineType<'_>,
255    > {
256        geo_traits::GeometryType::Rect(self)
257    }
258}
259
260// This is copied from rstar under the MIT/Apache 2 license
261// https://github.com/georust/rstar/blob/6c23af0f3acc0c4668ce6c368820e0fa986a65b4/rstar/src/algorithm/intersection_iterator.rs
262pub(crate) struct IntersectionIterator<'a, N, T1, T2>
263where
264    N: IndexableNum,
265    T1: RTreeIndex<N>,
266    T2: RTreeIndex<N>,
267{
268    left: &'a T1,
269    right: &'a T2,
270    todo_list: Vec<(usize, usize)>,
271    candidates: Vec<usize>,
272    phantom: PhantomData<N>,
273}
274
275impl<'a, N, T1, T2> IntersectionIterator<'a, N, T1, T2>
276where
277    N: IndexableNum,
278    T1: RTreeIndex<N>,
279    T2: RTreeIndex<N>,
280{
281    pub(crate) fn from_trees(root1: &'a T1, root2: &'a T2) -> Self {
282        let mut intersections = IntersectionIterator {
283            left: root1,
284            right: root2,
285            todo_list: Vec::new(),
286            candidates: Vec::new(),
287            phantom: PhantomData,
288        };
289        intersections.add_intersecting_children(&root1.root(), &root2.root());
290        intersections
291    }
292
293    #[allow(dead_code)]
294    pub(crate) fn new(root1: &'a Node<N, T1>, root2: &'a Node<N, T2>) -> Self {
295        let mut intersections = IntersectionIterator {
296            left: root1.tree,
297            right: root2.tree,
298            todo_list: Vec::new(),
299            candidates: Vec::new(),
300            phantom: PhantomData,
301        };
302        intersections.add_intersecting_children(root1, root2);
303        intersections
304    }
305
306    fn push_if_intersecting(&mut self, node1: &'_ Node<N, T1>, node2: &'_ Node<N, T2>) {
307        if node1.intersects(node2) {
308            self.todo_list.push((node1.pos, node2.pos));
309        }
310    }
311
312    fn add_intersecting_children(&mut self, parent1: &'_ Node<N, T1>, parent2: &'_ Node<N, T2>) {
313        if !parent1.intersects(parent2) {
314            return;
315        }
316
317        let children1 = parent1
318            .children_unchecked()
319            .filter(|c1| c1.intersects(parent2));
320
321        let mut children2 = take(&mut self.candidates);
322        children2.extend(
323            parent2
324                .children_unchecked()
325                .filter(|c2| c2.intersects(parent1))
326                .map(|c| c.pos),
327        );
328
329        for child1 in children1 {
330            for child2 in &children2 {
331                self.push_if_intersecting(&child1, &Node::new(self.right, *child2));
332            }
333        }
334
335        children2.clear();
336        self.candidates = children2;
337    }
338}
339
340impl<N, T1, T2> Iterator for IntersectionIterator<'_, N, T1, T2>
341where
342    N: IndexableNum,
343    T1: RTreeIndex<N>,
344    T2: RTreeIndex<N>,
345{
346    type Item = (u32, u32);
347
348    fn next(&mut self) -> Option<Self::Item> {
349        while let Some((left_index, right_index)) = self.todo_list.pop() {
350            let left = Node::new(self.left, left_index);
351            let right = Node::new(self.right, right_index);
352            match (left.is_leaf(), right.is_leaf()) {
353                (true, true) => {
354                    return Some((
355                        left.insertion_index_unchecked(),
356                        right.insertion_index_unchecked(),
357                    ))
358                }
359                (true, false) => right
360                    .children_unchecked()
361                    .for_each(|c| self.push_if_intersecting(&left, &c)),
362                (false, true) => left
363                    .children_unchecked()
364                    .for_each(|c| self.push_if_intersecting(&c, &right)),
365                (false, false) => self.add_intersecting_children(&left, &right),
366            }
367        }
368        None
369    }
370}
371
372#[cfg(test)]
373mod test {
374    use super::*;
375    use crate::test::flatbush_js_test_index;
376
377    #[test]
378    fn test_node() {
379        let tree = flatbush_js_test_index();
380
381        let top_box = tree.boxes_at_level(2).unwrap();
382
383        // Should only be one box
384        assert_eq!(top_box.len(), 4);
385
386        // Root node should match that one box in the top level
387        let root_node = tree.root();
388        assert_eq!(root_node.min_x(), top_box[0]);
389        assert_eq!(root_node.min_y(), top_box[1]);
390        assert_eq!(root_node.max_x(), top_box[2]);
391        assert_eq!(root_node.max_y(), top_box[3]);
392
393        assert!(root_node.is_parent());
394
395        let level_1_boxes = tree.boxes_at_level(1).unwrap();
396        let level_1 = root_node.children_unchecked().collect::<Vec<_>>();
397        assert_eq!(level_1.len(), level_1_boxes.len() / 4);
398    }
399}
400
401#[cfg(test)]
402mod test_issue_42 {
403    use std::collections::HashSet;
404
405    use crate::rtree::sort::HilbertSort;
406    use crate::rtree::{RTreeBuilder, RTreeIndex};
407    use geo::Polygon;
408    use geo::{BoundingRect, Geometry};
409    use geozero::geo_types::GeoWriter;
410    use geozero::geojson::read_geojson_fc;
411    use rstar::primitives::GeomWithData;
412    use rstar::{primitives::Rectangle, AABB};
413    use zip::ZipArchive;
414
415    // Find tree self-intersection canddiates using rstar
416    fn geo_contiguity(geom: &[Polygon]) -> HashSet<(usize, usize)> {
417        let to_insert = geom
418            .iter()
419            .enumerate()
420            .map(|(i, gi)| {
421                let rect = gi.bounding_rect().unwrap();
422                let aabb =
423                    AABB::from_corners([rect.min().x, rect.min().y], [rect.max().x, rect.max().y]);
424
425                GeomWithData::new(Rectangle::from_aabb(aabb), i)
426            })
427            .collect::<Vec<_>>();
428
429        let tree = rstar::RTree::bulk_load(to_insert);
430        let candidates = tree
431            .intersection_candidates_with_other_tree(&tree)
432            .map(|(left_candidate, right_candidate)| (left_candidate.data, right_candidate.data));
433
434        HashSet::from_iter(candidates)
435    }
436
437    // Find tree self-intersection canddiates using geo-index
438    fn geo_index_contiguity(geoms: &Vec<Polygon>, node_size: u16) -> HashSet<(usize, usize)> {
439        let mut tree_builder = RTreeBuilder::new_with_node_size(geoms.len() as _, node_size);
440        for geom in geoms {
441            tree_builder.add_rect(&geom.bounding_rect().unwrap());
442        }
443        let tree = tree_builder.finish::<HilbertSort>();
444
445        let candidates = tree
446            .intersection_candidates_with_other_tree(&tree)
447            .map(|(l, r)| (l as usize, r as usize));
448
449        HashSet::from_iter(candidates)
450    }
451
452    #[test]
453    fn test_repro_issue_42() {
454        let file = std::fs::File::open("fixtures/issue_42.geojson.zip").unwrap();
455        let mut zip_archive = ZipArchive::new(file).unwrap();
456        let zipped_file = zip_archive.by_name("guerry.geojson").unwrap();
457        let reader = std::io::BufReader::new(zipped_file);
458
459        let mut geo_writer = GeoWriter::new();
460        read_geojson_fc(reader, &mut geo_writer).unwrap();
461
462        let geoms = match geo_writer.take_geometry().unwrap() {
463            Geometry::GeometryCollection(gc) => gc.0,
464            _ => panic!(),
465        };
466
467        let mut polys = vec![];
468        for geom in geoms {
469            let poly = match geom {
470                Geometry::Polygon(poly) => poly,
471                _ => panic!(),
472            };
473            polys.push(poly);
474        }
475
476        let geo_index_self_intersection = geo_index_contiguity(&polys, 10);
477        let geo_self_intersection = geo_contiguity(&polys);
478
479        assert_eq!(
480            geo_index_self_intersection, geo_self_intersection,
481            "The two intersections should match!"
482        );
483    }
484}
geo_index/rtree/traversal.rs

geo_index/rtree/
traversal.rs
