From 81f89229b4b5a1b82814fb7acc1905c0a463a3c6 Mon Sep 17 00:00:00 2001
From: Tavian Barnes <tavianator@tavianator.com>
Date: Wed, 24 Feb 2021 14:23:35 -0500
Subject: kd: Add a non-consuming iter() to KdTree and FlatKdTree

---
 src/kd.rs | 74 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 72 insertions(+), 2 deletions(-)

(limited to 'src')

diff --git a/src/kd.rs b/src/kd.rs
index 5434ddd..d37321e 100644
--- a/src/kd.rs
+++ b/src/kd.rs
@@ -185,6 +185,11 @@ impl<T: Coordinates> KdTree<T> {
         }
     }
 
+    /// Iterate over the items stored in this tree.
+    pub fn iter(&self) -> Iter<T> {
+        (&self).into_iter()
+    }
+
     /// Rebalance this k-d tree.
     pub fn balance(&mut self) {
         let mut nodes = Vec::new();
@@ -265,7 +270,7 @@ impl<T> IntoIter<T> {
 impl<T> Iterator for IntoIter<T> {
     type Item = T;
 
-    fn next(&mut self) -> Option<T> {
+    fn next(&mut self) -> Option<Self::Item> {
         self.stack.pop().map(|node| {
             if let Some(left) = node.left {
                 self.stack.push(*left);
@@ -287,6 +292,45 @@ impl<T> IntoIterator for KdTree<T> {
     }
 }
 
+/// An iterator over the values in a k-d tree.
+#[derive(Debug)]
+pub struct Iter<'a, T> {
+    stack: Vec<&'a KdNode<T>>,
+}
+
+impl<'a, T> Iter<'a, T> {
+    fn new(node: &'a Option<KdNode<T>>) -> Self {
+        Self {
+            stack: node.as_ref().into_iter().collect(),
+        }
+    }
+}
+
+impl<'a, T> Iterator for Iter<'a, T> {
+    type Item = &'a T;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.stack.pop().map(|node| {
+            if let Some(left) = &node.left {
+                self.stack.push(left);
+            }
+            if let Some(right) = &node.right {
+                self.stack.push(right);
+            }
+            &node.item
+        })
+    }
+}
+
+impl<'a, T> IntoIterator for &'a KdTree<T> {
+    type Item = &'a T;
+    type IntoIter = Iter<'a, T>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        Iter::new(&self.root)
+    }
+}
+
 impl<K, V> NearestNeighbors<K, V> for KdTree<V>
 where
     K: KdProximity<V>,
@@ -411,6 +455,11 @@ impl<T: Coordinates> FlatKdTree<T> {
             nodes: FlatKdNode::balanced(items),
         }
     }
+
+    /// Iterate over the items stored in this tree.
+    pub fn iter(&self) -> FlatIter<T> {
+        (&self).into_iter()
+    }
 }
 
 impl<T: Coordinates> FromIterator<T> for FlatKdTree<T> {
@@ -426,7 +475,7 @@ pub struct FlatIntoIter<T>(std::vec::IntoIter<FlatKdNode<T>>);
 impl<T> Iterator for FlatIntoIter<T> {
     type Item = T;
 
-    fn next(&mut self) -> Option<T> {
+    fn next(&mut self) -> Option<Self::Item> {
         self.0.next().map(|n| n.item)
     }
 }
@@ -440,6 +489,27 @@ impl<T> IntoIterator for FlatKdTree<T> {
     }
 }
 
+/// An iterator over the values in a flat k-d tree.
+#[derive(Debug)]
+pub struct FlatIter<'a, T>(std::slice::Iter<'a, FlatKdNode<T>>);
+
+impl<'a, T> Iterator for FlatIter<'a, T> {
+    type Item = &'a T;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.0.next().map(|n| &n.item)
+    }
+}
+
+impl<'a, T> IntoIterator for &'a FlatKdTree<T> {
+    type Item = &'a T;
+    type IntoIter = FlatIter<'a, T>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        FlatIter(self.nodes.iter())
+    }
+}
+
 impl<K, V> NearestNeighbors<K, V> for FlatKdTree<V>
 where
     K: KdProximity<V>,
-- 
cgit v1.2.3