use crate::algorithm::bulk_load;
use crate::algorithm::iterators::*;
use crate::algorithm::nearest_neighbor;
use crate::algorithm::nearest_neighbor::NearestNeighborDistance2Iterator;
use crate::algorithm::nearest_neighbor::NearestNeighborIterator;
use crate::algorithm::removal;
use crate::algorithm::selection_functions::*;
use crate::envelope::Envelope;
use crate::node::ParentNode;
use crate::object::{PointDistance, RTreeObject};
use crate::params::{verify_parameters, DefaultParams, InsertionStrategy, RTreeParams};
use crate::Point;
use core::ops::ControlFlow;
#[cfg(not(test))]
use alloc::vec::Vec;
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};
impl<T, Params> Default for RTree<T, Params>
where
T: RTreeObject,
Params: RTreeParams,
{
fn default() -> Self {
Self::new_with_params()
}
}
#[derive(Clone)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[cfg_attr(
feature = "serde",
serde(bound(
serialize = "T: Serialize, T::Envelope: Serialize",
deserialize = "T: Deserialize<'de>, T::Envelope: Deserialize<'de>"
))
)]
pub struct RTree<T, Params = DefaultParams>
where
Params: RTreeParams,
T: RTreeObject,
{
root: ParentNode<T>,
size: usize,
_params: ::core::marker::PhantomData<Params>,
}
struct DebugHelper<'a, T, Params>
where
T: RTreeObject + ::core::fmt::Debug + 'a,
Params: RTreeParams + 'a,
{
rtree: &'a RTree<T, Params>,
}
impl<'a, T, Params> ::core::fmt::Debug for DebugHelper<'a, T, Params>
where
T: RTreeObject + ::core::fmt::Debug,
Params: RTreeParams,
{
fn fmt(&self, formatter: &mut ::core::fmt::Formatter) -> ::core::fmt::Result {
formatter.debug_set().entries(self.rtree.iter()).finish()
}
}
impl<T, Params> ::core::fmt::Debug for RTree<T, Params>
where
Params: RTreeParams,
T: RTreeObject + ::core::fmt::Debug,
{
fn fmt(&self, formatter: &mut ::core::fmt::Formatter) -> ::core::fmt::Result {
formatter
.debug_struct("RTree")
.field("size", &self.size)
.field("items", &DebugHelper { rtree: self })
.finish()
}
}
impl<T> RTree<T>
where
T: RTreeObject,
{
pub fn new() -> Self {
Self::new_with_params()
}
pub fn bulk_load(elements: Vec<T>) -> Self {
Self::bulk_load_with_params(elements)
}
}
impl<T, Params> RTree<T, Params>
where
Params: RTreeParams,
T: RTreeObject,
{
pub fn new_with_params() -> Self {
verify_parameters::<T, Params>();
RTree {
root: ParentNode::new_root::<Params>(),
size: 0,
_params: Default::default(),
}
}
pub fn bulk_load_with_params(elements: Vec<T>) -> Self {
Self::new_from_bulk_loading(elements, bulk_load::bulk_load_sequential::<_, Params>)
}
pub fn size(&self) -> usize {
self.size
}
pub(crate) fn size_mut(&mut self) -> &mut usize {
&mut self.size
}
pub fn iter(&self) -> RTreeIterator<T> {
RTreeIterator::new(&self.root, SelectAllFunc)
}
pub fn iter_mut(&mut self) -> RTreeIteratorMut<T> {
RTreeIteratorMut::new(&mut self.root, SelectAllFunc)
}
pub fn locate_in_envelope(&self, envelope: &T::Envelope) -> LocateInEnvelope<T> {
LocateInEnvelope::new(&self.root, SelectInEnvelopeFunction::new(envelope.clone()))
}
pub fn locate_in_envelope_mut(&mut self, envelope: &T::Envelope) -> LocateInEnvelopeMut<T> {
LocateInEnvelopeMut::new(
&mut self.root,
SelectInEnvelopeFunction::new(envelope.clone()),
)
}
pub fn locate_in_envelope_int<'a, V, B>(
&'a self,
envelope: &T::Envelope,
mut visitor: V,
) -> ControlFlow<B>
where
V: FnMut(&'a T) -> ControlFlow<B>,
{
select_nodes(
self.root(),
&SelectInEnvelopeFunction::new(envelope.clone()),
&mut visitor,
)
}
pub fn locate_in_envelope_int_mut<'a, V, B>(
&'a mut self,
envelope: &T::Envelope,
mut visitor: V,
) -> ControlFlow<B>
where
V: FnMut(&'a mut T) -> ControlFlow<B>,
{
select_nodes_mut(
self.root_mut(),
&SelectInEnvelopeFunction::new(envelope.clone()),
&mut visitor,
)
}
pub fn drain(&mut self) -> DrainIterator<T, SelectAllFunc, Params> {
self.drain_with_selection_function(SelectAllFunc)
}
pub fn drain_in_envelope(
&mut self,
envelope: T::Envelope,
) -> DrainIterator<T, SelectInEnvelopeFunction<T>, Params> {
let sel = SelectInEnvelopeFunction::new(envelope);
self.drain_with_selection_function(sel)
}
pub fn locate_in_envelope_intersecting(
&self,
envelope: &T::Envelope,
) -> LocateInEnvelopeIntersecting<T> {
LocateInEnvelopeIntersecting::new(
&self.root,
SelectInEnvelopeFuncIntersecting::new(envelope.clone()),
)
}
pub fn locate_in_envelope_intersecting_mut(
&mut self,
envelope: &T::Envelope,
) -> LocateInEnvelopeIntersectingMut<T> {
LocateInEnvelopeIntersectingMut::new(
&mut self.root,
SelectInEnvelopeFuncIntersecting::new(envelope.clone()),
)
}
pub fn locate_in_envelope_intersecting_int<'a, V, B>(
&'a self,
envelope: &T::Envelope,
mut visitor: V,
) -> ControlFlow<B>
where
V: FnMut(&'a T) -> ControlFlow<B>,
{
select_nodes(
self.root(),
&SelectInEnvelopeFuncIntersecting::new(envelope.clone()),
&mut visitor,
)
}
pub fn locate_in_envelope_intersecting_int_mut<'a, V, B>(
&'a mut self,
envelope: &T::Envelope,
mut visitor: V,
) -> ControlFlow<B>
where
V: FnMut(&'a mut T) -> ControlFlow<B>,
{
select_nodes_mut(
self.root_mut(),
&SelectInEnvelopeFuncIntersecting::new(envelope.clone()),
&mut visitor,
)
}
pub fn locate_with_selection_function<S: SelectionFunction<T>>(
&self,
selection_function: S,
) -> SelectionIterator<T, S> {
SelectionIterator::new(&self.root, selection_function)
}
pub fn locate_with_selection_function_mut<S: SelectionFunction<T>>(
&mut self,
selection_function: S,
) -> SelectionIteratorMut<T, S> {
SelectionIteratorMut::new(&mut self.root, selection_function)
}
pub fn intersection_candidates_with_other_tree<'a, U>(
&'a self,
other: &'a RTree<U>,
) -> IntersectionIterator<'a, T, U>
where
U: RTreeObject<Envelope = T::Envelope>,
{
IntersectionIterator::new(self.root(), other.root())
}
pub fn root(&self) -> &ParentNode<T> {
&self.root
}
pub(crate) fn root_mut(&mut self) -> &mut ParentNode<T> {
&mut self.root
}
fn new_from_bulk_loading(
elements: Vec<T>,
root_loader: impl Fn(Vec<T>) -> ParentNode<T>,
) -> Self {
verify_parameters::<T, Params>();
let size = elements.len();
let root = if size == 0 {
ParentNode::new_root::<Params>()
} else {
root_loader(elements)
};
RTree {
root,
size,
_params: Default::default(),
}
}
pub fn remove_with_selection_function<F>(&mut self, function: F) -> Option<T>
where
F: SelectionFunction<T>,
{
removal::DrainIterator::new(self, function).take(1).last()
}
pub fn drain_with_selection_function<F>(&mut self, function: F) -> DrainIterator<T, F, Params>
where
F: SelectionFunction<T>,
{
removal::DrainIterator::new(self, function)
}
pub fn drain_in_envelope_intersecting(
&mut self,
envelope: T::Envelope,
) -> DrainIterator<T, SelectInEnvelopeFuncIntersecting<T>, Params> {
let selection_function = SelectInEnvelopeFuncIntersecting::new(envelope);
self.drain_with_selection_function(selection_function)
}
}
impl<T, Params> RTree<T, Params>
where
Params: RTreeParams,
T: PointDistance,
{
pub fn locate_at_point(&self, point: &<T::Envelope as Envelope>::Point) -> Option<&T> {
self.locate_all_at_point(point).next()
}
pub fn locate_at_point_mut(
&mut self,
point: &<T::Envelope as Envelope>::Point,
) -> Option<&mut T> {
self.locate_all_at_point_mut(point).next()
}
pub fn locate_at_point_int(&self, point: &<T::Envelope as Envelope>::Point) -> Option<&T> {
match self.locate_all_at_point_int(point, ControlFlow::Break) {
ControlFlow::Break(node) => Some(node),
ControlFlow::Continue(()) => None,
}
}
pub fn locate_at_point_int_mut(
&mut self,
point: &<T::Envelope as Envelope>::Point,
) -> Option<&mut T> {
match self.locate_all_at_point_int_mut(point, ControlFlow::Break) {
ControlFlow::Break(node) => Some(node),
ControlFlow::Continue(()) => None,
}
}
pub fn locate_all_at_point(
&self,
point: &<T::Envelope as Envelope>::Point,
) -> LocateAllAtPoint<T> {
LocateAllAtPoint::new(&self.root, SelectAtPointFunction::new(point.clone()))
}
pub fn locate_all_at_point_mut(
&mut self,
point: &<T::Envelope as Envelope>::Point,
) -> LocateAllAtPointMut<T> {
LocateAllAtPointMut::new(&mut self.root, SelectAtPointFunction::new(point.clone()))
}
pub fn locate_all_at_point_int<'a, V, B>(
&'a self,
point: &<T::Envelope as Envelope>::Point,
mut visitor: V,
) -> ControlFlow<B>
where
V: FnMut(&'a T) -> ControlFlow<B>,
{
select_nodes(
&self.root,
&SelectAtPointFunction::new(point.clone()),
&mut visitor,
)
}
pub fn locate_all_at_point_int_mut<'a, V, B>(
&'a mut self,
point: &<T::Envelope as Envelope>::Point,
mut visitor: V,
) -> ControlFlow<B>
where
V: FnMut(&'a mut T) -> ControlFlow<B>,
{
select_nodes_mut(
&mut self.root,
&SelectAtPointFunction::new(point.clone()),
&mut visitor,
)
}
pub fn remove_at_point(&mut self, point: &<T::Envelope as Envelope>::Point) -> Option<T> {
let removal_function = SelectAtPointFunction::new(point.clone());
self.remove_with_selection_function(removal_function)
}
}
impl<T, Params> RTree<T, Params>
where
Params: RTreeParams,
T: RTreeObject + PartialEq,
{
pub fn contains(&self, t: &T) -> bool {
self.locate_in_envelope(&t.envelope()).any(|e| e == t)
}
pub fn remove(&mut self, t: &T) -> Option<T> {
let removal_function = SelectEqualsFunction::new(t);
self.remove_with_selection_function(removal_function)
}
}
impl<T, Params> RTree<T, Params>
where
Params: RTreeParams,
T: PointDistance,
{
pub fn nearest_neighbor(&self, query_point: &<T::Envelope as Envelope>::Point) -> Option<&T> {
if self.size > 0 {
nearest_neighbor::nearest_neighbor(&self.root, query_point.clone())
.or_else(|| self.nearest_neighbor_iter(query_point).next())
} else {
None
}
}
pub fn nearest_neighbors(&self, query_point: &<T::Envelope as Envelope>::Point) -> Vec<&T> {
nearest_neighbor::nearest_neighbors(&self.root, query_point.clone())
}
pub fn locate_within_distance(
&self,
query_point: <T::Envelope as Envelope>::Point,
max_squared_radius: <<T::Envelope as Envelope>::Point as Point>::Scalar,
) -> LocateWithinDistanceIterator<T> {
let selection_function = SelectWithinDistanceFunction::new(query_point, max_squared_radius);
LocateWithinDistanceIterator::new(self.root(), selection_function)
}
pub fn drain_within_distance(
&mut self,
query_point: <T::Envelope as Envelope>::Point,
max_squared_radius: <<T::Envelope as Envelope>::Point as Point>::Scalar,
) -> DrainIterator<T, SelectWithinDistanceFunction<T>, Params> {
let selection_function = SelectWithinDistanceFunction::new(query_point, max_squared_radius);
self.drain_with_selection_function(selection_function)
}
pub fn nearest_neighbor_iter(
&self,
query_point: &<T::Envelope as Envelope>::Point,
) -> NearestNeighborIterator<T> {
nearest_neighbor::NearestNeighborIterator::new(&self.root, query_point.clone())
}
#[deprecated(note = "Please use nearest_neighbor_iter_with_distance_2 instead")]
pub fn nearest_neighbor_iter_with_distance(
&self,
query_point: &<T::Envelope as Envelope>::Point,
) -> NearestNeighborDistance2Iterator<T> {
nearest_neighbor::NearestNeighborDistance2Iterator::new(&self.root, query_point.clone())
}
pub fn nearest_neighbor_iter_with_distance_2(
&self,
query_point: &<T::Envelope as Envelope>::Point,
) -> NearestNeighborDistance2Iterator<T> {
nearest_neighbor::NearestNeighborDistance2Iterator::new(&self.root, query_point.clone())
}
pub fn pop_nearest_neighbor(
&mut self,
query_point: &<T::Envelope as Envelope>::Point,
) -> Option<T> {
if let Some(neighbor) = self.nearest_neighbor(query_point) {
let removal_function = SelectByAddressFunction::new(neighbor.envelope(), neighbor);
self.remove_with_selection_function(removal_function)
} else {
None
}
}
}
impl<T, Params> RTree<T, Params>
where
T: RTreeObject,
Params: RTreeParams,
{
pub fn insert(&mut self, t: T) {
Params::DefaultInsertionStrategy::insert(self, t);
self.size += 1;
}
}
impl<T, Params> IntoIterator for RTree<T, Params>
where
T: RTreeObject,
Params: RTreeParams,
{
type IntoIter = IntoIter<T>;
type Item = T;
fn into_iter(self) -> Self::IntoIter {
IntoIter::new(self.root)
}
}
impl<'a, T, Params> IntoIterator for &'a RTree<T, Params>
where
T: RTreeObject,
Params: RTreeParams,
{
type IntoIter = RTreeIterator<'a, T>;
type Item = &'a T;
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
impl<'a, T, Params> IntoIterator for &'a mut RTree<T, Params>
where
T: RTreeObject,
Params: RTreeParams,
{
type IntoIter = RTreeIteratorMut<'a, T>;
type Item = &'a mut T;
fn into_iter(self) -> Self::IntoIter {
self.iter_mut()
}
}
#[cfg(test)]
mod test {
use super::RTree;
use crate::algorithm::rstar::RStarInsertionStrategy;
use crate::params::RTreeParams;
use crate::test_utilities::{create_random_points, SEED_1};
use crate::DefaultParams;
struct TestParams;
impl RTreeParams for TestParams {
const MIN_SIZE: usize = 10;
const MAX_SIZE: usize = 20;
const REINSERTION_COUNT: usize = 1;
type DefaultInsertionStrategy = RStarInsertionStrategy;
}
#[test]
fn test_remove_capacity() {
pub struct WeirdParams;
impl RTreeParams for WeirdParams {
const MIN_SIZE: usize = 1;
const MAX_SIZE: usize = 10;
const REINSERTION_COUNT: usize = 1;
type DefaultInsertionStrategy = RStarInsertionStrategy;
}
let mut items: Vec<[f32; 2]> = Vec::new();
for i in 0..2 {
items.push([i as f32, i as f32]);
}
let mut tree: RTree<_, WeirdParams> = RTree::bulk_load_with_params(items);
assert_eq!(tree.remove(&[1.0, 1.0]).unwrap(), [1.0, 1.0]);
}
#[test]
fn test_create_rtree_with_parameters() {
let tree: RTree<[f32; 2], TestParams> = RTree::new_with_params();
assert_eq!(tree.size(), 0);
}
#[test]
fn test_insert_single() {
let mut tree: RTree<_> = RTree::new();
tree.insert([0.02f32, 0.4f32]);
assert_eq!(tree.size(), 1);
assert!(tree.contains(&[0.02, 0.4]));
assert!(!tree.contains(&[0.3, 0.2]));
}
#[test]
fn test_insert_many() {
const NUM_POINTS: usize = 1000;
let points = create_random_points(NUM_POINTS, SEED_1);
let mut tree = RTree::new();
for p in &points {
tree.insert(*p);
tree.root.sanity_check::<DefaultParams>(true);
}
assert_eq!(tree.size(), NUM_POINTS);
for p in &points {
assert!(tree.contains(p));
}
}
#[test]
fn test_fmt_debug() {
let tree = RTree::bulk_load(vec![[0, 1], [0, 1]]);
let debug: String = format!("{:?}", tree);
assert_eq!(debug, "RTree { size: 2, items: {[0, 1], [0, 1]} }");
}
#[test]
fn test_default() {
let tree: RTree<[f32; 2]> = Default::default();
assert_eq!(tree.size(), 0);
}
#[cfg(feature = "serde")]
#[test]
fn test_serialization() {
use crate::test_utilities::create_random_integers;
use serde_json;
const SIZE: usize = 20;
let points = create_random_integers::<[i32; 2]>(SIZE, SEED_1);
let tree = RTree::bulk_load(points.clone());
let json = serde_json::to_string(&tree).expect("Serializing tree failed");
let parsed: RTree<[i32; 2]> =
serde_json::from_str(&json).expect("Deserializing tree failed");
assert_eq!(parsed.size(), SIZE);
for point in &points {
assert!(parsed.contains(point));
}
}
#[test]
fn test_bulk_load_crash() {
let bulk_nodes = vec![
[570.0, 1080.0, 89.0],
[30.0, 1080.0, 627.0],
[1916.0, 1080.0, 68.0],
[274.0, 1080.0, 790.0],
[476.0, 1080.0, 895.0],
[1557.0, 1080.0, 250.0],
[1546.0, 1080.0, 883.0],
[1512.0, 1080.0, 610.0],
[1729.0, 1080.0, 358.0],
[1841.0, 1080.0, 434.0],
[1752.0, 1080.0, 696.0],
[1674.0, 1080.0, 705.0],
[136.0, 1080.0, 22.0],
[1593.0, 1080.0, 71.0],
[586.0, 1080.0, 272.0],
[348.0, 1080.0, 373.0],
[502.0, 1080.0, 2.0],
[1488.0, 1080.0, 1072.0],
[31.0, 1080.0, 526.0],
[1695.0, 1080.0, 559.0],
[1663.0, 1080.0, 298.0],
[316.0, 1080.0, 417.0],
[1348.0, 1080.0, 731.0],
[784.0, 1080.0, 126.0],
[225.0, 1080.0, 847.0],
[79.0, 1080.0, 819.0],
[320.0, 1080.0, 504.0],
[1714.0, 1080.0, 1026.0],
[264.0, 1080.0, 229.0],
[108.0, 1080.0, 158.0],
[1665.0, 1080.0, 604.0],
[496.0, 1080.0, 231.0],
[1813.0, 1080.0, 865.0],
[1200.0, 1080.0, 326.0],
[1661.0, 1080.0, 818.0],
[135.0, 1080.0, 229.0],
[424.0, 1080.0, 1016.0],
[1708.0, 1080.0, 791.0],
[1626.0, 1080.0, 682.0],
[442.0, 1080.0, 895.0],
];
let nodes = vec![
[1916.0, 1060.0, 68.0],
[1664.0, 1060.0, 298.0],
[1594.0, 1060.0, 71.0],
[225.0, 1060.0, 846.0],
[1841.0, 1060.0, 434.0],
[502.0, 1060.0, 2.0],
[1625.5852, 1060.0122, 682.0],
[1348.5273, 1060.0029, 731.08124],
[316.36127, 1060.0298, 418.24515],
[1729.3253, 1060.0023, 358.50134],
];
let mut tree = RTree::bulk_load(bulk_nodes);
for node in nodes {
tree.insert(node);
tree.root().sanity_check::<DefaultParams>(false);
}
}
}