Add AccountTree and PartialAccountTree

CHANGELOG.md

@@ -7,6 +7,7 @@
 - Added getter for proof security level in `ProvenBatch` and `ProvenBlock` (#1259).
 - [BREAKING] Replaced the `ProvenBatch::new_unchecked` with the `ProvenBatch::new` method to initialize the struct with validations (#1260).
 - Added a retry strategy for worker's health check (#1255).
+- Add `AccountTree` and `PartialAccountTree` wrappers and enforce ID prefix uniqueness (#1254).
 
 ## 0.8.1 (2025-03-26) - `miden-objects` and `miden-tx` crates only.
 

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,17 +1,27 @@
 use anyhow::Context;
 use assert_matches::assert_matches;
-use miden_crypto::merkle::MerkleError;
+use miden_crypto::{EMPTY_WORD, Felt, FieldElement};
+use miden_lib::transaction::TransactionKernel;
 use miden_objects::{
-    NullifierTreeError,
+    AccountTreeError, Digest, NullifierTreeError,
+    account::{
+        Account, AccountBuilder, AccountId, AccountIdAnchor, StorageSlot,
+        delta::AccountUpdateDetails,
+    },
     batch::ProvenBatch,
-    block::{BlockInputs, ProposedBlock},
+    block::{BlockInputs, BlockNumber, ProposedBlock},
+    testing::account_component::AccountMockComponent,
+    transaction::{ProvenTransaction, ProvenTransactionBuilder},
+    vm::ExecutionProof,
 };
+use miden_tx::testing::{MockChain, TransactionContextBuilder};
+use winterfell::Proof;
 
 use crate::{
     LocalBlockProver, ProvenBlockError,
     tests::utils::{
-        TestSetup, generate_batch, generate_executed_tx_with_authenticated_notes,
-        generate_tracked_note, setup_chain,
+        ProvenTransactionExt, TestSetup, generate_batch,
+        generate_executed_tx_with_authenticated_notes, generate_tracked_note, setup_chain,
     },
 };
 
@@ -186,8 +196,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 +250,198 @@ fn proven_block_fails_on_nullifier_tree_root_mismatch() -> anyhow::Result<()> {
 
     Ok(())
 }
+
+/// Tests that creating an account when an existing account with the same account ID prefix exists,
+/// results in an error.
+#[test]
+fn proven_block_fails_on_creating_account_with_existing_account_id_prefix() -> anyhow::Result<()> {
+    // Construct a new account.
+    // --------------------------------------------------------------------------------------------
+
+    let mut mock_chain = MockChain::new();
+    let anchor_block = mock_chain.block_header(0);
+    let (account, seed) = AccountBuilder::new([5; 32])
+        .anchor(AccountIdAnchor::try_from(&anchor_block)?)
+        .with_component(
+            AccountMockComponent::new_with_slots(
+                TransactionKernel::testing_assembler(),
+                vec![StorageSlot::Value([5u32.into(); 4])],
+            )
+            .unwrap(),
+        )
+        .build()
+        .context("failed to build account")?;
+
+    let new_id = account.id();
+
+    // Construct a second account whose ID matches the prefix of the first and insert it into the
+    // chain, as if that account already existed. That way we can check if the block prover errors
+    // when we attempt to create the first account.
+    // --------------------------------------------------------------------------------------------
+
+    // Set some bits on the hash part of the suffix to make the account id distinct from the
+    // original one, but their prefix is still the same.
+    let existing_id = AccountId::try_from(u128::from(new_id) | 0xffff00)
+        .context("failed to convert account ID")?;
+
+    assert_eq!(
+        new_id.prefix(),
+        existing_id.prefix(),
+        "test requires that prefixes are the same"
+    );
+    assert_ne!(
+        new_id.suffix(),
+        existing_id.suffix(),
+        "test should work if suffixes are different, so we want to ensure it"
+    );
+    assert_eq!(account.init_commitment(), miden_objects::Digest::from(EMPTY_WORD));
+
+    let existing_account =
+        Account::mock(existing_id.into(), Felt::ZERO, TransactionKernel::testing_assembler());
+    mock_chain.add_pending_account(existing_account.clone());
+    mock_chain.seal_next_block();
+
+    // Execute the account-creating transaction.
+    // --------------------------------------------------------------------------------------------
+
+    let tx_inputs = mock_chain.get_transaction_inputs(account.clone(), Some(seed), &[], &[]);
+    let tx_context = TransactionContextBuilder::new(account)
+        .account_seed(Some(seed))
+        .tx_inputs(tx_inputs)
+        .build();
+    let tx = tx_context.execute().context("failed to execute account creating tx")?;
+    let tx =
+        ProvenTransaction::from_executed_transaction_mocked(tx, &mock_chain.latest_block_header());
+
+    let batch = generate_batch(&mut mock_chain, vec![tx]);
+    let batches = [batch];
+
+    let block_inputs = mock_chain.get_block_inputs(batches.iter());
+    // Sanity check: The mock chain account tree root should match the previous block header's
+    // account tree root.
+    assert_eq!(mock_chain.accounts().root(), block_inputs.prev_block_header().account_root());
+    assert_eq!(mock_chain.accounts().num_accounts(), 1);
+
+    // Sanity check: The block inputs should contain an account witness whose ID matches the
+    // existing ID.
+    assert_eq!(block_inputs.account_witnesses().len(), 1);
+    let witness = block_inputs
+        .account_witnesses()
+        .get(&new_id)
+        .context("block inputs did not contain witness for id")?;
+
+    // The witness should be for the **existing** account, because that's the one that exists in
+    // the tree and is therefore in the same SMT leaf that we would insert the new ID into.
+    assert_eq!(witness.id(), existing_id);
+    assert_eq!(witness.state_commitment(), existing_account.commitment());
+
+    let block = mock_chain.propose_block(batches).context("failed to propose block")?;
+
+    let err = LocalBlockProver::new(0).prove_without_batch_verification(block).unwrap_err();
+
+    // This should fail when we try to _insert_ the same two prefixes into the partial tree.
+    assert_matches!(
+        err,
+        ProvenBlockError::AccountIdPrefixDuplicate {
+            source: AccountTreeError::DuplicateIdPrefix { duplicate_prefix }
+        } if duplicate_prefix == new_id.prefix()
+    );
+
+    Ok(())
+}
+
+/// Tests that creating two accounts in the same block whose ID prefixes match, results in an error.
+#[test]
+fn proven_block_fails_on_creating_account_with_duplicate_account_id_prefix() -> anyhow::Result<()> {
+    // Construct a new account.
+    // --------------------------------------------------------------------------------------------
+
+    let mut mock_chain = MockChain::new();
+    let anchor_block = mock_chain.block_header(0);
+    let (account, _) = AccountBuilder::new([5; 32])
+        .anchor(AccountIdAnchor::try_from(&anchor_block)?)
+        .with_component(
+            AccountMockComponent::new_with_slots(
+                TransactionKernel::testing_assembler(),
+                vec![StorageSlot::Value([5u32.into(); 4])],
+            )
+            .unwrap(),
+        )
+        .build()
+        .context("failed to build account")?;
+
+    let id0 = account.id();
+
+    // Construct a second account whose ID matches the prefix of the first.
+    // --------------------------------------------------------------------------------------------
+
+    // Set some bits on the hash part of the suffix to make the account id distinct from the
+    // original one, but their prefix is still the same.
+    let id1 =
+        AccountId::try_from(u128::from(id0) | 0xffff00).context("failed to convert account ID")?;
+
+    assert_eq!(id0.prefix(), id1.prefix(), "test requires that prefixes are the same");
+    assert_ne!(
+        id0.suffix(),
+        id1.suffix(),
+        "test should work if suffixes are different, so we want to ensure it"
+    );
+
+    // Build two mocked proven transactions, each of which creates a new account and both share the
+    // same ID prefix but not the suffix.
+    // --------------------------------------------------------------------------------------------
+
+    let [tx0, tx1] =
+        [(id0, [0, 0, 0, 1u32]), (id1, [0, 0, 0, 2u32])].map(|(id, final_state_comm)| {
+            ProvenTransactionBuilder::new(
+                id,
+                Digest::default(),
+                Digest::from(final_state_comm),
+                anchor_block.block_num(),
+                anchor_block.commitment(),
+                BlockNumber::from(u32::MAX),
+                ExecutionProof::new(Proof::new_dummy(), Default::default()),
+            )
+            .account_update_details(AccountUpdateDetails::Private)
+            .build()
+            .unwrap()
+        });
+
+    // Build a batch from these transactions and attempt to prove a block.
+    // --------------------------------------------------------------------------------------------
+
+    let batch = generate_batch(&mut mock_chain, vec![tx0, tx1]);
+    let batches = [batch];
+
+    // Sanity check: The block inputs should contain two account witnesses that point to the same
+    // empty entry.
+    let block_inputs = mock_chain.get_block_inputs(batches.iter());
+    assert_eq!(block_inputs.account_witnesses().len(), 2);
+    let witness0 = block_inputs
+        .account_witnesses()
+        .get(&id0)
+        .context("block inputs did not contain witness for id0")?;
+    let witness1 = block_inputs
+        .account_witnesses()
+        .get(&id1)
+        .context("block inputs did not contain witness for id1")?;
+    assert_eq!(witness0.id(), id0);
+    assert_eq!(witness1.id(), id1);
+
+    assert_eq!(witness0.state_commitment(), Digest::default());
+    assert_eq!(witness1.state_commitment(), Digest::default());
+
+    let block = mock_chain.propose_block(batches).context("failed to propose block")?;
+
+    let err = LocalBlockProver::new(0).prove_without_batch_verification(block).unwrap_err();
+
+    // This should fail when we try to _track_ the same two prefixes in the partial tree.
+    assert_matches!(
+        err,
+        ProvenBlockError::AccountWitnessTracking {
+            source: AccountTreeError::DuplicateIdPrefix { duplicate_prefix }
+        } if duplicate_prefix == id0.prefix()
+    );
+
+    Ok(())
+}