1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
use alloc::{boxed::Box, vec::Vec};
use super::{IN_MEMORY_DEPTH, LargeSmt, SmtStorageReader, is_empty_parent};
use crate::{
Word,
hash::poseidon2::Poseidon2,
merkle::{InnerNodeInfo, smt::large::subtree::Subtree},
};
// ITERATORS
// ================================================================================================
enum InnerNodeIteratorState<'a> {
InMemory {
current_index: usize,
large_smt_in_memory_nodes: &'a [Word],
},
Subtree {
subtree_iter: Box<dyn Iterator<Item = Subtree> + 'a>,
current_subtree_node_iter: Option<Box<dyn Iterator<Item = InnerNodeInfo> + 'a>>,
},
Done,
}
pub struct LargeSmtInnerNodeIterator<'a, S: SmtStorageReader> {
large_smt: &'a LargeSmt<S>,
state: InnerNodeIteratorState<'a>,
}
impl<'a, S: SmtStorageReader> LargeSmtInnerNodeIterator<'a, S> {
pub(super) fn new(large_smt: &'a LargeSmt<S>) -> Self {
// in-memory nodes should never be empty
Self {
large_smt,
state: InnerNodeIteratorState::InMemory {
current_index: 0,
large_smt_in_memory_nodes: &large_smt.in_memory_nodes,
},
}
}
}
impl<S: SmtStorageReader> Iterator for LargeSmtInnerNodeIterator<'_, S> {
type Item = InnerNodeInfo;
/// Returns the next inner node info in the tree.
///
/// The iterator operates in three phases:
/// 1. InMemory: Iterates through the in-memory nodes (depths 0-IN_MEMORY_DEPTH-1)
/// 2. Subtree: Iterates through nodes in storage subtrees (depths IN_MEMORY_DEPTH-SMT_DEPTH)
/// 3. Done: No more nodes to iterate
fn next(&mut self) -> Option<Self::Item> {
loop {
match &mut self.state {
// Phase 1: Process in-memory nodes (depths 0-23)
InnerNodeIteratorState::InMemory { current_index, large_smt_in_memory_nodes } => {
// Iterate through nodes at depths 0 to IN_MEMORY_DEPTH-1
// Start at index 1 (root), max node index is (1 << IN_MEMORY_DEPTH) - 1
if *current_index == 0 {
*current_index = 1;
}
let max_node_idx = (1 << IN_MEMORY_DEPTH) - 1;
while *current_index <= max_node_idx {
let node_idx = *current_index;
*current_index += 1;
// Get children from flat layout: left at 2*i, right at 2*i+1
let left = large_smt_in_memory_nodes[node_idx * 2];
let right = large_smt_in_memory_nodes[node_idx * 2 + 1];
// Skip empty nodes
let depth = node_idx.ilog2() as u8;
let child_depth = depth + 1;
if !is_empty_parent(left, right, child_depth) {
return Some(InnerNodeInfo {
value: Poseidon2::merge(&[left, right]),
left,
right,
});
}
}
// All in-memory nodes processed. Transition to Phase 2: Subtree iteration
match self.large_smt.storage.iter_subtrees() {
Ok(subtree_iter) => {
self.state = InnerNodeIteratorState::Subtree {
subtree_iter,
current_subtree_node_iter: None,
};
continue; // Start processing subtrees immediately
},
Err(_e) => {
// Storage error occurred - we should propagate this properly
// For now, transition to Done state to avoid infinite loops
self.state = InnerNodeIteratorState::Done;
return None;
},
}
},
// Phase 2: Process storage subtrees (depths 25-64)
InnerNodeIteratorState::Subtree { subtree_iter, current_subtree_node_iter } => {
loop {
// First, try to get the next node from current subtree
if let Some(node_iter) = current_subtree_node_iter
&& let Some(info) = node_iter.as_mut().next()
{
return Some(info);
}
// Current subtree exhausted, move to next subtree
match subtree_iter.next() {
Some(next_subtree) => {
// Collect is necessary here because iter_inner_node_info returns
// an iterator borrowing from next_subtree, which would outlive
// the subtree itself. We need to eagerly evaluate to owned data.
#[expect(clippy::needless_collect)]
let infos: Vec<InnerNodeInfo> =
next_subtree.iter_inner_node_info().collect();
*current_subtree_node_iter = Some(Box::new(infos.into_iter()));
},
None => {
self.state = InnerNodeIteratorState::Done;
return None; // All subtrees processed
},
}
}
},
InnerNodeIteratorState::Done => {
return None; // Iteration finished.
},
}
}
}
}