Add iterator over ProgTree elements

src/prog_tree.rs

@@ -1,4 +1,4 @@
-use crate::{Arc, Error, Tree, Value, builder::Builder, utils::opt_packing_factor};
+use crate::{Arc, Error, Tree, Value, builder::Builder, iter::Iter, utils::{Length, opt_packing_factor}};
 use educe::Educe;
 use ethereum_hashing::hash32_concat;
 use parking_lot::RwLock;
@@ -167,4 +167,114 @@ impl<T: Value + Send + Sync> ProgTree<T> {
             }
         }
     }
+
+    /// Create an iterator over all elements in the progressive tree.
+    ///
+    /// The iterator traverses elements in order by visiting each binary subtree
+    /// (right child) at increasing progressive depths:
+    /// 1. All elements in the right child at the root level
+    /// 2. All elements in the right child of the first left node
+    /// 3. All elements in the right child of the second left node
+    ///
+    /// And so on, following the progressive tree structure as defined in EIP-7916.
+    pub fn iter(&self, length: usize) -> ProgTreeIter<'_, T> {
+        ProgTreeIter::new(self, length)
+    }
+}
+
+/// Iterator over elements in a progressive tree.
+///
+/// The iterator traverses each binary subtree (right child) in sequence by following
+/// the left spine of the progressive tree structure.
+#[derive(Debug)]
+pub struct ProgTreeIter<'a, T: Value> {
+    /// Current progressive node being traversed.
+    current_prog_node: Option<&'a ProgTree<T>>,
+    /// Current iterator over a binary subtree (Tree).
+    current_iter: Option<Iter<'a, T>>,
+    /// Progressive depth for calculating the next subtree depth.
+    prog_depth: u32,
+    /// Total number of elements to iterate.
+    length: usize,
+    /// Number of elements already yielded.
+    yielded: usize,
 }
+
+impl<'a, T: Value> ProgTreeIter<'a, T> {
+    fn new(root: &'a ProgTree<T>, length: usize) -> Self {
+        let mut iter = Self {
+            current_prog_node: Some(root),
+            current_iter: None,
+            prog_depth: 0,
+            length,
+            yielded: 0,
+        };
+
+        // Initialize by setting up the iterator for the first right child
+        iter.advance_to_next_subtree();
+        iter
+    }
+
+    /// Advance to the next binary subtree by moving to the left child and
+    /// setting up an iterator for its right child.
+    fn advance_to_next_subtree(&mut self) {
+        match self.current_prog_node {
+            None | Some(ProgTree::ProgZero) => {
+                // No more subtrees
+                self.current_iter = None;
+                self.current_prog_node = None;
+            }
+            Some(ProgTree::ProgNode { left, right, .. }) => {
+                self.prog_depth += 1;
+
+                // Calculate the depth and length for this binary subtree
+                let binary_depth = ProgTree::<T>::prog_depth_to_binary_depth(self.prog_depth);
+                let remaining = self.length.saturating_sub(self.yielded);
+                let capacity = ProgTree::<T>::capacity_at_depth(self.prog_depth);
+                let subtree_length = remaining.min(capacity);
+
+                // Create an iterator for the right subtree
+                self.current_iter = Some(Iter::from_index(
+                    0,
+                    right,
+                    binary_depth,
+                    Length(subtree_length),
+                ));
+
+                // Move to the left child for the next iteration
+                self.current_prog_node = Some(left);
+            }
+        }
+    }
+}
+
+impl<'a, T: Value> Iterator for ProgTreeIter<'a, T> {
+    type Item = &'a T;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            // Try to get the next item from the current binary tree iterator
+            if let Some(iter) = &mut self.current_iter
+                && let Some(value) = iter.next()
+            {
+                self.yielded += 1;
+                return Some(value);
+            }
+
+            // Current subtree exhausted, move to the next one
+            if self.current_prog_node.is_some() {
+                self.advance_to_next_subtree();
+            } else {
+                // No more subtrees to iterate
+                return None;
+            }
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let remaining = self.length.saturating_sub(self.yielded);
+        (remaining, Some(remaining))
+    }
+}
+
+impl<T: Value> ExactSizeIterator for ProgTreeIter<'_, T> {}

src/tests/prog.rs

@@ -26,3 +26,76 @@ fn wow_u64() {
 
     println!("{tree:#?}");
 }
+
+#[test]
+fn prog_tree_iterator() {
+    let mut tree = ProgTree::<u64>::empty();
+
+    // Build a tree with 65 elements
+    for i in 1..=65 {
+        tree = tree.push(i, i as usize - 1).unwrap();
+    }
+
+    // Iterate and collect all elements
+    let collected: Vec<_> = tree.iter(65).copied().collect();
+
+    // Verify we got all 65 elements in order
+    assert_eq!(collected.len(), 65);
+    for (i, &value) in collected.iter().enumerate() {
+        assert_eq!(value, (i + 1) as u64, "Element at index {} should be {}", i, i + 1);
+    }
+}
+
+#[test]
+fn prog_tree_iterator_empty() {
+    let tree = ProgTree::<u64>::empty();
+    let collected: Vec<_> = tree.iter(0).collect();
+    assert_eq!(collected.len(), 0);
+}
+
+#[test]
+fn prog_tree_iterator_small() {
+    let mut tree = ProgTree::<u64>::empty();
+
+    // Build a small tree with just 4 elements (one packed leaf)
+    for i in 1..=4 {
+        tree = tree.push(i, i as usize - 1).unwrap();
+    }
+
+    let collected: Vec<_> = tree.iter(4).copied().collect();
+    assert_eq!(collected, vec![1, 2, 3, 4]);
+}
+
+#[test]
+fn prog_tree_iterator_exact_size() {
+    let mut tree = ProgTree::<u64>::empty();
+
+    for i in 1..=20 {
+        tree = tree.push(i, i as usize - 1).unwrap();
+    }
+
+    let iter = tree.iter(20);
+    assert_eq!(iter.len(), 20);
+
+    let collected: Vec<_> = iter.copied().collect();
+    assert_eq!(collected.len(), 20);
+}
+
+#[test]
+fn prog_tree_iterator_hash256() {
+    let mut tree = ProgTree::<Hash256>::empty();
+
+    // Build a tree with non-packed values
+    for i in 1..=10 {
+        let hash = Hash256::repeat_byte(i as u8);
+        tree = tree.push(hash, i - 1).unwrap();
+    }
+
+    let collected: Vec<_> = tree.iter(10).collect();
+    assert_eq!(collected.len(), 10);
+
+    // Verify order
+    for (i, hash) in collected.iter().enumerate() {
+        assert_eq!(**hash, Hash256::repeat_byte((i + 1) as u8));
+    }
+}