Factor out MAST root calls, improve mock account code

miden-lib/src/transaction/mod.rs

@@ -1,7 +1,7 @@
 use alloc::{string::ToString, sync::Arc, vec::Vec};
 
 use miden_objects::{
-    accounts::AccountId,
+    accounts::{AccountCode, AccountId},
     assembly::{Assembler, DefaultSourceManager, KernelLibrary},
     transaction::{
         OutputNote, OutputNotes, TransactionArgs, TransactionInputs, TransactionOutputs,
@@ -267,7 +267,7 @@ impl TransactionKernel {
     /// The `kernel` library is added separately because even though the library (`api.masm`) and
     /// the kernel binary (`main.masm`) include this code, it is not exposed explicitly. By adding
     /// it separately, we can expose procedures from `/lib` and test them individually.
-    pub fn assembler_testing() -> Assembler {
+    pub fn testing_assembler() -> Assembler {
         let source_manager = Arc::new(DefaultSourceManager::default());
         let kernel_library = Self::kernel_as_library();
 
@@ -279,4 +279,12 @@ impl TransactionKernel {
             .with_library(kernel_library)
             .expect("failed to load kernel library (/lib)")
     }
+
+    /// Returns the testing assembler, and additionally contains the library for
+    /// [AccountCode::mock_library()], which is a mock wallet used in tests.
+    pub fn testing_assembler_with_mock_account() -> Assembler {
+        let assembler = Self::testing_assembler();
+        let library = AccountCode::mock_library(assembler.clone());
+        assembler.with_library(library).expect("failed to add mock account code")
+    }
 }

miden-tx/src/testing/executor.rs

@@ -37,7 +37,7 @@ impl<H: Host> CodeExecutor<H> {
 
     /// Compiles and runs the desired code in the host and returns the [Process] state
     pub fn run(self, code: &str) -> Result<Process<H>, ExecutionError> {
-        let program = TransactionKernel::assembler_testing().assemble_program(code).unwrap();
+        let program = TransactionKernel::testing_assembler().assemble_program(code).unwrap();
         self.execute_program(program)
     }
 

miden-tx/src/testing/mock_chain.rs

@@ -301,15 +301,15 @@ impl MockChain {
 
     pub fn add_new_wallet(&mut self, auth_method: Auth, assets: Vec<Asset>) -> Account {
         let account_builder = AccountBuilder::new(ChaCha20Rng::from_entropy())
-            .default_code(TransactionKernel::assembler_testing())
+            .default_code(TransactionKernel::testing_assembler())
             .nonce(Felt::ZERO)
             .add_assets(assets);
         self.add_from_account_builder(auth_method, account_builder)
     }
 
     pub fn add_existing_wallet(&mut self, auth_method: Auth, assets: Vec<Asset>) -> Account {
         let account_builder = AccountBuilder::new(ChaCha20Rng::from_entropy())
-            .default_code(TransactionKernel::assembler_testing())
+            .default_code(TransactionKernel::testing_assembler())
             .nonce(Felt::ONE)
             .add_assets(assets);
         self.add_from_account_builder(auth_method, account_builder)
@@ -331,7 +331,7 @@ impl MockChain {
         let faucet_metadata = SlotItem::new_value(1, 0, metadata);
 
         let account_builder = AccountBuilder::new(ChaCha20Rng::from_entropy())
-            .default_code(TransactionKernel::assembler_testing())
+            .default_code(TransactionKernel::testing_assembler())
             .nonce(Felt::ZERO)
             .account_type(AccountType::FungibleFaucet)
             .add_storage_item(faucet_metadata);
@@ -357,7 +357,7 @@ impl MockChain {
         let faucet_metadata = SlotItem::new_value(1, 0, metadata);
 
         let account_builder = AccountBuilder::new(ChaCha20Rng::from_entropy())
-            .default_code(TransactionKernel::assembler_testing())
+            .default_code(TransactionKernel::testing_assembler())
             .nonce(Felt::ONE)
             .account_type(AccountType::FungibleFaucet)
             .add_storage_item(faucet_metadata);

miden-tx/src/testing/tx_context/builder.rs

@@ -14,10 +14,9 @@ use miden_objects::{
         Account, AccountId,
     },
     assembly::Assembler,
-    assets::{Asset, FungibleAsset},
+    assets::{Asset, FungibleAsset, NonFungibleAsset},
     notes::{Note, NoteExecutionHint, NoteId, NoteType},
     testing::{
-        account_code::ACCOUNT_ADD_ASSET_TO_NOTE_MAST_ROOT,
         constants::{
             CONSUMED_ASSET_1_AMOUNT, CONSUMED_ASSET_2_AMOUNT, CONSUMED_ASSET_3_AMOUNT,
             NON_FUNGIBLE_ASSET_DATA_2,
@@ -53,7 +52,7 @@ pub struct TransactionContextBuilder {
 impl TransactionContextBuilder {
     pub fn new(account: Account) -> Self {
         Self {
-            assembler: TransactionKernel::assembler_testing(),
+            assembler: TransactionKernel::testing_assembler(),
             account,
             account_seed: None,
             input_notes: Vec::new(),
@@ -69,13 +68,15 @@ impl TransactionContextBuilder {
     }
 
     pub fn with_standard_account(nonce: Felt) -> Self {
-        let assembler = TransactionKernel::assembler_testing();
+        // Build standard account with normal assembler because the testing one already contains it
         let account = Account::mock(
             ACCOUNT_ID_REGULAR_ACCOUNT_UPDATABLE_CODE_ON_CHAIN,
             nonce,
-            assembler.clone(),
+            TransactionKernel::testing_assembler(),
         );
 
+        let assembler = TransactionKernel::testing_assembler_with_mock_account();
+
         Self {
             assembler: assembler.clone(),
             account,
@@ -93,9 +94,13 @@ impl TransactionContextBuilder {
     }
 
     pub fn with_fungible_faucet(acct_id: u64, nonce: Felt, initial_balance: Felt) -> Self {
-        let assembler = TransactionKernel::assembler_testing();
-        let account =
-            Account::mock_fungible_faucet(acct_id, nonce, initial_balance, assembler.clone());
+        let account = Account::mock_fungible_faucet(
+            acct_id,
+            nonce,
+            initial_balance,
+            TransactionKernel::testing_assembler(),
+        );
+        let assembler = TransactionKernel::testing_assembler_with_mock_account();
 
         Self {
             assembler,
@@ -114,13 +119,13 @@ impl TransactionContextBuilder {
     }
 
     pub fn with_non_fungible_faucet(acct_id: u64, nonce: Felt, empty_reserved_slot: bool) -> Self {
-        let assembler = TransactionKernel::assembler_testing();
         let account = Account::mock_non_fungible_faucet(
             acct_id,
             nonce,
             empty_reserved_slot,
-            assembler.clone(),
+            TransactionKernel::testing_assembler(),
         );
+        let assembler = TransactionKernel::testing_assembler_with_mock_account();
 
         Self {
             assembler,
@@ -215,6 +220,7 @@ impl TransactionContextBuilder {
         let var_name = format!(
             "
             use.miden::contracts::wallets::basic->wallet
+            use.test::account
 
             begin
                 # NOTE
@@ -227,7 +233,7 @@ impl TransactionContextBuilder {
                 call.wallet::create_note
 
                 push.{asset}
-                call.{ACCOUNT_ADD_ASSET_TO_NOTE_MAST_ROOT}
+                call.account::add_asset_to_note
                 dropw dropw dropw
             end
             ",
@@ -260,6 +266,7 @@ impl TransactionContextBuilder {
         let code = format!(
             "
             use.miden::contracts::wallets::basic->wallet
+            use.test::account
 
             begin
 
@@ -272,11 +279,10 @@ impl TransactionContextBuilder {
                 push.{aux0}
                 push.{tag0}
 
-
                 call.wallet::create_note
 
                 push.{asset0}
-                call.{ACCOUNT_ADD_ASSET_TO_NOTE_MAST_ROOT}
+                call.account::add_asset_to_note
                 dropw dropw dropw
 
                 # NOTE 1
@@ -289,7 +295,7 @@ impl TransactionContextBuilder {
                 call.wallet::create_note
 
                 push.{asset1}
-                call.{ACCOUNT_ADD_ASSET_TO_NOTE_MAST_ROOT}
+                call.account::add_asset_to_note
                 dropw dropw dropw
             end
             ",
@@ -441,7 +447,7 @@ impl TransactionContextBuilder {
             FungibleAsset::new(faucet_id_2, CONSUMED_ASSET_2_AMOUNT).unwrap().into();
         let fungible_asset_3: Asset =
             FungibleAsset::new(faucet_id_3, CONSUMED_ASSET_3_AMOUNT).unwrap().into();
-        let nonfungible_asset_1: Asset = Asset::mock_non_fungible(
+        let nonfungible_asset_1: Asset = NonFungibleAsset::mock(
             ACCOUNT_ID_NON_FUNGIBLE_FAUCET_ON_CHAIN,
             &NON_FUNGIBLE_ASSET_DATA_2,
         );
@@ -526,7 +532,7 @@ impl TransactionContextBuilder {
             FungibleAsset::new(faucet_id_2, CONSUMED_ASSET_2_AMOUNT).unwrap().into();
         let fungible_asset_3: Asset =
             FungibleAsset::new(faucet_id_3, CONSUMED_ASSET_3_AMOUNT).unwrap().into();
-        let nonfungible_asset_1: Asset = Asset::mock_non_fungible(
+        let nonfungible_asset_1: Asset = NonFungibleAsset::mock(
             ACCOUNT_ID_NON_FUNGIBLE_FAUCET_ON_CHAIN,
             &NON_FUNGIBLE_ASSET_DATA_2,
         );

miden-tx/src/tests/kernel_tests/test_account.rs

@@ -451,7 +451,7 @@ fn test_set_map_item() {
 #[test]
 fn test_storage_offset() {
     // setup assembler
-    let assembler = TransactionKernel::assembler_testing();
+    let assembler = TransactionKernel::testing_assembler();
 
     // The following code will execute the following logic that will be asserted during the test:
     //

miden-tx/src/tests/kernel_tests/test_asset.rs

@@ -2,7 +2,7 @@ use miden_objects::{
     accounts::account_id::testing::{
         ACCOUNT_ID_FUNGIBLE_FAUCET_ON_CHAIN, ACCOUNT_ID_NON_FUNGIBLE_FAUCET_ON_CHAIN,
     },
-    assets::Asset,
+    assets::NonFungibleAsset,
     testing::{
         constants::{
             FUNGIBLE_ASSET_AMOUNT, FUNGIBLE_FAUCET_INITIAL_BALANCE, NON_FUNGIBLE_ASSET_DATA,
@@ -62,7 +62,7 @@ fn test_create_non_fungible_asset_succeeds() {
     .build();
 
     let non_fungible_asset =
-        Asset::mock_non_fungible(ACCOUNT_ID_NON_FUNGIBLE_FAUCET_ON_CHAIN, &NON_FUNGIBLE_ASSET_DATA);
+        NonFungibleAsset::mock(ACCOUNT_ID_NON_FUNGIBLE_FAUCET_ON_CHAIN, &NON_FUNGIBLE_ASSET_DATA);
 
     let code = format!(
         "
@@ -95,7 +95,7 @@ fn test_validate_non_fungible_asset() {
     .build();
 
     let non_fungible_asset =
-        Asset::mock_non_fungible(ACCOUNT_ID_NON_FUNGIBLE_FAUCET_ON_CHAIN, &[1, 2, 3]);
+        NonFungibleAsset::mock(ACCOUNT_ID_NON_FUNGIBLE_FAUCET_ON_CHAIN, &[1, 2, 3]);
     let encoded = Word::from(non_fungible_asset);
 
     let code = format!(