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}