From 9759d4d4fcf6070d642274787b130300cf5f4fa9 Mon Sep 17 00:00:00 2001
From: Andrew Morris <voltrevo@gmail.com>
Date: Thu, 17 Aug 2023 12:43:06 +1000
Subject: [PATCH] Added (failing) into_iter

---
 radix_tree/src/lib.rs        | 27 +++++++++++
 radix_tree/src/radix_tree.rs | 90 ++++++++++++++++++++++++++++++++++++
 2 files changed, 117 insertions(+)

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::<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;
+    // }
+  }
 }
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<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()
+      }
+    }
+  }
+}