Skip to content

Commit

Permalink
Implement an owning iterator and IntoIter for RTree
Browse files Browse the repository at this point in the history
  • Loading branch information
adamreichold committed Dec 24, 2023
1 parent a8904d2 commit da2bd1f
Show file tree
Hide file tree
Showing 3 changed files with 63 additions and 3 deletions.
2 changes: 1 addition & 1 deletion rstar/src/algorithm/iterators.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ use crate::RTree;
use smallvec::SmallVec;

pub use super::intersection_iterator::IntersectionIterator;
pub use super::removal::DrainIterator;
pub use super::removal::{DrainIterator, IntoIter};

/// Iterator returned by [`RTree::locate_all_at_point`].
pub type LocateAllAtPoint<'a, T> = SelectionIterator<'a, T, SelectAtPointFunction<T>>;
Expand Down
55 changes: 55 additions & 0 deletions rstar/src/algorithm/removal.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,45 @@ use alloc::{vec, vec::Vec};
#[allow(unused_imports)] // Import is required when building without std
use num_traits::Float;

/// Iterator returned by `impl IntoIter for RTree`.
///
/// Consumes the whole tree and yields all leaf objects.
pub struct IntoIter<T>
where
T: RTreeObject,
{
node_stack: Vec<RTreeNode<T>>,
}

impl<T> IntoIter<T>
where
T: RTreeObject,
{
pub fn new(root: ParentNode<T>) -> Self {
Self {
node_stack: vec![RTreeNode::Parent(root)],
}
}
}

impl<T> Iterator for IntoIter<T>
where
T: RTreeObject,
{
type Item = T;

fn next(&mut self) -> Option<Self::Item> {
while let Some(node) = self.node_stack.pop() {
match node {
RTreeNode::Leaf(object) => return Some(object),
RTreeNode::Parent(parent) => self.node_stack.extend(parent.children),
}
}

None
}
}

/// Iterator returned by `RTree::drain_*` methods.
///
/// Draining iterator that removes elements of the tree selected by a
Expand Down Expand Up @@ -338,4 +377,20 @@ mod test {
assert_eq!(sel_count, 0);
assert_eq!(tree.size(), 1000 - 80 - 326);
}

#[test]
fn test_into_iter() {
const SIZE: usize = 100;
let mut points = create_random_points(SIZE, SEED_1);
let tree = RTree::bulk_load(points.clone());

let mut vec = tree.into_iter().collect::<Vec<_>>();

assert_eq!(vec.len(), points.len());

points.sort_unstable_by(|lhs, rhs| lhs.partial_cmp(rhs).unwrap());
vec.sort_unstable_by(|lhs, rhs| lhs.partial_cmp(rhs).unwrap());

assert_eq!(points, vec);
}
}
9 changes: 7 additions & 2 deletions rstar/src/rtree.rs
Original file line number Diff line number Diff line change
Expand Up @@ -834,12 +834,17 @@ where
}
}

impl<T, Params> RTree<T, Params>
impl<T, Params> IntoIterator for RTree<T, Params>
where
T: RTreeObject,
<T::Envelope as Envelope>::Point: Point,
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>
Expand Down

0 comments on commit da2bd1f

Please sign in to comment.