starknet_patricia_storage: split storage trait (#13234)

Author: nimrod-starkware

crates/starknet_committer/src/db/index_db/test_utils.rs

@@ -13,7 +13,7 @@ use starknet_patricia::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndice
 use starknet_patricia::patricia_merkle_tree::updated_skeleton_tree::hash_function::TreeHashFunction;
 use starknet_patricia_storage::db_object::{DBObject, EmptyKeyContext, HasStaticPrefix};
 use starknet_patricia_storage::map_storage::MapStorage;
-use starknet_patricia_storage::storage_trait::{DbHashMap, DbValue, Storage};
+use starknet_patricia_storage::storage_trait::{DbHashMap, DbValue, ReadOnlyStorage, Storage};
 
 use crate::block_committer::input::try_node_index_into_contract_address;
 use crate::db::db_layout::DbLayout;

crates/starknet_patricia_storage/src/aerospike_storage.rs

@@ -35,6 +35,7 @@ use crate::storage_trait::{
     EmptyStorageConfig,
     NoStats,
     PatriciaStorageResult,
+    ReadOnlyStorage,
     Storage,
 };
 
@@ -128,10 +129,7 @@ impl AerospikeStorage {
     }
 }
 
-impl Storage for AerospikeStorage {
-    type Stats = NoStats;
-    type Config = EmptyStorageConfig;
-
+impl ReadOnlyStorage for AerospikeStorage {
     async fn get(&mut self, key: &DbKey) -> PatriciaStorageResult<Option<DbValue>> {
         let record = self
             .client
@@ -140,17 +138,6 @@ impl Storage for AerospikeStorage {
         self.extract_value(&record)
     }
 
-    async fn set(&mut self, key: DbKey, value: DbValue) -> PatriciaStorageResult<()> {
-        Ok(self
-            .client
-            .put(
-                &self.config.write_policy,
-                &self.get_key(key)?,
-                &[as_bin!(&self.config.bin_name, value.0)],
-            )
-            .await?)
-    }
-
     async fn mget(&mut self, keys: &[&DbKey]) -> PatriciaStorageResult<Vec<Option<DbValue>>> {
         let mut ops = Vec::new();
         for key in keys.iter() {
@@ -174,6 +161,22 @@ impl Storage for AerospikeStorage {
             })
             .collect::<Result<_, _>>()
     }
+}
+
+impl Storage for AerospikeStorage {
+    type Stats = NoStats;
+    type Config = EmptyStorageConfig;
+
+    async fn set(&mut self, key: DbKey, value: DbValue) -> PatriciaStorageResult<()> {
+        Ok(self
+            .client
+            .put(
+                &self.config.write_policy,
+                &self.get_key(key)?,
+                &[as_bin!(&self.config.bin_name, value.0)],
+            )
+            .await?)
+    }
 
     async fn mset(&mut self, key_to_value: DbHashMap) -> PatriciaStorageResult<()> {
         let keys_and_bins: Vec<(DbKey, Bin)> = key_to_value

crates/starknet_patricia_storage/src/map_storage.rs

@@ -19,6 +19,7 @@ use crate::storage_trait::{
     NoStats,
     NullStorage,
     PatriciaStorageResult,
+    ReadOnlyStorage,
     Storage,
     StorageConfigTrait,
     StorageStats,
@@ -35,6 +36,16 @@ pub struct BorrowedStorage<'a, S: Storage> {
     pub storage: &'a mut S,
 }
 
+impl ReadOnlyStorage for MapStorage {
+    async fn get(&mut self, key: &DbKey) -> PatriciaStorageResult<Option<DbValue>> {
+        Ok(self.0.get(key).cloned())
+    }
+
+    async fn mget(&mut self, keys: &[&DbKey]) -> PatriciaStorageResult<Vec<Option<DbValue>>> {
+        Ok(keys.iter().map(|key| self.0.get(key).cloned()).collect())
+    }
+}
+
 impl Storage for MapStorage {
     type Stats = NoStats;
     type Config = EmptyStorageConfig;
@@ -67,14 +78,6 @@ impl Storage for MapStorage {
         Ok(())
     }
 
-    async fn get(&mut self, key: &DbKey) -> PatriciaStorageResult<Option<DbValue>> {
-        Ok(self.0.get(key).cloned())
-    }
-
-    async fn mget(&mut self, keys: &[&DbKey]) -> PatriciaStorageResult<Vec<Option<DbValue>>> {
-        Ok(keys.iter().map(|key| self.0.get(key).cloned()).collect())
-    }
-
     fn get_stats(&self) -> PatriciaStorageResult<Self::Stats> {
         Ok(NoStats)
     }
@@ -251,10 +254,7 @@ impl<S: Storage> CachedStorage<S> {
     }
 }
 
-impl<S: Storage> Storage for CachedStorage<S> {
-    type Stats = CachedStorageStats<S::Stats>;
-    type Config = CachedStorageConfig<S::Config>;
-
+impl<S: Storage> ReadOnlyStorage for CachedStorage<S> {
     async fn get(&mut self, key: &DbKey) -> PatriciaStorageResult<Option<DbValue>> {
         self.reads += 1;
         if let Some(cached_value) = self.cache.get(key) {
@@ -267,13 +267,6 @@ impl<S: Storage> Storage for CachedStorage<S> {
         Ok(storage_value)
     }
 
-    async fn set(&mut self, key: DbKey, value: DbValue) -> PatriciaStorageResult<()> {
-        self.writes += 1;
-        self.storage.set(key.clone(), value.clone()).await?;
-        self.update_cached_value(&key, &value);
-        Ok(())
-    }
-
     async fn mget(&mut self, keys: &[&DbKey]) -> PatriciaStorageResult<Vec<Option<DbValue>>> {
         let mut values = vec![None; keys.len()]; // The None values are placeholders.
         let mut keys_to_fetch = Vec::new();
@@ -302,6 +295,18 @@ impl<S: Storage> Storage for CachedStorage<S> {
 
         Ok(values)
     }
+}
+
+impl<S: Storage> Storage for CachedStorage<S> {
+    type Stats = CachedStorageStats<S::Stats>;
+    type Config = CachedStorageConfig<S::Config>;
+
+    async fn set(&mut self, key: DbKey, value: DbValue) -> PatriciaStorageResult<()> {
+        self.writes += 1;
+        self.storage.set(key.clone(), value.clone()).await?;
+        self.update_cached_value(&key, &value);
+        Ok(())
+    }
 
     async fn mset(&mut self, key_to_value: DbHashMap) -> PatriciaStorageResult<()> {
         self.writes += u128::try_from(key_to_value.len()).expect("usize should fit in u128");

crates/starknet_patricia_storage/src/mdbx_storage.rs

@@ -23,6 +23,7 @@ use crate::storage_trait::{
     DbValue,
     EmptyStorageConfig,
     PatriciaStorageResult,
+    ReadOnlyStorage,
     Storage,
     StorageStats,
 };
@@ -103,24 +104,13 @@ impl MdbxStorage {
     }
 }
 
-impl Storage for MdbxStorage {
-    type Stats = MdbxStorageStats;
-    type Config = EmptyStorageConfig;
-
+impl ReadOnlyStorage for MdbxStorage {
     async fn get(&mut self, key: &DbKey) -> PatriciaStorageResult<Option<DbValue>> {
         let txn = self.db.begin_ro_txn()?;
         let table = txn.open_table(None)?;
         Ok(txn.get(&table, &key.0)?.map(DbValue))
     }
 
-    async fn set(&mut self, key: DbKey, value: DbValue) -> PatriciaStorageResult<()> {
-        let txn = self.db.begin_rw_txn()?;
-        let table = txn.open_table(None)?;
-        txn.put(&table, key.0, value.0, WriteFlags::UPSERT)?;
-        txn.commit()?;
-        Ok(())
-    }
-
     async fn mget(&mut self, keys: &[&DbKey]) -> PatriciaStorageResult<Vec<Option<DbValue>>> {
         let txn = self.db.begin_ro_txn()?;
         let table = txn.open_table(None)?;
@@ -130,6 +120,19 @@ impl Storage for MdbxStorage {
         }
         Ok(res)
     }
+}
+
+impl Storage for MdbxStorage {
+    type Stats = MdbxStorageStats;
+    type Config = EmptyStorageConfig;
+
+    async fn set(&mut self, key: DbKey, value: DbValue) -> PatriciaStorageResult<()> {
+        let txn = self.db.begin_rw_txn()?;
+        let table = txn.open_table(None)?;
+        txn.put(&table, key.0, value.0, WriteFlags::UPSERT)?;
+        txn.commit()?;
+        Ok(())
+    }
 
     async fn mset(&mut self, key_to_value: DbHashMap) -> PatriciaStorageResult<()> {
         let txn = self.db.begin_rw_txn()?;

crates/starknet_patricia_storage/src/rocksdb_storage.rs

@@ -28,6 +28,7 @@ use crate::storage_trait::{
     DbValue,
     PatriciaStorageError,
     PatriciaStorageResult,
+    ReadOnlyStorage,
     Storage,
     StorageConfigTrait,
     StorageStats,
@@ -275,18 +276,11 @@ impl StorageStats for RocksDbStats {
     }
 }
 
-impl Storage for RocksDbStorage {
-    type Stats = RocksDbStats;
-    type Config = RocksDbStorageConfig;
-
+impl ReadOnlyStorage for RocksDbStorage {
     async fn get(&mut self, key: &DbKey) -> PatriciaStorageResult<Option<DbValue>> {
         Ok(self.db.get(&key.0)?.map(DbValue))
     }
 
-    async fn set(&mut self, key: DbKey, value: DbValue) -> PatriciaStorageResult<()> {
-        Ok(self.db.put_opt(&key.0, &value.0, &self.options.write_options)?)
-    }
-
     async fn mget(&mut self, keys: &[&DbKey]) -> PatriciaStorageResult<Vec<Option<DbValue>>> {
         let raw_keys = keys.iter().map(|k| &k.0);
         let res = self
@@ -297,6 +291,15 @@ impl Storage for RocksDbStorage {
             .collect::<Result<_, _>>()?;
         Ok(res)
     }
+}
+
+impl Storage for RocksDbStorage {
+    type Stats = RocksDbStats;
+    type Config = RocksDbStorageConfig;
+
+    async fn set(&mut self, key: DbKey, value: DbValue) -> PatriciaStorageResult<()> {
+        Ok(self.db.put_opt(&key.0, &value.0, &self.options.write_options)?)
+    }
 
     async fn mset(&mut self, key_to_value: DbHashMap) -> PatriciaStorageResult<()> {
         let mut batch = WriteBatch::default();

crates/starknet_patricia_storage/src/short_key_storage.rs

@@ -11,6 +11,7 @@ use crate::storage_trait::{
     DbValue,
     EmptyStorageConfig,
     PatriciaStorageResult,
+    ReadOnlyStorage,
     Storage,
 };
 
@@ -47,25 +48,27 @@ macro_rules! define_short_key_storage {
             }
         }
 
-        impl<S: Storage> Storage for $name<S> {
-            type Stats = S::Stats;
-            type Config = EmptyStorageConfig;
-
+        impl<S: Storage> ReadOnlyStorage for $name<S> {
             async fn get(&mut self, key: &DbKey) -> PatriciaStorageResult<Option<DbValue>> {
                 self.storage.get(&Self::small_key(key)).await
             }
 
-            async fn set(&mut self, key: DbKey, value: DbValue) -> PatriciaStorageResult<()> {
-                self.storage.set(Self::small_key(&key), value).await
-            }
-
             async 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()).await
             }
+        }
+
+        impl<S: Storage> Storage for $name<S> {
+            type Stats = S::Stats;
+            type Config = EmptyStorageConfig;
+
+            async fn set(&mut self, key: DbKey, value: DbValue) -> PatriciaStorageResult<()> {
+                self.storage.set(Self::small_key(&key), value).await
+            }
 
             async fn mset(&mut self, key_to_value: DbHashMap) -> PatriciaStorageResult<()> {
                 self.storage.mset(

crates/starknet_patricia_storage/src/storage_trait.rs

@@ -99,13 +99,11 @@ pub trait StorageConfigTrait:
     Clone + Debug + Serialize + PartialEq + Validate + SerializeConfig + Default + Send + Sync
 {
 }
-/// A trait for the storage. Does not assume concurrent access is possible - see [AsyncStorage].
-pub trait Storage: Send + Sync {
-    type Stats: StorageStats;
-    type Config: StorageConfigTrait;
 
+/// A read-only view of a storage. Does not assume concurrent access is possible.
+pub trait ReadOnlyStorage: Send + Sync {
     /// Returns value from storage, if it exists.
-    /// Uses a mutable &self to allow changes in the internal state of the storage (e.g.,
+    /// Uses a mutable `&self` to allow changes in the internal state of the storage (e.g.,
     /// for caching).
     // Use explicit desugaring of `async fn` to allow adding trait bounds to the return type, see
     // https://blog.rust-lang.org/2023/12/21/async-fn-rpit-in-traits.html#async-fn-in-public-traits
@@ -115,16 +113,6 @@ pub trait Storage: Send + Sync {
         key: &DbKey,
     ) -> impl Future<Output = PatriciaStorageResult<Option<DbValue>>> + Send;
 
-    /// Sets value in storage. If key already exists, its value is overwritten.
-    // Use explicit desugaring of `async fn` to allow adding trait bounds to the return type, see
-    // https://blog.rust-lang.org/2023/12/21/async-fn-rpit-in-traits.html#async-fn-in-public-traits
-    // for details.
-    fn set(
-        &mut self,
-        key: DbKey,
-        value: DbValue,
-    ) -> impl Future<Output = PatriciaStorageResult<()>> + Send;
-
     /// Returns values from storage in same order of given keys. Value is None for keys that do not
     /// exist.
     // Use explicit desugaring of `async fn` to allow adding trait bounds to the return type, see
@@ -134,6 +122,22 @@ pub trait Storage: Send + Sync {
         &mut self,
         keys: &[&DbKey],
     ) -> impl Future<Output = PatriciaStorageResult<Vec<Option<DbValue>>>> + Send;
+}
+
+/// A trait for the storage. Extends [ReadOnlyStorage] with write operations.
+pub trait Storage: ReadOnlyStorage {
+    type Stats: StorageStats;
+    type Config: StorageConfigTrait;
+
+    /// Sets value in storage. If key already exists, its value is overwritten.
+    // Use explicit desugaring of `async fn` to allow adding trait bounds to the return type, see
+    // https://blog.rust-lang.org/2023/12/21/async-fn-rpit-in-traits.html#async-fn-in-public-traits
+    // for details.
+    fn set(
+        &mut self,
+        key: DbKey,
+        value: DbValue,
+    ) -> impl Future<Output = PatriciaStorageResult<()>> + Send;
 
     /// Sets values in storage.
     // Use explicit desugaring of `async fn` to allow adding trait bounds to the return type, see
@@ -198,21 +202,23 @@ impl StorageConfigTrait for EmptyStorageConfig {}
 #[derive(Clone)]
 pub struct NullStorage;
 
-impl Storage for NullStorage {
-    type Stats = NoStats;
-    type Config = EmptyStorageConfig;
-
+impl ReadOnlyStorage for NullStorage {
     async fn get(&mut self, _key: &DbKey) -> PatriciaStorageResult<Option<DbValue>> {
         Ok(None)
     }
 
-    async fn set(&mut self, _key: DbKey, _value: DbValue) -> PatriciaStorageResult<()> {
-        Ok(())
-    }
-
     async fn mget(&mut self, keys: &[&DbKey]) -> PatriciaStorageResult<Vec<Option<DbValue>>> {
         Ok(vec![None; keys.len()])
     }
+}
+
+impl Storage for NullStorage {
+    type Stats = NoStats;
+    type Config = EmptyStorageConfig;
+
+    async fn set(&mut self, _key: DbKey, _value: DbValue) -> PatriciaStorageResult<()> {
+        Ok(())
+    }
 
     async fn mset(&mut self, _key_to_value: DbHashMap) -> PatriciaStorageResult<()> {
         Ok(())