Fix build after SMT breaking changes

block-producer/src/batch_builder/batch.rs

@@ -17,7 +17,7 @@ use crate::{SharedProvenTx, CREATED_NOTES_SMT_DEPTH, MAX_NUM_CREATED_NOTES_PER_B
 pub struct TransactionBatch {
     updated_accounts: BTreeMap<AccountId, AccountStates>,
     produced_nullifiers: Vec<Digest>,
-    created_notes_smt: SimpleSmt,
+    created_notes_smt: SimpleSmt<CREATED_NOTES_SMT_DEPTH>,
     /// The notes stored `created_notes_smt`
     created_notes: Vec<NoteEnvelope>,
 }
@@ -64,8 +64,7 @@ impl TransactionBatch {
             // TODO: document under what circumstances SMT creating can fail
             (
                 created_notes.clone(),
-                SimpleSmt::with_contiguous_leaves(
-                    CREATED_NOTES_SMT_DEPTH,
+                SimpleSmt::<CREATED_NOTES_SMT_DEPTH>::with_contiguous_leaves(
                     created_notes.into_iter().flat_map(|note_envelope| {
                         [note_envelope.note_id().into(), note_envelope.metadata().into()]
                     }),

block-producer/src/block_builder/prover/tests.rs

@@ -464,17 +464,14 @@ async fn test_compute_note_root_success() {
 
     // The first 2 txs were put in the first batch; the 3rd was put in the second. It will lie in
     // the second subtree of depth 12
-    let notes_smt = SimpleSmt::with_leaves(
-        CREATED_NOTES_TREE_DEPTH,
-        vec![
-            (0u64, notes_created[0].note_id().into()),
-            (1u64, notes_created[0].metadata().into()),
-            (2u64, notes_created[1].note_id().into()),
-            (3u64, notes_created[1].metadata().into()),
-            (2u64.pow(13), notes_created[2].note_id().into()),
-            (2u64.pow(13) + 1, notes_created[2].metadata().into()),
-        ],
-    )
+    let notes_smt = SimpleSmt::<CREATED_NOTES_TREE_DEPTH>::with_leaves(vec![
+        (0u64, notes_created[0].note_id().into()),
+        (1u64, notes_created[0].metadata().into()),
+        (2u64, notes_created[1].note_id().into()),
+        (3u64, notes_created[1].metadata().into()),
+        (2u64.pow(13), notes_created[2].note_id().into()),
+        (2u64.pow(13) + 1, notes_created[2].metadata().into()),
+    ])
     .unwrap();
 
     // Compare roots

block-producer/src/test_utils/block.rs

@@ -2,7 +2,8 @@ use std::{collections::BTreeMap, sync::Arc};
 
 use miden_node_proto::domain::BlockInputs;
 use miden_objects::{
-    accounts::AccountId, crypto::merkle::Mmr, notes::NoteEnvelope, BlockHeader, Digest, ONE, ZERO,
+    accounts::AccountId, crypto::merkle::Mmr, notes::NoteEnvelope, BlockHeader, Digest,
+    ACCOUNT_TREE_DEPTH, ONE, ZERO,
 };
 use miden_vm::crypto::SimpleSmt;
 
@@ -27,10 +28,8 @@ pub async fn build_expected_block_header(
         batches.iter().flat_map(|batch| batch.updated_accounts()).collect();
     let new_account_root = {
         let mut store_accounts = store.accounts.read().await.clone();
-        for (account_id, new_account_state) in updated_accounts.iter() {
-            store_accounts
-                .update_leaf(u64::from(*account_id), new_account_state.into())
-                .unwrap();
+        for &(account_id, new_account_state) in updated_accounts.iter() {
+            store_accounts.insert(account_id.into(), new_account_state.into());
         }
 
         store_accounts.root()
@@ -95,7 +94,7 @@ pub async fn build_actual_block_header(
 
 #[derive(Debug)]
 pub struct MockBlockBuilder {
-    store_accounts: SimpleSmt,
+    store_accounts: SimpleSmt<ACCOUNT_TREE_DEPTH>,
     store_chain_mmr: Mmr,
     last_block_header: BlockHeader,
 
@@ -121,10 +120,8 @@ impl MockBlockBuilder {
         mut self,
         updated_accounts: Vec<(AccountId, Digest)>,
     ) -> Self {
-        for (account_id, new_account_state) in updated_accounts.iter() {
-            self.store_accounts
-                .update_leaf(u64::from(*account_id), new_account_state.into())
-                .unwrap();
+        for &(account_id, new_account_state) in updated_accounts.iter() {
+            self.store_accounts.insert(account_id.into(), new_account_state.into());
         }
 
         self.updated_accounts = Some(updated_accounts);

block-producer/src/test_utils/store.rs

@@ -1,6 +1,7 @@
 use async_trait::async_trait;
+use miden_crypto::merkle::ValuePath;
 use miden_node_proto::domain::{AccountInputRecord, BlockInputs};
-use miden_objects::{crypto::merkle::Mmr, BlockHeader, EMPTY_WORD, ONE, ZERO};
+use miden_objects::{crypto::merkle::Mmr, BlockHeader, ACCOUNT_TREE_DEPTH, EMPTY_WORD, ONE, ZERO};
 use miden_vm::crypto::SimpleSmt;
 
 use super::*;
@@ -10,12 +11,10 @@ use crate::{
     SharedProvenTx,
 };
 
-const ACCOUNT_SMT_DEPTH: u8 = 64;
-
 /// Builds a [`MockStoreSuccess`]
 #[derive(Debug, Default)]
 pub struct MockStoreSuccessBuilder {
-    accounts: Option<SimpleSmt>,
+    accounts: Option<SimpleSmt<ACCOUNT_TREE_DEPTH>>,
     consumed_nullifiers: Option<BTreeSet<Digest>>,
     chain_mmr: Option<Mmr>,
 }
@@ -35,7 +34,7 @@ impl MockStoreSuccessBuilder {
             let accounts =
                 accounts.into_iter().map(|(account_id, hash)| (account_id.into(), hash.into()));
 
-            SimpleSmt::with_leaves(ACCOUNT_SMT_DEPTH, accounts).unwrap()
+            SimpleSmt::<ACCOUNT_TREE_DEPTH>::with_leaves(accounts).unwrap()
         };
 
         self.accounts = Some(accounts_smt);
@@ -62,7 +61,7 @@ impl MockStoreSuccessBuilder {
     }
 
     pub fn build(self) -> MockStoreSuccess {
-        let accounts_smt = self.accounts.unwrap_or(SimpleSmt::new(ACCOUNT_SMT_DEPTH).unwrap());
+        let accounts_smt = self.accounts.unwrap_or(SimpleSmt::<ACCOUNT_TREE_DEPTH>::new().unwrap());
         let chain_mmr = self.chain_mmr.unwrap_or_default();
 
         let initial_block_header = BlockHeader::new(
@@ -93,7 +92,7 @@ impl MockStoreSuccessBuilder {
 
 pub struct MockStoreSuccess {
     /// Map account id -> account hash
-    pub accounts: Arc<RwLock<SimpleSmt>>,
+    pub accounts: Arc<RwLock<SimpleSmt<ACCOUNT_TREE_DEPTH>>>,
 
     /// Stores the nullifiers of the notes that were consumed
     pub consumed_nullifiers: Arc<RwLock<BTreeSet<Digest>>>,
@@ -128,7 +127,7 @@ impl ApplyBlock for MockStoreSuccess {
 
         // update accounts
         for &(account_id, account_hash) in block.updated_accounts.iter() {
-            locked_accounts.update_leaf(account_id.into(), account_hash.into()).unwrap();
+            locked_accounts.insert(account_id.into(), account_hash.into());
         }
         debug_assert_eq!(locked_accounts.root(), block.header.account_root());
 
@@ -164,7 +163,7 @@ impl Store for MockStoreSuccess {
         let locked_consumed_nullifiers = self.consumed_nullifiers.read().await;
 
         let account_hash = {
-            let account_hash = locked_accounts.get_leaf(proven_tx.account_id().into()).unwrap();
+            let account_hash = locked_accounts.get_leaf(&proven_tx.account_id().into());
 
             if account_hash == EMPTY_WORD {
                 None
@@ -202,12 +201,14 @@ impl Store for MockStoreSuccess {
 
             updated_accounts
                 .map(|&account_id| {
-                    let account_hash = locked_accounts.get_leaf(account_id.into()).unwrap();
-                    let proof = locked_accounts.get_leaf_path(account_id.into()).unwrap();
+                    let ValuePath {
+                        value: account_hash,
+                        path: proof,
+                    } = locked_accounts.open(&account_id.into());
 
                     AccountInputRecord {
                         account_id,
-                        account_hash: account_hash.into(),
+                        account_hash,
                         proof,
                     }
                 })

store/src/db/tests.rs

@@ -1,4 +1,4 @@
-use miden_crypto::{hash::rpo::RpoDigest, merkle::NodeIndex, StarkField};
+use miden_crypto::{hash::rpo::RpoDigest, merkle::LeafIndex, StarkField};
 use miden_node_proto::{
     block_header::BlockHeader as ProtobufBlockHeader,
     digest::Digest as ProtobufDigest,
@@ -273,9 +273,9 @@ fn test_notes() {
     let tag = 5;
     let note_hash = num_to_rpo_digest(3);
     let values = [(note_index as u64, *note_hash)];
-    let notes_db = SimpleSmt::with_leaves(NOTE_LEAF_DEPTH, values.iter().cloned()).unwrap();
-    let idx = NodeIndex::new(NOTE_LEAF_DEPTH, note_index as u64).unwrap();
-    let merkle_path = notes_db.get_path(idx).unwrap();
+    let notes_db = SimpleSmt::<NOTE_LEAF_DEPTH>::with_leaves(values.iter().cloned()).unwrap();
+    let idx = LeafIndex::<NOTE_LEAF_DEPTH>::new(note_index as u64).unwrap();
+    let merkle_path = notes_db.open(&idx).path;
 
     let merkle_path: Vec<ProtobufDigest> =
         merkle_path.nodes().iter().map(|n| (*n).into()).collect();

store/src/genesis.rs

@@ -3,11 +3,9 @@ use miden_objects::{
     accounts::Account,
     notes::NOTE_LEAF_DEPTH,
     utils::serde::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable},
-    BlockHeader, Digest,
+    BlockHeader, Digest, ACCOUNT_TREE_DEPTH,
 };
 
-use crate::state::ACCOUNT_DB_DEPTH;
-
 pub const GENESIS_BLOCK_NUM: u32 = 0;
 
 /// Represents the state at genesis, which will be used to derive the genesis block.
@@ -31,9 +29,10 @@ impl GenesisState {
     }
 
     /// Returns the block header and the account SMT
-    pub fn into_block_parts(self) -> Result<(BlockHeader, SimpleSmt), MerkleError> {
-        let account_smt = SimpleSmt::with_leaves(
-            ACCOUNT_DB_DEPTH,
+    pub fn into_block_parts(
+        self
+    ) -> Result<(BlockHeader, SimpleSmt<ACCOUNT_TREE_DEPTH>), MerkleError> {
+        let account_smt: SimpleSmt<ACCOUNT_TREE_DEPTH> = SimpleSmt::with_leaves(
             self.accounts
                 .into_iter()
                 .map(|account| (account.id().into(), account.hash().into())),

store/src/state.rs

@@ -8,7 +8,8 @@ use anyhow::{anyhow, bail, Result};
 use miden_crypto::{
     hash::rpo::RpoDigest,
     merkle::{
-        MerkleError, MerklePath, Mmr, MmrDelta, MmrPeaks, SimpleSmt, TieredSmt, TieredSmtProof,
+        LeafIndex, MerkleError, MerklePath, Mmr, MmrDelta, MmrPeaks, SimpleSmt, TieredSmt,
+        TieredSmtProof, ValuePath,
     },
     Felt, FieldElement, Word, EMPTY_WORD,
 };
@@ -26,7 +27,7 @@ use miden_node_proto::{
 };
 use miden_objects::{
     notes::{NoteMetadata, NOTE_LEAF_DEPTH},
-    BlockHeader,
+    BlockHeader, ACCOUNT_TREE_DEPTH,
 };
 use tokio::{
     sync::{oneshot, Mutex, RwLock},
@@ -41,19 +42,14 @@ use crate::{
     COMPONENT,
 };
 
-// CONSTANTS
-// ================================================================================================
-
-pub(crate) const ACCOUNT_DB_DEPTH: u8 = 64;
-
 // STRUCTURES
 // ================================================================================================
 
 /// Container for state that needs to be updated atomically.
 struct InnerState {
     nullifier_tree: TieredSmt,
     chain_mmr: Mmr,
-    account_tree: SimpleSmt,
+    account_tree: SimpleSmt<ACCOUNT_TREE_DEPTH>,
 }
 
 /// The rollup state
@@ -252,7 +248,8 @@ impl State {
             // update account tree
             let mut account_tree = inner.account_tree.clone();
             for (account_id, account_hash) in accounts.iter() {
-                account_tree.update_leaf(*account_id, account_hash.try_into()?)?;
+                account_tree
+                    .insert(LeafIndex::new_max_depth(*account_id), account_hash.try_into()?);
             }
 
             if account_tree.root() != new_block.account_root() {
@@ -272,11 +269,13 @@ impl State {
                 .map(|note| {
                     // Safety: This should never happen, the note_tree is created directly form
                     // this list of notes
-                    let merkle_path =
-                        note_tree.get_leaf_path(note.note_index as u64).map_err(|_| {
+                    let leaf_index: LeafIndex<NOTE_LEAF_DEPTH> =
+                        LeafIndex::new(note.note_index as u64).map_err(|_| {
                             anyhow!("Internal server error, unable to created proof for note")
                         })?;
 
+                    let merkle_path = note_tree.open(&leaf_index).path;
+
                     Ok(Note {
                         block_num: new_block.block_num(),
                         note_hash: note.note_hash.clone(),
@@ -397,11 +396,13 @@ impl State {
             .iter()
             .cloned()
             .map(|account_id| {
-                let account_hash = inner.account_tree.get_leaf(account_id)?;
-                let merkle_path = inner.account_tree.get_leaf_path(account_id)?;
+                let ValuePath {
+                    value: account_hash,
+                    path: merkle_path,
+                } = inner.account_tree.open(&LeafIndex::new_max_depth(account_id));
                 Ok(AccountStateWithProof {
                     account_id,
-                    account_hash,
+                    account_hash: account_hash.into(),
                     merkle_path,
                 })
             })
@@ -420,7 +421,11 @@ impl State {
 
         let account = AccountState {
             account_id,
-            account_hash: inner.account_tree.get_leaf(account_id)?,
+            account_hash: inner
+                .account_tree
+                .open(&LeafIndex::new_max_depth(account_id))
+                .value
+                .into(),
         };
 
         let nullifier_blocks = nullifiers
@@ -465,7 +470,9 @@ fn block_to_nullifier_data(block: BlockNumber) -> Word {
 }
 
 /// Creates a [SimpleSmt] tree from the `notes`.
-pub fn build_notes_tree(notes: &[NoteCreated]) -> Result<SimpleSmt, anyhow::Error> {
+pub fn build_notes_tree(
+    notes: &[NoteCreated]
+) -> Result<SimpleSmt<NOTE_LEAF_DEPTH>, anyhow::Error> {
     // TODO: create SimpleSmt without this allocation
     let mut entries: Vec<(u64, Word)> = Vec::with_capacity(notes.len() * 2);
 
@@ -478,7 +485,7 @@ pub fn build_notes_tree(notes: &[NoteCreated]) -> Result<SimpleSmt, anyhow::Erro
         entries.push((index + 1, note_metadata.into()));
     }
 
-    Ok(SimpleSmt::with_leaves(NOTE_LEAF_DEPTH, entries)?)
+    Ok(SimpleSmt::with_leaves(entries)?)
 }
 
 #[instrument(skip(db))]
@@ -516,15 +523,15 @@ async fn load_mmr(db: &mut Db) -> Result<Mmr> {
 }
 
 #[instrument(skip(db))]
-async fn load_accounts(db: &mut Db) -> Result<SimpleSmt> {
+async fn load_accounts(db: &mut Db) -> Result<SimpleSmt<ACCOUNT_TREE_DEPTH>> {
     let account_data: Result<Vec<(u64, Word)>> = db
         .select_account_hashes()
         .await?
         .into_iter()
         .map(|(id, account_hash)| Ok((id, account_hash.try_into()?)))
         .collect();
 
-    let smt = SimpleSmt::with_leaves(ACCOUNT_DB_DEPTH, account_data?)?;
+    let smt = SimpleSmt::with_leaves(account_data?)?;
 
     Ok(smt)
 }