feat: remove asset_witnesses field from TransactionInputs

CHANGELOG.md

@@ -17,6 +17,7 @@
 
 ### Changes
 
+- [BREAKING] Removed `asset_witnesses` field from `TransactionInputs` ([#2274](https://github.com/0xMiden/miden-base/pull/2274)).
 - No longer pad the note inputs on insertion into advice map ([#2232](https://github.com/0xMiden/miden-base/pull/2232)).
 - Added proc-macro `WordWrapper` to ease implementation of `Word`-wrapping types ([#2071](https://github.com/0xMiden/miden-base/pull/2108)).
 - [BREAKING] Added `BlockBody` and `BlockProof` structs in preparation for validator signatures and deferred block proving ([#2012](https://github.com/0xMiden/miden-base/pull/2012)).

crates/miden-protocol/src/transaction/inputs/mod.rs

@@ -116,7 +116,14 @@ impl TransactionInputs {
     }
 
     /// Replaces the transaction inputs and assigns the given asset witnesses.
+    ///
+    /// The asset witness data is added to both the internal advice inputs and the transaction
+    /// arguments' advice inputs, so it can be retrieved later via
+    /// [`Self::read_vault_asset_witnesses`].
     pub fn with_asset_witnesses(mut self, witnesses: Vec<AssetWitness>) -> Self {
+        let mut tx_adv = TransactionAdviceInputs::default();
+        witnesses.iter().for_each(|witness| tx_adv.add_asset_witness(witness.clone()));
+        self.extend_advice_inputs(tx_adv.into_advice_inputs());
         self.asset_witnesses = witnesses;
         self
     }
@@ -127,12 +134,6 @@ impl TransactionInputs {
         self
     }
 
-    /// Replaces the transaction inputs and assigns the given transaction arguments.
-    pub fn with_tx_args(mut self, tx_args: TransactionArgs) -> Self {
-        self.set_tx_args_inner(tx_args);
-        self
-    }
-
     /// Replaces the transaction inputs and assigns the given foreign account slot names.
     pub fn with_foreign_account_slot_names(
         mut self,
@@ -142,6 +143,12 @@ impl TransactionInputs {
         self
     }
 
+    /// Replaces the transaction inputs and assigns the given transaction arguments.
+    pub fn with_tx_args(mut self, tx_args: TransactionArgs) -> Self {
+        self.set_tx_args_inner(tx_args);
+        self
+    }
+
     /// Replaces the transaction inputs and assigns the given advice inputs.
     pub fn with_advice_inputs(mut self, advice_inputs: AdviceInputs) -> Self {
         self.set_advice_inputs(advice_inputs);
@@ -165,6 +172,15 @@ impl TransactionInputs {
         self.tx_args.extend_advice_inputs(self.advice_inputs.clone());
     }
 
+    /// Extends the advice inputs with the provided inputs.
+    ///
+    /// This extends both the internal advice inputs and the transaction arguments' advice inputs,
+    /// ensuring that `self.advice_inputs` is always a subset of `self.tx_args.advice_inputs()`.
+    pub fn extend_advice_inputs(&mut self, advice_inputs: AdviceInputs) {
+        self.advice_inputs.extend(advice_inputs.clone());
+        self.tx_args.extend_advice_inputs(advice_inputs);
+    }
+
     /// Updates the transaction arguments of the inputs.
     #[cfg(feature = "testing")]
     pub fn set_tx_args(&mut self, tx_args: TransactionArgs) {
@@ -412,8 +428,8 @@ impl TransactionInputs {
 
     /// Replaces the current tx_args with the provided value.
     ///
-    /// This also appends advice inputs from these transaction inputs to the advice inputs of the
-    /// tx args.
+    /// This also merges the advice inputs from both the transaction inputs and the tx args,
+    /// ensuring that `self.advice_inputs` is always a subset of `self.tx_args.advice_inputs()`.
     fn set_tx_args_inner(&mut self, tx_args: TransactionArgs) {
         self.tx_args = tx_args;
         self.tx_args.extend_advice_inputs(self.advice_inputs.clone());

crates/miden-protocol/src/transaction/kernel/advice_inputs.rs

@@ -106,16 +106,6 @@ impl TransactionAdviceInputs {
         .into_iter()
     }
 
-    // PUBLIC UTILITIES
-    // --------------------------------------------------------------------------------------------
-
-    /// Returns the advice map key where:
-    /// - the seed for native accounts is stored.
-    /// - the account header for foreign accounts is stored.
-    pub fn account_id_map_key(id: AccountId) -> Word {
-        Word::from([id.suffix(), id.prefix().as_felt(), ZERO, ZERO])
-    }
-
     // MUTATORS
     // --------------------------------------------------------------------------------------------
 
@@ -309,7 +299,7 @@ impl TransactionAdviceInputs {
     }
 
     /// Adds an asset witness to the advice inputs.
-    fn add_asset_witness(&mut self, witness: AssetWitness) {
+    pub(crate) fn add_asset_witness(&mut self, witness: AssetWitness) {
         self.extend_merkle_store(witness.authenticated_nodes());
 
         let smt_proof = SmtProof::from(witness);
@@ -421,6 +411,16 @@ impl TransactionAdviceInputs {
     fn extend_merkle_store(&mut self, iter: impl Iterator<Item = InnerNodeInfo>) {
         self.0.store.extend(iter);
     }
+
+    // PUBLIC UTILITIES
+    // --------------------------------------------------------------------------------------------
+
+    /// Returns the advice map key where:
+    /// - the seed for native accounts is stored.
+    /// - the account header for foreign accounts is stored.
+    pub fn account_id_map_key(id: AccountId) -> Word {
+        Word::from([id.suffix(), id.prefix().as_felt(), ZERO, ZERO])
+    }
 }
 
 // CONVERSIONS

crates/miden-testing/tests/scripts/fee.rs

@@ -1,5 +1,9 @@
 use anyhow::Context;
+use miden_protocol::account::AccountType;
 use miden_protocol::asset::FungibleAsset;
+use miden_protocol::crypto::merkle::smt::Smt;
+use miden_protocol::note::NoteType;
+use miden_protocol::testing::account_id::{ACCOUNT_ID_SENDER, AccountIdBuilder};
 use miden_protocol::{self, Felt, Word};
 use miden_testing::{Auth, MockChain};
 
@@ -49,3 +53,56 @@ async fn prove_account_creation_with_fees() -> anyhow::Result<()> {
 
     Ok(())
 }
+
+/// Test that reexecuting a transaction with no authenticator and the tx inputs from a first
+/// successful execution is possible.
+///
+/// The test setup is a non-empty account vault with a native asset that is sufficient to cover the
+/// fee of the transaction.
+///
+/// When re-executing, the initial_fee_asset_balance is extracted from the asset witness, which is
+/// done differently during execution and re-execution. Since we cannot assert this balance matches
+/// directly, this test ensures that the witness fetched during re-execution isn't 0.
+#[tokio::test]
+async fn tx_reexecution_fetches_correct_fee_asset_witness() -> anyhow::Result<()> {
+    let native_asset_id = AccountIdBuilder::new()
+        .account_type(AccountType::FungibleFaucet)
+        .build_with_seed([4; 32]);
+    let mut builder =
+        MockChain::builder().verification_base_fee(50).native_asset_id(native_asset_id);
+    // Use basic auth so the tx requires a signature for successful execution.
+    // Setup a non-empty vault with a non-zero fee asset. Without this, the test has no point.
+    let account = builder.add_existing_mock_account_with_assets(
+        Auth::BasicAuth,
+        [FungibleAsset::new(native_asset_id, 10_000)?.into()],
+    )?;
+    assert_ne!(
+        account.vault().root(),
+        Smt::default().root(),
+        "account should have a non-default vault root for this test"
+    );
+
+    let note = builder.add_p2id_note(
+        ACCOUNT_ID_SENDER.try_into()?,
+        account.id(),
+        &[FungibleAsset::mock(3)],
+        NoteType::Public,
+    )?;
+    let chain = builder.build()?;
+
+    let tx = chain
+        .build_tx_context(account.id(), &[note.id()], &[])?
+        .build()?
+        .execute()
+        .await?;
+
+    let _reexecuted_tx = chain
+        .build_tx_context(account.id(), &[note.id()], &[])?
+        .authenticator(None)
+        .tx_inputs(tx.tx_inputs().clone())
+        .build()?
+        .execute()
+        .await?;
+
+    Ok(())
+}

crates/miden-tx/src/executor/mod.rs

@@ -1,5 +1,6 @@
 use alloc::collections::BTreeSet;
 use alloc::sync::Arc;
+use alloc::vec::Vec;
 
 use miden_processor::fast::FastProcessor;
 use miden_processor::{AdviceInputs, ExecutionError, StackInputs};
@@ -9,6 +10,7 @@ use miden_protocol::assembly::DefaultSourceManager;
 use miden_protocol::assembly::debuginfo::SourceManagerSync;
 use miden_protocol::asset::{Asset, AssetVaultKey};
 use miden_protocol::block::BlockNumber;
+use miden_protocol::crypto::merkle::NodeIndex;
 use miden_protocol::transaction::{
     ExecutedTransaction,
     InputNote,
@@ -40,6 +42,25 @@ pub use notes_checker::{
     NoteConsumptionInfo,
 };
 
+// PREPARED TRANSACTION INPUTS
+// ================================================================================================
+
+/// Wrapper for transaction inputs and the initial fee asset balance.
+///
+/// This struct is returned by [`TransactionExecutor::prepare_tx_inputs`] and contains the
+/// transaction inputs along with the initial fee asset balance, which is needed for proper
+/// transaction execution.
+pub(crate) struct PreparedTransactionInputs {
+    pub inputs: TransactionInputs,
+    pub initial_fee_balance: u64,
+}
+
+impl PreparedTransactionInputs {
+    fn new(inputs: TransactionInputs, initial_fee_balance: u64) -> Self {
+        Self { inputs, initial_fee_balance }
+    }
+}
+
 // TRANSACTION EXECUTOR
 // ================================================================================================
 
@@ -179,9 +200,10 @@ where
         notes: InputNotes<InputNote>,
         tx_args: TransactionArgs,
     ) -> Result<ExecutedTransaction, TransactionExecutorError> {
-        let tx_inputs = self.prepare_tx_inputs(account_id, block_ref, notes, tx_args).await?;
+        let prepared = self.prepare_tx_inputs(account_id, block_ref, notes, tx_args).await?;
 
-        let (mut host, stack_inputs, advice_inputs) = self.prepare_transaction(&tx_inputs).await?;
+        let (mut host, stack_inputs, advice_inputs) =
+            self.prepare_transaction(&prepared.inputs, prepared.initial_fee_balance).await?;
 
         let processor = FastProcessor::new_debug(stack_inputs.as_slice(), advice_inputs);
         let output = processor
@@ -199,7 +221,7 @@ where
             ..Default::default()
         };
 
-        build_executed_transaction(advice_inputs, tx_inputs, stack_outputs, host)
+        build_executed_transaction(advice_inputs, prepared.inputs, stack_outputs, host)
     }
 
     // SCRIPT EXECUTION
@@ -224,9 +246,10 @@ where
         tx_args.extend_advice_inputs(advice_inputs);
 
         let notes = InputNotes::default();
-        let tx_inputs = self.prepare_tx_inputs(account_id, block_ref, notes, tx_args).await?;
+        let prepared = self.prepare_tx_inputs(account_id, block_ref, notes, tx_args).await?;
 
-        let (mut host, stack_inputs, advice_inputs) = self.prepare_transaction(&tx_inputs).await?;
+        let (mut host, stack_inputs, advice_inputs) =
+            self.prepare_transaction(&prepared.inputs, prepared.initial_fee_balance).await?;
 
         let processor =
             FastProcessor::new_with_advice_inputs(stack_inputs.as_slice(), advice_inputs);
@@ -247,13 +270,15 @@ where
     // This method has a one-to-many call relationship with the `prepare_transaction` method. This
     // method needs to be called only once in order to allow many transactions to be prepared based
     // on the transaction inputs returned by this method.
+    //
+    // Returns the prepared transaction inputs containing both the inputs and initial fee balance.
     async fn prepare_tx_inputs(
         &self,
         account_id: AccountId,
         block_ref: BlockNumber,
         input_notes: InputNotes<InputNote>,
         tx_args: TransactionArgs,
-    ) -> Result<TransactionInputs, TransactionExecutorError> {
+    ) -> Result<PreparedTransactionInputs, TransactionExecutorError> {
         let (mut asset_vault_keys, mut ref_blocks) = validate_input_notes(&input_notes, block_ref)?;
         ref_blocks.insert(block_ref);
 
@@ -270,19 +295,49 @@ where
                 .expect("fee asset should be a fungible asset");
         asset_vault_keys.insert(fee_asset_vault_key);
 
-        // Fetch the witnesses for all asset vault keys.
-        let asset_witnesses = self
-            .data_store
-            .get_vault_asset_witnesses(account_id, account.vault().root(), asset_vault_keys)
-            .await
-            .map_err(TransactionExecutorError::FetchAssetWitnessFailed)?;
+        // Filter out asset vault keys whose merkle paths are already present in the advice inputs.
+        // This optimization avoids redundant DataStore calls during re-execution.
+        let vault_root = account.vault().root();
+        let advice_store = &tx_args.advice_inputs().store;
+        asset_vault_keys.retain(|key| {
+            let node_index = NodeIndex::from(key.to_leaf_index());
+            !advice_store.has_path(vault_root, node_index)
+        });
+
+        // Fetch the witnesses for asset vault keys that still need them.
+        let asset_witnesses = if !asset_vault_keys.is_empty() {
+            self.data_store
+                .get_vault_asset_witnesses(account_id, vault_root, asset_vault_keys)
+                .await
+                .map_err(TransactionExecutorError::FetchAssetWitnessFailed)?
+        } else {
+            Vec::new()
+        };
 
         let tx_inputs = TransactionInputs::new(account, block_header, blockchain, input_notes)
             .map_err(TransactionExecutorError::InvalidTransactionInputs)?
             .with_tx_args(tx_args)
             .with_asset_witnesses(asset_witnesses);
 
-        Ok(tx_inputs)
+        // Read the fee asset witness back from the tx inputs.
+        // This way, whether the witness was just added or was already present doesn't matter, which
+        // abstracts over execution and re-execution cases.
+        let fee_asset_witnesses = tx_inputs
+            .read_vault_asset_witnesses(vault_root, [fee_asset_vault_key].into())
+            .unwrap_or_default();
+
+        // Calculate the initial fee asset balance.
+        let fee_asset_witness =
+            fee_asset_witnesses.iter().find_map(|witness| witness.find(fee_asset_vault_key));
+        let initial_fee_asset_balance = match fee_asset_witness {
+            Some(Asset::Fungible(fee_asset)) => fee_asset.amount(),
+            Some(Asset::NonFungible(_)) => {
+                return Err(TransactionExecutorError::FeeAssetMustBeFungible);
+            },
+            None => 0,
+        };
+
+        Ok(PreparedTransactionInputs::new(tx_inputs, initial_fee_asset_balance))
     }
 
     /// Prepares the data needed for transaction execution.
@@ -292,6 +347,7 @@ where
     async fn prepare_transaction(
         &self,
         tx_inputs: &TransactionInputs,
+        initial_fee_asset_balance: u64,
     ) -> Result<
         (TransactionExecutorHost<'store, 'auth, STORE, AUTH>, StackInputs, AdviceInputs),
         TransactionExecutorError,
@@ -317,26 +373,6 @@ where
         let account_procedure_index_map =
             AccountProcedureIndexMap::new([tx_inputs.account().code()]);
 
-        let initial_fee_asset_balance = {
-            let native_asset_id = tx_inputs.block_header().fee_parameters().native_asset_id();
-            let fee_asset_vault_key = AssetVaultKey::from_account_id(native_asset_id)
-                .expect("fee asset should be a fungible asset");
-
-            let fee_asset_witness = tx_inputs
-                .asset_witnesses()
-                .iter()
-                .find_map(|witness| witness.find(fee_asset_vault_key));
-
-            match fee_asset_witness {
-                Some(Asset::Fungible(fee_asset)) => fee_asset.amount(),
-                Some(Asset::NonFungible(_)) => {
-                    return Err(TransactionExecutorError::FeeAssetMustBeFungible);
-                },
-                // If the witness does not contain the asset, its balance is zero.
-                None => 0,
-            }
-        };
-
         let host = TransactionExecutorHost::new(
             tx_inputs.account(),
             input_notes.clone(),
@@ -367,7 +403,6 @@ fn build_executed_transaction<STORE: DataStore + Sync, AUTH: TransactionAuthenti
 ) -> Result<ExecutedTransaction, TransactionExecutorError> {
     // Note that the account delta does not contain the removed transaction fee, so it is the
     // "pre-fee" delta of the transaction.
-
     let (
         pre_fee_account_delta,
         _input_notes,