Add AccountTree and PartialAccountTree

CHANGELOG.md

@@ -5,6 +5,7 @@
 - [BREAKING] Add `TransactionHeader` and include it in batches and blocks (#1247).
 - [BREAKING] Hash keys in storage maps before insertion into the SMT (#1250).
 - Added getter for proof security level in `ProvenBatch` and `ProvenBlock` (#1259).
+- Add `AccountTree` and `PartialAccountTree` wrappers (#1254).
 
 ## 0.8.1 (2025-03-26)
 

crates/miden-block-prover/src/errors.rs

@@ -1,19 +1,16 @@
-use miden_crypto::merkle::MerkleError;
-use miden_objects::{Digest, NullifierTreeError, account::AccountId};
+use miden_objects::{AccountTreeError, Digest, NullifierTreeError};
 use thiserror::Error;
 
 #[derive(Debug, Error)]
 pub enum ProvenBlockError {
     #[error("nullifier witness has a different root than the current nullifier tree root")]
     NullifierWitnessRootMismatch(#[source] NullifierTreeError),
 
-    #[error(
-        "account witness for account {account_id} has a different root than the current account tree root"
-    )]
-    AccountWitnessRootMismatch {
-        account_id: AccountId,
-        source: MerkleError,
-    },
+    #[error("failed to track account witness")]
+    AccountWitnessTracking { source: AccountTreeError },
+
+    #[error("account ID prefix already exists in the tree")]
+    AccountIdPrefixDuplicate { source: AccountTreeError },
 
     #[error(
         "account tree root of the previous block header is {prev_block_account_root} but the root of the partial tree computed from account witnesses is {stale_account_root}, indicating that the witnesses are stale"

crates/miden-block-prover/src/local_block_prover.rs

@@ -1,14 +1,13 @@
 use std::{collections::BTreeMap, vec::Vec};
 
-use miden_crypto::merkle::{LeafIndex, PartialMerkleTree};
 use miden_lib::transaction::TransactionKernel;
 use miden_objects::{
-    Digest, Word,
+    Digest,
     account::AccountId,
     block::{
         AccountUpdateWitness, BlockAccountUpdate, BlockHeader, BlockNoteIndex, BlockNoteTree,
-        BlockNumber, NullifierWitness, OutputNoteBatch, PartialNullifierTree, ProposedBlock,
-        ProvenBlock,
+        BlockNumber, NullifierWitness, OutputNoteBatch, PartialAccountTree, PartialNullifierTree,
+        ProposedBlock, ProvenBlock,
     },
     note::Nullifier,
     transaction::ChainMmr,
@@ -89,7 +88,7 @@ impl LocalBlockProver {
 
         let (
             batches,
-            mut account_updated_witnesses,
+            account_updated_witnesses,
             output_note_batches,
             created_nullifiers,
             chain_mmr,
@@ -115,7 +114,7 @@ impl LocalBlockProver {
         // --------------------------------------------------------------------------------------------
 
         let new_account_root =
-            compute_account_root(&mut account_updated_witnesses, &prev_block_header)?;
+            compute_account_root(&account_updated_witnesses, &prev_block_header)?;
 
         // Insert the previous block header into the block chain MMR to get the new chain
         // commitment.
@@ -208,7 +207,7 @@ fn compute_nullifiers(
     // its corresponding nullifier witness, so we don't have to check again whether they match.
     for witness in created_nullifiers.into_values() {
         partial_nullifier_tree
-            .add_nullifier_witness(witness)
+            .track_nullifier_witness(witness)
             .map_err(ProvenBlockError::NullifierWitnessRootMismatch)?;
     }
 
@@ -248,34 +247,21 @@ fn compute_chain_commitment(mut chain_mmr: ChainMmr, prev_block_header: BlockHea
 /// It uses a PartialMerkleTree for now while we use a SimpleSmt for the account tree. Once that is
 /// updated to an Smt, we can use a PartialSmt instead.
 fn compute_account_root(
-    updated_accounts: &mut Vec<(AccountId, AccountUpdateWitness)>,
+    updated_accounts: &[(AccountId, AccountUpdateWitness)],
     prev_block_header: &BlockHeader,
 ) -> Result<Digest, ProvenBlockError> {
     // If no accounts were updated, the account tree root is unchanged.
     if updated_accounts.is_empty() {
         return Ok(prev_block_header.account_root());
     }
 
-    let mut partial_account_tree = PartialMerkleTree::new();
-
     // First reconstruct the current account tree from the provided merkle paths.
-    for (account_id, witness) in updated_accounts.iter_mut() {
-        let account_leaf_index = LeafIndex::from(*account_id);
-        // Shouldn't the value in PartialMerkleTree::add_path be a Word instead of a Digest?
-        // PartialMerkleTree::update_leaf (below) takes a Word as a value, so this seems
-        // inconsistent.
-        partial_account_tree
-            .add_path(
-                account_leaf_index.value(),
-                witness.initial_state_commitment(),
-                // We don't need the merkle path later, so we can take it out.
-                core::mem::take(witness.initial_state_proof_mut()),
-            )
-            .map_err(|source| ProvenBlockError::AccountWitnessRootMismatch {
-                account_id: *account_id,
-                source,
-            })?;
-    }
+    // If a witness points to a leaf where multiple account IDs share the same prefix, this will
+    // return an error.
+    let mut partial_account_tree = PartialAccountTree::with_witnesses(
+        updated_accounts.iter().map(|(_, update_witness)| update_witness.to_witness()),
+    )
+    .map_err(|source| ProvenBlockError::AccountWitnessTracking { source })?;
 
     // Check the account tree root in the previous block header matches the reconstructed tree's
     // root.
@@ -288,13 +274,14 @@ fn compute_account_root(
 
     // Second, update the account tree by inserting the new final account state commitments to
     // compute the new root of the account tree.
-    // TODO: Move this loop into a method on `PartialAccountTree` once it is added.
-    for (account_id, witness) in updated_accounts {
-        let account_leaf_index = LeafIndex::from(*account_id);
-        partial_account_tree
-            .update_leaf(account_leaf_index.value(), Word::from(witness.final_state_commitment()))
-            .expect("every account leaf should have been inserted in the first loop");
-    }
+    // If an account ID's prefix already exists in the tree, this will return an error.
+    // Note that we have inserted all witnesses that we want to update into the partial account
+    // tree, so we should not run into the untracked key error.
+    partial_account_tree
+        .upsert_state_commitments(updated_accounts.iter().map(|(account_id, update_witness)| {
+            (*account_id, update_witness.final_state_commitment())
+        }))
+        .map_err(|source| ProvenBlockError::AccountIdPrefixDuplicate { source })?;
 
     Ok(partial_account_tree.root())
 }

crates/miden-block-prover/src/tests/proven_block_error.rs

@@ -1,8 +1,7 @@
 use anyhow::Context;
 use assert_matches::assert_matches;
-use miden_crypto::merkle::MerkleError;
 use miden_objects::{
-    NullifierTreeError,
+    AccountTreeError, NullifierTreeError,
     batch::ProvenBatch,
     block::{BlockInputs, ProposedBlock},
 };
@@ -186,8 +185,8 @@ fn proven_block_fails_on_account_tree_root_mismatch() -> anyhow::Result<()> {
 
     assert_matches!(
         error,
-        ProvenBlockError::AccountWitnessRootMismatch {
-            source: MerkleError::ConflictingRoots { .. },
+        ProvenBlockError::AccountWitnessTracking {
+            source: AccountTreeError::TreeRootConflict { .. },
             ..
         }
     );
@@ -240,3 +239,5 @@ fn proven_block_fails_on_nullifier_tree_root_mismatch() -> anyhow::Result<()> {
 
     Ok(())
 }
+
+// TODO: Add test to make sure duplicate account ID prefixes result in an error.

crates/miden-block-prover/src/tests/proven_block_success.rs

@@ -1,7 +1,7 @@
 use std::{collections::BTreeMap, vec::Vec};
 
 use anyhow::Context;
-use miden_crypto::merkle::{LeafIndex, SimpleSmt, Smt};
+use miden_crypto::merkle::{SimpleSmt, Smt};
 use miden_objects::{
     ACCOUNT_TREE_DEPTH, Felt, FieldElement, MIN_PROOF_SECURITY_LEVEL,
     batch::BatchNoteTree,
@@ -104,13 +104,14 @@ fn proven_block_success() -> anyhow::Result<()> {
         );
     }
 
-    // Compute expected account root on the full SimpleSmt.
+    // Compute expected account root on the full account tree.
     // --------------------------------------------------------------------------------------------
 
     let mut expected_account_tree = chain.accounts().clone();
     for (account_id, witness) in proposed_block.updated_accounts() {
         expected_account_tree
-            .insert(LeafIndex::from(*account_id), *witness.final_state_commitment());
+            .insert(*account_id, witness.final_state_commitment())
+            .context("failed to insert account id into account tree")?;
     }
 
     // Prove block.

crates/miden-lib/asm/kernels/transaction/lib/account.masm

@@ -61,8 +61,8 @@ const.ERR_FOREIGN_ACCOUNT_ID_IS_ZERO=0x00020181
 # Maximum allowed number of foreign account to be loaded (64) was exceeded.
 const.ERR_FOREIGN_ACCOUNT_MAX_NUMBER_EXCEEDED=0x00020183
 
-# State of the current foreign account is invalid.
-const.ERR_FOREIGN_ACCOUNT_INVALID=0x00020182
+# Commitment of the foreign account in the advice provider does not match the commitment in the account tree.
+const.ERR_FOREIGN_ACCOUNT_INVALID_COMMITMENT=0x00020182
 
 # Unknown version in account ID.
 const.ERR_ACCOUNT_ID_UNKNOWN_VERSION=0x00020146
@@ -1122,23 +1122,21 @@ export.validate_current_foreign_account
     # => [ACCOUNT_DB_ROOT]
 
     # get the current account ID
-    exec.memory::get_account_id swap drop
-    # => [account_id_prefix, ACCOUNT_DB_ROOT]
+    push.0.0 exec.memory::get_account_id
+    # => [account_id_prefix, account_id_suffix, 0, 0, ACCOUNT_DB_ROOT]
 
-    # push the depth of the account database tree
-    push.ACCOUNT_TREE_DEPTH
-    # => [depth, account_id_prefix, ACCOUNT_DB_ROOT]
-
-    # get the foreign account commitment
-    exec.get_current_commitment
-    # => [FOREIGN_ACCOUNT_COMMITMENT, depth, account_id_prefix, ACCOUNT_DB_ROOT]
+    # retrieve the commitment of the foreign account from the current account tree
+    # this would abort if the proof for the commitment was invalid for the account root,
+    # so this implicitly verifies its correctness
+    exec.smt::get
+    # => [FOREIGN_ACCOUNT_COMMITMENT, ACCOUNT_DB_ROOT]
 
-    # verify that the account database has the hash of the current foreign account
-    mtree_verify.err=ERR_FOREIGN_ACCOUNT_INVALID
-    # => [FOREIGN_ACCOUNT_COMMITMENT, depth, account_id_prefix, ACCOUNT_DB_ROOT]
+    # get the foreign account's commitment from memory and compare with the verified commitment
+    exec.get_current_commitment assert_eqw.err=ERR_FOREIGN_ACCOUNT_INVALID_COMMITMENT
+    # => [ACCOUNT_DB_ROOT]
 
     # clean the stack
-    dropw drop drop dropw
+    dropw
 end
 
 #! Hashes the provided map key before using it as the key in an SMT.

crates/miden-lib/src/errors/tx_kernel_errors.rs

@@ -144,8 +144,8 @@ pub const ERR_ACCOUNT_STACK_UNDERFLOW: u32 = 0x20156;
 pub const ERR_FOREIGN_ACCOUNT_CONTEXT_AGAINST_NATIVE_ACCOUNT: u32 = 0x20180;
 /// ID of the provided foreign account equals zero.
 pub const ERR_FOREIGN_ACCOUNT_ID_IS_ZERO: u32 = 0x20181;
-/// State of the current foreign account is invalid.
-pub const ERR_FOREIGN_ACCOUNT_INVALID: u32 = 0x20182;
+/// Commitment of the foreign account in the advice provider does not match the commitment in the account tree.
+pub const ERR_FOREIGN_ACCOUNT_INVALID_COMMITMENT: u32 = 0x20182;
 /// Maximum allowed number of foreign account to be loaded (64) was exceeded.
 pub const ERR_FOREIGN_ACCOUNT_MAX_NUMBER_EXCEEDED: u32 = 0x20183;
 
@@ -275,7 +275,7 @@ pub const TX_KERNEL_ERRORS: [(u32, &str); 88] = [
 
     (ERR_FOREIGN_ACCOUNT_CONTEXT_AGAINST_NATIVE_ACCOUNT, "Creation of a foreign context against the native account is forbidden"),
     (ERR_FOREIGN_ACCOUNT_ID_IS_ZERO, "ID of the provided foreign account equals zero."),
-    (ERR_FOREIGN_ACCOUNT_INVALID, "State of the current foreign account is invalid."),
+    (ERR_FOREIGN_ACCOUNT_INVALID_COMMITMENT, "Commitment of the foreign account in the advice provider does not match the commitment in the account tree."),
     (ERR_FOREIGN_ACCOUNT_MAX_NUMBER_EXCEEDED, "Maximum allowed number of foreign account to be loaded (64) was exceeded."),
 
     (ERR_FAUCET_BURN_CANNOT_EXCEED_EXISTING_TOTAL_SUPPLY, "Asset amount to burn can not exceed the existing total supply"),

crates/miden-lib/src/transaction/mod.rs

@@ -4,8 +4,8 @@ use miden_objects::{
     Digest, EMPTY_WORD, Felt, TransactionOutputError, ZERO,
     account::{AccountCode, AccountHeader, AccountId, AccountStorageHeader},
     assembly::{Assembler, DefaultSourceManager, KernelLibrary},
-    block::BlockNumber,
-    crypto::merkle::{MerkleError, MerklePath},
+    block::{AccountWitness, BlockNumber},
+    crypto::merkle::MerkleError,
     transaction::{
         OutputNote, OutputNotes, TransactionArgs, TransactionInputs, TransactionOutputs,
     },
@@ -202,7 +202,7 @@ impl TransactionKernel {
         account_header: &AccountHeader,
         account_code: &AccountCode,
         storage_header: &AccountStorageHeader,
-        merkle_path: &MerklePath,
+        account_witness: &AccountWitness,
     ) -> Result<(), MerkleError> {
         let account_id = account_header.id();
         let storage_root = account_header.storage_commitment();
@@ -221,11 +221,18 @@ impl TransactionKernel {
             (code_root, account_code.as_elements()),
         ]);
 
-        // Extend the advice inputs with Merkle store data
+        let account_leaf = account_witness.as_proof().leaf();
+        let account_leaf_hash = account_leaf.hash();
+
+        // extend the merkle store and map with account witnesses merkle path
         advice_inputs.extend_merkle_store(
-            // The prefix is the index in the account tree.
-            merkle_path.inner_nodes(account_id.prefix().as_u64(), account_header.commitment())?,
+            account_witness
+                .as_proof()
+                .path()
+                .inner_nodes(account_id.prefix().as_u64(), account_leaf_hash)?,
         );
+        // populate advice map with the account's leaf
+        advice_inputs.extend_map([(account_leaf_hash, account_leaf.to_elements())]);
 
         Ok(())
     }

crates/miden-lib/src/transaction/procedures/kernel_v0.rs

@@ -72,7 +72,7 @@ pub const KERNEL0_PROCEDURES: [Digest; 36] = [
     // tx_get_block_timestamp
     digest!("0x786863e6dbcd5026619afd3831b7dcbf824cda54950b0e0724ebf9d9370ec723"),
     // tx_start_foreign_context
-    digest!("0x6e7ab409b5661e1d03dc73bd7b6e32ca7785a983e49ab2d056e9c786620e5c20"),
+    digest!("0xf7075576799b6c2b88bca3f36ed7c3b2560ed4bc0e630f1e30db0be84e746c6e"),
     // tx_end_foreign_context
     digest!("0x90a107168d81c1c0c23890e61fb7910a64b4711afd0bf8c3098d74737e4853ba"),
     // tx_get_expiration_delta

crates/miden-objects/src/account/account_id/mod.rs

@@ -26,14 +26,14 @@ use alloc::string::{String, ToString};
 use core::fmt;
 
 pub use id_version::AccountIdVersion;
-use miden_crypto::{merkle::LeafIndex, utils::hex_to_bytes};
+use miden_crypto::utils::hex_to_bytes;
 use vm_core::{
     Felt, Word,
     utils::{ByteReader, Deserializable, Serializable},
 };
 use vm_processor::{DeserializationError, Digest};
 
-use crate::{ACCOUNT_TREE_DEPTH, AccountError, errors::AccountIdError};
+use crate::{AccountError, errors::AccountIdError};
 
 /// The identifier of an [`Account`](crate::account::Account).
 ///
@@ -398,15 +398,6 @@ impl From<AccountId> for u128 {
     }
 }
 
-/// Account IDs are used as indexes in the account database, which is a tree of depth 64.
-impl From<AccountId> for LeafIndex<ACCOUNT_TREE_DEPTH> {
-    fn from(id: AccountId) -> Self {
-        match id {
-            AccountId::V0(account_id) => account_id.into(),
-        }
-    }
-}
-
 // CONVERSIONS TO ACCOUNT ID
 // ================================================================================================
 

crates/miden-objects/src/account/account_id/v0/mod.rs

@@ -6,7 +6,7 @@ use alloc::{
 use core::fmt;
 
 use bech32::{Bech32m, primitives::decode::CheckedHrpstring};
-use miden_crypto::{merkle::LeafIndex, utils::hex_to_bytes};
+use miden_crypto::utils::hex_to_bytes;
 pub use prefix::AccountIdPrefixV0;
 use vm_core::{
     Felt, Word,
@@ -15,7 +15,7 @@ use vm_core::{
 use vm_processor::{DeserializationError, Digest};
 
 use crate::{
-    ACCOUNT_TREE_DEPTH, AccountError, Hasher,
+    AccountError, Hasher,
     account::{
         AccountIdAnchor, AccountIdVersion, AccountStorageMode, AccountType,
         account_id::{
@@ -327,13 +327,6 @@ impl From<AccountIdV0> for u128 {
     }
 }
 
-/// Account IDs are used as indexes in the account database, which is a tree of depth 64.
-impl From<AccountIdV0> for LeafIndex<ACCOUNT_TREE_DEPTH> {
-    fn from(id: AccountIdV0) -> Self {
-        LeafIndex::new_max_depth(id.prefix().as_u64())
-    }
-}
-
 // CONVERSIONS TO ACCOUNT ID
 // ================================================================================================