refactor: introduce `StorageMapKey` and `HashedStorageMapKey` Word wrappers by Farukest · Pull Request #2431 · 0xMiden/protocol

Farukest · 2026-02-11T14:33:00Z

Summary

Introduces type-safe Word wrappers for storage map keys to prevent accidentally mixing up raw and hashed keys:

StorageMapKey(Word) — wraps a raw, user-chosen key with hash() -> HashedStorageMapKey and to_leaf_index() methods.
HashedStorageMapKey(Word) — wraps a hashed key with to_leaf_index() method.

Both types use #[derive(WordWrapper)] from miden-protocol-macros.

The utility functions StorageMap::hash_key, StorageMap::map_key_to_leaf_index and StorageMap::hashed_map_key_to_leaf_index are removed and replaced by methods on the new types. All internal and external call sites are updated accordingly.

…appers

PhilippGackstatter · 2026-02-16T18:53:00Z

crates/miden-protocol/src/account/storage/map/key.rs

+/// underlying SMT. Wrapping the hashed key in a distinct type prevents accidentally using a raw
+/// key where a hashed key is expected and vice-versa.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, WordWrapper)]
+pub struct HashedStorageMapKey(Word);


Despite my initial suggestion, I'd now name this to StorageMapKeyHash, since the naming pattern we usually have is to append things at the end, e.g. StorageSlot and StorageSlotContent, StorageSlotName, etc.

PhilippGackstatter · 2026-02-17T08:33:50Z

crates/miden-protocol/src/account/storage/map/key.rs

+    /// Returns the leaf index in the SMT for this raw key.
+    ///
+    /// This first hashes the key, then converts it to a leaf index.
+    pub fn to_leaf_index(&self) -> LeafIndex<SMT_DEPTH> {
+        self.hash().to_leaf_index()
+    }


I'd remove this implementation to make it clearer that only a StorageMapKeyHash can be converted into a LeafIndex - i.e. not StorageMapKey but StorageMapKeyHash are used as the key in an SMT.

PhilippGackstatter · 2026-02-17T08:38:31Z

crates/miden-protocol/src/account/storage/map/mod.rs

@@ -140,8 +144,8 @@ impl StorageMap {
    ///
    /// Conceptually, an opening is a Merkle path to the leaf, as well as the leaf itself.
    pub fn open(&self, raw_key: &Word) -> StorageMapWitness {


Suggested change

pub fn open(&self, raw_key: &Word) -> StorageMapWitness {

pub fn open(&self, key: impl Into<StorageMapKeyHash>) -> StorageMapWitness {

I think we should update all the StorageMap APIs that take a key to take it in this form. So, any type that can be converted into a StorageMapKeyHash. This should provide good type safety: If you provide a StorageMapKey, it will be converted (and hashed) into a StorageMapKeyHash while you can also provide a StorageMapKeyHash if that's all you have. This way we allow using both types of keys while users don't have to think about those details.

Maybe this will also require refactoring StorageMap, PartialStorageMap and StorageMapWitness to use BTreeMap<StorageMapKey, Word> internally, but take a look.

PhilippGackstatter · 2026-02-17T08:39:18Z

crates/miden-protocol/src/account/storage/map/partial.rs

@@ -98,7 +98,7 @@ impl PartialStorageMap {
    /// - [`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<Word> {


I would do the same impl Into<StorageMapKeyHash> change for PartialStorageMap APIs.

Farukest · 2026-02-17T18:51:35Z

All done @PhilippGackstatter 👍

PhilippGackstatter

Looks good!

I left a few more comments, after which I think we can probably merge.

I think there are more places that we should eventually refactor, but it will be better to do this in a follow-up PR to keep this PR more reasonable in scope.

Follow-ups

StorageMap::with_entries might also benefit from changing to (StorageMapKey, Word) constructor for type safety, e.g. the API is clearer on which of the two words is the key and the value. Also, it will avoid any doubt about whether the provided key needs to be the raw key or the hashed key.

We should also make use of this type in these places which will make the data store api much clearer:

https://github.com/0xMiden/miden-base/blob/753fb1c9a24bad4dc891bc784a635ba5089c2c99/crates/miden-tx/src/host/tx_event.rs#L86-L93

https://github.com/0xMiden/miden-base/blob/753fb1c9a24bad4dc891bc784a635ba5089c2c99/crates/miden-tx/src/executor/data_store.rs#L66-L71

PhilippGackstatter · 2026-02-18T09:25:27Z

crates/miden-protocol/src/account/storage/map/mod.rs

+            .iter()
+            .find(|(k, _)| StorageMapKey::from_raw(**k).hash() == key_hash)
+            .map(|(k, v)| (*k, *v))
+            .unwrap_or_default();


I think we should be able to keep self.entries.get here?

I'm not sure that we can use get here anymore, since we don't have the original raw key here, only its hash, while the entries map uses raw keys.

But I'm wondering whether we need to have key: impl Into<StorageMapKeyHash> here in the first place: I agree with you that it will be much more convenient from the user standpoint, but will the hashed key be ever used by... well, user? It seems like this is an internal value only used to balance the tree, but will there be any cases when user will operate with hashed keys directly? If not, maybe it makes sense to keep the raw key (StorageMapKey in our case) as a parameter of this procedure? Since it will allow us just to get the value using self.entries.get method, saving us the burden of getting the raw key.

Ah, true, thank you! We actually need to pass the raw key to StorageMapWitness, and so if we only have the hash, this won't work. So indeed we should take key: &StorageMapKey as an input to open. Technically for get we could use StorageMapKeyHash without problems, but it may be better to require StorageMapKey for API uniformity. My bad for suggesting StorageMapKeyHash!

PhilippGackstatter · 2026-02-18T09:33:01Z

crates/miden-protocol/src/account/storage/map/partial.rs

+    /// Accepts either a [`StorageMapKey`] (which will be hashed) or a [`StorageMapKeyHash`]
+    /// directly.
+    pub fn get(&self, key: impl Into<StorageMapKeyHash>) -> Option<Word> {
+        let hash_word: Word = key.into().into();


Suggested change

let hash_word: Word = key.into().into();

let hash_word = key.into().as_word();

Nit: This might be a bit easier to read.

PhilippGackstatter · 2026-02-18T09:41:16Z

crates/miden-protocol/src/account/storage/map/key.rs

+impl From<Word> for StorageMapKey {
+    fn from(word: Word) -> Self {
+        Self(word)
+    }
+}


This is technically correct, but we previously decided not to implement From<Word> for these kinds of word-wrapper types, so that the "provenance" of the StorageMapKey becomes clearer. E.g. not every Word is semantically a StorageMapKey that is valid for a given map, and you need to use the deliberately less convenient from_raw API to vouch that the key is valid.

PhilippGackstatter · 2026-02-18T09:42:05Z

crates/miden-protocol/src/account/storage/map/witness.rs


        for raw_key in raw_keys.into_iter() {
-            let hashed_map_key = StorageMap::hash_key(raw_key);
+            let hashed_map_key: Word = StorageMapKey::from_raw(raw_key).hash().into();


I would also change the constructors of StorageMapWitness to use StorageMapKey, e.g.:

StorageMapWitness::new

StorageMapWitness::new_unchecked

So that we don't need to handle raw_key here. We can also rename raw_key to use key: StorageMapKey, since it's clear from the type now that it's the "raw key".

Fumuran

Looks good, thank you! In addition to Philipp's comments, I left just one small optional nit inline.

Fumuran · 2026-02-18T13:27:41Z

crates/miden-protocol/src/account/storage/map/mod.rs

+        let hash_word: Word = key_hash.into();
+        let smt_proof = self.smt.open(&hash_word);


optional nit: we can optimize this a bit

Suggested change

let hash_word: Word = key_hash.into();

let smt_proof = self.smt.open(&hash_word);

let smt_proof = self.smt.open(&key_hash.into());

Farukest · 2026-02-18T15:11:55Z

All done @PhilippGackstatter @Fumuran let me know if I missed something :)

refactor: introduce StorageMapKey and HashedStorageMapKey Word wr…

435b677

…appers

Farukest force-pushed the feat/issue-1864-storage-map-key-wrappers branch from 6082b2c to 435b677 Compare February 11, 2026 14:36

Farukest added 2 commits February 14, 2026 08:45

Merge branch 'next' into feat/issue-1864-storage-map-key-wrappers

bca117d

Merge branch 'next' into feat/issue-1864-storage-map-key-wrappers

8f4d015

bobbinth requested a review from PhilippGackstatter February 15, 2026 23:28

Fumuran self-requested a review February 16, 2026 18:46

PhilippGackstatter reviewed Feb 16, 2026

View reviewed changes

PhilippGackstatter reviewed Feb 17, 2026

View reviewed changes

Farukest force-pushed the feat/issue-1864-storage-map-key-wrappers branch from 35e637f to e44465f Compare February 17, 2026 21:50

PhilippGackstatter reviewed Feb 18, 2026

View reviewed changes

Fumuran reviewed Feb 18, 2026

View reviewed changes

Farukest force-pushed the feat/issue-1864-storage-map-key-wrappers branch 2 times, most recently from 1c9ce8c to 3ac810b Compare February 18, 2026 15:11

Farukest force-pushed the feat/issue-1864-storage-map-key-wrappers branch from 6bb353f to 5d7b47c Compare February 18, 2026 15:17

refactor: address review feedback for storage map key wrappers

e3fb5c2

Farukest force-pushed the feat/issue-1864-storage-map-key-wrappers branch from 5d7b47c to e3fb5c2 Compare February 18, 2026 16:00

Farukest added 2 commits February 18, 2026 19:00

Merge branch 'next' into feat/issue-1864-storage-map-key-wrappers

724be92

Merge branch 'next' into feat/issue-1864-storage-map-key-wrappers

2fe6e76

	pub fn open(&self, raw_key: &Word) -> StorageMapWitness {
	pub fn open(&self, key: impl Into<StorageMapKeyHash>) -> StorageMapWitness {

	let hash_word: Word = key.into().into();
	let hash_word = key.into().as_word();

		let hash_word: Word = key_hash.into();
		let smt_proof = self.smt.open(&hash_word);

	let hash_word: Word = key_hash.into();
	let smt_proof = self.smt.open(&hash_word);
	let smt_proof = self.smt.open(&key_hash.into());

Conversation

Farukest commented Feb 11, 2026

Summary

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Farukest commented Feb 17, 2026

Uh oh!

PhilippGackstatter left a comment

Choose a reason for hiding this comment

Follow-ups

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Fumuran left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Farukest commented Feb 18, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants