Merge pull request #376 from 0xPolygonMiden/bobbin-chain-mmr

Convert `ChainMmr` to use `PartialMmr`

miden-lib/src/tests/test_prologue.rs

@@ -193,18 +193,10 @@ fn chain_mmr_memory_assertions<A: AdviceProvider>(
     // The number of leaves should be stored at the CHAIN_MMR_NUM_LEAVES_PTR
     assert_eq!(
         process.get_mem_value(ContextId::root(), CHAIN_MMR_NUM_LEAVES_PTR).unwrap()[0],
-        Felt::new(inputs.block_chain().mmr().forest() as u64)
+        Felt::new(inputs.block_chain().chain_length() as u64)
     );
 
-    for (i, peak) in inputs
-        .block_chain()
-        .mmr()
-        .peaks(inputs.block_chain().mmr().forest())
-        .unwrap()
-        .peaks()
-        .iter()
-        .enumerate()
-    {
+    for (i, peak) in inputs.block_chain().peaks().peaks().iter().enumerate() {
         // The peaks should be stored at the CHAIN_MMR_PEAKS_PTR
         let i: u32 = i.try_into().expect(
             "Number of peaks is log2(number_of_leaves), this value won't be larger than 2**32",

miden-tx/src/lib.rs

@@ -3,9 +3,9 @@ use miden_objects::{
     accounts::{Account, AccountCode, AccountId},
     assembly::CodeBlock,
     notes::{NoteOrigin, NoteScript},
-    transaction::{PreparedTransaction, TransactionResult},
+    transaction::{ChainMmr, PreparedTransaction, TransactionResult},
     utils::collections::BTreeMap,
-    AccountError, BlockHeader, ChainMmr, Digest, Hasher, TransactionResultError,
+    AccountError, BlockHeader, Digest, Hasher, TransactionResultError,
 };
 use vm_core::Program;
 use vm_processor::{ExecutionError, RecAdviceProvider};

miden-tx/tests/common/mod.rs

@@ -6,7 +6,8 @@ use miden_objects::{
     assets::{Asset, FungibleAsset},
     crypto::{dsa::rpo_falcon512::KeyPair, utils::Serializable},
     notes::{Note, NoteOrigin, NoteScript, RecordedNote},
-    BlockHeader, ChainMmr, Felt, Word,
+    transaction::ChainMmr,
+    BlockHeader, Felt, Word,
 };
 use miden_tx::{DataStore, DataStoreError};
 use mock::{

mock/src/lib.rs

@@ -2,8 +2,8 @@ use miden_lib::{assembler::assembler, memory};
 use miden_objects::{
     accounts::Account,
     notes::{Note, NoteVault, RecordedNote},
-    transaction::{PreparedTransaction, TransactionScript},
-    BlockHeader, ChainMmr, Felt, StarkField,
+    transaction::{ChainMmr, PreparedTransaction, TransactionScript},
+    BlockHeader, Felt, StarkField,
 };
 use std::{fs::File, io::Read, path::PathBuf};
 use vm_processor::{

mock/src/mock/chain.rs

@@ -10,10 +10,11 @@ use core::fmt;
 use miden_objects::{
     accounts::{Account, AccountId, AccountType, SlotItem},
     assets::Asset,
-    crypto::merkle::{NodeIndex, SimpleSmt, TieredSmt},
+    crypto::merkle::{Mmr, NodeIndex, PartialMmr, SimpleSmt, TieredSmt},
     notes::{Note, NoteInclusionProof, RecordedNote, NOTE_LEAF_DEPTH, NOTE_TREE_DEPTH},
-    utils::collections::Vec,
-    BlockHeader, ChainMmr, Digest, Felt, FieldElement, StarkField, Word,
+    transaction::ChainMmr,
+    utils::collections::{BTreeMap, Vec},
+    BlockHeader, Digest, Felt, FieldElement, StarkField, Word,
 };
 use rand::{Rng, SeedableRng};
 
@@ -121,7 +122,7 @@ impl<R: Rng> Objects<R> {
 #[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
 pub struct MockChain<R> {
     /// An append-only structure used to represent the history of blocks produced for this chain.
-    chain: ChainMmr,
+    chain: Mmr,
 
     /// History of produced blocks.
     blocks: Vec<BlockHeader>,
@@ -193,7 +194,7 @@ impl<R: Rng + SeedableRng> MockChain<R> {
         let account_rng = R::from_rng(&mut rng).expect("rng seeding failed");
         let account_id_builder = AccountIdBuilder::new(account_rng);
         Self {
-            chain: ChainMmr::default(),
+            chain: Mmr::default(),
             blocks: vec![],
             nullifiers: TieredSmt::default(),
             accounts: SimpleSmt::new(64).expect("depth too big for SimpleSmt"),
@@ -459,7 +460,7 @@ impl<R: Rng + SeedableRng> MockChain<R> {
         let notes = self.pending_objects.build_notes_tree();
 
         let previous = self.blocks.last();
-        let peaks = self.chain.mmr().peaks(self.chain.mmr().forest()).unwrap();
+        let peaks = self.chain.peaks(self.chain.forest()).unwrap();
         let chain_root: Digest = peaks.hash_peaks();
         let account_root = self.accounts.root();
         let prev_hash = previous.map_or(Digest::default(), |header| header.hash());
@@ -488,7 +489,7 @@ impl<R: Rng + SeedableRng> MockChain<R> {
         );
 
         self.blocks.push(header);
-        self.chain.mmr_mut().add(header.hash());
+        self.chain.add(header.hash());
         self.objects.update_with(&mut self.pending_objects, header, &notes);
 
         header
@@ -531,13 +532,13 @@ impl<R: Rng + SeedableRng> MockChain<R> {
     // ACCESSORS
     // ----------------------------------------------------------------------------------------
 
-    /// Get a reference to the latest [ChainMmr].
-    pub fn chain(&self) -> &ChainMmr {
-        &self.chain
+    /// Get the latest [ChainMmr].
+    pub fn chain(&self) -> ChainMmr {
+        mmr_to_chain_mmr(&self.chain)
     }
 
     /// Get a reference to [BlockHeader] with `block_number`.
-    pub fn blockheader(&self, block_number: usize) -> &BlockHeader {
+    pub fn block_header(&self, block_number: usize) -> &BlockHeader {
         &self.blocks[block_number]
     }
 
@@ -623,10 +624,11 @@ pub fn mock_chain_data(consumed_notes: Vec<Note>) -> (ChainMmr, Vec<RecordedNote
     ];
 
     // instantiate and populate MMR
-    let mut chain_mmr = ChainMmr::default();
+    let mut mmr = Mmr::default();
     for block_header in block_chain.iter() {
-        chain_mmr.mmr_mut().add(block_header.hash())
+        mmr.add(block_header.hash())
     }
+    let chain_mmr = mmr_to_chain_mmr(&mmr);
 
     // set origin for consumed notes using chain and block data
     let recorded_notes = consumed_notes
@@ -651,3 +653,22 @@ pub fn mock_chain_data(consumed_notes: Vec<Note>) -> (ChainMmr, Vec<RecordedNote
 
     (chain_mmr, recorded_notes)
 }
+
+// HELPER FUNCTIONS
+// ================================================================================================
+
+/// Converts the MMR into partial MMR by copying all leaves from MMR to partial MMR.
+fn mmr_to_chain_mmr(mmr: &Mmr) -> ChainMmr {
+    let num_leaves = mmr.forest();
+    let mut partial_mmr = PartialMmr::from_peaks(mmr.peaks(mmr.forest()).unwrap());
+    let mut headers = BTreeMap::new();
+
+    for i in 0..num_leaves {
+        let node = mmr.get(i).unwrap();
+        let path = mmr.open(i, mmr.forest()).unwrap().merkle_path;
+        partial_mmr.add(i, node, &path).unwrap();
+        headers.insert(i as u32, node);
+    }
+
+    ChainMmr::new(partial_mmr, headers).unwrap()
+}

mock/src/mock/transaction.rs

@@ -12,9 +12,9 @@ use miden_lib::assembler::assembler;
 use miden_objects::{
     accounts::Account,
     notes::{Note, RecordedNote},
-    transaction::ExecutedTransaction,
+    transaction::{ChainMmr, ExecutedTransaction},
     utils::collections::Vec,
-    BlockHeader, ChainMmr, Felt, FieldElement,
+    BlockHeader, Felt, FieldElement,
 };
 use vm_processor::AdviceInputs;
 
@@ -59,12 +59,8 @@ pub fn mock_inputs(
     let (chain_mmr, recorded_notes) = mock_chain_data(consumed_notes);
 
     // Block header
-    let block_header = mock_block_header(
-        4,
-        Some(chain_mmr.mmr().peaks(chain_mmr.mmr().forest()).unwrap().hash_peaks()),
-        None,
-        &[account.clone()],
-    );
+    let block_header =
+        mock_block_header(4, Some(chain_mmr.peaks().hash_peaks()), None, &[account.clone()]);
 
     // Transaction inputs
     (account, block_header, chain_mmr, recorded_notes, auxiliary_data)
@@ -118,12 +114,8 @@ pub fn mock_inputs_with_existing(
     let (chain_mmr, recorded_notes) = mock_chain_data(consumed_notes);
 
     // Block header
-    let block_header = mock_block_header(
-        4,
-        Some(chain_mmr.mmr().peaks(chain_mmr.mmr().forest()).unwrap().hash_peaks()),
-        None,
-        &[account.clone()],
-    );
+    let block_header =
+        mock_block_header(4, Some(chain_mmr.peaks().hash_peaks()), None, &[account.clone()]);
 
     // Transaction inputs
     (account, block_header, chain_mmr, recorded_notes, auxiliary_data)
@@ -157,7 +149,7 @@ pub fn mock_executed_tx(asset_preservation: AssetPreservationStatus) -> Executed
     // Block header
     let block_header = mock_block_header(
         4,
-        Some(chain_mmr.mmr().peaks(chain_mmr.mmr().forest()).unwrap().hash_peaks()),
+        Some(chain_mmr.peaks().hash_peaks()),
         None,
         &[initial_account.clone()],
     );

objects/src/errors.rs

@@ -259,6 +259,35 @@ impl fmt::Display for NoteError {
 #[cfg(feature = "std")]
 impl std::error::Error for NoteError {}
 
+// CHAIN MMR ERROR
+// ================================================================================================
+
+#[derive(Debug, Clone)]
+pub enum ChainMmrError {
+    BlockNumTooBig {
+        chain_length: usize,
+        block_num: usize,
+    },
+}
+
+impl ChainMmrError {
+    pub fn block_num_too_big(chain_length: usize, block_num: usize) -> Self {
+        Self::BlockNumTooBig {
+            chain_length,
+            block_num,
+        }
+    }
+}
+
+impl fmt::Display for ChainMmrError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{:?}", self)
+    }
+}
+
+#[cfg(feature = "std")]
+impl std::error::Error for ChainMmrError {}
+
 // TRANSACTION SCRIPT ERROR
 // ================================================================================================
 

objects/src/lib.rs

@@ -16,15 +16,12 @@ pub mod notes;
 pub mod block;
 pub use block::BlockHeader;
 
-pub mod chain;
-pub use chain::ChainMmr;
-
 pub mod transaction;
 
 mod errors;
 pub use errors::{
-    AccountDeltaError, AccountError, AssetError, ExecutedTransactionError, NoteError,
-    PreparedTransactionError, TransactionResultError, TransactionScriptError,
+    AccountDeltaError, AccountError, AssetError, ChainMmrError, ExecutedTransactionError,
+    NoteError, PreparedTransactionError, TransactionResultError, TransactionScriptError,
     TransactionWitnessError,
 };
 

objects/src/transaction/chain_mmr.rs

@@ -0,0 +1,74 @@
+use super::{AdviceInputsBuilder, Digest, ToAdviceInputs};
+use crate::{
+    crypto::merkle::{MmrPeaks, PartialMmr},
+    utils::collections::{BTreeMap, Vec},
+    ChainMmrError,
+};
+
+// CHAIN MMR
+// ================================================================================================
+
+/// A struct that represents the chain Merkle Mountain Range (MMR).
+///
+/// The MMR allows for efficient authentication of input notes during transaction execution.
+/// Authentication is achieved by providing inclusion proofs for the notes consumed in the
+/// transaction against the chain MMR root associated with the latest block known at the time
+/// of transaction execution.
+///
+/// [ChainMmr] represents a partial view into the actual MMR and contains authentication paths
+/// for a limited set of blocks. The intent is to include only the blocks relevant for execution
+/// of a specific transaction (i.e., the blocks corresponding to all input notes).
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct ChainMmr {
+    /// Partial view of the Chain MMR with authentication paths for the blocks listed below.
+    mmr: PartialMmr,
+    /// A list of `(block_num, block_hash)` tuples for all blocks for which the partial MMR
+    /// contains authentication paths.
+    blocks: Vec<(usize, Digest)>,
+}
+
+impl ChainMmr {
+    // CONSTRUCTOR
+    // --------------------------------------------------------------------------------------------
+    /// Returns a new [ChainMmr] instantiated from the provided partial MMR and a map mapping
+    /// block_num |-> block_hash.
+    ///
+    /// # Errors
+    /// Returns an error if maximum block_num is greater than the chain length implied by the
+    /// provided partial MMR.
+    pub fn new(mmr: PartialMmr, blocks: BTreeMap<u32, Digest>) -> Result<Self, ChainMmrError> {
+        let chain_length = mmr.forest();
+        let max_block_num = blocks.keys().next_back().cloned().unwrap_or_default() as usize;
+        if max_block_num >= chain_length {
+            return Err(ChainMmrError::block_num_too_big(chain_length, max_block_num));
+        }
+
+        let blocks = blocks.into_iter().map(|(key, val)| (key as usize, val)).collect();
+        Ok(Self { mmr, blocks })
+    }
+
+    // PUBLIC ACCESSORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Returns peaks of this MMR.
+    pub fn peaks(&self) -> MmrPeaks {
+        self.mmr.peaks()
+    }
+
+    /// Returns total number of blocks contain in the chain described by this MMR.
+    pub fn chain_length(&self) -> usize {
+        self.mmr.forest()
+    }
+}
+
+impl ToAdviceInputs for &ChainMmr {
+    fn to_advice_inputs<T: AdviceInputsBuilder>(&self, target: &mut T) {
+        // Add the Mmr nodes to the merkle store
+        target.add_merkle_nodes(self.mmr.inner_nodes(self.blocks.iter()));
+
+        // Extract Mmr accumulator
+        let peaks = self.mmr.peaks();
+
+        peaks.to_advice_inputs(target);
+    }
+}