feat: Introduce Word wrapper for storage map keys

crates/miden-lib/src/account/auth/ecdsa_k256_keccak_acl.rs

@@ -159,7 +159,7 @@ impl From<AuthEcdsaK256KeccakAcl> 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,15 +242,15 @@ 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);
             }
         } else {
             // 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());
         }

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<AuthEcdsaK256KeccakMultisig> 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<AuthEcdsaK256KeccakMultisig> 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));
     }

crates/miden-lib/src/account/auth/rpo_falcon_512_acl.rs

@@ -160,7 +160,7 @@ impl From<AuthRpoFalcon512Acl> 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,15 +243,15 @@ 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);
             }
         } else {
             // 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());
         }

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<AuthRpoFalcon512Multisig> 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<AuthRpoFalcon512Multisig> 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));
     }

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,23 +588,29 @@ 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::<Result<Vec<(Word, Word)>, _>>()?;
+                    .collect::<Result<Vec<(StorageMapKey, Word)>, _>>()?;
 
                 StorageMap::with_entries(entries).map_err(|err| {
                     AccountComponentTemplateError::StorageMapHasDuplicateKeys(Box::new(err))
                 })
             },
             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))
                     });
                 }

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"),
         }

crates/miden-objects/src/account/delta/storage.rs

@@ -381,7 +381,7 @@ impl From<StorageMap> for StorageMapDelta {
         StorageMapDelta::new(
             map.into_entries()
                 .into_iter()
-                .map(|(key, value)| (LexicographicWord::new(key), value))
+                .map(|(key, value)| (LexicographicWord::from(key), value))
                 .collect(),
         )
     }

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<Account> 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);

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)
+    }
+}