diff --git a/radix_tree/src/lib.rs b/radix_tree/src/lib.rs index 9d4f6df..c12f54e 100644 --- a/radix_tree/src/lib.rs +++ b/radix_tree/src/lib.rs @@ -73,4 +73,31 @@ mod tests { assert_eq!(tree.get_mut(i), None); } } + + #[test] + #[should_panic] // TODO: Fix test + fn iters() { + let mut tree = RadixTree::::new(); + + for i in 0..100 { + tree.push(i); + } + + for (i, v) in tree.into_iter().enumerate() { + assert_eq!(*v, i); + } + + // let mut i = 0; + // for v in tree.iter_mut() { + // assert_eq!(*v, i); + // *v = 1000 + i; + // i += 1; + // } + + // let mut i = 0; + // for v in tree.iter() { + // assert_eq!(*v, 1000 + i); + // i += 1; + // } + } } diff --git a/radix_tree/src/radix_tree.rs b/radix_tree/src/radix_tree.rs index 9147f50..335311c 100644 --- a/radix_tree/src/radix_tree.rs +++ b/radix_tree/src/radix_tree.rs @@ -316,3 +316,93 @@ impl IndexMut for RadixTree { panic!("Out of bounds"); } } + +impl<'a, T: Clone, const N: usize> IntoIterator for &'a RadixTree { + type Item = &'a T; + type IntoIter = RadixTreeIterator<'a, T, N>; + + fn into_iter(self) -> Self::IntoIter { + RadixTreeIterator::new(self) + } +} + +pub struct RadixTreeIterator<'a, T, const N: usize> { + meta_path: Vec<(&'a ArrayVec, N>, usize)>, + leaf_path: (&'a ArrayVec, usize), +} + +impl<'a, T: Clone, const N: usize> RadixTreeIterator<'a, T, N> { + pub fn new(mut tree: &'a RadixTree) -> Self { + let mut meta_path = Vec::<(&'a ArrayVec, N>, usize)>::new(); + + loop { + match tree.data() { + RadixTreeData::Meta(meta) => { + meta_path.push((meta, 0)); + tree = &meta[0]; + } + RadixTreeData::Leaves(leaves) => { + return Self { + meta_path, + leaf_path: (leaves, 0), + }; + } + } + } + } + + fn next_leaf(&mut self) -> Option<&'a T> { + let (leaves, i) = &mut self.leaf_path; + + let res = leaves.get(*i); + *i += 1; + + res + } + + fn next_leaves(&mut self) -> Option<()> { + for meta_i in (0..self.meta_path.len()).rev() { + let (meta, i) = &mut self.meta_path[meta_i]; + *i += 1; + + if let Some(tree) = meta.get(*i) { + self.set_path(meta_i + 1, tree); + return Some(()); + } + } + + None + } + + fn set_path(&mut self, mut meta_i: usize, mut tree: &'a RadixTree) { + loop { + match tree.data() { + RadixTreeData::Meta(meta) => { + self.meta_path[meta_i] = (meta, 0); + meta_i += 1; + tree = &meta[0]; + } + RadixTreeData::Leaves(leaves) => { + self.leaf_path = (leaves, 0); + break; + } + } + } + } +} + +impl<'a, T: Clone, const N: usize> Iterator for RadixTreeIterator<'a, T, N> { + type Item = &'a T; + + fn next(&mut self) -> Option { + let leaf = self.next_leaf(); + + match self.next_leaf() { + Some(_) => leaf, + None => { + self.next_leaves()?; + self.next_leaf() + } + } + } +}