From 81f89229b4b5a1b82814fb7acc1905c0a463a3c6 Mon Sep 17 00:00:00 2001 From: Tavian Barnes 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 KdTree { } } + /// Iterate over the items stored in this tree. + pub fn iter(&self) -> Iter { + (&self).into_iter() + } + /// Rebalance this k-d tree. pub fn balance(&mut self) { let mut nodes = Vec::new(); @@ -265,7 +270,7 @@ impl IntoIter { impl Iterator for IntoIter { type Item = T; - fn next(&mut self) -> Option { + fn next(&mut self) -> Option { self.stack.pop().map(|node| { if let Some(left) = node.left { self.stack.push(*left); @@ -287,6 +292,45 @@ impl IntoIterator for KdTree { } } +/// An iterator over the values in a k-d tree. +#[derive(Debug)] +pub struct Iter<'a, T> { + stack: Vec<&'a KdNode>, +} + +impl<'a, T> Iter<'a, T> { + fn new(node: &'a Option>) -> 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.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 { + type Item = &'a T; + type IntoIter = Iter<'a, T>; + + fn into_iter(self) -> Self::IntoIter { + Iter::new(&self.root) + } +} + impl NearestNeighbors for KdTree where K: KdProximity, @@ -411,6 +455,11 @@ impl FlatKdTree { nodes: FlatKdNode::balanced(items), } } + + /// Iterate over the items stored in this tree. + pub fn iter(&self) -> FlatIter { + (&self).into_iter() + } } impl FromIterator for FlatKdTree { @@ -426,7 +475,7 @@ pub struct FlatIntoIter(std::vec::IntoIter>); impl Iterator for FlatIntoIter { type Item = T; - fn next(&mut self) -> Option { + fn next(&mut self) -> Option { self.0.next().map(|n| n.item) } } @@ -440,6 +489,27 @@ impl IntoIterator for FlatKdTree { } } +/// An iterator over the values in a flat k-d tree. +#[derive(Debug)] +pub struct FlatIter<'a, T>(std::slice::Iter<'a, FlatKdNode>); + +impl<'a, T> Iterator for FlatIter<'a, T> { + type Item = &'a T; + + fn next(&mut self) -> Option { + self.0.next().map(|n| &n.item) + } +} + +impl<'a, T> IntoIterator for &'a FlatKdTree { + type Item = &'a T; + type IntoIter = FlatIter<'a, T>; + + fn into_iter(self) -> Self::IntoIter { + FlatIter(self.nodes.iter()) + } +} + impl NearestNeighbors for FlatKdTree where K: KdProximity, -- cgit v1.2.3