Skip to content

Commit

Permalink
Added (failing) into_iter
Browse files Browse the repository at this point in the history
  • Loading branch information
voltrevo committed Aug 17, 2023
1 parent 0eb7a51 commit 9759d4d
Show file tree
Hide file tree
Showing 2 changed files with 117 additions and 0 deletions.
27 changes: 27 additions & 0 deletions radix_tree/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::<usize, 4>::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;
// }
}
}
90 changes: 90 additions & 0 deletions radix_tree/src/radix_tree.rs
Original file line number Diff line number Diff line change
Expand Up @@ -316,3 +316,93 @@ impl<T: Clone, const N: usize> IndexMut<usize> for RadixTree<T, N> {
panic!("Out of bounds");
}
}

impl<'a, T: Clone, const N: usize> IntoIterator for &'a RadixTree<T, N> {
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<RadixTree<T, N>, N>, usize)>,
leaf_path: (&'a ArrayVec<T, N>, usize),
}

impl<'a, T: Clone, const N: usize> RadixTreeIterator<'a, T, N> {
pub fn new(mut tree: &'a RadixTree<T, N>) -> Self {
let mut meta_path = Vec::<(&'a ArrayVec<RadixTree<T, N>, 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<T, N>) {
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<Self::Item> {
let leaf = self.next_leaf();

match self.next_leaf() {
Some(_) => leaf,
None => {
self.next_leaves()?;
self.next_leaf()
}
}
}
}

0 comments on commit 9759d4d

Please sign in to comment.