Use AccountTree type wrapper

Cargo.toml

@@ -35,17 +35,17 @@ assert_matches            = { version = "1.5" }
 http                      = { version = "1.3" }
 itertools                 = { version = "0.14" }
 miden-air                 = { version = "0.13" }
-miden-lib                 = { git = "https://github.com/0xPolygonMiden/miden-base.git", branch = "next" }
+miden-lib                 = { git = "https://github.com/0xPolygonMiden/miden-base.git", branch = "pgackst-account-tree" }
 miden-node-block-producer = { path = "crates/block-producer", version = "0.9" }
 miden-node-proto          = { path = "crates/proto", version = "0.9" }
 miden-node-proto-build    = { path = "proto", version = "0.9" }
 miden-node-rpc            = { path = "crates/rpc", version = "0.9" }
 miden-node-store          = { path = "crates/store", version = "0.9" }
 miden-node-test-macro     = { path = "crates/test-macro" }
 miden-node-utils          = { path = "crates/utils", version = "0.9" }
-miden-objects             = { git = "https://github.com/0xPolygonMiden/miden-base.git", branch = "next" }
+miden-objects             = { git = "https://github.com/0xPolygonMiden/miden-base.git", branch = "pgackst-account-tree" }
 miden-processor           = { version = "0.13" }
-miden-tx                  = { git = "https://github.com/0xPolygonMiden/miden-base.git", branch = "next" }
+miden-tx                  = { git = "https://github.com/0xPolygonMiden/miden-base.git", branch = "pgackst-account-tree" }
 prost                     = { version = "0.13" }
 rand                      = { version = "0.9" }
 thiserror                 = { version = "2.0", default-features = false }

bin/stress-test/Cargo.toml

@@ -19,7 +19,7 @@ workspace = true
 anyhow = { workspace = true }
 clap = { version = "4.5", features = ["derive", "string"] }
 miden-air = { workspace = true }
-miden-block-prover = { git = "https://github.com/0xPolygonMiden/miden-base.git", branch = "next", features = [
+miden-block-prover = { git = "https://github.com/0xPolygonMiden/miden-base.git", branch = "pgackst-account-tree", features = [
   "testing",
 ] }
 miden-lib = { workspace = true }

crates/block-producer/Cargo.toml

@@ -21,18 +21,18 @@ tracing-forest = ["miden-node-utils/tracing-forest"]
 async-trait = { version = "0.1" }
 futures = { version = "0.3" }
 itertools = { workspace = true }
-miden-block-prover = { git = "https://github.com/0xPolygonMiden/miden-base.git", branch = "next" }
+miden-block-prover = { git = "https://github.com/0xPolygonMiden/miden-base.git", branch = "pgackst-account-tree" }
 miden-lib = { workspace = true }
 miden-node-proto = { workspace = true }
 miden-node-utils = { workspace = true }
 miden-objects = { workspace = true }
 miden-processor = { workspace = true }
-miden-proving-service-client = { git = "https://github.com/0xPolygonMiden/miden-base.git", branch = "next", features = [
+miden-proving-service-client = { git = "https://github.com/0xPolygonMiden/miden-base.git", branch = "pgackst-account-tree", features = [
   "batch-prover",
   "block-prover",
 ] }
 miden-tx = { workspace = true }
-miden-tx-batch-prover = { git = "https://github.com/0xPolygonMiden/miden-base.git", branch = "next" }
+miden-tx-batch-prover = { git = "https://github.com/0xPolygonMiden/miden-base.git", branch = "pgackst-account-tree" }
 rand = { version = "0.9" }
 thiserror = { workspace = true }
 tokio = { workspace = true, features = ["macros", "net", "rt-multi-thread", "sync", "time"] }

crates/block-producer/src/test_utils/block.rs

@@ -1,11 +1,11 @@
 use miden_objects::{
-    ACCOUNT_TREE_DEPTH, Digest,
+    Digest,
     batch::ProvenBatch,
     block::{
-        BlockAccountUpdate, BlockHeader, BlockNoteIndex, BlockNoteTree, OutputNoteBatch,
-        ProvenBlock,
+        AccountTree, BlockAccountUpdate, BlockHeader, BlockNoteIndex, BlockNoteTree,
+        OutputNoteBatch, ProvenBlock,
     },
-    crypto::merkle::{Mmr, SimpleSmt},
+    crypto::merkle::Mmr,
     note::Nullifier,
     transaction::{OrderedTransactionHeaders, OutputNote},
 };
@@ -34,7 +34,9 @@ pub async fn build_expected_block_header(
     let new_account_root = {
         let mut store_accounts = store.accounts.read().await.clone();
         for (&account_id, update) in updated_accounts {
-            store_accounts.insert(account_id.into(), update.final_state_commitment().into());
+            store_accounts
+                .insert(account_id, update.final_state_commitment())
+                .expect("error: duplicate account ID prefix");
         }
 
         store_accounts.root()
@@ -71,7 +73,7 @@ pub async fn build_expected_block_header(
 
 #[derive(Debug)]
 pub struct MockBlockBuilder {
-    store_accounts: SimpleSmt<ACCOUNT_TREE_DEPTH>,
+    store_accounts: AccountTree,
     store_chain_mmr: Mmr,
     last_block_header: BlockHeader,
 
@@ -105,7 +107,8 @@ impl MockBlockBuilder {
     pub fn account_updates(mut self, updated_accounts: Vec<BlockAccountUpdate>) -> Self {
         for update in &updated_accounts {
             self.store_accounts
-                .insert(update.account_id().into(), update.final_state_commitment().into());
+                .insert(update.account_id(), update.final_state_commitment())
+                .expect("error: duplicate account ID prefix");
         }
 
         self.updated_accounts = Some(updated_accounts);

crates/block-producer/src/test_utils/store.rs

@@ -5,11 +5,11 @@ use std::{
 };
 
 use miden_objects::{
-    ACCOUNT_TREE_DEPTH, Digest, EMPTY_WORD, ZERO,
+    Digest, EMPTY_WORD, Word, ZERO,
     account::AccountId,
     batch::ProvenBatch,
-    block::{BlockHeader, BlockNumber, OutputNoteBatch, ProvenBlock},
-    crypto::merkle::{Mmr, SimpleSmt, Smt},
+    block::{AccountTree, BlockHeader, BlockNumber, OutputNoteBatch, ProvenBlock},
+    crypto::merkle::{Mmr, Smt},
     note::{NoteId, NoteInclusionProof},
     transaction::ProvenTransaction,
 };
@@ -26,7 +26,7 @@ use crate::{
 /// Builds a [`MockStoreSuccess`]
 #[derive(Debug)]
 pub struct MockStoreSuccessBuilder {
-    accounts: Option<SimpleSmt<ACCOUNT_TREE_DEPTH>>,
+    accounts: Option<AccountTree>,
     notes: Option<Vec<OutputNoteBatch>>,
     produced_nullifiers: Option<BTreeSet<Digest>>,
     chain_mmr: Option<Mmr>,
@@ -36,16 +36,13 @@ pub struct MockStoreSuccessBuilder {
 impl MockStoreSuccessBuilder {
     pub fn from_batches<'a>(batches_iter: impl Iterator<Item = &'a ProvenBatch> + Clone) -> Self {
         let accounts_smt = {
-            let accounts = batches_iter
-                .clone()
-                .flat_map(|batch| {
-                    batch
-                        .account_updates()
-                        .iter()
-                        .map(|(account_id, update)| (account_id, update.initial_state_commitment()))
-                })
-                .map(|(account_id, commitment)| (account_id.prefix().into(), commitment.into()));
-            SimpleSmt::<ACCOUNT_TREE_DEPTH>::with_leaves(accounts).unwrap()
+            let accounts = batches_iter.clone().flat_map(|batch| {
+                batch
+                    .account_updates()
+                    .iter()
+                    .map(|(account_id, update)| (*account_id, update.initial_state_commitment()))
+            });
+            AccountTree::with_entries(accounts).unwrap()
         };
 
         Self {
@@ -58,12 +55,7 @@ impl MockStoreSuccessBuilder {
     }
 
     pub fn from_accounts(accounts: impl Iterator<Item = (AccountId, Digest)>) -> Self {
-        let accounts_smt = {
-            let accounts = accounts
-                .map(|(account_id, commitment)| (account_id.prefix().into(), commitment.into()));
-
-            SimpleSmt::<ACCOUNT_TREE_DEPTH>::with_leaves(accounts).unwrap()
-        };
+        let accounts_smt = AccountTree::with_entries(accounts).unwrap();
 
         Self {
             accounts: Some(accounts_smt),
@@ -107,7 +99,7 @@ impl MockStoreSuccessBuilder {
 
     pub fn build(self) -> MockStoreSuccess {
         let block_num = self.block_num.unwrap_or(1.into());
-        let accounts_smt = self.accounts.unwrap_or(SimpleSmt::new().unwrap());
+        let accounts_smt = self.accounts.unwrap_or_default();
         let notes = self.notes.unwrap_or_default();
         let block_note_tree = note_created_smt_from_note_batches(notes.iter());
         let note_root = block_note_tree.root();
@@ -168,7 +160,7 @@ impl MockStoreSuccessBuilder {
 
 pub struct MockStoreSuccess {
     /// Map account id -> account commitment
-    pub accounts: Arc<RwLock<SimpleSmt<ACCOUNT_TREE_DEPTH>>>,
+    pub accounts: Arc<RwLock<AccountTree>>,
 
     /// Stores the nullifiers of the notes that were consumed
     pub produced_nullifiers: Arc<RwLock<Smt>>,
@@ -202,7 +194,8 @@ impl MockStoreSuccess {
         // update accounts
         for update in block.updated_accounts() {
             locked_accounts
-                .insert(update.account_id().into(), update.final_state_commitment().into());
+                .insert(update.account_id(), update.final_state_commitment())
+                .expect("TODO: what should we do with this error?");
         }
         let header = block.header();
         debug_assert_eq!(locked_accounts.root(), header.account_root());
@@ -254,12 +247,12 @@ impl MockStoreSuccess {
         let locked_produced_nullifiers = self.produced_nullifiers.read().await;
 
         let account_commitment = {
-            let account_commitment = locked_accounts.get_leaf(&proven_tx.account_id().into());
+            let account_commitment = locked_accounts.get(proven_tx.account_id());
 
-            if account_commitment == EMPTY_WORD {
+            if Word::from(account_commitment) == EMPTY_WORD {
                 None
             } else {
-                Some(account_commitment.into())
+                Some(account_commitment)
             }
         };
 

crates/proto/src/domain/account.rs

@@ -5,7 +5,7 @@ use miden_objects::{
     Digest,
     account::{Account, AccountHeader, AccountId},
     block::BlockNumber,
-    crypto::{hash::rpo::RpoDigest, merkle::MerklePath},
+    crypto::{hash::rpo::RpoDigest, merkle::SmtProof},
     utils::{Deserializable, Serializable},
 };
 
@@ -155,16 +155,14 @@ impl TryInto<StorageMapKeysProof> for proto::requests::get_account_proofs_reques
 #[derive(Clone, Debug)]
 pub struct AccountWitnessRecord {
     pub account_id: AccountId,
-    pub initial_state_commitment: Digest,
-    pub proof: MerklePath,
+    pub proof: SmtProof,
 }
 
 impl From<AccountWitnessRecord> for proto::responses::AccountWitness {
     fn from(from: AccountWitnessRecord) -> Self {
         Self {
             account_id: Some(from.account_id.into()),
-            initial_state_commitment: Some(from.initial_state_commitment.into()),
-            proof: Some(Into::into(&from.proof)),
+            proof: Some(from.proof.into()),
         }
     }
 }
@@ -180,15 +178,8 @@ impl TryFrom<proto::responses::AccountWitness> for AccountWitnessRecord {
                 .account_id
                 .ok_or(proto::responses::AccountWitness::missing_field(stringify!(account_id)))?
                 .try_into()?,
-            initial_state_commitment: account_witness_record
-                .initial_state_commitment
-                .ok_or(proto::responses::AccountWitness::missing_field(stringify!(
-                    account_commitment
-                )))?
-                .try_into()?,
             proof: account_witness_record
                 .proof
-                .as_ref()
                 .ok_or(proto::responses::AccountWitness::missing_field(stringify!(proof)))?
                 .try_into()?,
         })

crates/proto/src/domain/block.rs

@@ -112,13 +112,8 @@ impl From<BlockInputs> for GetBlockInputsResponse {
             account_witnesses: account_witnesses
                 .into_iter()
                 .map(|(id, witness)| {
-                    let (initial_state_commitment, proof) = witness.into_parts();
-                    AccountWitnessRecord {
-                        account_id: id,
-                        initial_state_commitment,
-                        proof,
-                    }
-                    .into()
+                    let proof = witness.into_proof();
+                    AccountWitnessRecord { account_id: id, proof }.into()
                 })
                 .collect(),
             nullifier_witnesses: nullifier_witnesses
@@ -153,10 +148,8 @@ impl TryFrom<GetBlockInputsResponse> for BlockInputs {
                 let witness_record: AccountWitnessRecord = entry.try_into()?;
                 Ok((
                     witness_record.account_id,
-                    AccountWitness::new(
-                        witness_record.initial_state_commitment,
-                        witness_record.proof,
-                    ),
+                    AccountWitness::new(witness_record.proof)
+                        .expect("account witness record should contain a valid account witness"),
                 ))
             })
             .collect::<Result<BTreeMap<_, _>, ConversionError>>()?;

crates/proto/src/generated/responses.rs

@@ -81,19 +81,15 @@ pub struct SyncNoteResponse {
     #[prost(message, repeated, tag = "4")]
     pub notes: ::prost::alloc::vec::Vec<super::note::NoteSyncRecord>,
 }
-/// An account returned as a response to the `GetBlockInputs`.
+/// An account witness returned as a response to the `GetBlockInputs`.
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct AccountWitness {
     /// The account ID.
     #[prost(message, optional, tag = "1")]
     pub account_id: ::core::option::Option<super::account::AccountId>,
-    /// The latest account state commitment used as the initial state of the requested block.
-    /// This will be the zero digest if the account doesn't exist.
+    /// The SMT proof to verify the account's state commitment in the account tree.
     #[prost(message, optional, tag = "2")]
-    pub initial_state_commitment: ::core::option::Option<super::digest::Digest>,
-    /// Merkle path to verify the account's inclusion in the account tree.
-    #[prost(message, optional, tag = "3")]
-    pub proof: ::core::option::Option<super::merkle::MerklePath>,
+    pub proof: ::core::option::Option<super::smt::SmtOpening>,
 }
 /// A nullifier returned as a response to the `GetBlockInputs`.
 #[derive(Clone, PartialEq, ::prost::Message)]
@@ -220,7 +216,7 @@ pub struct GetAccountStateDeltaResponse {
 /// Represents the result of getting account proofs.
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct GetAccountProofsResponse {
-    /// Block number at which the state of the account was returned.
+    /// Block number at which the state of the accounts is returned.
     #[prost(fixed32, tag = "1")]
     pub block_num: u32,
     /// List of account state infos for the requested account keys.
@@ -233,14 +229,11 @@ pub struct AccountProofsResponse {
     /// Account ID.
     #[prost(message, optional, tag = "1")]
     pub account_id: ::core::option::Option<super::account::AccountId>,
-    /// Account commitment.
+    /// Proof of inclusion for the account's current state commitment in the account tree.
     #[prost(message, optional, tag = "2")]
-    pub account_commitment: ::core::option::Option<super::digest::Digest>,
-    /// Authentication path from the `account_root` of the block header to the account.
-    #[prost(message, optional, tag = "3")]
-    pub account_proof: ::core::option::Option<super::merkle::MerklePath>,
+    pub proof: ::core::option::Option<super::smt::SmtOpening>,
     /// State header for public accounts. Filled only if `include_headers` flag is set to `true`.
-    #[prost(message, optional, tag = "4")]
+    #[prost(message, optional, tag = "3")]
     pub state_header: ::core::option::Option<AccountStateHeader>,
 }
 /// State header for public accounts.

crates/store/src/errors.rs

@@ -2,7 +2,7 @@ use std::io;
 
 use deadpool::managed::PoolError;
 use miden_objects::{
-    AccountDeltaError, AccountError, NoteError,
+    AccountDeltaError, AccountError, AccountTreeError, NoteError,
     account::AccountId,
     block::BlockNumber,
     crypto::{
@@ -108,7 +108,7 @@ pub enum StateInitializationError {
     #[error("failed to create nullifier tree")]
     FailedToCreateNullifierTree(#[from] NullifierTreeError),
     #[error("failed to create accounts tree")]
-    FailedToCreateAccountsTree(#[from] MerkleError),
+    FailedToCreateAccountsTree(#[source] AccountTreeError),
 }
 
 #[derive(Debug, Error)]