rstar/
rtree.rs

1use crate::algorithm::bulk_load;
2use crate::algorithm::iterators::*;
3use crate::algorithm::nearest_neighbor;
4use crate::algorithm::nearest_neighbor::NearestNeighborDistance2Iterator;
5use crate::algorithm::nearest_neighbor::NearestNeighborIterator;
6use crate::algorithm::removal;
7use crate::algorithm::selection_functions::*;
8use crate::envelope::Envelope;
9use crate::node::ParentNode;
10use crate::object::{PointDistance, RTreeObject};
11use crate::params::{verify_parameters, DefaultParams, InsertionStrategy, RTreeParams};
12use crate::Point;
13use core::ops::ControlFlow;
14
15#[cfg(not(test))]
16use alloc::vec::Vec;
17
18#[cfg(feature = "serde")]
19use serde::{Deserialize, Serialize};
20
21impl<T, Params> Default for RTree<T, Params>
22where
23    T: RTreeObject,
24    Params: RTreeParams,
25{
26    fn default() -> Self {
27        Self::new_with_params()
28    }
29}
30
31/// An n-dimensional r-tree data structure.
32///
33/// # R-Trees
34/// R-Trees are data structures containing multi-dimensional objects like points, rectangles
35/// or polygons. They are optimized for retrieving the nearest neighbor at any point.
36///
37/// R-trees can efficiently find answers to queries like "Find the nearest point of a polygon",
38/// "Find all police stations within a rectangle" or "Find the 10 nearest restaurants, sorted
39/// by their distances". Compared to a naive implementation for these scenarios that runs
40/// in `O(n)` for `n` inserted elements, r-trees reduce this time to `O(log(n))`.
41///
42/// However, creating an r-tree is time consuming
43/// and runs in `O(n * log(n))`. Thus, r-trees are suited best if many queries and only few
44/// insertions are made. `rstar` also supports [bulk loading](RTree::bulk_load),
45/// which cuts down the constant factors when creating an r-tree significantly compared to
46/// sequential insertions.
47///
48/// R-trees are also _dynamic_: points can be inserted and removed from an existing tree.
49///
50/// ## Partitioning heuristics
51/// The inserted objects are internally partitioned into several boxes which should have small
52/// overlap and volume. This is done heuristically. While the originally proposed heuristic focused
53/// on fast insertion operations, the resulting r-trees were often suboptimally structured. Another
54/// heuristic, called `R*-tree` (r-star-tree), was proposed to improve the tree structure at the cost of
55/// longer insertion operations and is currently the crate's only implemented
56/// [InsertionStrategy].
57///
58/// # Usage
59/// The items inserted into an r-tree must implement the [RTreeObject]
60/// trait. To support nearest neighbor queries, implement the [PointDistance]
61/// trait. Some useful geometric primitives that implement the above traits can be found in the
62/// [crate::primitives] module. Several primitives in the [`geo-types`](https://docs.rs/geo-types/) crate also
63/// implement these traits.
64///
65/// ## Example
66/// ```
67/// use rstar::RTree;
68///
69/// let mut tree = RTree::new();
70/// tree.insert([0.1, 0.0f32]);
71/// tree.insert([0.2, 0.1]);
72/// tree.insert([0.3, 0.0]);
73///
74/// assert_eq!(tree.nearest_neighbor(&[0.4, -0.1]), Some(&[0.3, 0.0]));
75/// tree.remove(&[0.3, 0.0]);
76/// assert_eq!(tree.nearest_neighbor(&[0.4, 0.3]), Some(&[0.2, 0.1]));
77///
78/// assert_eq!(tree.size(), 2);
79/// // &RTree implements IntoIterator!
80/// for point in &tree {
81///     println!("Tree contains a point {:?}", point);
82/// }
83/// ```
84///
85/// ## Supported point types
86/// All types implementing the [Point] trait can be used as underlying point type.
87/// By default, fixed size arrays can be used as points.
88///
89/// # Associating Data with Geometries
90/// Users wishing to store associated data with geometries can use [crate::primitives::GeomWithData].
91///
92/// # Runtime and Performance
93/// The runtime of query operations (e.g. `nearest neighbor` or `contains`) is usually
94/// `O(log(n))`, where `n` refers to the number of elements contained in the r-tree.
95/// A naive sequential algorithm would take `O(n)` time. However, r-trees incur higher
96/// build up times: inserting an element into an r-tree costs `O(log(n))` time.
97///
98/// Most of the selection methods, meaning those with names beginning with `locate_`,
99/// return iterators which are driven externally and can therefore be combined into
100/// more complex pipelines using the combinators defined on the [`Iterator`] trait.
101///
102/// This flexiblity does come at the cost of temporary heap allocations required
103/// to keep track of the iteration state. Alternative methods using internal iteration
104/// are provided to avoid this overhead, their names ending in `_int` or `_int_mut`.
105///
106/// They use a callback-based interface to pass the selected objects on to the caller
107/// thereby being able to use the stack to keep track of the state required for
108/// traversing the tree.
109///
110/// # Bulk loading
111/// In many scenarios, insertion is only carried out once for many points. In this case,
112/// [RTree::bulk_load] will be considerably faster. Its total run time
113/// is still `O(nlog(n))`, i.e. `O(log(n))` per element inserted, but the scaling
114/// factor is, on average, significantly improved compared with performing single
115/// insertion n times in a row. **Note the performance caveat
116/// related to the computation of the envelope**.
117///
118/// # Element distribution
119/// The tree's performance heavily relies on the spatial distribution of its elements.
120/// Best performance is achieved if:
121///  * No element is inserted more than once
122///  * The overlapping area of elements is as small as
123///    possible.
124///
125/// For the edge case that all elements are overlapping (e.g, one and the same element
126/// is contained `n` times), the performance of most operations usually degrades to `O(n)`.
127///
128/// # Type Parameters
129/// * `T`: The type of objects stored in the r-tree.
130/// * `Params`: Compile time parameters that change the r-tree's internal layout. Refer to the
131///   [RTreeParams] trait for more information.
132///
133/// # Defining methods generic over r-trees
134/// If a library defines a method that should be generic over the r-tree type signature, make
135/// sure to include both type parameters like this:
136/// ```
137/// # use rstar::{RTree,RTreeObject, RTreeParams};
138/// pub fn generic_rtree_function<T, Params>(tree: &mut RTree<T, Params>)
139/// where
140///   T: RTreeObject,
141///   Params: RTreeParams
142/// {
143///   // ...
144/// }
145/// ```
146/// Otherwise, any user of `generic_rtree_function` would be forced to use
147/// a tree with default parameters.
148///
149/// # (De)Serialization
150/// Enable the `serde` feature for [Serde](https://crates.io/crates/serde) support.
151///
152/// ## Further reading
153/// For more information refer to the [wikipedia article](https://en.wikipedia.org/wiki/R-tree).
154///
155#[derive(Clone)]
156#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
157#[cfg_attr(
158    feature = "serde",
159    serde(bound(
160        serialize = "T: Serialize, T::Envelope: Serialize",
161        deserialize = "T: Deserialize<'de>, T::Envelope: Deserialize<'de>"
162    ))
163)]
164pub struct RTree<T, Params = DefaultParams>
165where
166    Params: RTreeParams,
167    T: RTreeObject,
168{
169    root: ParentNode<T>,
170    size: usize,
171    _params: ::core::marker::PhantomData<Params>,
172}
173
174struct DebugHelper<'a, T, Params>
175where
176    T: RTreeObject + ::core::fmt::Debug + 'a,
177    Params: RTreeParams + 'a,
178{
179    rtree: &'a RTree<T, Params>,
180}
181
182impl<'a, T, Params> ::core::fmt::Debug for DebugHelper<'a, T, Params>
183where
184    T: RTreeObject + ::core::fmt::Debug,
185    Params: RTreeParams,
186{
187    fn fmt(&self, formatter: &mut ::core::fmt::Formatter) -> ::core::fmt::Result {
188        formatter.debug_set().entries(self.rtree.iter()).finish()
189    }
190}
191
192impl<T, Params> ::core::fmt::Debug for RTree<T, Params>
193where
194    Params: RTreeParams,
195    T: RTreeObject + ::core::fmt::Debug,
196{
197    fn fmt(&self, formatter: &mut ::core::fmt::Formatter) -> ::core::fmt::Result {
198        formatter
199            .debug_struct("RTree")
200            .field("size", &self.size)
201            .field("items", &DebugHelper { rtree: self })
202            .finish()
203    }
204}
205
206impl<T> RTree<T>
207where
208    T: RTreeObject,
209{
210    /// Creates a new, empty r-tree.
211    ///
212    /// The created r-tree is configured with [default parameters](DefaultParams).
213    pub fn new() -> Self {
214        Self::new_with_params()
215    }
216
217    /// Creates a new r-tree with some elements already inserted.
218    ///
219    /// This method should be the preferred way for creating r-trees. It both
220    /// runs faster and yields an r-tree with better internal structure that
221    /// improves query performance.
222    ///
223    /// This method implements the overlap minimizing top-down bulk loading algorithm (OMT)
224    /// as described in [this paper by Lee and Lee (2003)](http://ceur-ws.org/Vol-74/files/FORUM_18.pdf).
225    ///
226    /// # Runtime
227    /// Bulk loading runs in `O(n * log(n))`, where `n` is the number of loaded
228    /// elements.
229    ///
230    /// # Note
231    /// The envelope of each element will be accessed many times during loading. If that computation
232    /// is expensive, **consider memoizing it** using [`CachedEnvelope`][crate::primitives::CachedEnvelope].
233    pub fn bulk_load(elements: Vec<T>) -> Self {
234        Self::bulk_load_with_params(elements)
235    }
236}
237
238impl<T, Params> RTree<T, Params>
239where
240    Params: RTreeParams,
241    T: RTreeObject,
242{
243    /// Creates a new, empty r-tree.
244    ///
245    /// The tree's compile time parameters must be specified. Refer to the
246    /// [RTreeParams] trait for more information and a usage example.
247    pub fn new_with_params() -> Self {
248        verify_parameters::<T, Params>();
249        RTree {
250            root: ParentNode::new_root::<Params>(),
251            size: 0,
252            _params: Default::default(),
253        }
254    }
255
256    /// Creates a new r-tree with some given elements and configurable parameters.
257    ///
258    /// For more information refer to [RTree::bulk_load]
259    /// and [RTreeParams].
260    pub fn bulk_load_with_params(elements: Vec<T>) -> Self {
261        Self::new_from_bulk_loading(elements, bulk_load::bulk_load_sequential::<_, Params>)
262    }
263
264    /// Returns the number of objects in an r-tree.
265    ///
266    /// # Example
267    /// ```
268    /// use rstar::RTree;
269    ///
270    /// let mut tree = RTree::new();
271    /// assert_eq!(tree.size(), 0);
272    /// tree.insert([0.0, 1.0, 2.0]);
273    /// assert_eq!(tree.size(), 1);
274    /// tree.remove(&[0.0, 1.0, 2.0]);
275    /// assert_eq!(tree.size(), 0);
276    /// ```
277    pub fn size(&self) -> usize {
278        self.size
279    }
280
281    pub(crate) fn size_mut(&mut self) -> &mut usize {
282        &mut self.size
283    }
284
285    /// Returns an iterator over all elements contained in the tree.
286    ///
287    /// The order in which the elements are returned is not specified.
288    ///
289    /// # Example
290    /// ```
291    /// use rstar::RTree;
292    /// let tree = RTree::bulk_load(vec![(0.0, 0.1), (0.3, 0.2), (0.4, 0.2)]);
293    /// for point in tree.iter() {
294    ///     println!("This tree contains point {:?}", point);
295    /// }
296    /// ```
297    pub fn iter(&self) -> RTreeIterator<T> {
298        RTreeIterator::new(&self.root, SelectAllFunc)
299    }
300
301    /// Returns an iterator over all mutable elements contained in the tree.
302    ///
303    /// The order in which the elements are returned is not specified.
304    ///
305    /// *Note*: It is a logic error to change an inserted item's position or dimensions. This
306    /// method is primarily meant for own implementations of [RTreeObject]
307    /// which can contain arbitrary additional data.
308    /// If the position or location of an inserted object need to change, you will need to [RTree::remove]
309    /// and reinsert it.
310    ///
311    pub fn iter_mut(&mut self) -> RTreeIteratorMut<T> {
312        RTreeIteratorMut::new(&mut self.root, SelectAllFunc)
313    }
314
315    /// Returns all elements contained in an [Envelope].
316    ///
317    /// Usually, an envelope is an [axis aligned bounding box](crate::AABB). This
318    /// method can be used to retrieve all elements that are fully contained within an envelope.
319    ///
320    /// # Example
321    /// ```
322    /// use rstar::{RTree, AABB};
323    /// let mut tree = RTree::bulk_load(vec![
324    ///   [0.0, 0.0],
325    ///   [0.0, 1.0],
326    ///   [1.0, 1.0]
327    /// ]);
328    /// let half_unit_square = AABB::from_corners([0.0, 0.0], [0.5, 1.0]);
329    /// let unit_square = AABB::from_corners([0.0, 0.0], [1.0, 1.0]);
330    /// let elements_in_half_unit_square = tree.locate_in_envelope(&half_unit_square);
331    /// let elements_in_unit_square = tree.locate_in_envelope(&unit_square);
332    /// assert_eq!(elements_in_half_unit_square.count(), 2);
333    /// assert_eq!(elements_in_unit_square.count(), 3);
334    /// ```
335    pub fn locate_in_envelope(&self, envelope: &T::Envelope) -> LocateInEnvelope<T> {
336        LocateInEnvelope::new(&self.root, SelectInEnvelopeFunction::new(envelope.clone()))
337    }
338
339    /// Mutable variant of [locate_in_envelope](#method.locate_in_envelope).
340    pub fn locate_in_envelope_mut(&mut self, envelope: &T::Envelope) -> LocateInEnvelopeMut<T> {
341        LocateInEnvelopeMut::new(
342            &mut self.root,
343            SelectInEnvelopeFunction::new(envelope.clone()),
344        )
345    }
346
347    /// Variant of [`locate_in_envelope`][Self::locate_in_envelope] using internal iteration.
348    pub fn locate_in_envelope_int<'a, V, B>(
349        &'a self,
350        envelope: &T::Envelope,
351        mut visitor: V,
352    ) -> ControlFlow<B>
353    where
354        V: FnMut(&'a T) -> ControlFlow<B>,
355    {
356        select_nodes(
357            self.root(),
358            &SelectInEnvelopeFunction::new(envelope.clone()),
359            &mut visitor,
360        )
361    }
362
363    /// Mutable variant of [`locate_in_envelope_mut`][Self::locate_in_envelope_mut].
364    pub fn locate_in_envelope_int_mut<'a, V, B>(
365        &'a mut self,
366        envelope: &T::Envelope,
367        mut visitor: V,
368    ) -> ControlFlow<B>
369    where
370        V: FnMut(&'a mut T) -> ControlFlow<B>,
371    {
372        select_nodes_mut(
373            self.root_mut(),
374            &SelectInEnvelopeFunction::new(envelope.clone()),
375            &mut visitor,
376        )
377    }
378
379    /// Returns a draining iterator over all elements contained in the tree.
380    ///
381    /// The order in which the elements are returned is not specified.
382    ///
383    /// See
384    /// [drain_with_selection_function](#method.drain_with_selection_function)
385    /// for more information.
386    pub fn drain(&mut self) -> DrainIterator<T, SelectAllFunc, Params> {
387        self.drain_with_selection_function(SelectAllFunc)
388    }
389
390    /// Draining variant of [locate_in_envelope](#method.locate_in_envelope).
391    pub fn drain_in_envelope(
392        &mut self,
393        envelope: T::Envelope,
394    ) -> DrainIterator<T, SelectInEnvelopeFunction<T>, Params> {
395        let sel = SelectInEnvelopeFunction::new(envelope);
396        self.drain_with_selection_function(sel)
397    }
398
399    /// Returns all elements whose envelope intersects a given envelope.
400    ///
401    /// Any element fully contained within an envelope is also returned by this method. Two
402    /// envelopes that "touch" each other (e.g. by sharing only a common corner) are also
403    /// considered to intersect. Usually, an envelope is an [axis aligned bounding box](crate::AABB).
404    /// This method will return all elements whose AABB has some common area with
405    /// a given AABB.
406    ///
407    /// # Example
408    /// ```
409    /// use rstar::{RTree, AABB};
410    /// use rstar::primitives::Rectangle;
411    ///
412    /// let left_piece = AABB::from_corners([0.0, 0.0], [0.4, 1.0]);
413    /// let right_piece = AABB::from_corners([0.6, 0.0], [1.0, 1.0]);
414    /// let middle_piece = AABB::from_corners([0.25, 0.0], [0.75, 1.0]);
415    ///
416    /// let mut tree = RTree::<Rectangle<_>>::bulk_load(vec![
417    ///   left_piece.into(),
418    ///   right_piece.into(),
419    ///   middle_piece.into(),
420    /// ]);
421    ///
422    /// let elements_intersecting_left_piece = tree.locate_in_envelope_intersecting(&left_piece);
423    /// // The left piece should not intersect the right piece!
424    /// assert_eq!(elements_intersecting_left_piece.count(), 2);
425    /// let elements_intersecting_middle = tree.locate_in_envelope_intersecting(&middle_piece);
426    /// // Only the middle piece intersects all pieces within the tree
427    /// assert_eq!(elements_intersecting_middle.count(), 3);
428    ///
429    /// let large_piece = AABB::from_corners([-100., -100.], [100., 100.]);
430    /// let elements_intersecting_large_piece = tree.locate_in_envelope_intersecting(&large_piece);
431    /// // Any element that is fully contained should also be returned:
432    /// assert_eq!(elements_intersecting_large_piece.count(), 3);
433    /// ```
434    pub fn locate_in_envelope_intersecting(
435        &self,
436        envelope: &T::Envelope,
437    ) -> LocateInEnvelopeIntersecting<T> {
438        LocateInEnvelopeIntersecting::new(
439            &self.root,
440            SelectInEnvelopeFuncIntersecting::new(envelope.clone()),
441        )
442    }
443
444    /// Mutable variant of [locate_in_envelope_intersecting](#method.locate_in_envelope_intersecting)
445    pub fn locate_in_envelope_intersecting_mut(
446        &mut self,
447        envelope: &T::Envelope,
448    ) -> LocateInEnvelopeIntersectingMut<T> {
449        LocateInEnvelopeIntersectingMut::new(
450            &mut self.root,
451            SelectInEnvelopeFuncIntersecting::new(envelope.clone()),
452        )
453    }
454
455    /// Variant of [`locate_in_envelope_intersecting`][Self::locate_in_envelope_intersecting] using internal iteration.
456    pub fn locate_in_envelope_intersecting_int<'a, V, B>(
457        &'a self,
458        envelope: &T::Envelope,
459        mut visitor: V,
460    ) -> ControlFlow<B>
461    where
462        V: FnMut(&'a T) -> ControlFlow<B>,
463    {
464        select_nodes(
465            self.root(),
466            &SelectInEnvelopeFuncIntersecting::new(envelope.clone()),
467            &mut visitor,
468        )
469    }
470
471    /// Mutable variant of [`locate_in_envelope_intersecting_int`][Self::locate_in_envelope_intersecting_int].
472    pub fn locate_in_envelope_intersecting_int_mut<'a, V, B>(
473        &'a mut self,
474        envelope: &T::Envelope,
475        mut visitor: V,
476    ) -> ControlFlow<B>
477    where
478        V: FnMut(&'a mut T) -> ControlFlow<B>,
479    {
480        select_nodes_mut(
481            self.root_mut(),
482            &SelectInEnvelopeFuncIntersecting::new(envelope.clone()),
483            &mut visitor,
484        )
485    }
486
487    /// Locates elements in the r-tree defined by a selection function.
488    ///
489    /// Refer to the documentation of [`SelectionFunction`] for
490    /// more information.
491    ///
492    /// Usually, other `locate` methods should cover most common use cases. This method is only required
493    /// in more specific situations.
494    pub fn locate_with_selection_function<S: SelectionFunction<T>>(
495        &self,
496        selection_function: S,
497    ) -> SelectionIterator<T, S> {
498        SelectionIterator::new(&self.root, selection_function)
499    }
500
501    /// Mutable variant of [`locate_with_selection_function`](#method.locate_with_selection_function).
502    pub fn locate_with_selection_function_mut<S: SelectionFunction<T>>(
503        &mut self,
504        selection_function: S,
505    ) -> SelectionIteratorMut<T, S> {
506        SelectionIteratorMut::new(&mut self.root, selection_function)
507    }
508
509    /// Returns all possible intersecting objects of this and another tree.
510    ///
511    /// This will return all objects whose _envelopes_ intersect. No geometric intersection
512    /// checking is performed.
513    pub fn intersection_candidates_with_other_tree<'a, U>(
514        &'a self,
515        other: &'a RTree<U>,
516    ) -> IntersectionIterator<'a, T, U>
517    where
518        U: RTreeObject<Envelope = T::Envelope>,
519    {
520        IntersectionIterator::new(self.root(), other.root())
521    }
522
523    /// Returns the tree's root node.
524    ///
525    /// Usually, you will not need to call this method. However, for debugging purposes or for
526    /// advanced algorithms, knowledge about the tree's internal structure may be required.
527    /// For these cases, this method serves as an entry point.
528    pub fn root(&self) -> &ParentNode<T> {
529        &self.root
530    }
531
532    pub(crate) fn root_mut(&mut self) -> &mut ParentNode<T> {
533        &mut self.root
534    }
535
536    fn new_from_bulk_loading(
537        elements: Vec<T>,
538        root_loader: impl Fn(Vec<T>) -> ParentNode<T>,
539    ) -> Self {
540        verify_parameters::<T, Params>();
541        let size = elements.len();
542        let root = if size == 0 {
543            ParentNode::new_root::<Params>()
544        } else {
545            root_loader(elements)
546        };
547        RTree {
548            root,
549            size,
550            _params: Default::default(),
551        }
552    }
553
554    /// Removes and returns a single element from the tree. The element to remove is specified
555    /// by a [`SelectionFunction`].
556    ///
557    /// See also: [`RTree::remove`], [`RTree::remove_at_point`]
558    ///
559    pub fn remove_with_selection_function<F>(&mut self, function: F) -> Option<T>
560    where
561        F: SelectionFunction<T>,
562    {
563        removal::DrainIterator::new(self, function).take(1).last()
564    }
565
566    /// Drain elements selected by a [`SelectionFunction`]. Returns an
567    /// iterator that successively removes selected elements and returns
568    /// them. This is the most generic drain API, see also:
569    /// [`RTree::drain_in_envelope_intersecting`],
570    /// [`RTree::drain_within_distance`].
571    ///
572    /// # Remarks
573    ///
574    /// This API is similar to `Vec::drain_filter`, but stopping the
575    /// iteration would stop the removal. However, the returned iterator
576    /// must be properly dropped. Leaking this iterator leads to a leak
577    /// amplification, where all the elements in the tree are leaked.
578    pub fn drain_with_selection_function<F>(&mut self, function: F) -> DrainIterator<T, F, Params>
579    where
580        F: SelectionFunction<T>,
581    {
582        removal::DrainIterator::new(self, function)
583    }
584
585    /// Drains elements intersecting the `envelope`. Similar to
586    /// `locate_in_envelope_intersecting`, except the elements are removed
587    /// and returned via an iterator.
588    pub fn drain_in_envelope_intersecting(
589        &mut self,
590        envelope: T::Envelope,
591    ) -> DrainIterator<T, SelectInEnvelopeFuncIntersecting<T>, Params> {
592        let selection_function = SelectInEnvelopeFuncIntersecting::new(envelope);
593        self.drain_with_selection_function(selection_function)
594    }
595}
596
597impl<T, Params> RTree<T, Params>
598where
599    Params: RTreeParams,
600    T: PointDistance,
601{
602    /// Returns a single object that covers a given point.
603    ///
604    /// Method [contains_point](PointDistance::contains_point)
605    /// is used to determine if a tree element contains the given point.
606    ///
607    /// If multiple elements contain the given point, any of them is returned.
608    pub fn locate_at_point(&self, point: &<T::Envelope as Envelope>::Point) -> Option<&T> {
609        self.locate_all_at_point(point).next()
610    }
611
612    /// Mutable variant of [RTree::locate_at_point].
613    pub fn locate_at_point_mut(
614        &mut self,
615        point: &<T::Envelope as Envelope>::Point,
616    ) -> Option<&mut T> {
617        self.locate_all_at_point_mut(point).next()
618    }
619
620    /// Variant of [`locate_at_point`][Self::locate_at_point] using internal iteration.
621    pub fn locate_at_point_int(&self, point: &<T::Envelope as Envelope>::Point) -> Option<&T> {
622        match self.locate_all_at_point_int(point, ControlFlow::Break) {
623            ControlFlow::Break(node) => Some(node),
624            ControlFlow::Continue(()) => None,
625        }
626    }
627
628    /// Mutable variant of [`locate_at_point_int`][Self::locate_at_point_int].
629    pub fn locate_at_point_int_mut(
630        &mut self,
631        point: &<T::Envelope as Envelope>::Point,
632    ) -> Option<&mut T> {
633        match self.locate_all_at_point_int_mut(point, ControlFlow::Break) {
634            ControlFlow::Break(node) => Some(node),
635            ControlFlow::Continue(()) => None,
636        }
637    }
638
639    /// Locate all elements containing a given point.
640    ///
641    /// Method [PointDistance::contains_point] is used
642    /// to determine if a tree element contains the given point.
643    /// # Example
644    /// ```
645    /// use rstar::RTree;
646    /// use rstar::primitives::Rectangle;
647    ///
648    /// let tree = RTree::bulk_load(vec![
649    ///   Rectangle::from_corners([0.0, 0.0], [2.0, 2.0]),
650    ///   Rectangle::from_corners([1.0, 1.0], [3.0, 3.0])
651    /// ]);
652    ///
653    /// assert_eq!(tree.locate_all_at_point(&[1.5, 1.5]).count(), 2);
654    /// assert_eq!(tree.locate_all_at_point(&[0.0, 0.0]).count(), 1);
655    /// assert_eq!(tree.locate_all_at_point(&[-1., 0.0]).count(), 0);
656    /// ```
657    pub fn locate_all_at_point(
658        &self,
659        point: &<T::Envelope as Envelope>::Point,
660    ) -> LocateAllAtPoint<T> {
661        LocateAllAtPoint::new(&self.root, SelectAtPointFunction::new(point.clone()))
662    }
663
664    /// Mutable variant of [`locate_all_at_point`][Self::locate_all_at_point].
665    pub fn locate_all_at_point_mut(
666        &mut self,
667        point: &<T::Envelope as Envelope>::Point,
668    ) -> LocateAllAtPointMut<T> {
669        LocateAllAtPointMut::new(&mut self.root, SelectAtPointFunction::new(point.clone()))
670    }
671
672    /// Variant of [`locate_all_at_point`][Self::locate_all_at_point] using internal iteration.
673    pub fn locate_all_at_point_int<'a, V, B>(
674        &'a self,
675        point: &<T::Envelope as Envelope>::Point,
676        mut visitor: V,
677    ) -> ControlFlow<B>
678    where
679        V: FnMut(&'a T) -> ControlFlow<B>,
680    {
681        select_nodes(
682            &self.root,
683            &SelectAtPointFunction::new(point.clone()),
684            &mut visitor,
685        )
686    }
687
688    /// Mutable variant of [`locate_all_at_point_int`][Self::locate_all_at_point_int].
689    pub fn locate_all_at_point_int_mut<'a, V, B>(
690        &'a mut self,
691        point: &<T::Envelope as Envelope>::Point,
692        mut visitor: V,
693    ) -> ControlFlow<B>
694    where
695        V: FnMut(&'a mut T) -> ControlFlow<B>,
696    {
697        select_nodes_mut(
698            &mut self.root,
699            &SelectAtPointFunction::new(point.clone()),
700            &mut visitor,
701        )
702    }
703
704    /// Removes an element containing a given point.
705    ///
706    /// The removed element, if any, is returned. If multiple elements cover the given point,
707    /// only one of them is removed and returned.
708    ///
709    /// # Example
710    /// ```
711    /// use rstar::RTree;
712    /// use rstar::primitives::Rectangle;
713    ///
714    /// let mut tree = RTree::bulk_load(vec![
715    ///   Rectangle::from_corners([0.0, 0.0], [2.0, 2.0]),
716    ///   Rectangle::from_corners([1.0, 1.0], [3.0, 3.0])
717    /// ]);
718    ///
719    /// assert!(tree.remove_at_point(&[1.5, 1.5]).is_some());
720    /// assert!(tree.remove_at_point(&[1.5, 1.5]).is_some());
721    /// assert!(tree.remove_at_point(&[1.5, 1.5]).is_none());
722    ///```
723    pub fn remove_at_point(&mut self, point: &<T::Envelope as Envelope>::Point) -> Option<T> {
724        let removal_function = SelectAtPointFunction::new(point.clone());
725        self.remove_with_selection_function(removal_function)
726    }
727}
728
729impl<T, Params> RTree<T, Params>
730where
731    Params: RTreeParams,
732    T: RTreeObject + PartialEq,
733{
734    /// Returns `true` if a given element is equal (`==`) to an element in the
735    /// r-tree.
736    ///
737    /// This method will only work correctly if two equal elements also have the
738    /// same envelope.
739    ///
740    /// # Example
741    /// ```
742    /// use rstar::RTree;
743    ///
744    /// let mut tree = RTree::new();
745    /// assert!(!tree.contains(&[0.0, 2.0]));
746    /// tree.insert([0.0, 2.0]);
747    /// assert!(tree.contains(&[0.0, 2.0]));
748    /// ```
749    pub fn contains(&self, t: &T) -> bool {
750        self.locate_in_envelope(&t.envelope()).any(|e| e == t)
751    }
752
753    /// Removes and returns an element of the r-tree equal (`==`) to a given element.
754    ///
755    /// If multiple elements equal to the given elements are contained in the tree, only
756    /// one of them is removed and returned.
757    ///
758    /// This method will only work correctly if two equal elements also have the
759    /// same envelope.
760    ///
761    /// # Example
762    /// ```
763    /// use rstar::RTree;
764    ///
765    /// let mut tree = RTree::new();
766    /// tree.insert([0.0, 2.0]);
767    /// // The element can be inserted twice just fine
768    /// tree.insert([0.0, 2.0]);
769    /// assert!(tree.remove(&[0.0, 2.0]).is_some());
770    /// assert!(tree.remove(&[0.0, 2.0]).is_some());
771    /// assert!(tree.remove(&[0.0, 2.0]).is_none());
772    /// ```
773    pub fn remove(&mut self, t: &T) -> Option<T> {
774        let removal_function = SelectEqualsFunction::new(t);
775        self.remove_with_selection_function(removal_function)
776    }
777}
778
779impl<T, Params> RTree<T, Params>
780where
781    Params: RTreeParams,
782    T: PointDistance,
783{
784    /// Returns the nearest neighbor for a given point.
785    ///
786    /// The distance is calculated by calling
787    /// [PointDistance::distance_2]
788    ///
789    /// # Example
790    /// ```
791    /// use rstar::RTree;
792    /// let tree = RTree::bulk_load(vec![
793    ///   [0.0, 0.0],
794    ///   [0.0, 1.0],
795    /// ]);
796    /// assert_eq!(tree.nearest_neighbor(&[-1., 0.0]), Some(&[0.0, 0.0]));
797    /// assert_eq!(tree.nearest_neighbor(&[0.0, 2.0]), Some(&[0.0, 1.0]));
798    /// ```
799    pub fn nearest_neighbor(&self, query_point: &<T::Envelope as Envelope>::Point) -> Option<&T> {
800        if self.size > 0 {
801            // The single-nearest-neighbor retrieval may in rare cases return None due to
802            // rounding issues. The iterator will still work, though.
803            nearest_neighbor::nearest_neighbor(&self.root, query_point.clone())
804                .or_else(|| self.nearest_neighbor_iter(query_point).next())
805        } else {
806            None
807        }
808    }
809
810    /// Returns the nearest neighbors for a given point.
811    ///
812    /// The distance is calculated by calling
813    /// [PointDistance::distance_2]
814    ///
815    /// All returned values will have the exact same distance from the given query point.
816    /// Returns an empty `Vec` if the tree is empty.
817    ///
818    /// # Example
819    /// ```
820    /// use rstar::RTree;
821    /// let tree = RTree::bulk_load(vec![
822    ///   [0.0, 0.0],
823    ///   [0.0, 1.0],
824    ///   [1.0, 0.0],
825    /// ]);
826    ///
827    /// // A single nearest neighbor
828    /// assert_eq!(tree.nearest_neighbors(&[0.01, 0.01]), &[&[0.0, 0.0]]);
829    ///
830    /// // Two nearest neighbors
831    /// let nearest_two = tree.nearest_neighbors(&[1.0, 1.0]);
832    /// assert_eq!(nearest_two.len(), 2);
833    /// assert!(nearest_two.contains(&&[0.0, 1.0]));
834    /// assert!(nearest_two.contains(&&[1.0, 0.0]));
835    /// ```
836    pub fn nearest_neighbors(&self, query_point: &<T::Envelope as Envelope>::Point) -> Vec<&T> {
837        nearest_neighbor::nearest_neighbors(&self.root, query_point.clone())
838    }
839
840    /// Returns all elements of the tree within a certain distance.
841    ///
842    /// The elements may be returned in any order. Each returned element
843    /// will have a squared distance less or equal to the given squared distance.
844    ///
845    /// This method makes use of [PointDistance::distance_2_if_less_or_equal].
846    /// If performance is critical and the distance calculation to the object is fast,
847    /// overwriting this function may be beneficial.
848    pub fn locate_within_distance(
849        &self,
850        query_point: <T::Envelope as Envelope>::Point,
851        max_squared_radius: <<T::Envelope as Envelope>::Point as Point>::Scalar,
852    ) -> LocateWithinDistanceIterator<T> {
853        let selection_function = SelectWithinDistanceFunction::new(query_point, max_squared_radius);
854        LocateWithinDistanceIterator::new(self.root(), selection_function)
855    }
856
857    /// Drain all elements of the tree within a certain distance.
858    ///
859    /// Similar to [`RTree::locate_within_distance`], but removes and
860    /// returns the elements via an iterator.
861    pub fn drain_within_distance(
862        &mut self,
863        query_point: <T::Envelope as Envelope>::Point,
864        max_squared_radius: <<T::Envelope as Envelope>::Point as Point>::Scalar,
865    ) -> DrainIterator<T, SelectWithinDistanceFunction<T>, Params> {
866        let selection_function = SelectWithinDistanceFunction::new(query_point, max_squared_radius);
867        self.drain_with_selection_function(selection_function)
868    }
869
870    /// Returns all elements of the tree sorted by their distance to a given point.
871    ///
872    /// # Runtime
873    /// Every `next()` call runs in `O(log(n))`. Creating the iterator runs in
874    /// `O(log(n))`.
875    /// The [r-tree documentation](RTree) contains more information about
876    /// r-tree performance.
877    ///
878    /// # Example
879    /// ```
880    /// use rstar::RTree;
881    /// let tree = RTree::bulk_load(vec![
882    ///   [0.0, 0.0],
883    ///   [0.0, 1.0],
884    /// ]);
885    ///
886    /// let nearest_neighbors = tree.nearest_neighbor_iter(&[0.5, 0.0]).collect::<Vec<_>>();
887    /// assert_eq!(nearest_neighbors, vec![&[0.0, 0.0], &[0.0, 1.0]]);
888    /// ```
889    pub fn nearest_neighbor_iter(
890        &self,
891        query_point: &<T::Envelope as Envelope>::Point,
892    ) -> NearestNeighborIterator<T> {
893        nearest_neighbor::NearestNeighborIterator::new(&self.root, query_point.clone())
894    }
895
896    /// Returns `(element, distance^2)` tuples of the tree sorted by their distance to a given point.
897    ///
898    /// The distance is calculated by calling
899    /// [PointDistance::distance_2].
900    #[deprecated(note = "Please use nearest_neighbor_iter_with_distance_2 instead")]
901    pub fn nearest_neighbor_iter_with_distance(
902        &self,
903        query_point: &<T::Envelope as Envelope>::Point,
904    ) -> NearestNeighborDistance2Iterator<T> {
905        nearest_neighbor::NearestNeighborDistance2Iterator::new(&self.root, query_point.clone())
906    }
907
908    /// Returns `(element, distance^2)` tuples of the tree sorted by their distance to a given point.
909    ///
910    /// The distance is calculated by calling
911    /// [PointDistance::distance_2].
912    pub fn nearest_neighbor_iter_with_distance_2(
913        &self,
914        query_point: &<T::Envelope as Envelope>::Point,
915    ) -> NearestNeighborDistance2Iterator<T> {
916        nearest_neighbor::NearestNeighborDistance2Iterator::new(&self.root, query_point.clone())
917    }
918
919    /// Removes the nearest neighbor for a given point and returns it.
920    ///
921    /// The distance is calculated by calling
922    /// [PointDistance::distance_2].
923    ///
924    /// # Example
925    /// ```
926    /// use rstar::RTree;
927    /// let mut tree = RTree::bulk_load(vec![
928    ///   [0.0, 0.0],
929    ///   [0.0, 1.0],
930    /// ]);
931    /// assert_eq!(tree.pop_nearest_neighbor(&[0.0, 0.0]), Some([0.0, 0.0]));
932    /// assert_eq!(tree.pop_nearest_neighbor(&[0.0, 0.0]), Some([0.0, 1.0]));
933    /// assert_eq!(tree.pop_nearest_neighbor(&[0.0, 0.0]), None);
934    /// ```
935    pub fn pop_nearest_neighbor(
936        &mut self,
937        query_point: &<T::Envelope as Envelope>::Point,
938    ) -> Option<T> {
939        if let Some(neighbor) = self.nearest_neighbor(query_point) {
940            let removal_function = SelectByAddressFunction::new(neighbor.envelope(), neighbor);
941            self.remove_with_selection_function(removal_function)
942        } else {
943            None
944        }
945    }
946}
947
948impl<T, Params> RTree<T, Params>
949where
950    T: RTreeObject,
951    Params: RTreeParams,
952{
953    /// Inserts a new element into the r-tree.
954    ///
955    /// If the element is already present in the tree, it will now be present twice.
956    ///
957    /// # Runtime
958    /// This method runs in `O(log(n))`.
959    /// The [r-tree documentation](RTree) contains more information about
960    /// r-tree performance.
961    pub fn insert(&mut self, t: T) {
962        Params::DefaultInsertionStrategy::insert(self, t);
963        self.size += 1;
964    }
965}
966
967impl<T, Params> IntoIterator for RTree<T, Params>
968where
969    T: RTreeObject,
970    Params: RTreeParams,
971{
972    type IntoIter = IntoIter<T>;
973    type Item = T;
974
975    fn into_iter(self) -> Self::IntoIter {
976        IntoIter::new(self.root)
977    }
978}
979
980impl<'a, T, Params> IntoIterator for &'a RTree<T, Params>
981where
982    T: RTreeObject,
983    Params: RTreeParams,
984{
985    type IntoIter = RTreeIterator<'a, T>;
986    type Item = &'a T;
987
988    fn into_iter(self) -> Self::IntoIter {
989        self.iter()
990    }
991}
992
993impl<'a, T, Params> IntoIterator for &'a mut RTree<T, Params>
994where
995    T: RTreeObject,
996    Params: RTreeParams,
997{
998    type IntoIter = RTreeIteratorMut<'a, T>;
999    type Item = &'a mut T;
1000
1001    fn into_iter(self) -> Self::IntoIter {
1002        self.iter_mut()
1003    }
1004}
1005
1006#[cfg(test)]
1007mod test {
1008    use super::RTree;
1009    use crate::algorithm::rstar::RStarInsertionStrategy;
1010    use crate::params::RTreeParams;
1011    use crate::test_utilities::{create_random_points, SEED_1};
1012    use crate::DefaultParams;
1013
1014    struct TestParams;
1015    impl RTreeParams for TestParams {
1016        const MIN_SIZE: usize = 10;
1017        const MAX_SIZE: usize = 20;
1018        const REINSERTION_COUNT: usize = 1;
1019        type DefaultInsertionStrategy = RStarInsertionStrategy;
1020    }
1021
1022    #[test]
1023    fn test_remove_capacity() {
1024        pub struct WeirdParams;
1025
1026        impl RTreeParams for WeirdParams {
1027            const MIN_SIZE: usize = 1;
1028            const MAX_SIZE: usize = 10;
1029            const REINSERTION_COUNT: usize = 1;
1030            type DefaultInsertionStrategy = RStarInsertionStrategy;
1031        }
1032
1033        let mut items: Vec<[f32; 2]> = Vec::new();
1034        for i in 0..2 {
1035            items.push([i as f32, i as f32]);
1036        }
1037        let mut tree: RTree<_, WeirdParams> = RTree::bulk_load_with_params(items);
1038        assert_eq!(tree.remove(&[1.0, 1.0]).unwrap(), [1.0, 1.0]);
1039    }
1040
1041    #[test]
1042    fn test_create_rtree_with_parameters() {
1043        let tree: RTree<[f32; 2], TestParams> = RTree::new_with_params();
1044        assert_eq!(tree.size(), 0);
1045    }
1046
1047    #[test]
1048    fn test_insert_single() {
1049        let mut tree: RTree<_> = RTree::new();
1050        tree.insert([0.02f32, 0.4f32]);
1051        assert_eq!(tree.size(), 1);
1052        assert!(tree.contains(&[0.02, 0.4]));
1053        assert!(!tree.contains(&[0.3, 0.2]));
1054    }
1055
1056    #[test]
1057    fn test_insert_many() {
1058        const NUM_POINTS: usize = 1000;
1059        let points = create_random_points(NUM_POINTS, SEED_1);
1060        let mut tree = RTree::new();
1061        for p in &points {
1062            tree.insert(*p);
1063            tree.root.sanity_check::<DefaultParams>(true);
1064        }
1065        assert_eq!(tree.size(), NUM_POINTS);
1066        for p in &points {
1067            assert!(tree.contains(p));
1068        }
1069    }
1070
1071    #[test]
1072    fn test_fmt_debug() {
1073        let tree = RTree::bulk_load(vec![[0, 1], [0, 1]]);
1074        let debug: String = format!("{:?}", tree);
1075        assert_eq!(debug, "RTree { size: 2, items: {[0, 1], [0, 1]} }");
1076    }
1077
1078    #[test]
1079    fn test_default() {
1080        let tree: RTree<[f32; 2]> = Default::default();
1081        assert_eq!(tree.size(), 0);
1082    }
1083
1084    #[cfg(feature = "serde")]
1085    #[test]
1086    fn test_serialization() {
1087        use crate::test_utilities::create_random_integers;
1088
1089        use serde_json;
1090        const SIZE: usize = 20;
1091        let points = create_random_integers::<[i32; 2]>(SIZE, SEED_1);
1092        let tree = RTree::bulk_load(points.clone());
1093        let json = serde_json::to_string(&tree).expect("Serializing tree failed");
1094        let parsed: RTree<[i32; 2]> =
1095            serde_json::from_str(&json).expect("Deserializing tree failed");
1096        assert_eq!(parsed.size(), SIZE);
1097        for point in &points {
1098            assert!(parsed.contains(point));
1099        }
1100    }
1101
1102    #[test]
1103    fn test_bulk_load_crash() {
1104        let bulk_nodes = vec![
1105            [570.0, 1080.0, 89.0],
1106            [30.0, 1080.0, 627.0],
1107            [1916.0, 1080.0, 68.0],
1108            [274.0, 1080.0, 790.0],
1109            [476.0, 1080.0, 895.0],
1110            [1557.0, 1080.0, 250.0],
1111            [1546.0, 1080.0, 883.0],
1112            [1512.0, 1080.0, 610.0],
1113            [1729.0, 1080.0, 358.0],
1114            [1841.0, 1080.0, 434.0],
1115            [1752.0, 1080.0, 696.0],
1116            [1674.0, 1080.0, 705.0],
1117            [136.0, 1080.0, 22.0],
1118            [1593.0, 1080.0, 71.0],
1119            [586.0, 1080.0, 272.0],
1120            [348.0, 1080.0, 373.0],
1121            [502.0, 1080.0, 2.0],
1122            [1488.0, 1080.0, 1072.0],
1123            [31.0, 1080.0, 526.0],
1124            [1695.0, 1080.0, 559.0],
1125            [1663.0, 1080.0, 298.0],
1126            [316.0, 1080.0, 417.0],
1127            [1348.0, 1080.0, 731.0],
1128            [784.0, 1080.0, 126.0],
1129            [225.0, 1080.0, 847.0],
1130            [79.0, 1080.0, 819.0],
1131            [320.0, 1080.0, 504.0],
1132            [1714.0, 1080.0, 1026.0],
1133            [264.0, 1080.0, 229.0],
1134            [108.0, 1080.0, 158.0],
1135            [1665.0, 1080.0, 604.0],
1136            [496.0, 1080.0, 231.0],
1137            [1813.0, 1080.0, 865.0],
1138            [1200.0, 1080.0, 326.0],
1139            [1661.0, 1080.0, 818.0],
1140            [135.0, 1080.0, 229.0],
1141            [424.0, 1080.0, 1016.0],
1142            [1708.0, 1080.0, 791.0],
1143            [1626.0, 1080.0, 682.0],
1144            [442.0, 1080.0, 895.0],
1145        ];
1146
1147        let nodes = vec![
1148            [1916.0, 1060.0, 68.0],
1149            [1664.0, 1060.0, 298.0],
1150            [1594.0, 1060.0, 71.0],
1151            [225.0, 1060.0, 846.0],
1152            [1841.0, 1060.0, 434.0],
1153            [502.0, 1060.0, 2.0],
1154            [1625.5852, 1060.0122, 682.0],
1155            [1348.5273, 1060.0029, 731.08124],
1156            [316.36127, 1060.0298, 418.24515],
1157            [1729.3253, 1060.0023, 358.50134],
1158        ];
1159        let mut tree = RTree::bulk_load(bulk_nodes);
1160        for node in nodes {
1161            // Bulk loading will create nodes larger than Params::MAX_SIZE,
1162            // which is intentional and not harmful.
1163            tree.insert(node);
1164            tree.root().sanity_check::<DefaultParams>(false);
1165        }
1166    }
1167}