From cae4a32ebda00986b379d831ea4cd84f794b8236 Mon Sep 17 00:00:00 2001
From: Adam Reichold <adam.reichold@t-online.de>
Date: Sun, 24 Dec 2023 10:41:53 +0100
Subject: [PATCH] Implement an owning iterator and IntoIter for RTree

---
 rstar/CHANGELOG.md               |  1 +
 rstar/src/algorithm/iterators.rs |  2 +-
 rstar/src/algorithm/removal.rs   | 55 ++++++++++++++++++++++++++++++++
 rstar/src/rtree.rs               |  9 ++++--
 4 files changed, 64 insertions(+), 3 deletions(-)

diff --git a/rstar/CHANGELOG.md b/rstar/CHANGELOG.md
index 9d88cf7..141ffa6 100644
--- a/rstar/CHANGELOG.md
+++ b/rstar/CHANGELOG.md
@@ -2,6 +2,7 @@
 
 ## Added
 - Add optional support for the [mint](https://docs.rs/mint/0.5.9/mint/index.html) crate
+- Implemented `IntoIter` for `RTree`, i.e. added a owning iterator
 - Added cached envelope bulk load benchmark
 
 ## Changed
diff --git a/rstar/src/algorithm/iterators.rs b/rstar/src/algorithm/iterators.rs
index 782b0be..c802734 100644
--- a/rstar/src/algorithm/iterators.rs
+++ b/rstar/src/algorithm/iterators.rs
@@ -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>>;
diff --git a/rstar/src/algorithm/removal.rs b/rstar/src/algorithm/removal.rs
index 6b498a8..6788ae3 100644
--- a/rstar/src/algorithm/removal.rs
+++ b/rstar/src/algorithm/removal.rs
@@ -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(crate) 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
@@ -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);
+    }
 }
diff --git a/rstar/src/rtree.rs b/rstar/src/rtree.rs
index c2d1865..357057f 100644
--- a/rstar/src/rtree.rs
+++ b/rstar/src/rtree.rs
@@ -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>