From 0a755fdcceca078fd6e909c5e30f37fdabb6883b Mon Sep 17 00:00:00 2001 From: varun-doshi Date: Sat, 29 Nov 2025 23:59:51 +0530 Subject: [PATCH] feat: Introduce Word wrapper for storage map keys --- .../src/account/auth/ecdsa_k256_keccak_acl.rs | 6 +- .../auth/ecdsa_k256_keccak_multisig.rs | 16 ++- .../src/account/auth/rpo_falcon_512_acl.rs | 6 +- .../account/auth/rpo_falcon_512_multisig.rs | 16 ++- .../template/storage/entry_content.rs | 18 ++-- .../account/component/template/storage/mod.rs | 4 +- .../src/account/delta/storage.rs | 2 +- crates/miden-objects/src/account/mod.rs | 20 +++- .../storage/map/hashed_storage_map_key.rs | 37 +++++++ .../src/account/storage/map/mod.rs | 59 +++++----- .../src/account/storage/map/partial.rs | 17 +-- .../account/storage/map/storage_map_key.rs | 101 ++++++++++++++++++ .../src/account/storage/map/witness.rs | 20 ++-- .../miden-objects/src/account/storage/mod.rs | 6 +- .../src/account/storage/partial.rs | 12 +-- crates/miden-objects/src/errors.rs | 9 +- crates/miden-objects/src/testing/storage.rs | 17 ++- .../src/kernel_tests/tx/test_account.rs | 7 +- .../src/kernel_tests/tx/test_lazy_loading.rs | 11 +- .../miden-testing/src/tx_context/context.rs | 4 +- crates/miden-tx/src/executor/exec_host.rs | 2 +- .../src/host/storage_delta_tracker.rs | 2 +- crates/miden-tx/src/host/tx_event.rs | 14 +-- 23 files changed, 288 insertions(+), 118 deletions(-) create mode 100644 crates/miden-objects/src/account/storage/map/hashed_storage_map_key.rs create mode 100644 crates/miden-objects/src/account/storage/map/storage_map_key.rs diff --git a/crates/miden-lib/src/account/auth/ecdsa_k256_keccak_acl.rs b/crates/miden-lib/src/account/auth/ecdsa_k256_keccak_acl.rs index 186ebd026f..5dec5fa540 100644 --- a/crates/miden-lib/src/account/auth/ecdsa_k256_keccak_acl.rs +++ b/crates/miden-lib/src/account/auth/ecdsa_k256_keccak_acl.rs @@ -159,7 +159,7 @@ impl From for AccountComponent { .auth_trigger_procedures .iter() .enumerate() - .map(|(i, proc_root)| (Word::from([i as u32, 0, 0, 0]), *proc_root)); + .map(|(i, proc_root)| (Word::from([i as u32, 0, 0, 0]).into(), *proc_root)); // Safe to unwrap because we know that the map keys are unique. storage_slots.push(StorageSlot::Map(StorageMap::with_entries(map_entries).unwrap())); @@ -242,7 +242,7 @@ mod tests { for (i, expected_proc_root) in auth_trigger_procedures.iter().enumerate() { let proc_root = account .storage() - .get_map_item(2, Word::from([i as u32, 0, 0, 0])) + .get_map_item(2, Word::from([i as u32, 0, 0, 0]).into()) .expect("storage map access failed"); assert_eq!(proc_root, *expected_proc_root); } @@ -250,7 +250,7 @@ mod tests { // When no procedures, the map should return empty for key [0,0,0,0] let proc_root = account .storage() - .get_map_item(2, Word::empty()) + .get_map_item(2, Word::empty().into()) .expect("storage map access failed"); assert_eq!(proc_root, Word::empty()); } diff --git a/crates/miden-lib/src/account/auth/ecdsa_k256_keccak_multisig.rs b/crates/miden-lib/src/account/auth/ecdsa_k256_keccak_multisig.rs index 7bb653cb62..4e02151eef 100644 --- a/crates/miden-lib/src/account/auth/ecdsa_k256_keccak_multisig.rs +++ b/crates/miden-lib/src/account/auth/ecdsa_k256_keccak_multisig.rs @@ -2,7 +2,7 @@ use alloc::collections::BTreeSet; use alloc::vec::Vec; use miden_objects::account::auth::PublicKeyCommitment; -use miden_objects::account::{AccountComponent, StorageMap, StorageSlot}; +use miden_objects::account::{AccountComponent, StorageMap, StorageMapKey, StorageSlot}; use miden_objects::{AccountError, Word}; use crate::account::components::ecdsa_k256_keccak_multisig_library; @@ -125,7 +125,7 @@ impl From for AccountComponent { .approvers() .iter() .enumerate() - .map(|(i, pub_key)| (Word::from([i as u32, 0, 0, 0]), (*pub_key).into())); + .map(|(i, pub_key)| (Word::from([i as u32, 0, 0, 0]).into(), (*pub_key).into())); // Safe to unwrap because we know that the map keys are unique. storage_slots.push(StorageSlot::Map(StorageMap::with_entries(map_entries).unwrap())); @@ -136,11 +136,9 @@ impl From for AccountComponent { // Slot 3: A map which stores procedure thresholds (PROC_ROOT -> threshold) let proc_threshold_roots = StorageMap::with_entries( - multisig - .config - .proc_thresholds() - .iter() - .map(|(proc_root, threshold)| (*proc_root, Word::from([*threshold, 0, 0, 0]))), + multisig.config.proc_thresholds().iter().map(|(proc_root, threshold)| { + (StorageMapKey::from(*proc_root), Word::from([*threshold, 0, 0, 0])) + }), ) .unwrap(); storage_slots.push(StorageSlot::Map(proc_threshold_roots)); @@ -193,7 +191,7 @@ mod tests { for (i, expected_pub_key) in approvers.iter().enumerate() { let stored_pub_key = account .storage() - .get_map_item(1, Word::from([i as u32, 0, 0, 0])) + .get_map_item(1, Word::from([i as u32, 0, 0, 0]).into()) .expect("storage map access failed"); assert_eq!(stored_pub_key, Word::from(*expected_pub_key)); } @@ -224,7 +222,7 @@ mod tests { let stored_pub_key = account .storage() - .get_map_item(1, Word::from([0u32, 0, 0, 0])) + .get_map_item(1, Word::from([0u32, 0, 0, 0]).into()) .expect("storage map access failed"); assert_eq!(stored_pub_key, Word::from(pub_key)); } diff --git a/crates/miden-lib/src/account/auth/rpo_falcon_512_acl.rs b/crates/miden-lib/src/account/auth/rpo_falcon_512_acl.rs index f9982d5456..41a402402b 100644 --- a/crates/miden-lib/src/account/auth/rpo_falcon_512_acl.rs +++ b/crates/miden-lib/src/account/auth/rpo_falcon_512_acl.rs @@ -160,7 +160,7 @@ impl From for AccountComponent { .auth_trigger_procedures .iter() .enumerate() - .map(|(i, proc_root)| (Word::from([i as u32, 0, 0, 0]), *proc_root)); + .map(|(i, proc_root)| (Word::from([i as u32, 0, 0, 0]).into(), *proc_root)); // Safe to unwrap because we know that the map keys are unique. storage_slots.push(StorageSlot::Map(StorageMap::with_entries(map_entries).unwrap())); @@ -243,7 +243,7 @@ mod tests { for (i, expected_proc_root) in auth_trigger_procedures.iter().enumerate() { let proc_root = account .storage() - .get_map_item(2, Word::from([i as u32, 0, 0, 0])) + .get_map_item(2, Word::from([i as u32, 0, 0, 0]).into()) .expect("storage map access failed"); assert_eq!(proc_root, *expected_proc_root); } @@ -251,7 +251,7 @@ mod tests { // When no procedures, the map should return empty for key [0,0,0,0] let proc_root = account .storage() - .get_map_item(2, Word::empty()) + .get_map_item(2, Word::empty().into()) .expect("storage map access failed"); assert_eq!(proc_root, Word::empty()); } diff --git a/crates/miden-lib/src/account/auth/rpo_falcon_512_multisig.rs b/crates/miden-lib/src/account/auth/rpo_falcon_512_multisig.rs index 2d933d990a..5173f2aa63 100644 --- a/crates/miden-lib/src/account/auth/rpo_falcon_512_multisig.rs +++ b/crates/miden-lib/src/account/auth/rpo_falcon_512_multisig.rs @@ -2,7 +2,7 @@ use alloc::collections::BTreeSet; use alloc::vec::Vec; use miden_objects::account::auth::PublicKeyCommitment; -use miden_objects::account::{AccountComponent, StorageMap, StorageSlot}; +use miden_objects::account::{AccountComponent, StorageMap, StorageMapKey, StorageSlot}; use miden_objects::{AccountError, Word}; use crate::account::components::rpo_falcon_512_multisig_library; @@ -125,7 +125,7 @@ impl From for AccountComponent { .approvers() .iter() .enumerate() - .map(|(i, pub_key)| (Word::from([i as u32, 0, 0, 0]), (*pub_key).into())); + .map(|(i, pub_key)| (Word::from([i as u32, 0, 0, 0]).into(), (*pub_key).into())); // Safe to unwrap because we know that the map keys are unique. storage_slots.push(StorageSlot::Map(StorageMap::with_entries(map_entries).unwrap())); @@ -136,11 +136,9 @@ impl From for AccountComponent { // Slot 3: A map which stores procedure thresholds (PROC_ROOT -> threshold) let proc_threshold_roots = StorageMap::with_entries( - multisig - .config - .proc_thresholds() - .iter() - .map(|(proc_root, threshold)| (*proc_root, Word::from([*threshold, 0, 0, 0]))), + multisig.config.proc_thresholds().iter().map(|(proc_root, threshold)| { + (StorageMapKey::from(*proc_root), Word::from([*threshold, 0, 0, 0])) + }), ) .unwrap(); storage_slots.push(StorageSlot::Map(proc_threshold_roots)); @@ -193,7 +191,7 @@ mod tests { for (i, expected_pub_key) in approvers.iter().enumerate() { let stored_pub_key = account .storage() - .get_map_item(1, Word::from([i as u32, 0, 0, 0])) + .get_map_item(1, Word::from([i as u32, 0, 0, 0]).into()) .expect("storage map access failed"); assert_eq!(stored_pub_key, Word::from(*expected_pub_key)); } @@ -224,7 +222,7 @@ mod tests { let stored_pub_key = account .storage() - .get_map_item(1, Word::from([0u32, 0, 0, 0])) + .get_map_item(1, Word::from([0u32, 0, 0, 0]).into()) .expect("storage map access failed"); assert_eq!(stored_pub_key, Word::from(pub_key)); } diff --git a/crates/miden-objects/src/account/component/template/storage/entry_content.rs b/crates/miden-objects/src/account/component/template/storage/entry_content.rs index 54780435b3..1709a9645e 100644 --- a/crates/miden-objects/src/account/component/template/storage/entry_content.rs +++ b/crates/miden-objects/src/account/component/template/storage/entry_content.rs @@ -12,8 +12,8 @@ use super::{ StorageValueName, TemplateRequirementsIter, }; -use crate::account::StorageMap; use crate::account::component::template::AccountComponentTemplateError; +use crate::account::{StorageMap, StorageMapKey}; use crate::utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable}; use crate::{Felt, FieldElement, Word}; @@ -588,15 +588,17 @@ impl MapRepresentation { let entries = entries .iter() .map(|map_entry| { - let key = map_entry - .key() - .try_build_word(init_storage_data, identifier.name.clone())?; + let key = StorageMapKey::new_unchecked( + map_entry + .key() + .try_build_word(init_storage_data, identifier.name.clone())?, + ); let value = map_entry .value() .try_build_word(init_storage_data, identifier.name.clone())?; Ok((key, value)) }) - .collect::, _>>()?; + .collect::, _>>()?; StorageMap::with_entries(entries).map_err(|err| { AccountComponentTemplateError::StorageMapHasDuplicateKeys(Box::new(err)) @@ -604,7 +606,11 @@ impl MapRepresentation { }, MapRepresentation::Template { identifier } => { if let Some(entries) = init_storage_data.map_entries(&identifier.name) { - return StorageMap::with_entries(entries.clone()).map_err(|err| { + return StorageMap::with_entries(entries.into_iter().map(|map_entry| { + let key = StorageMapKey::new_unchecked(map_entry.0); + (key, map_entry.1) + })) + .map_err(|err| { AccountComponentTemplateError::StorageMapHasDuplicateKeys(Box::new(err)) }); } diff --git a/crates/miden-objects/src/account/component/template/storage/mod.rs b/crates/miden-objects/src/account/component/template/storage/mod.rs index e1380e025f..7fa2b8ca6b 100644 --- a/crates/miden-objects/src/account/component/template/storage/mod.rs +++ b/crates/miden-objects/src/account/component/template/storage/mod.rs @@ -746,7 +746,7 @@ mod tests { "0x0000000000000000000000000000000000000000000000000000000000000010", ) .unwrap(); - assert_eq!(storage_map.get(&main_key), main_value_expected); + assert_eq!(storage_map.get(&main_key.into()), main_value_expected); }, _ => panic!("expected map storage slot"), } @@ -917,7 +917,7 @@ mod tests { "0x0000000000000000000000000000000000000000000000000000000000000002", ) .unwrap(); - assert_eq!(storage_map.get(&second_key), second_value); + assert_eq!(storage_map.get(&second_key.into()), second_value); }, _ => panic!("expected map storage slot"), } diff --git a/crates/miden-objects/src/account/delta/storage.rs b/crates/miden-objects/src/account/delta/storage.rs index baeda54931..e763bf6507 100644 --- a/crates/miden-objects/src/account/delta/storage.rs +++ b/crates/miden-objects/src/account/delta/storage.rs @@ -381,7 +381,7 @@ impl From for StorageMapDelta { StorageMapDelta::new( map.into_entries() .into_iter() - .map(|(key, value)| (LexicographicWord::new(key), value)) + .map(|(key, value)| (LexicographicWord::from(key), value)) .collect(), ) } diff --git a/crates/miden-objects/src/account/mod.rs b/crates/miden-objects/src/account/mod.rs index 0918e43ce8..1595e4ccdd 100644 --- a/crates/miden-objects/src/account/mod.rs +++ b/crates/miden-objects/src/account/mod.rs @@ -70,6 +70,7 @@ pub use storage::{ PartialStorageMap, SlotName, StorageMap, + StorageMapKey, StorageMapWitness, StorageSlot, StorageSlotType, @@ -449,7 +450,7 @@ impl TryFrom for AccountDelta { storage_map .into_entries() .into_iter() - .map(|(key, value)| (LexicographicWord::from(key), value)) + .map(|(key, value)| (LexicographicWord::from(key.inner()), value)) .collect(), ); map_slots.insert(slot_idx, map_delta); @@ -639,6 +640,7 @@ mod tests { PartialAccount, StorageMap, StorageMapDelta, + StorageMapKey, StorageSlot, }; use crate::asset::{Asset, AssetVault, FungibleAsset, NonFungibleAsset}; @@ -700,7 +702,12 @@ mod tests { let storage_slot_value_1 = StorageSlot::Value(Word::from([5, 6, 7, 8u32])); let mut storage_map = StorageMap::with_entries([ ( - Word::new([Felt::new(101), Felt::new(102), Felt::new(103), Felt::new(104)]), + StorageMapKey::from(Word::new([ + Felt::new(101), + Felt::new(102), + Felt::new(103), + Felt::new(104), + ])), Word::from([ Felt::new(1_u64), Felt::new(2_u64), @@ -709,7 +716,12 @@ mod tests { ]), ), ( - Word::new([Felt::new(105), Felt::new(106), Felt::new(107), Felt::new(108)]), + StorageMapKey::from(Word::new([ + Felt::new(105), + Felt::new(106), + Felt::new(107), + Felt::new(108), + ])), Word::new([Felt::new(5_u64), Felt::new(6_u64), Felt::new(7_u64), Felt::new(8_u64)]), ), ]) @@ -730,7 +742,7 @@ mod tests { let updated_map = StorageMapDelta::from_iters([], [(new_map_entry.0, new_map_entry.1.into())]); - storage_map.insert(new_map_entry.0, new_map_entry.1.into()).unwrap(); + storage_map.insert(new_map_entry.0.into(), new_map_entry.1.into()).unwrap(); // build account delta let final_nonce = Felt::new(2); diff --git a/crates/miden-objects/src/account/storage/map/hashed_storage_map_key.rs b/crates/miden-objects/src/account/storage/map/hashed_storage_map_key.rs new file mode 100644 index 0000000000..b98033f9fa --- /dev/null +++ b/crates/miden-objects/src/account/storage/map/hashed_storage_map_key.rs @@ -0,0 +1,37 @@ +use core::fmt; + +use miden_crypto::Word; + +use crate::Felt; + +#[derive(Debug, PartialEq, Eq, Clone, Copy, PartialOrd, Ord)] +pub struct HashedStorageMapKey(Word); + +impl HashedStorageMapKey { + /// Creates a new [`HashedStorageMapKey`] from the given [`Word`] **without performing + /// validation**. + /// + /// ## Warning + /// + /// This function **does not check** whether the provided `Word` represents a valid + /// fungible or non-fungible asset key. + pub fn new_unchecked(value: Word) -> Self { + Self(value) + } + + pub fn inner(&self) -> Word { + self.0 + } + + /// Returns the leaf index of a map key. + pub fn hashed_map_key_to_leaf_index(&self) -> Felt { + // The third element in an SMT key is the index. + self.0[3] + } +} + +impl fmt::Display for HashedStorageMapKey { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.0) + } +} diff --git a/crates/miden-objects/src/account/storage/map/mod.rs b/crates/miden-objects/src/account/storage/map/mod.rs index 4b262f88ed..82f081969c 100644 --- a/crates/miden-objects/src/account/storage/map/mod.rs +++ b/crates/miden-objects/src/account/storage/map/mod.rs @@ -7,7 +7,13 @@ use super::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serial use crate::account::StorageMapDelta; use crate::crypto::merkle::{InnerNodeInfo, LeafIndex, SMT_DEPTH, Smt, SmtLeaf}; use crate::errors::StorageMapError; -use crate::{AccountError, Felt, Hasher}; +use crate::{AccountError, Felt}; + +mod storage_map_key; +pub use storage_map_key::StorageMapKey; + +mod hashed_storage_map_key; +pub use hashed_storage_map_key::HashedStorageMapKey; mod partial; pub use partial::PartialStorageMap; @@ -46,7 +52,7 @@ pub struct StorageMap { /// /// It is an invariant of this type that the map's entries are always consistent with the SMT's /// entries and vice-versa. - entries: BTreeMap, + entries: BTreeMap, } impl StorageMap { @@ -78,8 +84,8 @@ impl StorageMap { /// /// Returns an error if: /// - the provided entries contain multiple values for the same key. - pub fn with_entries>( - entries: impl IntoIterator, + pub fn with_entries>( + entries: impl IntoIterator, ) -> Result { let mut map = BTreeMap::new(); @@ -97,8 +103,8 @@ impl StorageMap { } /// Creates a new [`StorageMap`] from the given map. For internal use. - fn from_btree_map(entries: BTreeMap) -> Self { - let hashed_keys_iter = entries.iter().map(|(key, value)| (Self::hash_key(*key), *value)); + fn from_btree_map(entries: BTreeMap) -> Self { + let hashed_keys_iter = entries.iter().map(|(key, value)| (key.hash().inner(), *value)); let smt = Smt::with_entries(hashed_keys_iter) .expect("btree maps should not contain duplicate keys"); @@ -131,15 +137,15 @@ impl StorageMap { /// Returns the value corresponding to the key or [`Self::EMPTY_VALUE`] if the key is not /// associated with a value. - pub fn get(&self, raw_key: &Word) -> Word { + pub fn get(&self, raw_key: &StorageMapKey) -> Word { self.entries.get(raw_key).copied().unwrap_or_default() } /// Returns an opening of the leaf associated with raw key. /// /// Conceptually, an opening is a Merkle path to the leaf, as well as the leaf itself. - pub fn open(&self, raw_key: &Word) -> StorageMapWitness { - let hashed_map_key = Self::hash_key(*raw_key); + pub fn open(&self, raw_key: &StorageMapKey) -> StorageMapWitness { + let hashed_map_key = raw_key.hash().inner(); let smt_proof = self.smt.open(&hashed_map_key); let value = self.entries.get(raw_key).copied().unwrap_or_default(); @@ -159,7 +165,7 @@ impl StorageMap { /// Returns an iterator over the key-value pairs in this storage map. /// /// Note that the returned key is the raw map key. - pub fn entries(&self) -> impl Iterator { + pub fn entries(&self) -> impl Iterator { self.entries.iter() } @@ -175,14 +181,14 @@ impl StorageMap { /// [`Self::EMPTY_VALUE`] if no entry was previously present. /// /// If the provided `value` is [`Self::EMPTY_VALUE`] the entry will be removed. - pub fn insert(&mut self, raw_key: Word, value: Word) -> Result { + pub fn insert(&mut self, raw_key: StorageMapKey, value: Word) -> Result { if value == EMPTY_WORD { self.entries.remove(&raw_key); } else { self.entries.insert(raw_key, value); } - let hashed_key = Self::hash_key(raw_key); + let hashed_key = raw_key.hash().inner(); self.smt .insert(hashed_key, value) .map_err(AccountError::MaxNumStorageMapLeavesExceeded) @@ -192,28 +198,16 @@ impl StorageMap { pub fn apply_delta(&mut self, delta: &StorageMapDelta) -> Result { // apply the updated and cleared leaves to the storage map for (&key, &value) in delta.entries().iter() { - self.insert(key.into_inner(), value)?; + self.insert(StorageMapKey::new_unchecked(key.into_inner()), value)?; } Ok(self.root()) } /// Consumes the map and returns the underlying map of entries. - pub fn into_entries(self) -> BTreeMap { + pub fn into_entries(self) -> BTreeMap { self.entries } - - /// Hashes the given key to get the key of the SMT. - pub fn hash_key(raw_key: Word) -> Word { - Hasher::hash_elements(raw_key.as_elements()) - } - - // TODO: Replace with https://github.com/0xMiden/crypto/issues/515 once implemented. - /// Returns the leaf index of a map key. - pub fn hashed_map_key_to_leaf_index(hashed_map_key: Word) -> Felt { - // The third element in an SMT key is the index. - hashed_map_key[3] - } } impl Default for StorageMap { @@ -247,6 +241,7 @@ mod tests { use assert_matches::assert_matches; use super::{Deserializable, EMPTY_STORAGE_MAP_ROOT, Serializable, StorageMap, Word}; + use crate::account::StorageMapKey; use crate::errors::StorageMapError; #[test] @@ -257,9 +252,9 @@ mod tests { assert_eq!(storage_map_default, StorageMap::read_from_bytes(&bytes).unwrap()); // StorageMap with values - let storage_map_leaves_2: [(Word, Word); 2] = [ - (Word::from([101, 102, 103, 104u32]), Word::from([1, 2, 3, 4u32])), - (Word::from([105, 106, 107, 108u32]), Word::from([5, 6, 7, 8u32])), + let storage_map_leaves_2: [(StorageMapKey, Word); 2] = [ + (Word::from([101, 102, 103, 104u32]).into(), Word::from([1, 2, 3, 4u32])), + (Word::from([105, 106, 107, 108u32]).into(), Word::from([5, 6, 7, 8u32])), ]; let storage_map = StorageMap::with_entries(storage_map_leaves_2).unwrap(); assert_eq!(storage_map.num_entries(), 2); @@ -282,9 +277,9 @@ mod tests { #[test] fn account_storage_map_fails_on_duplicate_entries() { // StorageMap with values - let storage_map_leaves_2: [(Word, Word); 2] = [ - (Word::from([101, 102, 103, 104u32]), Word::from([1, 2, 3, 4u32])), - (Word::from([101, 102, 103, 104u32]), Word::from([5, 6, 7, 8u32])), + let storage_map_leaves_2: [(StorageMapKey, Word); 2] = [ + (Word::from([101, 102, 103, 104u32]).into(), Word::from([1, 2, 3, 4u32])), + (Word::from([101, 102, 103, 104u32]).into(), Word::from([5, 6, 7, 8u32])), ]; let error = StorageMap::with_entries(storage_map_leaves_2).unwrap_err(); diff --git a/crates/miden-objects/src/account/storage/map/partial.rs b/crates/miden-objects/src/account/storage/map/partial.rs index 92fd97e868..d608bec6f8 100644 --- a/crates/miden-objects/src/account/storage/map/partial.rs +++ b/crates/miden-objects/src/account/storage/map/partial.rs @@ -12,6 +12,7 @@ use miden_crypto::merkle::{ SmtProof, }; +use super::StorageMapKey; use crate::account::{StorageMap, StorageMapWitness}; use crate::utils::serde::{ByteReader, DeserializationError}; @@ -36,7 +37,7 @@ pub struct PartialStorageMap { /// /// It is an invariant of this type that the map's entries are always consistent with the /// partial SMT's entries and vice-versa. - entries: BTreeMap, + entries: BTreeMap, } impl PartialStorageMap { @@ -104,8 +105,8 @@ impl PartialStorageMap { /// - a non-empty [`Word`] if the key is tracked by this map and exists in it, /// - [`Word::empty`] if the key is tracked by this map and does not exist, /// - `None` if the key is not tracked by this map. - pub fn get(&self, raw_key: &Word) -> Option { - let hashed_key = StorageMap::hash_key(*raw_key); + pub fn get(&self, raw_key: &StorageMapKey) -> Option { + let hashed_key = raw_key.hash().inner(); // This returns an error if the key is not tracked which we map to a `None`. self.partial_smt.get_value(&hashed_key).ok() } @@ -118,8 +119,8 @@ impl PartialStorageMap { /// /// Returns an error if: /// - the key is not tracked by this partial storage map. - pub fn open(&self, raw_key: &Word) -> Result { - let hashed_key = StorageMap::hash_key(*raw_key); + pub fn open(&self, raw_key: &StorageMapKey) -> Result { + let hashed_key = raw_key.hash().inner(); let smt_proof = self.partial_smt.open(&hashed_key)?; let value = self.entries.get(raw_key).copied().unwrap_or_default(); @@ -139,7 +140,7 @@ impl PartialStorageMap { /// Returns an iterator over the key-value pairs in this storage map. /// /// Note that the returned key is the raw map key. - pub fn entries(&self) -> impl Iterator { + pub fn entries(&self) -> impl Iterator { self.entries.iter() } @@ -174,8 +175,8 @@ impl Deserializable for PartialStorageMap { let num_entries: usize = source.read()?; for _ in 0..num_entries { - let key: Word = source.read()?; - let hashed_map_key = StorageMap::hash_key(key); + let key: StorageMapKey = StorageMapKey::new_unchecked(source.read()?); + let hashed_map_key = key.hash().inner(); let value = partial_smt.get_value(&hashed_map_key).map_err(|err| { DeserializationError::InvalidValue(format!( "failed to find map key {key} in partial SMT: {err}" diff --git a/crates/miden-objects/src/account/storage/map/storage_map_key.rs b/crates/miden-objects/src/account/storage/map/storage_map_key.rs new file mode 100644 index 0000000000..8b199919cb --- /dev/null +++ b/crates/miden-objects/src/account/storage/map/storage_map_key.rs @@ -0,0 +1,101 @@ +use core::fmt; + +use miden_crypto::Word; +use miden_crypto::word::LexicographicWord; + +use super::{ + ByteReader, + ByteWriter, + Deserializable, + DeserializationError, + Felt, + HashedStorageMapKey, + Serializable, +}; +use crate::Hasher; +use crate::asset::AssetVaultKey; + +#[derive(Debug, PartialEq, Eq, Clone, Copy, PartialOrd, Ord)] +pub struct StorageMapKey(Word); + +impl StorageMapKey { + /// Creates a new [`StorageMapKey`] from the given [`Word`] **without performing validation**. + /// + /// ## Warning + /// + /// This function **does not check** whether the provided `Word` represents a valid + /// fungible or non-fungible asset key. + pub const fn new_unchecked(value: Word) -> Self { + Self(value) + } + + pub fn inner(&self) -> Word { + self.0 + } + + /// Hashes the given key to get the key of the SMT. + pub fn hash(&self) -> HashedStorageMapKey { + HashedStorageMapKey::new_unchecked(Hasher::hash_elements(self.0.as_elements())) + } + + pub fn as_elements(&self) -> &[Felt] { + self.0.as_elements() + } +} + +impl fmt::Display for StorageMapKey { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.0) + } +} + +// SERIALIZATION +// ================================================================================================ + +impl Serializable for StorageMapKey { + fn write_into(&self, target: &mut W) { + self.0.write_into(target); + } + + fn get_size_hint(&self) -> usize { + self.0.get_size_hint() + } +} + +impl Deserializable for StorageMapKey { + fn read_from(source: &mut R) -> Result { + let word = source.read()?; + Ok(Self::new_unchecked(word)) + } +} + +impl From for StorageMapKey { + fn from(value: Word) -> Self { + Self::new_unchecked(value) + } +} + +impl From for Word { + fn from(value: StorageMapKey) -> Self { + value.0 + } +} + +impl From for LexicographicWord { + fn from(value: StorageMapKey) -> Self { + LexicographicWord::from(value.0) + } +} + +impl From for StorageMapKey { + fn from(vault_key: AssetVaultKey) -> Self { + let vault_key_word: Word = vault_key.into(); + StorageMapKey::from(vault_key_word) + } +} + +impl From<[Felt; 4]> for StorageMapKey { + fn from(value: [Felt; 4]) -> Self { + StorageMapKey::new_unchecked(Word::from(value)) + } +} diff --git a/crates/miden-objects/src/account/storage/map/witness.rs b/crates/miden-objects/src/account/storage/map/witness.rs index 401d8345a4..662174f51a 100644 --- a/crates/miden-objects/src/account/storage/map/witness.rs +++ b/crates/miden-objects/src/account/storage/map/witness.rs @@ -2,8 +2,8 @@ use alloc::collections::BTreeMap; use miden_crypto::merkle::{InnerNodeInfo, SmtProof}; +use super::StorageMapKey; use crate::Word; -use crate::account::StorageMap; use crate::errors::StorageMapError; /// A witness of an asset in a [`StorageMap`](super::StorageMap). @@ -25,7 +25,7 @@ pub struct StorageMapWitness { /// /// It is an invariant of this type that the map's entries are always consistent with the SMT's /// entries and vice-versa. - entries: BTreeMap, + entries: BTreeMap, } impl StorageMapWitness { @@ -40,12 +40,12 @@ impl StorageMapWitness { /// - Any of the map keys is not contained in the proof. pub fn new( proof: SmtProof, - raw_keys: impl IntoIterator, + raw_keys: impl IntoIterator, ) -> Result { let mut entries = BTreeMap::new(); for raw_key in raw_keys.into_iter() { - let hashed_map_key = StorageMap::hash_key(raw_key); + let hashed_map_key = raw_key.hash().inner(); let value = proof.get(&hashed_map_key).ok_or(StorageMapError::MissingKey { raw_key })?; entries.insert(raw_key, value); @@ -62,7 +62,7 @@ impl StorageMapWitness { /// details. pub fn new_unchecked( proof: SmtProof, - raw_key_values: impl IntoIterator, + raw_key_values: impl IntoIterator, ) -> Self { Self { proof, @@ -82,15 +82,15 @@ impl StorageMapWitness { /// - a non-empty [`Word`] if the key is tracked by this witness and exists in it, /// - [`Word::empty`] if the key is tracked by this witness and does not exist, /// - `None` if the key is not tracked by this witness. - pub fn get(&self, raw_key: &Word) -> Option { - let hashed_key = StorageMap::hash_key(*raw_key); + pub fn get(&self, raw_key: &StorageMapKey) -> Option { + let hashed_key = raw_key.hash().inner(); self.proof.get(&hashed_key) } /// Returns an iterator over the key-value pairs in this witness. /// /// Note that the returned key is the raw map key. - pub fn entries(&self) -> impl Iterator { + pub fn entries(&self) -> impl Iterator { self.entries.iter() } @@ -119,7 +119,7 @@ mod tests { #[test] fn creating_witness_fails_on_missing_key() { // Create a storage map with one key-value pair - let key1 = Word::from([1, 2, 3, 4u32]); + let key1 = StorageMapKey::from(Word::from([1, 2, 3, 4u32])); let value1 = Word::from([10, 20, 30, 40u32]); let entries = [(key1, value1)]; let storage_map = StorageMap::with_entries(entries).unwrap(); @@ -128,7 +128,7 @@ mod tests { let proof = storage_map.open(&key1).into(); // Try to create a witness for a different key that's not in the proof - let missing_key = Word::from([5, 6, 7, 8u32]); + let missing_key = StorageMapKey::from(Word::from([5, 6, 7, 8u32])); let result = StorageMapWitness::new(proof, [missing_key]); assert_matches!(result, Err(StorageMapError::MissingKey { raw_key }) => { diff --git a/crates/miden-objects/src/account/storage/mod.rs b/crates/miden-objects/src/account/storage/mod.rs index 168f12ee85..0436b27f5f 100644 --- a/crates/miden-objects/src/account/storage/mod.rs +++ b/crates/miden-objects/src/account/storage/mod.rs @@ -19,7 +19,7 @@ mod slot; pub use slot::{SlotName, StorageSlot, StorageSlotType}; mod map; -pub use map::{PartialStorageMap, StorageMap, StorageMapWitness}; +pub use map::{PartialStorageMap, StorageMap, StorageMapKey, StorageMapWitness}; mod header; pub use header::{AccountStorageHeader, StorageSlotHeader}; @@ -147,7 +147,7 @@ impl AccountStorage { /// # Errors: /// - If the index is out of bounds /// - If the [StorageSlot] is not [StorageSlotType::Map] - pub fn get_map_item(&self, index: u8, key: Word) -> Result { + pub fn get_map_item(&self, index: u8, key: StorageMapKey) -> Result { match self.slots.get(index as usize).ok_or(AccountError::StorageIndexOutOfBounds { slots_len: self.slots.len() as u8, index, @@ -243,7 +243,7 @@ impl AccountStorage { pub fn set_map_item( &mut self, index: u8, - key: Word, + key: StorageMapKey, value: Word, ) -> Result<(Word, Word), AccountError> { // check if index is in bounds diff --git a/crates/miden-objects/src/account/storage/partial.rs b/crates/miden-objects/src/account/storage/partial.rs index 53c90828a7..fa538c7750 100644 --- a/crates/miden-objects/src/account/storage/partial.rs +++ b/crates/miden-objects/src/account/storage/partial.rs @@ -171,23 +171,23 @@ mod tests { let map_key_absent: Word = [9u64, 12, 18, 3].try_into()?; let mut map_1 = StorageMap::new(); - map_1.insert(map_key_absent, Word::try_from([1u64, 2, 3, 2])?).unwrap(); - map_1.insert(map_key_present, Word::try_from([5u64, 4, 3, 2])?).unwrap(); - assert_eq!(map_1.get(&map_key_present), [5u64, 4, 3, 2].try_into()?); + map_1.insert(map_key_absent.into(), Word::try_from([1u64, 2, 3, 2])?).unwrap(); + map_1.insert(map_key_present.into(), Word::try_from([5u64, 4, 3, 2])?).unwrap(); + assert_eq!(map_1.get(&map_key_present.into()), [5u64, 4, 3, 2].try_into()?); let storage = AccountStorage::new(vec![StorageSlot::Map(map_1.clone())]).unwrap(); // Create partial storage with validation of one map key let storage_header = AccountStorageHeader::from(&storage); - let witness = map_1.open(&map_key_present); + let witness = map_1.open(&map_key_present.into()); let partial_storage = PartialStorage::new(storage_header, [PartialStorageMap::with_witnesses([witness])?]) .context("creating partial storage")?; let retrieved_map = partial_storage.maps.get(&partial_storage.header.slot(0)?.1).unwrap(); - assert!(retrieved_map.open(&map_key_absent).is_err()); - assert!(retrieved_map.open(&map_key_present).is_ok()); + assert!(retrieved_map.open(&map_key_absent.into()).is_err()); + assert!(retrieved_map.open(&map_key_present.into()).is_ok()); Ok(()) } } diff --git a/crates/miden-objects/src/errors.rs b/crates/miden-objects/src/errors.rs index a11d45d421..8342098704 100644 --- a/crates/miden-objects/src/errors.rs +++ b/crates/miden-objects/src/errors.rs @@ -23,6 +23,7 @@ use crate::account::{ AccountStorage, AccountType, SlotName, + StorageMapKey, StorageValueName, StorageValueNameError, TemplateTypeError, @@ -407,9 +408,13 @@ pub enum AccountDeltaError { #[derive(Debug, Error)] pub enum StorageMapError { #[error("map entries contain key {key} twice with values {value0} and {value1}")] - DuplicateKey { key: Word, value0: Word, value1: Word }, + DuplicateKey { + key: StorageMapKey, + value0: Word, + value1: Word, + }, #[error("map key {raw_key} is not present in provided SMT proof")] - MissingKey { raw_key: Word }, + MissingKey { raw_key: StorageMapKey }, } // BATCH ACCOUNT UPDATE ERROR diff --git a/crates/miden-objects/src/testing/storage.rs b/crates/miden-objects/src/testing/storage.rs index 85f8f673f5..eb5735bb96 100644 --- a/crates/miden-objects/src/testing/storage.rs +++ b/crates/miden-objects/src/testing/storage.rs @@ -11,6 +11,7 @@ use crate::account::{ AccountStorageDelta, StorageMap, StorageMapDelta, + StorageMapKey, StorageSlot, }; use crate::note::NoteAssets; @@ -85,13 +86,23 @@ pub const STORAGE_VALUE_0: Word = Word::new([Felt::new(1), Felt::new(2), Felt::new(3), Felt::new(4)]); pub const STORAGE_VALUE_1: Word = Word::new([Felt::new(5), Felt::new(6), Felt::new(7), Felt::new(8)]); -pub const STORAGE_LEAVES_2: [(Word, Word); 2] = [ +pub const STORAGE_LEAVES_2: [(StorageMapKey, Word); 2] = [ ( - Word::new([Felt::new(101), Felt::new(102), Felt::new(103), Felt::new(104)]), + StorageMapKey::new_unchecked(Word::new([ + Felt::new(101), + Felt::new(102), + Felt::new(103), + Felt::new(104), + ])), Word::new([Felt::new(1_u64), Felt::new(2_u64), Felt::new(3_u64), Felt::new(4_u64)]), ), ( - Word::new([Felt::new(105), Felt::new(106), Felt::new(107), Felt::new(108)]), + StorageMapKey::new_unchecked(Word::new([ + Felt::new(105), + Felt::new(106), + Felt::new(107), + Felt::new(108), + ])), Word::new([Felt::new(5_u64), Felt::new(6_u64), Felt::new(7_u64), Felt::new(8_u64)]), ), ]; diff --git a/crates/miden-testing/src/kernel_tests/tx/test_account.rs b/crates/miden-testing/src/kernel_tests/tx/test_account.rs index a4f5a4398a..54b30eadcf 100644 --- a/crates/miden-testing/src/kernel_tests/tx/test_account.rs +++ b/crates/miden-testing/src/kernel_tests/tx/test_account.rs @@ -30,6 +30,7 @@ use miden_objects::account::{ AccountStorageMode, AccountType, StorageMap, + StorageMapKey, StorageSlot, }; use miden_objects::assembly::diagnostics::{IntoDiagnostic, NamedSource, Report, WrapErr, miette}; @@ -77,7 +78,7 @@ pub async fn compute_commitment() -> miette::Result<()> { let mut account_clone = account.clone(); let key = Word::from([1, 2, 3, 4u32]); let value = Word::from([2, 3, 4, 5u32]); - account_clone.storage_mut().set_map_item(2, key, value).unwrap(); + account_clone.storage_mut().set_map_item(2, key.into(), value).unwrap(); let expected_commitment = account_clone.commitment(); let tx_script = format!( @@ -603,7 +604,7 @@ async fn test_set_map_item() -> miette::Result<()> { let exec_output = &tx_context.execute_code(&code).await.unwrap(); let mut new_storage_map = AccountStorage::mock_map(); - new_storage_map.insert(new_key, new_value).unwrap(); + new_storage_map.insert(new_key.into(), new_value).unwrap(); assert_eq!( new_storage_map.root(), @@ -1009,7 +1010,7 @@ async fn prove_account_creation_with_non_empty_storage() -> anyhow::Result<()> { let slot1 = StorageSlot::Value(Word::from([10, 20, 30, 40u32])); let mut map_entries = Vec::new(); for _ in 0..10 { - map_entries.push((rand_value::(), rand_value::())); + map_entries.push((StorageMapKey::from(rand_value::()), rand_value::())); } let map_slot = StorageSlot::Map(StorageMap::with_entries(map_entries.clone())?); diff --git a/crates/miden-testing/src/kernel_tests/tx/test_lazy_loading.rs b/crates/miden-testing/src/kernel_tests/tx/test_lazy_loading.rs index 83829fcd5e..a7417b4288 100644 --- a/crates/miden-testing/src/kernel_tests/tx/test_lazy_loading.rs +++ b/crates/miden-testing/src/kernel_tests/tx/test_lazy_loading.rs @@ -172,7 +172,8 @@ async fn setting_map_item_with_lazy_loading_succeeds() -> anyhow::Result<()> { let non_existent_key = Word::from([5, 5, 5, 5u32]); assert!( - mock_map.open(&non_existent_key).get(&non_existent_key).unwrap() == Word::empty(), + mock_map.open(&non_existent_key.into()).get(&non_existent_key.into()).unwrap() + == Word::empty(), "test setup requires that the non existent key does not exist" ); @@ -218,7 +219,10 @@ async fn setting_map_item_with_lazy_loading_succeeds() -> anyhow::Result<()> { .await?; let map_delta = tx.account_delta().storage().maps().get(&map_index).unwrap(); - assert_eq!(map_delta.entries().get(&LexicographicWord::new(existing_key)).unwrap(), &value0); + assert_eq!( + map_delta.entries().get(&LexicographicWord::new(existing_key.inner())).unwrap(), + &value0 + ); assert_eq!( map_delta.entries().get(&LexicographicWord::new(non_existent_key)).unwrap(), &value1 @@ -236,7 +240,8 @@ async fn getting_map_item_with_lazy_loading_succeeds() -> anyhow::Result<()> { let non_existent_key = Word::from([5, 5, 5, 5u32]); assert!( - mock_map.open(&non_existent_key).get(&non_existent_key).unwrap() == Word::empty(), + mock_map.open(&non_existent_key.into()).get(&non_existent_key.into()).unwrap() + == Word::empty(), "test setup requires that the non existent key does not exist" ); diff --git a/crates/miden-testing/src/tx_context/context.rs b/crates/miden-testing/src/tx_context/context.rs index 9289a9e598..834a7fbae3 100644 --- a/crates/miden-testing/src/tx_context/context.rs +++ b/crates/miden-testing/src/tx_context/context.rs @@ -292,7 +292,7 @@ impl DataStore for TransactionContext { )) })?; - Ok(storage_map.open(&map_key)) + Ok(storage_map.open(&map_key.into())) } else { let (foreign_account, _witness) = self .foreign_account_inputs @@ -322,7 +322,7 @@ impl DataStore for TransactionContext { )) })?; - Ok(map.open(&map_key)) + Ok(map.open(&map_key.into())) } } } diff --git a/crates/miden-tx/src/executor/exec_host.rs b/crates/miden-tx/src/executor/exec_host.rs index 0a7eafbf41..570a63790b 100644 --- a/crates/miden-tx/src/executor/exec_host.rs +++ b/crates/miden-tx/src/executor/exec_host.rs @@ -546,7 +546,7 @@ where self.on_account_storage_map_witness_requested( active_account_id, map_root, - map_key, + map_key.into(), ) .await }, diff --git a/crates/miden-tx/src/host/storage_delta_tracker.rs b/crates/miden-tx/src/host/storage_delta_tracker.rs index 81da62182b..c42590bee9 100644 --- a/crates/miden-tx/src/host/storage_delta_tracker.rs +++ b/crates/miden-tx/src/host/storage_delta_tracker.rs @@ -78,7 +78,7 @@ impl StorageDeltaTracker { storage_map.entries().for_each(|(key, value)| { storage_delta_tracker.set_map_item( slot_idx, - *key, + (*key).into(), Word::empty(), *value, ); diff --git a/crates/miden-tx/src/host/tx_event.rs b/crates/miden-tx/src/host/tx_event.rs index d233965c45..8edb48b2a6 100644 --- a/crates/miden-tx/src/host/tx_event.rs +++ b/crates/miden-tx/src/host/tx_event.rs @@ -1,7 +1,7 @@ use alloc::vec::Vec; use miden_lib::transaction::{EventId, TransactionEventId}; -use miden_objects::account::{AccountId, StorageMap, StorageSlotType}; +use miden_objects::account::{AccountId, StorageMap, StorageMapKey, StorageSlotType}; use miden_objects::asset::{Asset, AssetVault, AssetVaultKey, FungibleAsset}; use miden_objects::note::{NoteId, NoteInputs, NoteMetadata, NoteRecipient, NoteScript}; use miden_objects::transaction::TransactionSummary; @@ -75,7 +75,7 @@ pub(crate) enum TransactionEvent { /// The root of the storage map for which a witness is requested. map_root: Word, /// The raw map key for which a witness is requested. - map_key: Word, + map_key: StorageMapKey, }, /// The data necessary to request an asset witness from the data store. @@ -293,7 +293,7 @@ impl TransactionEvent { process, slot_index, current_map_root, - map_key, + map_key.into(), )? }, @@ -308,7 +308,7 @@ impl TransactionEvent { process, slot_index, current_map_root, - map_key, + map_key.into(), )? }, @@ -602,11 +602,11 @@ fn on_account_storage_map_item_accessed<'store, STORE>( process: &ProcessState, slot_index: Felt, current_map_root: Word, - map_key: Word, + map_key: StorageMapKey, ) -> Result, TransactionKernelError> { let active_account_id = process.get_active_account_id()?; - let hashed_map_key = StorageMap::hash_key(map_key); - let leaf_index = StorageMap::hashed_map_key_to_leaf_index(hashed_map_key); + let hashed_map_key = map_key.hash(); + let leaf_index = hashed_map_key.hashed_map_key_to_leaf_index(); let slot_index = u8::try_from(slot_index).map_err(|err| { TransactionKernelError::other(format!("failed to convert slot index into u8: {err}"))