WIP: implement multi-asset fee PoC

crates/miden-protocol/asm/kernels/transaction/lib/epilogue.masm

@@ -22,6 +22,8 @@ const ERR_EPILOGUE_AUTH_PROCEDURE_CALLED_FROM_WRONG_CONTEXT="auth procedure has
 
 const ERR_EPILOGUE_NONCE_CANNOT_BE_0="nonce cannot be 0 after an account-creating transaction"
 
+const ERR_EPILOGUE_FEE_PROC_EXCEEDED_BUDGET="fee procedure exceeded its cycle budget"
+
 # CONSTANTS
 # =================================================================================================
 
@@ -54,6 +56,9 @@ const NUM_POST_COMPUTE_FEE_CYCLES=608
 # The number of cycles the epilogue is estimated to take after compute_fee has been executed.
 const ESTIMATED_AFTER_COMPUTE_FEE_CYCLES=NUM_POST_COMPUTE_FEE_CYCLES+SMT_SET_ADDITIONAL_CYCLES
 
+# Maximum number of cycles the user-defined fee procedure is allowed to consume.
+const MAX_FEE_PROC_CYCLES=2000
+
 # OUTPUT NOTES PROCEDURES
 # =================================================================================================
 
@@ -291,8 +296,6 @@ proc create_native_fee_asset
     exec.memory::get_native_asset_id
     # => [native_asset_id_suffix, native_asset_id_prefix, fee_amount]
 
-    # assume the fee asset does not have callbacks
-    # this should be addressed more holistically with a fee construction refactor
     push.0
     # => [enable_callbacks, native_asset_id_suffix, native_asset_id_prefix, fee_amount]
 
@@ -301,23 +304,65 @@ proc create_native_fee_asset
     # => [FEE_ASSET_KEY, FEE_ASSET_VALUE]
 end
 
-#! Computes the fee of this transaction and removes the asset from the native account's vault.
+# TODO(multi-asset-fees): Wire up execute_fee_procedure once existing auth components
+# bundle a @fee_script procedure. The dispatch code below is complete but not called
+# from compute_and_remove_fee yet because accounts without @fee_script would crash.
+# The open question is how a user-space fee procedure obtains the native faucet ID
+# (storage slot vs. FEE_ARGS vs. syscall).
+
+#! Executes the account's fee procedure to convert computation units into a fee asset.
 #!
-#! Note that this does not have to account for the fee asset in the output vault explicitly,
-#! because the fee asset is removed from the account vault after build_output_vault and because
-#! it is not added to an output note. Effectively, the fee asset bypasses the asset preservation
-#! check. That's okay, because the logic is entirely determined by the transaction kernel.
+#! The fee procedure is at index 1 in the account's procedure list (placed there by the
+#! AccountProcedureBuilder when it sees the @fee_script attribute).
 #!
-#! Inputs:  []
-#! Outputs: [native_asset_id_suffix, native_asset_id_prefix, fee_amount]
+#! Inputs:  [computation_units]
+#! Outputs: [FEE_ASSET_KEY, FEE_ASSET_VALUE]
 #!
 #! Where:
-#! - fee_amount is the computed fee amount of the transaction in the native asset.
-#! - native_asset_id_{prefix,suffix} are the prefix and suffix felts of the faucet that issues the
-#!   native asset.
+#! - computation_units is the fee amount computed by compute_fee.
+#! - FEE_ASSET_KEY is the asset vault key of the fee asset chosen by the fee procedure.
+#! - FEE_ASSET_VALUE is the value of the fee asset.
 #!
 #! Panics if:
-#! - the account vault contains less than the computed fee.
+#! - the fee procedure exceeds MAX_FEE_PROC_CYCLES.
+proc execute_fee_procedure
+    # build the dyncall input: [computation_units, FEE_ARGS, pad(11)]
+    padw padw push.0.0.0
+    # => [pad(11), computation_units]
+
+    exec.memory::get_fee_args
+    # => [FEE_ARGS, pad(11), computation_units]
+
+    movup.15
+    # => [computation_units, FEE_ARGS, pad(11)]
+
+    # fee procedure is at index 1 within the account procedures section.
+    push.1 exec.memory::get_account_procedure_ptr
+    # => [fee_procedure_ptr, computation_units, FEE_ARGS, pad(11)]
+
+    padw dup.4 mem_loadw_le
+    # => [FEE_PROC_ROOT, fee_procedure_ptr, computation_units, FEE_ARGS, pad(11)]
+
+    dyncall
+    # => [OUTPUT_3, OUTPUT_2, OUTPUT_1, OUTPUT_0]
+
+    # TODO: enforce MAX_FEE_PROC_CYCLES budget via a memory slot (movup.16 is out of range)
+
+    # the fee procedure returns [FEE_ASSET_KEY, FEE_ASSET_VALUE, pad(8)]
+    swapw.2 dropw swapw.2 dropw
+    # => [FEE_ASSET_KEY, FEE_ASSET_VALUE]
+end
+
+#! Computes the fee of this transaction and removes the asset from the native account's vault.
+#!
+#! TODO(multi-asset-fees): once existing auth components bundle a @fee_script procedure,
+#! switch from create_native_fee_asset to execute_fee_procedure so accounts can pay in any asset.
+#!
+#! Note that this deliberately does not use account::remove_asset_from_vault, because that
+#! procedure modifies the vault delta.
+#!
+#! Inputs:  []
+#! Outputs: [native_asset_id_suffix, native_asset_id_prefix, fee_amount]
 proc compute_and_remove_fee
     # compute the fee the tx needs to pay
     exec.compute_fee dup
@@ -337,14 +382,6 @@ proc compute_and_remove_fee
     movdn.9 movdn.9
     # => [FEE_ASSET_KEY, FEE_ASSET_VALUE, native_asset_id_suffix, native_asset_id_prefix, fee_amount]
 
-    # remove the fee from the native account's vault
-    # note that this deliberately does not use account::remove_asset_from_vault, because that
-    # procedure modifies the vault delta of the account. That is undesirable because it does not
-    # take a constant number of cycles, which makes it much harder to calculate the number of
-    # cycles the kernel takes after computing the fee. It is also unnecessary, because the delta
-    # commitment has already been computed and so any modifications done to the delta at this point
-    # are essentially ignored.
-
     # fetch the vault root ptr
     exec.memory::get_account_vault_root_ptr movdn.8
     # => [FEE_ASSET_KEY, FEE_ASSET_VALUE, account_vault_root_ptr, native_asset_id_suffix, native_asset_id_prefix, fee_amount]
@@ -378,7 +415,7 @@ end
 #!
 #! Where:
 #! - OUTPUT_NOTES_COMMITMENT is the commitment of the output notes.
-#! - ACCOUNT_UPDATE_COMMITMENT is the hash of the the final account commitment and account
+#! - ACCOUNT_UPDATE_COMMITMENT is the hash of the final account commitment and account
 #!   delta commitment.
 #! - fee_amount is the computed fee amount of the transaction denominated in the native asset.
 #! - native_asset_id_{prefix,suffix} are the prefix and suffix felts of the faucet that issues the

crates/miden-protocol/asm/kernels/transaction/lib/memory.masm

@@ -112,6 +112,9 @@ const TX_SCRIPT_ARGS_PTR=432
 # The memory address at which the auth procedure arguments are stored.
 const AUTH_ARGS_PTR=436
 
+# The memory address at which the fee procedure arguments are stored.
+const FEE_ARGS_PTR=440
+
 # GLOBAL BLOCK DATA
 # -------------------------------------------------------------------------------------------------
 
@@ -706,6 +709,28 @@ pub proc set_auth_args
     mem_storew_le.AUTH_ARGS_PTR
 end
 
+#! Returns the fee procedure arguments.
+#!
+#! Inputs:  []
+#! Outputs: [FEE_ARGS]
+#!
+#! Where:
+#! - FEE_ARGS is the argument passed to the fee procedure.
+pub proc get_fee_args
+    padw mem_loadw_le.FEE_ARGS_PTR
+end
+
+#! Sets the fee procedure arguments.
+#!
+#! Inputs:  [FEE_ARGS]
+#! Outputs: [FEE_ARGS]
+#!
+#! Where:
+#! - FEE_ARGS is the argument passed to the fee procedure.
+pub proc set_fee_args
+    mem_storew_le.FEE_ARGS_PTR
+end
+
 # BLOCK DATA
 # -------------------------------------------------------------------------------------------------
 

crates/miden-protocol/asm/kernels/transaction/lib/prologue.masm

@@ -1064,6 +1064,25 @@ proc process_auth_procedure_data
     # => []
 end
 
+#! Saves the fee procedure args to memory.
+#!
+#! Inputs:
+#!   Operand stack: []
+#!   Advice stack: [FEE_ARGS]
+#! Outputs:
+#!   Operand stack: []
+#!   Advice stack: []
+#!
+#! Where:
+#! - FEE_ARGS is the argument passed to the fee procedure.
+proc process_fee_procedure_data
+    padw adv_loadw
+    # => [FEE_ARGS]
+
+    exec.memory::set_fee_args dropw
+    # => []
+end
+
 # TRANSACTION PROLOGUE
 # =================================================================================================
 
@@ -1101,6 +1120,7 @@ end
 #!     TX_SCRIPT_ROOT,
 #!     TX_SCRIPT_ARGS,
 #!     AUTH_ARGS,
+#!     FEE_ARGS,
 #!   ]
 #!   Advice map: {
 #!      PARTIAL_BLOCKCHAIN_COMMITMENT: [MMR_PEAKS],
@@ -1162,6 +1182,7 @@ pub proc prepare_transaction
     exec.process_input_notes_data
     exec.process_tx_script_data
     exec.process_auth_procedure_data
+    exec.process_fee_procedure_data
     # => []
 
     push.MAX_BLOCK_NUM exec.memory::set_expiration_block_num

crates/miden-protocol/src/account/code/mod.rs

@@ -336,7 +336,7 @@ impl PrettyPrint for AccountCode {
 
 /// A helper type for building the set of account procedures from account components.
 ///
-/// In particular, this ensures that the auth procedure ends up at index 0.
+/// Ensures the auth procedure ends up at index 0 and the fee procedure at index 1.
 struct AccountProcedureBuilder {
     procedures: Vec<AccountProcedureRoot>,
 }
@@ -346,18 +346,27 @@ impl AccountProcedureBuilder {
         Self { procedures: Vec::new() }
     }
 
-    /// This method must be called before add_component is called.
+    /// Adds the first (auth + fee) component. Must be called before `add_component`.
+    ///
+    /// The component must contain exactly one `@auth_script` procedure and exactly one
+    /// `@fee_script` procedure. The auth procedure is placed at index 0 and the fee procedure
+    /// at index 1.
     fn add_auth_component(&mut self, component: &AccountComponent) -> Result<(), AccountError> {
         let mut auth_proc_count = 0;
+        let mut fee_proc_count = 0;
 
-        for (proc_root, is_auth) in component.procedures() {
+        for (proc_root, is_auth, is_fee) in component.procedures() {
             self.add_procedure(proc_root);
 
             if is_auth {
                 let auth_proc_idx = self.procedures.len() - 1;
                 self.procedures.swap(0, auth_proc_idx);
                 auth_proc_count += 1;
             }
+
+            if is_fee {
+                fee_proc_count += 1;
+            }
         }
 
         if auth_proc_count == 0 {
@@ -366,14 +375,38 @@ impl AccountProcedureBuilder {
             return Err(AccountError::AccountComponentMultipleAuthProcedures);
         }
 
+        if fee_proc_count > 1 {
+            return Err(AccountError::AccountComponentMultipleFeeProcedures);
+        }
+
+        // Place the fee procedure at index 1 if one exists. We need a second pass since the
+        // auth swap may have moved things around.
+        if fee_proc_count == 1 {
+            let mut fee_proc_idx = None;
+            for (idx, proc_root) in self.procedures.iter().enumerate() {
+                for (comp_root, _, is_fee) in component.procedures() {
+                    if is_fee && proc_root == &comp_root {
+                        fee_proc_idx = Some(idx);
+                    }
+                }
+            }
+
+            if let Some(idx) = fee_proc_idx {
+                self.procedures.swap(1, idx);
+            }
+        }
+
         Ok(())
     }
 
     fn add_component(&mut self, component: &AccountComponent) -> Result<(), AccountError> {
-        for (proc_root, is_auth) in component.procedures() {
+        for (proc_root, is_auth, is_fee) in component.procedures() {
             if is_auth {
                 return Err(AccountError::AccountCodeMultipleAuthComponents);
             }
+            if is_fee {
+                return Err(AccountError::AccountCodeMultipleFeeComponents);
+            }
             self.add_procedure(proc_root);
         }
 

crates/miden-protocol/src/account/component/mod.rs

@@ -21,6 +21,9 @@ use crate::{MastForest, Word};
 /// The attribute name used to mark the authentication procedure in an account component.
 const AUTH_SCRIPT_ATTRIBUTE: &str = "auth_script";
 
+/// The attribute name used to mark the fee procedure in an account component.
+const FEE_SCRIPT_ATTRIBUTE: &str = "fee_script";
+
 // ACCOUNT COMPONENT
 // ================================================================================================
 
@@ -189,12 +192,12 @@ impl AccountComponent {
         self.metadata.supported_types().contains(&account_type)
     }
 
-    /// Returns an iterator over ([`AccountProcedureRoot`], is_auth) for all procedures in this
-    /// component.
+    /// Returns an iterator over ([`AccountProcedureRoot`], is_auth, is_fee) for all procedures in
+    /// this component.
     ///
     /// A procedure is considered an authentication procedure if it has the `@auth_script`
-    /// attribute.
-    pub fn procedures(&self) -> impl Iterator<Item = (AccountProcedureRoot, bool)> + '_ {
+    /// attribute, and a fee procedure if it has the `@fee_script` attribute.
+    pub fn procedures(&self) -> impl Iterator<Item = (AccountProcedureRoot, bool, bool)> + '_ {
         let library = self.code.as_library();
         library.exports().filter_map(|export| {
             export.as_procedure().map(|proc_export| {
@@ -204,7 +207,8 @@ impl AccountComponent {
                     .expect("export node not in the forest")
                     .digest();
                 let is_auth = proc_export.attributes.has(AUTH_SCRIPT_ATTRIBUTE);
-                (AccountProcedureRoot::from_raw(digest), is_auth)
+                let is_fee = proc_export.attributes.has(FEE_SCRIPT_ATTRIBUTE);
+                (AccountProcedureRoot::from_raw(digest), is_auth, is_fee)
             })
         })
     }

crates/miden-protocol/src/batch/proposed_batch.rs

@@ -440,7 +440,7 @@ mod tests {
     use crate::Word;
     use crate::account::delta::AccountUpdateDetails;
     use crate::account::{AccountIdVersion, AccountStorageMode, AccountType};
-    use crate::asset::FungibleAsset;
+    use crate::asset::{Asset, FungibleAsset};
     use crate::transaction::{InputNoteCommitment, OutputNote, ProvenTransaction, TxAccountUpdate};
 
     #[test]
@@ -497,7 +497,7 @@ mod tests {
             Vec::<OutputNote>::new(),
             block_num,
             block_ref,
-            FungibleAsset::mock(100).unwrap_fungible(),
+            Asset::from(FungibleAsset::mock(100).unwrap_fungible()),
             expiration_block_num,
             proof,
         )

crates/miden-protocol/src/errors/mod.rs

@@ -122,6 +122,12 @@ pub enum AccountError {
     AccountComponentMastForestMergeError(#[source] MastForestError),
     #[error("account component contains multiple authentication procedures")]
     AccountComponentMultipleAuthProcedures,
+    #[error("account component contains multiple fee procedures")]
+    AccountComponentMultipleFeeProcedures,
+    #[error("account code does not contain a fee component")]
+    AccountCodeNoFeeComponent,
+    #[error("account code contains multiple fee components")]
+    AccountCodeMultipleFeeComponents,
     #[error("failed to update asset vault")]
     AssetVaultUpdateError(#[source] AssetVaultError),
     #[error("account build error: {0}")]

crates/miden-protocol/src/transaction/executed_tx.rs

@@ -14,7 +14,7 @@ use super::{
     TransactionOutputs,
 };
 use crate::account::PartialAccount;
-use crate::asset::FungibleAsset;
+use crate::asset::Asset;
 use crate::block::{BlockHeader, BlockNumber};
 use crate::transaction::TransactionInputs;
 use crate::utils::serde::{
@@ -117,7 +117,7 @@ impl ExecutedTransaction {
     }
 
     /// Returns the fee of the transaction.
-    pub fn fee(&self) -> FungibleAsset {
+    pub fn fee(&self) -> Asset {
         self.tx_outputs.fee()
     }
 

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

@@ -203,6 +203,9 @@ impl TransactionAdviceInputs {
 
         // --- auth procedure args --------------------------------------------
         self.extend_stack(tx_args.auth_args());
+
+        // --- fee procedure args ---------------------------------------------
+        self.extend_stack(tx_args.fee_args());
     }
 
     // BLOCKCHAIN INJECTIONS