From 8ce732fd3cd97e78b11a468a62e081f0a20118e4 Mon Sep 17 00:00:00 2001
From: Dori Medini <dori@starkware.co>
Date: Wed, 29 Oct 2025 15:48:52 +0200
Subject: [PATCH] starknet_patricia_storage: add short-key storage

---
 Cargo.lock                                    |   2 +
 crates/starknet_patricia_storage/Cargo.toml   |   2 +
 crates/starknet_patricia_storage/src/lib.rs   |   1 +
 .../src/short_key_storage.rs                  | 106 ++++++++++++++++++
 4 files changed, 111 insertions(+)
 create mode 100644 crates/starknet_patricia_storage/src/short_key_storage.rs

diff --git a/Cargo.lock b/Cargo.lock
index 657abfda572..59bfbd4d8c5 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -12651,6 +12651,8 @@ dependencies = [
 name = "starknet_patricia_storage"
 version = "0.17.0-rc.0"
 dependencies = [
+ "blake2",
+ "digest 0.10.7",
  "hex",
  "libmdbx",
  "lru 0.12.5",
diff --git a/crates/starknet_patricia_storage/Cargo.toml b/crates/starknet_patricia_storage/Cargo.toml
index 1f522f900f7..a711125f2f6 100644
--- a/crates/starknet_patricia_storage/Cargo.toml
+++ b/crates/starknet_patricia_storage/Cargo.toml
@@ -13,6 +13,8 @@ testing = []
 workspace = true
 
 [dependencies]
+blake2.workspace = true
+digest.workspace = true
 hex.workspace = true
 libmdbx.workspace = true
 lru.workspace = true
diff --git a/crates/starknet_patricia_storage/src/lib.rs b/crates/starknet_patricia_storage/src/lib.rs
index 4ef92af19cc..2f93013a347 100644
--- a/crates/starknet_patricia_storage/src/lib.rs
+++ b/crates/starknet_patricia_storage/src/lib.rs
@@ -4,4 +4,5 @@ pub mod map_storage;
 #[cfg(test)]
 pub mod map_storage_test;
 pub mod mdbx_storage;
+pub mod short_key_storage;
 pub mod storage_trait;
diff --git a/crates/starknet_patricia_storage/src/short_key_storage.rs b/crates/starknet_patricia_storage/src/short_key_storage.rs
new file mode 100644
index 00000000000..87d30703c42
--- /dev/null
+++ b/crates/starknet_patricia_storage/src/short_key_storage.rs
@@ -0,0 +1,106 @@
+use std::marker::PhantomData;
+
+use blake2::digest::consts::{
+    U16,
+    U17,
+    U18,
+    U19,
+    U20,
+    U21,
+    U22,
+    U23,
+    U24,
+    U25,
+    U26,
+    U27,
+    U28,
+    U29,
+    U30,
+    U31,
+    U32,
+};
+use blake2::Blake2s;
+use digest::Digest;
+
+use crate::storage_trait::{DbHashMap, DbKey, DbValue, PatriciaStorageResult, Storage};
+
+#[macro_export]
+macro_rules! define_short_key_storage {
+    ($( ( $sizes:ty, $names:ident ) ),+ $(,)?) => {
+        $(
+            $crate::define_short_key_storage!($sizes, $names);
+        )+
+    };
+
+    ($size:ty, $name:ident) => {
+        /// A storage that hashes (using blake) each key to a $size - byte key.
+        pub struct $name<S: Storage> {
+            storage: S,
+            _n_bytes: PhantomData<$size>,
+        }
+
+        impl<S: Storage> $name<S> {
+            pub fn new(storage: S) -> Self {
+                Self { storage, _n_bytes: PhantomData }
+            }
+
+            pub fn small_key(key: &DbKey) -> DbKey {
+                let mut hasher = Blake2s::<$size>::new();
+                hasher.update(key.0.as_slice());
+                let result = hasher.finalize();
+                DbKey(result.as_slice().to_vec())
+            }
+        }
+
+        impl<S: Storage> Storage for $name<S> {
+            fn get(&mut self, key: &DbKey) -> PatriciaStorageResult<Option<DbValue>> {
+                self.storage.get(&Self::small_key(key))
+            }
+
+            fn set(&mut self, key: DbKey, value: DbValue) -> PatriciaStorageResult<Option<DbValue>> {
+                self.storage.set(Self::small_key(&key), value)
+            }
+
+            fn mget(&mut self, keys: &[&DbKey]) -> PatriciaStorageResult<Vec<Option<DbValue>>> {
+                let small_keys = keys
+                    .iter()
+                    .map(|key| Self::small_key(key))
+                    .collect::<Vec<_>>();
+                self.storage.mget(small_keys.iter().collect::<Vec<&DbKey>>().as_slice())
+            }
+
+            fn mset(&mut self, key_to_value: DbHashMap) -> PatriciaStorageResult<()> {
+                self.storage.mset(
+                    key_to_value
+                        .into_iter()
+                        .map(|(key, value)| (Self::small_key(&key), value))
+                        .collect()
+                )
+            }
+
+            fn delete(&mut self, key: &DbKey) -> PatriciaStorageResult<Option<DbValue>> {
+                self.storage.delete(&Self::small_key(key))
+            }
+        }
+    };
+}
+
+define_short_key_storage!(
+    (U16, ShortKeyStorage16),
+    (U17, ShortKeyStorage17),
+    (U18, ShortKeyStorage18),
+    (U19, ShortKeyStorage19),
+    (U20, ShortKeyStorage20),
+    (U21, ShortKeyStorage21),
+    (U22, ShortKeyStorage22),
+    (U23, ShortKeyStorage23),
+    (U24, ShortKeyStorage24),
+    (U25, ShortKeyStorage25),
+    (U26, ShortKeyStorage26),
+    (U27, ShortKeyStorage27),
+    (U28, ShortKeyStorage28),
+    (U29, ShortKeyStorage29),
+    (U30, ShortKeyStorage30),
+    (U31, ShortKeyStorage31),
+    (U32, ShortKeyStorage32)
+);