From b1ac058c81a28926e4b97420f9dd95a63a877c55 Mon Sep 17 00:00:00 2001 From: Tavian Barnes Date: Wed, 24 Feb 2021 14:24:26 -0500 Subject: vp: Add a non-consuming iter() to VpTree and FlatVpTree --- src/vp.rs | 92 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 92 insertions(+) (limited to 'src') diff --git a/src/vp.rs b/src/vp.rs index 67a094f..bf91d57 100644 --- a/src/vp.rs +++ b/src/vp.rs @@ -201,6 +201,11 @@ impl VpTree { } } + /// Iterate over the items stored in this tree. + pub fn iter(&self) -> Iter { + (&self).into_iter() + } + /// Rebalance this VP tree. pub fn balance(&mut self) { let mut nodes = Vec::new(); @@ -327,6 +332,56 @@ impl IntoIterator for VpTree { } } +/// An iterator over the values in a VP tree. +pub struct Iter<'a, T: Proximity> { + stack: Vec<&'a VpNode>, +} + +impl<'a, T: Proximity> Iter<'a, T> { + fn new(node: &'a Option>) -> Self { + Self { + stack: node.as_ref().into_iter().collect(), + } + } +} + +impl<'a, T> Debug for Iter<'a, T> +where + T: Proximity + Debug, + DistanceValue: Debug, +{ + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + f.debug_struct("Iter") + .field("stack", &self.stack) + .finish() + } +} + +impl<'a, T: Proximity> Iterator for Iter<'a, T> { + type Item = &'a T; + + fn next(&mut self) -> Option<&'a T> { + self.stack.pop().map(|node| { + if let Some(inside) = &node.inside { + self.stack.push(inside); + } + if let Some(outside) = &node.outside { + self.stack.push(outside); + } + &node.item + }) + } +} + +impl<'a, T: Proximity> IntoIterator for &'a VpTree { + type Item = &'a T; + type IntoIter = Iter<'a, T>; + + fn into_iter(self) -> Self::IntoIter { + Iter::new(&self.root) + } +} + impl NearestNeighbors for VpTree where K: Proximity, @@ -452,6 +507,11 @@ impl FlatVpTree { nodes: FlatVpNode::balanced(items), } } + + /// Iterate over the items stored in this tree. + pub fn iter(&self) -> FlatIter { + (&self).into_iter() + } } impl Debug for FlatVpTree @@ -504,6 +564,38 @@ impl IntoIterator for FlatVpTree { } } +/// An iterator over the values in a flat VP tree. +pub struct FlatIter<'a, T: Proximity>(std::slice::Iter<'a, FlatVpNode>); + +impl<'a, T> Debug for FlatIter<'a, T> +where + T: Proximity + Debug, + DistanceValue: Debug, +{ + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + f.debug_tuple("FlatIter") + .field(&self.0) + .finish() + } +} + +impl<'a, T: Proximity> Iterator for FlatIter<'a, T> { + type Item = &'a T; + + fn next(&mut self) -> Option<&'a T> { + self.0.next().map(|n| &n.item) + } +} + +impl<'a, T: Proximity> IntoIterator for &'a FlatVpTree { + type Item = &'a T; + type IntoIter = FlatIter<'a, T>; + + fn into_iter(self) -> Self::IntoIter { + FlatIter(self.nodes.iter()) + } +} + impl NearestNeighbors for FlatVpTree where K: Proximity, -- cgit v1.2.3