starknet_committer_cli: add option for key size

Author: dorimedini-starkware

crates/starknet_committer_cli/src/main.rs

@@ -7,6 +7,7 @@ use clap::{Args, Parser, Subcommand};
 use starknet_committer_cli::commands::run_storage_benchmark;
 use starknet_patricia_storage::map_storage::{CachedStorage, MapStorage};
 use starknet_patricia_storage::mdbx_storage::MdbxStorage;
+use starknet_patricia_storage::short_key_storage::wrap_storage_or_panic;
 use starknet_patricia_storage::storage_trait::Storage;
 use tracing::info;
 use tracing::level_filters::LevelFilter;
@@ -43,6 +44,10 @@ struct StorageArgs {
     /// checkpointing is ignored.
     #[clap(long, default_value = "cached-mdbx")]
     storage_type: StorageType,
+    /// If not none, wraps the storage in a key-shrinking storage.
+    /// See `short_key_storage.rs` for more details and allowed values.
+    #[clap(long, default_value = None)]
+    key_size: Option<u8>,
     /// If using cached storage, the size of the cache.
     #[clap(long, default_value = "1000000")]
     cache_size: usize,
@@ -85,6 +90,7 @@ pub async fn run_committer_cli(
             n_iterations,
             n_diffs,
             storage_type,
+            key_size,
             cache_size,
             checkpoint_interval,
             log_level,
@@ -102,13 +108,18 @@ pub async fn run_committer_cli(
 
             let (mut storage, checkpoint_dir): (Box<dyn Storage>, Option<&str>) = match storage_type
             {
-                StorageType::MapStorage => (Box::new(MapStorage::default()), None),
+                StorageType::MapStorage => {
+                    (wrap_storage_or_panic(key_size, MapStorage::default()), None)
+                }
                 StorageType::Mdbx => {
                     let storage_path = storage_path
                         .unwrap_or_else(|| format!("{data_path}/storage/{storage_type:?}"));
                     fs::create_dir_all(&storage_path).expect("Failed to create storage directory.");
                     (
-                        Box::new(MdbxStorage::open(Path::new(&storage_path)).unwrap()),
+                        wrap_storage_or_panic(
+                            key_size,
+                            MdbxStorage::open(Path::new(&storage_path)).unwrap(),
+                        ),
                         Some(&checkpoint_dir),
                     )
                 }
@@ -117,10 +128,13 @@ pub async fn run_committer_cli(
                         .unwrap_or_else(|| format!("{data_path}/storage/{storage_type:?}"));
                     fs::create_dir_all(&storage_path).expect("Failed to create storage directory.");
                     (
-                        Box::new(CachedStorage::new(
-                            MdbxStorage::open(Path::new(&storage_path)).unwrap(),
-                            NonZeroUsize::new(cache_size).unwrap(),
-                        )),
+                        wrap_storage_or_panic(
+                            key_size,
+                            CachedStorage::new(
+                                MdbxStorage::open(Path::new(&storage_path)).unwrap(),
+                                NonZeroUsize::new(cache_size).unwrap(),
+                            ),
+                        ),
                         Some(&checkpoint_dir),
                     )
                 }

crates/starknet_patricia_storage/src/short_key_storage.rs

@@ -26,17 +26,33 @@ use crate::storage_trait::{DbHashMap, DbKey, DbValue, PatriciaStorageResult, Sto
 
 #[macro_export]
 macro_rules! define_short_key_storage {
-    ($( ( $sizes:ty, $names:ident ) ),+ $(,)?) => {
+    ($( ( $sizes:expr, $sizes_type:ty, $names:ident ) ),+ $(,)?) => {
         $(
-            $crate::define_short_key_storage!($sizes, $names);
+            $crate::define_short_key_storage!($sizes, $sizes_type, $names);
         )+
+
+        /// Wrap an existing storage implementation in a boxed short key storage implementation.
+        /// If no size is given, boxes and returns the original storage.
+        /// Panics if the size is not within the allowed range.
+        pub fn wrap_storage_or_panic<S: Storage + 'static>(
+            size: Option<u8>,
+            storage: S,
+        ) -> Box<dyn Storage> {
+            let Some(size) = size else {
+                return Box::new(storage);
+            };
+            match size {
+                $( $sizes => Box::new($names::new(storage)), )+
+                size => panic!("Invalid key size {size}."),
+            }
+        }
     };
 
-    ($size:ty, $name:ident) => {
+    ($size:expr, $size_type:ty, $name:ident) => {
         /// A storage that hashes each key to a $size - byte key.
         pub struct $name<S: Storage> {
             storage: S,
-            _bytes: PhantomData<$size>,
+            _bytes: PhantomData<$size_type>,
         }
 
         impl<S: Storage> $name<S> {
@@ -45,7 +61,7 @@ macro_rules! define_short_key_storage {
             }
 
             pub fn small_key(key: &DbKey) -> DbKey {
-                let mut hasher = Blake2s::<$size>::new();
+                let mut hasher = Blake2s::<$size_type>::new();
                 hasher.update(key.0.as_slice());
                 let result = hasher.finalize();
                 DbKey(result.as_slice().to_vec())
@@ -81,26 +97,27 @@ macro_rules! define_short_key_storage {
             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)
+    (16, U16, ShortKeyStorage16),
+    (17, U17, ShortKeyStorage17),
+    (18, U18, ShortKeyStorage18),
+    (19, U19, ShortKeyStorage19),
+    (20, U20, ShortKeyStorage20),
+    (21, U21, ShortKeyStorage21),
+    (22, U22, ShortKeyStorage22),
+    (23, U23, ShortKeyStorage23),
+    (24, U24, ShortKeyStorage24),
+    (25, U25, ShortKeyStorage25),
+    (26, U26, ShortKeyStorage26),
+    (27, U27, ShortKeyStorage27),
+    (28, U28, ShortKeyStorage28),
+    (29, U29, ShortKeyStorage29),
+    (30, U30, ShortKeyStorage30),
+    (31, U31, ShortKeyStorage31),
+    (32, U32, ShortKeyStorage32)
 );