feat: Add Mmr::nodes_from and Mmr::from_nodes_unchecked

CHANGELOG.md

@@ -4,6 +4,8 @@
 #### Features
 
 - [BREAKING] Added `trace`, `trace.CONST`, and `trace.event("...")` assembly as syntactic sugar for emitting optional read-only trace events. This adds variants `Trace` and `TraceImm` to the public enum `miden_assembly_syntax::ast::Instruction` ([#3478](https://github.com/0xMiden/miden-vm/pull/3478)).
+- Added `Mmr::nodes_from(start)`, returning the MMR's nodes at indices `start..` in insertion (postorder) order ([#3585](https://github.com/0xMiden/miden-vm/pull/3585)).
+- [BREAKING] Added `Mmr::from_nodes_unchecked(forest, nodes)`, constructing an MMR from its complete postorder node array without recomputing hashes ([#3585](https://github.com/0xMiden/miden-vm/pull/3585)).
 
 #### Changes
 

crates/crypto/src/merkle/mmr/error.rs

@@ -14,6 +14,8 @@ pub enum MmrError {
     ForestOutOfBounds(usize, usize),
     #[error("mmr forest size {requested} exceeds maximum {max}")]
     ForestSizeExceeded { requested: usize, max: usize },
+    #[error("mmr node count {actual} does not match forest node count {expected}")]
+    InvalidNodeCount { expected: usize, actual: usize },
     #[error("mmr peak does not match the computed merkle root of the provided authentication path")]
     PeakPathMismatch,
     #[error("requested peak index is {peak_idx} but the number of peaks is {peaks_len}")]

crates/crypto/src/merkle/mmr/full.rs

@@ -10,8 +10,8 @@
 //! depths, i.e. as part of adding a new element to the forest the trees with same depth are
 //! merged, creating a new tree with depth d+1, this process is continued until the property is
 //! reestablished.
-use alloc::{sync::Arc, vec::Vec};
-use core::ops::Index;
+use alloc::{string::ToString, sync::Arc, vec::Vec};
+use core::{iter::FusedIterator, ops::Index, slice};
 
 use super::{
     super::{InnerNodeInfo, MerklePath},
@@ -83,8 +83,24 @@ impl NodeStore {
     }
 
     /// Returns an iterator over all nodes in the store, in insertion (postorder) order.
-    pub fn iter(&self) -> impl Iterator<Item = &Word> {
-        self.chunks.iter().flat_map(|chunk| chunk.iter())
+    pub fn iter(&self) -> MmrNodeIter<'_> {
+        self.iter_from(0)
+    }
+
+    /// Returns an iterator over the nodes at indices `start..`, in insertion (postorder) order.
+    ///
+    /// Skips directly to the chunk containing `start` instead of walking from the front. Returns
+    /// an empty iterator if `start >= self.len()`.
+    pub fn iter_from(&self, start: usize) -> MmrNodeIter<'_> {
+        let first_chunk = start / NODE_CHUNK_CAPACITY;
+        let offset = start % NODE_CHUNK_CAPACITY;
+        match self.chunks.get(first_chunk) {
+            Some(chunk) => MmrNodeIter {
+                current: chunk[offset.min(chunk.len())..].iter(),
+                chunks: &self.chunks[first_chunk + 1..],
+            },
+            None => MmrNodeIter { current: [].iter(), chunks: &[] },
+        }
     }
 }
 
@@ -165,6 +181,37 @@ impl Mmr {
         Self::try_from_iter_with_limit(values, Forest::MAX_LEAVES)
     }
 
+    /// Constructs an MMR from its forest and complete node array, in insertion (postorder) order,
+    /// e.g. as previously obtained from `mmr.nodes_from(0).copied()` (see [Mmr::nodes_from]).
+    ///
+    /// The only validation performed is structural: the node count must match `forest`. The
+    /// nodes are otherwise taken verbatim — no hashes are recomputed or verified.
+    /// Comparing the result's [Mmr::peaks] against a trusted commitment checks the accumulator
+    /// state, but because the peaks are read from the stored nodes rather than recomputed, it does
+    /// not validate any non-peak nodes. The nodes must therefore come from a trusted source, e.g.
+    /// the caller's own previously validated state.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the number of nodes does not match the node count of `forest`.
+    pub fn from_nodes_unchecked(
+        forest: Forest,
+        nodes: impl IntoIterator<Item = Word>,
+    ) -> Result<Self, MmrError> {
+        Self::from_store(forest, nodes.into_iter().collect())
+    }
+
+    /// Constructs an MMR from its forest and node store, validating the node count.
+    fn from_store(forest: Forest, nodes: NodeStore) -> Result<Self, MmrError> {
+        if nodes.len() != forest.num_nodes() {
+            return Err(MmrError::InvalidNodeCount {
+                expected: forest.num_nodes(),
+                actual: nodes.len(),
+            });
+        }
+        Ok(Self { forest, nodes })
+    }
+
     pub(crate) fn try_from_iter_with_limit<T: IntoIterator<Item = Word>>(
         values: T,
         max_leaves: usize,
@@ -194,6 +241,19 @@ impl Mmr {
         self.forest
     }
 
+    /// Returns an iterator over the MMR's nodes at indices `start..`, in insertion (postorder)
+    /// order. Returns an empty iterator if `start` is greater than or equal to the total node
+    /// count, which is given by `self.forest().num_nodes()`.
+    ///
+    /// The node buffer is strictly append-only, so a consumer that has persisted the first
+    /// `start` nodes can incrementally sync by appending only the nodes returned here.
+    ///
+    /// Positioning is cheap: the iterator starts directly at `start` without walking the
+    /// preceding nodes, and it knows its exact remaining length.
+    pub fn nodes_from(&self, start: usize) -> impl ExactSizeIterator<Item = &Word> + Clone {
+        self.nodes.iter_from(start)
+    }
+
     // FUNCTIONALITY
     // ============================================================================================
 
@@ -492,14 +552,19 @@ impl Deserializable for Mmr {
     fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
         let forest = Forest::read_from(source)?;
         let count = source.read_usize()?;
+        // Reject a forest/count mismatch before reading the nodes, so malformed input fails fast.
         if count != forest.num_nodes() {
-            return Err(DeserializationError::InvalidValue(alloc::format!(
-                "MMR node count {count} does not match forest node count {}",
-                forest.num_nodes()
-            )));
+            return Err(DeserializationError::InvalidValue(
+                MmrError::InvalidNodeCount {
+                    expected: forest.num_nodes(),
+                    actual: count,
+                }
+                .to_string(),
+            ));
         }
         let nodes = source.read_many_iter(count)?.collect::<Result<NodeStore, _>>()?;
-        Ok(Self { forest, nodes })
+        Self::from_store(forest, nodes)
+            .map_err(|err| DeserializationError::InvalidValue(err.to_string()))
     }
 }
 
@@ -553,6 +618,84 @@ impl<'de> serde::Deserialize<'de> for NodeStore {
 // ITERATOR
 // ===============================================================================================
 
+/// Iterator over a suffix of the [Mmr]'s node buffer, in insertion (postorder) order.
+///
+/// Underlies [Mmr::nodes_from], which returns it opaquely. Positioning is cheap: [Iterator::nth]
+/// (and therefore [Iterator::skip]) jumps over whole chunks instead of advancing one node at a
+/// time, and the iterator knows its exact remaining length ([ExactSizeIterator]).
+#[derive(Clone, Debug)]
+pub(super) struct MmrNodeIter<'a> {
+    /// Remainder of the chunk currently being yielded.
+    current: slice::Iter<'a, Word>,
+    /// Chunks after the current one; every chunk except the last is full.
+    chunks: &'a [Arc<Vec<Word>>],
+}
+
+impl<'a> MmrNodeIter<'a> {
+    /// Advances `current` to the next chunk, or returns `None` if no chunks remain.
+    ///
+    /// On `None` the iterator is left exhausted even if `current` still held items, so a caller
+    /// skipping past the end (see [Iterator::nth]) doesn't leave the skipped items behind.
+    fn advance_chunk(&mut self) -> Option<()> {
+        match self.chunks.split_first() {
+            Some((chunk, rest)) => {
+                self.current = chunk.iter();
+                self.chunks = rest;
+                Some(())
+            },
+            None => {
+                self.current = [].iter();
+                None
+            },
+        }
+    }
+}
+
+impl<'a> Iterator for MmrNodeIter<'a> {
+    type Item = &'a Word;
+
+    fn next(&mut self) -> Option<&'a Word> {
+        loop {
+            if let Some(node) = self.current.next() {
+                return Some(node);
+            }
+            self.advance_chunk()?;
+        }
+    }
+
+    fn nth(&mut self, mut n: usize) -> Option<&'a Word> {
+        loop {
+            let len = self.current.len();
+            if n < len {
+                return self.current.nth(n);
+            }
+            n -= len;
+            self.advance_chunk()?;
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let len = self.len();
+        (len, Some(len))
+    }
+
+    fn count(self) -> usize {
+        self.len()
+    }
+}
+
+impl ExactSizeIterator for MmrNodeIter<'_> {
+    fn len(&self) -> usize {
+        self.current.len()
+            + match self.chunks.split_last() {
+                Some((last, full)) => full.len() * NODE_CHUNK_CAPACITY + last.len(),
+                None => 0,
+            }
+    }
+}
+
+impl FusedIterator for MmrNodeIter<'_> {}
+
 /// Yields inner nodes of the [Mmr].
 pub struct MmrNodes<'a> {
     /// [Mmr] being yielded, when its `forest` value is matched, the iterations is finished.
@@ -638,7 +781,7 @@ mod tests {
     use super::{super::nodes_from_mask, NODE_CHUNK_CAPACITY};
     use crate::{
         Felt, Word, ZERO,
-        merkle::mmr::{Forest, Mmr},
+        merkle::mmr::{Forest, Mmr, MmrError},
         utils::{Deserializable, DeserializationError, Serializable},
     };
 
@@ -772,6 +915,147 @@ mod tests {
         assert!(!Arc::ptr_eq(orig_chunks.last().unwrap(), clone_chunks.last().unwrap()));
     }
 
+    #[test]
+    fn test_nodes_from() {
+        // Span multiple chunks to cover chunk boundaries.
+        let mmr = Mmr::try_from_iter(leaves(2 * NODE_CHUNK_CAPACITY as u64)).unwrap();
+        let num_nodes = mmr.forest().num_nodes();
+        let all: Vec<Word> = mmr.nodes_from(0).copied().collect();
+        assert_eq!(all.len(), num_nodes);
+
+        // Starts at chunk boundaries, mid-chunk, and in the last (partial) chunk.
+        for start in [
+            0,
+            1,
+            NODE_CHUNK_CAPACITY - 1,
+            NODE_CHUNK_CAPACITY,
+            NODE_CHUNK_CAPACITY + 1,
+            num_nodes - 1,
+            num_nodes,
+            num_nodes + 1,
+        ] {
+            let suffix: Vec<Word> = mmr.nodes_from(start).copied().collect();
+            assert_eq!(suffix, all[start.min(num_nodes)..]);
+        }
+    }
+
+    #[test]
+    fn test_node_iter_skip_matches_nodes_from() {
+        // Span multiple chunks so skips cross chunk boundaries.
+        let mmr = Mmr::try_from_iter(leaves(2 * NODE_CHUNK_CAPACITY as u64)).unwrap();
+        let num_nodes = mmr.forest().num_nodes();
+        let all: Vec<Word> = mmr.nodes_from(0).copied().collect();
+
+        for start in [
+            0,
+            1,
+            NODE_CHUNK_CAPACITY - 1,
+            NODE_CHUNK_CAPACITY,
+            NODE_CHUNK_CAPACITY + 1,
+            num_nodes - 1,
+            num_nodes,
+            num_nodes + 1,
+        ] {
+            let skipped: Vec<Word> = mmr.nodes_from(0).skip(start).copied().collect();
+            assert_eq!(skipped, all[start.min(num_nodes)..]);
+        }
+
+        // `nth` positions across a chunk boundary and resumes in order.
+        let mut iter = mmr.nodes_from(0);
+        assert_eq!(iter.nth(NODE_CHUNK_CAPACITY + 1), Some(&all[NODE_CHUNK_CAPACITY + 1]));
+        assert_eq!(iter.next(), Some(&all[NODE_CHUNK_CAPACITY + 2]));
+
+        // `nth` past the end exhausts the iterator.
+        let mut iter = mmr.nodes_from(0);
+        assert_eq!(iter.nth(num_nodes), None);
+        assert_eq!(iter.next(), None);
+    }
+
+    #[test]
+    fn test_node_iter_len() {
+        let mmr = Mmr::try_from_iter(leaves(2 * NODE_CHUNK_CAPACITY as u64)).unwrap();
+        let num_nodes = mmr.forest().num_nodes();
+
+        for start in [0, 1, NODE_CHUNK_CAPACITY, num_nodes - 1, num_nodes, num_nodes + 1] {
+            let iter = mmr.nodes_from(start);
+            assert_eq!(iter.len(), num_nodes.saturating_sub(start));
+            assert_eq!(iter.size_hint(), (iter.len(), Some(iter.len())));
+        }
+
+        // The length stays exact as the iterator advances, including across chunks.
+        let mut iter = mmr.nodes_from(0);
+        iter.next();
+        assert_eq!(iter.len(), num_nodes - 1);
+        iter.nth(NODE_CHUNK_CAPACITY);
+        assert_eq!(iter.len(), num_nodes - NODE_CHUNK_CAPACITY - 2);
+        assert_eq!(iter.count(), num_nodes - NODE_CHUNK_CAPACITY - 2);
+    }
+
+    #[test]
+    fn test_nodes_from_incremental_persistence() {
+        // Simulate a flat-file consumer: persist all nodes, grow the MMR, append only the new
+        // nodes, and verify the result matches a full dump of the final state.
+        let initial_leaves = NODE_CHUNK_CAPACITY as u64 / 2;
+        let final_leaves = 2 * NODE_CHUNK_CAPACITY as u64;
+
+        let mut mmr = Mmr::try_from_iter(leaves(initial_leaves)).unwrap();
+        let mut persisted: Vec<Word> = mmr.nodes_from(0).copied().collect();
+
+        for leaf in leaves(final_leaves).skip(initial_leaves as usize) {
+            mmr.add(leaf).unwrap();
+        }
+        persisted.extend(mmr.nodes_from(persisted.len()).copied());
+
+        assert_eq!(persisted.len(), mmr.forest().num_nodes());
+        assert!(mmr.nodes == persisted.as_slice());
+    }
+
+    #[test]
+    fn test_from_nodes_unchecked_round_trip() {
+        // Sizes: empty forest, single-chunk, and multi-chunk with a partial last chunk.
+        let multi_chunk = NODE_CHUNK_CAPACITY as u64 + NODE_CHUNK_CAPACITY as u64 / 2;
+        for num_leaves in [0, 1, 8, multi_chunk] {
+            let mmr = Mmr::try_from_iter(leaves(num_leaves)).unwrap();
+            let rebuilt =
+                Mmr::from_nodes_unchecked(mmr.forest(), mmr.nodes_from(0).copied()).unwrap();
+            assert_eq!(mmr.forest, rebuilt.forest);
+            assert_eq!(mmr.nodes, rebuilt.nodes);
+            assert_eq!(mmr.peaks(), rebuilt.peaks());
+
+            // Openings from the rebuilt MMR still verify against its peaks.
+            let peaks = rebuilt.peaks();
+            for pos in [0, num_leaves.saturating_sub(1) as usize] {
+                if num_leaves > 0 {
+                    let proof = rebuilt.open(pos).unwrap();
+                    let leaf = rebuilt.get(pos).unwrap();
+                    peaks.verify(leaf, proof).unwrap();
+                }
+            }
+        }
+    }
+
+    #[test]
+    fn test_from_nodes_unchecked_rejects_count_mismatch() {
+        let mmr = Mmr::try_from_iter(leaves(8)).unwrap();
+        let nodes: Vec<Word> = mmr.nodes_from(0).copied().collect();
+
+        let too_few =
+            Mmr::from_nodes_unchecked(mmr.forest(), nodes.iter().copied().take(nodes.len() - 1));
+        assert!(matches!(
+            too_few,
+            Err(MmrError::InvalidNodeCount { expected, actual })
+                if expected == nodes.len() && actual == nodes.len() - 1
+        ));
+
+        let too_many =
+            Mmr::from_nodes_unchecked(mmr.forest(), nodes.iter().copied().chain([Word::empty()]));
+        assert!(matches!(
+            too_many,
+            Err(MmrError::InvalidNodeCount { expected, actual })
+                if expected == nodes.len() && actual == nodes.len() + 1
+        ));
+    }
+
     #[test]
     fn test_nodes_from_mask_at_max_leaves() {
         let expected = (Forest::MAX_LEAVES as u128)