perf: avoid double-hashing the asset vault key in fungible add/remove

CHANGELOG.md

@@ -33,6 +33,7 @@
 - Optimized `rbac::grant_role_internal` and `rbac::revoke_role_internal` by removing the redundant membership read and rearranging the stack ([#3090](https://github.com/0xMiden/protocol/pull/3090)).
 - [BREAKING] Refactored `TokenPolicyManager` by adding `invoke_send_policy` / `invoke_receive_policy` wrappers (stored in the protocol reserved asset callback slots) that read the active policy root from the new `active_send_policy_proc_root` / `active_receive_policy_proc_root` storage slots ([#3047](https://github.com/0xMiden/protocol/pull/3047)).
 - [BREAKING] Flipped `AuthSingleSigAcl` ACL to an exempt list: every called procedure now requires a signature unless its root is in `exempt_procedures` [#3065](https://github.com/0xMiden/protocol/pull/3065).
+- [BREAKING] Changed `asset_vault::peek_asset` to accept a pre-hashed `ASSET_KEY_HASH` instead of a raw `ASSET_KEY`; fungible add/remove now hash the vault key once internally, eliminating a redundant `poseidon2::hash` per operation ([#3073](https://github.com/0xMiden/protocol/pull/3073)).
 - Added regression tests ensuring a `TokenPolicyManager` with only reserved send/receive policies installs the protocol reserved asset callback slots, so `has_callbacks` is correct from creation and minted assets carry the callback flag ([#3091](https://github.com/0xMiden/protocol/pull/3091)).
 - Fixed the `TokenPolicyManager` `get_mint_policy` / `get_burn_policy` / `get_send_policy` / `get_receive_policy` getters to align the 16-felt `call` ABI. ([#3114](https://github.com/0xMiden/protocol/pull/3114)).
 - Added a definition of the Miden operator on the architecture overview page and linked it from the note lifecycle ([#3017](https://github.com/0xMiden/protocol/pull/3017)).

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

@@ -45,7 +45,7 @@ pub proc get_asset
     # => [ASSET_VALUE]
 end
 
-#! Returns the _peeked_ asset associated with the provided asset vault key.
+#! Returns the _peeked_ asset associated with the provided hashed asset vault key.
 #!
 #! WARNING: Peeked means the asset is loaded from the advice provider, which is susceptible to
 #! manipulation from a malicious host. Therefore this should only be used when the inclusion of the
@@ -61,18 +61,14 @@ end
 #! merkle paths from the merkle store, since this is only possible for the account vault. Ensure
 #! that the merkle paths are present prior to calling.
 #!
-#! Inputs:  [ASSET_KEY, vault_root_ptr]
+#! Inputs:  [ASSET_KEY_HASH, vault_root_ptr]
 #! Outputs: [ASSET_VALUE]
 #!
 #! Where:
 #! - vault_root_ptr is a pointer to the memory location at which the vault root is stored.
-#! - ASSET_KEY is the asset vault key of the asset to fetch.
+#! - ASSET_KEY_HASH is the hashed asset vault key of the asset to fetch (see hash_asset_key).
 #! - ASSET_VALUE is the retrieved asset.
 pub proc peek_asset
-    # hash the asset vault key before using it as the SMT key
-    exec.hash_asset_key
-    # => [ASSET_KEY_HASH, vault_root_ptr]
-
     # load the asset vault root from memory
     padw movup.8 mem_loadw_le
     # => [ASSET_VAULT_ROOT, ASSET_KEY_HASH]
@@ -130,41 +126,45 @@ pub proc add_fungible_asset
     movup.8 loc_store.0
     # => [ASSET_KEY, ASSET_VALUE]
 
+    # hash the asset vault key once; it is used as the SMT key for both the peek and the set below
+    exec.hash_asset_key
+    # => [ASSET_KEY_HASH, ASSET_VALUE]
+
     dupw loc_load.0 movdn.4
-    # => [ASSET_KEY, vault_root_ptr, ASSET_KEY, ASSET_VALUE]
+    # => [ASSET_KEY_HASH, vault_root_ptr, ASSET_KEY_HASH, ASSET_VALUE]
 
     exec.peek_asset
-    # => [INITIAL_ASSET_VALUE, ASSET_KEY, ASSET_VALUE]
+    # => [INITIAL_ASSET_VALUE, ASSET_KEY_HASH, ASSET_VALUE]
 
     # since we have peeked the value, we need to later assert that the actual value matches this
     # one, so we'll keep a copy for later
     swapw dupw.1
-    # => [INITIAL_ASSET_VALUE, ASSET_KEY, INITIAL_ASSET_VALUE, ASSET_VALUE]
+    # => [INITIAL_ASSET_VALUE, ASSET_KEY_HASH, INITIAL_ASSET_VALUE, ASSET_VALUE]
 
     movupw.3
-    # => [ASSET_VALUE, INITIAL_ASSET_VALUE, ASSET_KEY, INITIAL_ASSET_VALUE]
+    # => [ASSET_VALUE, INITIAL_ASSET_VALUE, ASSET_KEY_HASH, INITIAL_ASSET_VALUE]
 
     # Merge the assets.
     # ---------------------------------------------------------------------------------------------
 
     exec.fungible_asset::merge
-    # => [FINAL_ASSET_VALUE, ASSET_KEY, INITIAL_ASSET_VALUE]
+    # => [FINAL_ASSET_VALUE, ASSET_KEY_HASH, INITIAL_ASSET_VALUE]
 
     swapw
-    # => [ASSET_KEY, FINAL_ASSET_VALUE, INITIAL_ASSET_VALUE]
+    # => [ASSET_KEY_HASH, FINAL_ASSET_VALUE, INITIAL_ASSET_VALUE]
 
     # Insert the merged asset.
     # ---------------------------------------------------------------------------------------------
 
     # load the vault root
     padw loc_load.0 mem_loadw_le
-    # => [VAULT_ROOT, ASSET_KEY, FINAL_ASSET_VALUE, INITIAL_ASSET_VALUE]
+    # => [VAULT_ROOT, ASSET_KEY_HASH, FINAL_ASSET_VALUE, INITIAL_ASSET_VALUE]
 
     swapw dupw.2
-    # => [FINAL_ASSET_VALUE, ASSET_KEY, VAULT_ROOT, FINAL_ASSET_VALUE, INITIAL_ASSET_VALUE]
+    # => [FINAL_ASSET_VALUE, ASSET_KEY_HASH, VAULT_ROOT, FINAL_ASSET_VALUE, INITIAL_ASSET_VALUE]
 
-    # hash the asset key and update the asset in the vault
-    exec.set_asset
+    # update the asset in the vault
+    exec.smt::set
     # => [PREV_VAULT_VALUE, NEW_VAULT_ROOT, FINAL_ASSET_VALUE, INITIAL_ASSET_VALUE]
 
     # assert PREV_VAULT_VALUE = INITIAL_ASSET_VALUE to make sure peek_asset returned the correct asset
@@ -194,20 +194,24 @@ end
 #! Panics if:
 #! - the vault already contains the same non-fungible asset.
 pub proc add_non_fungible_asset
+    # hash the asset vault key before using it as the SMT key
+    exec.hash_asset_key
+    # => [ASSET_KEY_HASH, ASSET_VALUE, vault_root_ptr]
+
     # Load VAULT_ROOT and insert asset.
     # ---------------------------------------------------------------------------------------------
 
     padw dup.12
-    # => [vault_root_ptr, pad(4), ASSET_KEY, ASSET_VALUE, vault_root_ptr]
+    # => [vault_root_ptr, pad(4), ASSET_KEY_HASH, ASSET_VALUE, vault_root_ptr]
 
     mem_loadw_le swapw
-    # => [ASSET_KEY, VAULT_ROOT, ASSET_VALUE, vault_root_ptr]
+    # => [ASSET_KEY_HASH, VAULT_ROOT, ASSET_VALUE, vault_root_ptr]
 
     dupw.2
-    # => [ASSET_VALUE, ASSET_KEY, VAULT_ROOT, ASSET_VALUE, vault_root_ptr]
+    # => [ASSET_VALUE, ASSET_KEY_HASH, VAULT_ROOT, ASSET_VALUE, vault_root_ptr]
 
-    # hash the asset key and insert the asset into the vault
-    exec.set_asset
+    # insert the asset into the vault
+    exec.smt::set
     # => [OLD_VAL, VAULT_ROOT', ASSET_VALUE, vault_root_ptr]
 
     # assert old value was empty
@@ -296,38 +300,42 @@ end
 #! - the amount of the asset in the vault is less than the amount to be removed.
 @locals(4)
 pub proc remove_fungible_asset
+    # hash the asset vault key once; it is used as the SMT key for both the peek and the set below
+    exec.hash_asset_key
+    # => [ASSET_KEY_HASH, ASSET_VALUE, vault_root_ptr]
+
     dupw movdnw.2
-    # => [ASSET_KEY, ASSET_VALUE, ASSET_KEY, vault_root_ptr]
+    # => [ASSET_KEY_HASH, ASSET_VALUE, ASSET_KEY_HASH, vault_root_ptr]
 
     dup.12 movdn.4
-    # => [ASSET_KEY, vault_root_ptr, ASSET_VALUE, ASSET_KEY, vault_root_ptr]
+    # => [ASSET_KEY_HASH, vault_root_ptr, ASSET_VALUE, ASSET_KEY_HASH, vault_root_ptr]
 
     exec.peek_asset
-    # => [INITIAL_ASSET_VALUE, ASSET_VALUE, ASSET_KEY, vault_root_ptr]
+    # => [INITIAL_ASSET_VALUE, ASSET_VALUE, ASSET_KEY_HASH, vault_root_ptr]
 
     movdnw.2
-    # => [ASSET_VALUE, ASSET_KEY, INITIAL_ASSET_VALUE, vault_root_ptr]
+    # => [ASSET_VALUE, ASSET_KEY_HASH, INITIAL_ASSET_VALUE, vault_root_ptr]
 
     dupw.2 swapw
-    # => [ASSET_VALUE, INITIAL_ASSET_VALUE, ASSET_KEY, INITIAL_ASSET_VALUE, vault_root_ptr]
+    # => [ASSET_VALUE, INITIAL_ASSET_VALUE, ASSET_KEY_HASH, INITIAL_ASSET_VALUE, vault_root_ptr]
 
     # compute FINAL_ASSET_VALUE = INITIAL_ASSET_VALUE - ASSET_VALUE
     exec.fungible_asset::split
-    # => [FINAL_ASSET_VALUE, ASSET_KEY, INITIAL_ASSET_VALUE, vault_root_ptr]
+    # => [FINAL_ASSET_VALUE, ASSET_KEY_HASH, INITIAL_ASSET_VALUE, vault_root_ptr]
 
     # store FINAL_ASSET_VALUE so we can return it at the end
     loc_storew_le.0
-    # => [FINAL_ASSET_VALUE, ASSET_KEY, INITIAL_ASSET_VALUE, vault_root_ptr]
+    # => [FINAL_ASSET_VALUE, ASSET_KEY_HASH, INITIAL_ASSET_VALUE, vault_root_ptr]
 
     dup.12 padw movup.4 mem_loadw_le
-    # => [VAULT_ROOT, FINAL_ASSET_VALUE, ASSET_KEY, INITIAL_ASSET_VALUE, vault_root_ptr]
+    # => [VAULT_ROOT, FINAL_ASSET_VALUE, ASSET_KEY_HASH, INITIAL_ASSET_VALUE, vault_root_ptr]
 
     movdnw.2
-    # => [FINAL_ASSET_VALUE, ASSET_KEY, VAULT_ROOT, INITIAL_ASSET_VALUE, vault_root_ptr]
+    # => [FINAL_ASSET_VALUE, ASSET_KEY_HASH, VAULT_ROOT, INITIAL_ASSET_VALUE, vault_root_ptr]
 
-    # hash the asset key and update the asset in the vault; the old value is asserted below to be
-    # equivalent to the peeked value provided via peek_asset
-    exec.set_asset
+    # update the asset in the vault; the old value is asserted below to be equivalent to the
+    # peeked value provided via peek_asset
+    exec.smt::set
     # => [OLD_VALUE, NEW_VAULT_ROOT, INITIAL_ASSET_VALUE, vault_root_ptr]
 
     dupw.2
@@ -371,16 +379,20 @@ end
 #! Panics if:
 #! - the non-fungible asset is not found in the vault.
 pub proc remove_non_fungible_asset
+    # hash the asset vault key before using it as the SMT key
+    exec.hash_asset_key
+    # => [ASSET_KEY_HASH, ASSET_VALUE, vault_root_ptr]
+
     # load vault root
     padw dup.12 mem_loadw_le
-    # => [VAULT_ROOT, ASSET_KEY, ASSET_VALUE, vault_root_ptr]
+    # => [VAULT_ROOT, ASSET_KEY_HASH, ASSET_VALUE, vault_root_ptr]
 
-    # prepare insertion of an EMPTY_WORD into the vault at the asset key to remove the asset
+    # prepare insertion of an EMPTY_WORD into the vault at the hashed asset key to remove the asset
     swapw padw
-    # => [EMPTY_WORD, ASSET_KEY, VAULT_ROOT, ASSET_VALUE, vault_root_ptr]
+    # => [EMPTY_WORD, ASSET_KEY_HASH, VAULT_ROOT, ASSET_VALUE, vault_root_ptr]
 
-    # hash the asset key and insert the empty word into the vault to remove the asset
-    exec.set_asset
+    # insert the empty word into the vault to remove the asset
+    exec.smt::set
     # => [REMOVED_ASSET_VALUE, NEW_VAULT_ROOT, ASSET_VALUE, vault_root_ptr]
 
     # dup ASSET_VALUE so it survives the assert; the assert proves it equals REMOVED_ASSET_VALUE
@@ -451,23 +463,3 @@ proc hash_asset_key
     exec.poseidon2::hash
     # => [ASSET_KEY_HASH]
 end
-
-#! Hashes the raw asset vault key and writes ASSET_VALUE into the asset vault SMT at the hashed key,
-#! returning the previous value stored there.
-#!
-#! Inputs:  [ASSET_VALUE, ASSET_KEY, VAULT_ROOT]
-#! Outputs: [OLD_VALUE, NEW_VAULT_ROOT]
-#!
-#! Where:
-#! - ASSET_KEY is the raw (unhashed) asset vault key.
-#! - ASSET_VALUE is the value to write into the vault at the hashed key.
-#! - VAULT_ROOT is the current root of the asset vault SMT.
-#! - OLD_VALUE is the value previously stored at the hashed key.
-#! - NEW_VAULT_ROOT is the root of the asset vault SMT after the write.
-proc set_asset
-    swapw exec.hash_asset_key swapw
-    # => [ASSET_VALUE, ASSET_KEY_HASH, VAULT_ROOT]
-
-    exec.smt::set
-    # => [OLD_VALUE, NEW_VAULT_ROOT]
-end

crates/miden-testing/src/kernel_tests/tx/test_asset_vault.rs

@@ -78,6 +78,7 @@ async fn peek_asset_returns_correct_asset() -> anyhow::Result<()> {
         use $kernel::prologue
         use $kernel::memory
         use $kernel::asset_vault
+        use miden::core::crypto::hashes::poseidon2
 
         begin
             exec.prologue::prepare_transaction
@@ -87,10 +88,14 @@ async fn peek_asset_returns_correct_asset() -> anyhow::Result<()> {
             # => [ASSET_KEY, account_vault_root_ptr]
 
             # emit an event to fetch the merkle path for the asset since peek_asset does not do
-            # that
+            # that. the event handler expects the raw ASSET_KEY, so it is emitted before hashing.
             emit.event("miden::protocol::account::vault_before_get_asset")
             # => [ASSET_KEY, account_vault_root_ptr]
 
+            # hash the asset vault key into the SMT key that peek_asset expects
+            exec.poseidon2::hash
+            # => [ASSET_KEY_HASH, account_vault_root_ptr]
+
             exec.asset_vault::peek_asset
             # => [PEEKED_ASSET_VALUE]