bugfix: allow account code to change

Changes to the account code were forbidden in the Rust code, because the
account delta nonce validation assumed the nonce would only change with
changes to the vault or storage.

Currently the account delta doesn't track changes to the account code,
so as an intermediary solution this changes the check to allow code
changes anytime. The correct solution would be to introduce the code
change to the account delta too, and introduce that as another case in
which the delta must change, and then forbid changes otherwise.

miden-lib/asm/miden/account.masm

@@ -86,22 +86,19 @@ end
 
 #! Sets an item in the account storage. Panics if the index is out of bounds.
 #!
-#! Stack: [index, V']
+#! Stack: [index, V', 0, 0, 0]
 #! Output: [R', V]
 #!
 #! - index is the index of the item to set.
 #! - V' is the value to set.
 #! - V is the previous value of the item.
 #! - R' is the new storage root.
 export.set_item
-    push.0 movdn.5 push.0 movdn.5 push.0 movdn.5
-    # => [index, V', 0, 0, 0]
-
     syscall.set_account_item
     # => [R', V]
 end
 
-#! Gets a map item from the account storage. Panics if 
+#! Gets a map item from the account storage. Panics if
 #! - the index for the map is out of bounds, means >255
 #! - the slot item at index is not a map
 #!
@@ -120,7 +117,7 @@ export.get_map_item
     # => [VALUE, 0]
 end
 
-#! Sets a map item in the account storage. Panics if 
+#! Sets a map item in the account storage. Panics if
 #! - the index for the map is out of bounds, means >255
 #! - the slot item at index is not a map
 #!
@@ -134,11 +131,10 @@ end
 #! - OLD_MAP_VALUE is the old value at KEY.
 export.set_map_item
     syscall.set_account_map_item
-    # => [OLD_MAP_ROOT, OLD_MAP_VALUE]
-
-    # prepare stack for return
-    push.0 movdn.9
     # => [OLD_MAP_ROOT, OLD_MAP_VALUE, 0]
+
+    movup.8 drop
+    # => [OLD_MAP_ROOT, OLD_MAP_VALUE]
 end
 
 #! Sets the code of the account the transaction is being executed against. This procedure can only

miden-lib/asm/miden/contracts/wallets/basic.masm

@@ -73,3 +73,27 @@ export.send_asset.1
     # prepare the stack for return - stack has 6 elements too many
     movupw.3 dropw swap drop swap drop
 end
+
+#! Increments the account nonce by the provided value.
+#!
+#! Stack: [value]
+#! Output: []
+#!
+#! - value is the value to increment the nonce by. value can be at most 2^32 - 1 otherwise this
+#!   procedure panics.
+export.incr_nonce
+    exec.account::incr_nonce
+    # => []
+end
+
+#! Sets the code of the account the transaction is being executed against. This procedure can only
+#! executed on regular accounts with updatable code. Otherwise, this procedure fails.
+#!
+#! Stack: [CODE_ROOT]
+#! Output: []
+#!
+#! - CODE_ROOT is the hash of the code to set.
+export.set_code
+    exec.account::set_code
+    # => []
+end

miden-lib/asm/miden/kernels/tx/memory.masm

@@ -633,8 +633,12 @@ end
 #! Where:
 #! - acct_nonce is the account nonce.
 export.set_acct_nonce
+    # load the data stored in the same word as the nonce
     padw push.ACCT_ID_AND_NONCE_PTR mem_loadw
+    # => [DATA, acct_nonce]
+
     drop movup.3 push.ACCT_ID_AND_NONCE_PTR mem_storew dropw
+    # => []
 end
 
 #! Sets the code root of the account.

miden-lib/src/transaction/errors.rs

@@ -98,7 +98,7 @@ impl fmt::Display for TransactionKernelError {
                 )
             },
             TransactionKernelError::UnknownAccountProcedure(proc_root) => {
-                write!(f, "account procedure with root {proc_root} is not in the advice provider")
+                write!(f, "Could not find account procedure with digest {proc_root} is in the advice provider, check the transactinon's account code")
             },
         }
     }

miden-tx/src/testing/tx_context.rs

@@ -51,6 +51,9 @@ pub struct TransactionContext {
 }
 
 /// Data required to compile the [miden_objects::transaction::TransactionScript] under the [TransactionContext].
+///
+/// This allows the account code to be loaded in the [Assembler] so that the transaction script can
+/// the account's procedures.
 pub enum ScriptAndInputs {
     Empty,
     Some {

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

@@ -10,15 +10,19 @@ use miden_objects::{
         AccountId, AccountStorage, AccountType, StorageSlotType,
     },
     crypto::{hash::rpo::RpoDigest, merkle::LeafIndex},
-    testing::{prepare_word, storage::STORAGE_LEAVES_2},
+    testing::{
+        account_code::{ACCOUNT_INCR_NONCE_MAST_ROOT, ACCOUNT_SET_CODE_MAST_ROOT},
+        prepare_word,
+        storage::STORAGE_LEAVES_2,
+    },
     transaction::TransactionArgs,
 };
-use vm_processor::{Felt, MemAdviceProvider};
+use vm_processor::{AdviceMap, Felt, MemAdviceProvider};
 
 use super::{ProcessState, StackInputs, Word, ONE, ZERO};
 use crate::{
-    testing::{executor::CodeExecutor, TransactionContextBuilder},
-    tests::kernel_tests::{output_notes_data_procedure, read_root_mem_value},
+    testing::{executor::CodeExecutor, ScriptAndInputs, TransactionContextBuilder},
+    tests::kernel_tests::read_root_mem_value,
 };
 
 // ACCOUNT CODE TESTS
@@ -65,37 +69,33 @@ pub fn test_set_code_succeeds() {
     .with_mock_notes_preserved()
     .build();
 
-    let output_notes_data_procedure =
-        output_notes_data_procedure(tx_context.expected_output_notes());
-
-    let code = format!(
+    let tx_script = format!(
         "
         use.miden::account
-        use.miden::kernels::tx::prologue
-        use.miden::kernels::tx::epilogue
 
-        {output_notes_data_procedure}
         begin
-            exec.prologue::prepare_transaction
-
             push.0.1.2.3
-            exec.account::set_code
-
-            exec.create_mock_notes
+            call.{ACCOUNT_SET_CODE_MAST_ROOT}
+            dropw
 
             push.1
-            exec.account::incr_nonce
-
-            exec.epilogue::finalize_transaction
+            call.{ACCOUNT_INCR_NONCE_MAST_ROOT}
+            drop
         end
         "
     );
 
-    let process = tx_context.execute_with_custom_main(&code, TransactionArgs::default()).unwrap();
+    let executed_transaction = tx_context
+        .execute_transaction(
+            None,
+            AdviceMap::default(),
+            ScriptAndInputs::new(tx_script, vec![], vec![]),
+        )
+        .unwrap();
 
     assert_eq!(
-        read_root_mem_value(&process, ACCT_CODE_ROOT_PTR),
-        [ZERO, ONE, Felt::new(2), Felt::new(3)],
+        executed_transaction.final_account().code_root(),
+        RpoDigest::from([ZERO, ONE, Felt::new(2), Felt::new(3)]),
         "the code root must change after the epilogue",
     );
 }