starknet_patricia_storage: move is_gatherable to ReadOnlyStorage

Author: ArielElp

crates/starknet_patricia_storage/src/aerospike_storage.rs

@@ -173,6 +173,10 @@ impl ReadOnlyStorage for AerospikeStorage {
     async fn mget_mut(&mut self, keys: &[&DbKey]) -> PatriciaStorageResult<Vec<Option<DbValue>>> {
         ImmutableReadOnlyStorage::mget(self, keys).await
     }
+
+    fn as_gatherable_storage(&mut self) -> Option<&mut impl GatherableStorage> {
+        Some(self)
+    }
 }
 
 impl Storage for AerospikeStorage {
@@ -256,10 +260,6 @@ impl Storage for AerospikeStorage {
     fn get_async_self(&self) -> Option<impl AsyncStorage> {
         Some(self.clone())
     }
-
-    fn as_gatherable_storage(&mut self) -> Option<&mut impl GatherableStorage> {
-        Some(self)
-    }
 }
 
 impl GatherableStorage for AerospikeStorage {}

crates/starknet_patricia_storage/src/map_storage.rs

@@ -60,6 +60,10 @@ impl ReadOnlyStorage for MapStorage {
     async fn mget_mut(&mut self, keys: &[&DbKey]) -> PatriciaStorageResult<Vec<Option<DbValue>>> {
         ImmutableReadOnlyStorage::mget(self, keys).await
     }
+
+    fn as_gatherable_storage(&mut self) -> Option<&mut impl GatherableStorage> {
+        Some(self)
+    }
 }
 
 impl Storage for MapStorage {
@@ -102,10 +106,6 @@ impl Storage for MapStorage {
         // Need a concrete Option type.
         None::<NullStorage>
     }
-
-    fn as_gatherable_storage(&mut self) -> Option<&mut impl GatherableStorage> {
-        Some(self)
-    }
 }
 
 impl GatherableStorage for MapStorage {}
@@ -303,7 +303,7 @@ impl<S: Storage + ImmutableReadOnlyStorage + 'static> GatherableStorage for Cach
 }
 
 // TODO(Nimrod): Find a way to share the implementation with `ImmutableReadOnlyStorage`.
-impl<S: Storage> ReadOnlyStorage for CachedStorage<S> {
+impl<S: Storage + ImmutableReadOnlyStorage + 'static> ReadOnlyStorage for CachedStorage<S> {
     async fn get_mut(&mut self, key: &DbKey) -> PatriciaStorageResult<Option<DbValue>> {
         self.reads.fetch_add(1, Ordering::Relaxed);
         if let Some(cached_value) = self.cache.get(key) {
@@ -346,6 +346,10 @@ impl<S: Storage> ReadOnlyStorage for CachedStorage<S> {
 
         Ok(values)
     }
+
+    fn as_gatherable_storage(&mut self) -> Option<&mut impl GatherableStorage> {
+        Some(self)
+    }
 }
 
 impl<S: Storage + ImmutableReadOnlyStorage + 'static> Storage for CachedStorage<S> {
@@ -420,8 +424,4 @@ impl<S: Storage + ImmutableReadOnlyStorage + 'static> Storage for CachedStorage<
         // Need a concrete Option type.
         None::<NullStorage>
     }
-
-    fn as_gatherable_storage(&mut self) -> Option<&mut impl GatherableStorage> {
-        Some(self)
-    }
 }

crates/starknet_patricia_storage/src/mdbx_storage.rs

@@ -132,6 +132,10 @@ impl ReadOnlyStorage for MdbxStorage {
     async fn mget_mut(&mut self, keys: &[&DbKey]) -> PatriciaStorageResult<Vec<Option<DbValue>>> {
         ImmutableReadOnlyStorage::mget(self, keys).await
     }
+
+    fn as_gatherable_storage(&mut self) -> Option<&mut impl GatherableStorage> {
+        Some(self)
+    }
 }
 
 impl Storage for MdbxStorage {
@@ -191,10 +195,6 @@ impl Storage for MdbxStorage {
     fn get_async_self(&self) -> Option<impl AsyncStorage> {
         Some(self.clone())
     }
-
-    fn as_gatherable_storage(&mut self) -> Option<&mut impl GatherableStorage> {
-        Some(self)
-    }
 }
 
 impl GatherableStorage for MdbxStorage {}

crates/starknet_patricia_storage/src/reads_collector_storage.rs

@@ -2,7 +2,9 @@ use crate::storage_trait::{
     DbHashMap,
     DbKey,
     DbValue,
+    GatherableStorage,
     ImmutableReadOnlyStorage,
+    NullStorage,
     PatriciaStorageResult,
     ReadOnlyStorage,
 };
@@ -43,4 +45,8 @@ impl<'a, S: ImmutableReadOnlyStorage> ReadOnlyStorage for ReadsCollectorStorage<
         }
         Ok(values)
     }
+
+    fn as_gatherable_storage(&mut self) -> Option<&mut impl GatherableStorage> {
+        None::<&mut NullStorage>
+    }
 }

crates/starknet_patricia_storage/src/rocksdb_storage.rs

@@ -331,6 +331,10 @@ impl ReadOnlyStorage for RocksDbStorage {
     async fn mget_mut(&mut self, keys: &[&DbKey]) -> PatriciaStorageResult<Vec<Option<DbValue>>> {
         ImmutableReadOnlyStorage::mget(self, keys).await
     }
+
+    fn as_gatherable_storage(&mut self) -> Option<&mut impl GatherableStorage> {
+        Some(self)
+    }
 }
 
 impl Storage for RocksDbStorage {
@@ -376,10 +380,6 @@ impl Storage for RocksDbStorage {
     fn get_async_self(&self) -> Option<impl AsyncStorage> {
         Some(self.clone())
     }
-
-    fn as_gatherable_storage(&mut self) -> Option<&mut impl GatherableStorage> {
-        Some(self)
-    }
 }
 
 impl GatherableStorage for RocksDbStorage {}

crates/starknet_patricia_storage/src/short_key_storage.rs

@@ -82,6 +82,10 @@ macro_rules! define_short_key_storage {
                     .mget_mut(small_keys.iter().collect::<Vec<&DbKey>>().as_slice())
                     .await
             }
+
+            fn as_gatherable_storage(&mut self) -> Option<&mut impl GatherableStorage> {
+                None::<&mut NullStorage>
+            }
         }
 
         impl<S: Storage> Storage for $name<S> {
@@ -123,10 +127,6 @@ macro_rules! define_short_key_storage {
             fn get_async_self(&self) -> Option<impl AsyncStorage> {
                 Some($name::new(self.storage.get_async_self()?))
             }
-
-            fn as_gatherable_storage(&mut self) -> Option<&mut impl GatherableStorage> {
-                None::<&mut NullStorage>
-            }
         }
 
         impl<S: AsyncStorage> Clone for $name<S> {

crates/starknet_patricia_storage/src/storage_trait.rs

@@ -197,6 +197,10 @@ pub trait ReadOnlyStorage: Send + Sync {
         &mut self,
         keys: &[&DbKey],
     ) -> impl Future<Output = PatriciaStorageResult<Vec<Option<DbValue>>>> + Send;
+
+    /// If the storage supports concurrent task execution via [GatherableStorage::gather], returns
+    /// a mutable reference to it. Returns `None` otherwise.
+    fn as_gatherable_storage(&mut self) -> Option<&mut impl GatherableStorage>;
 }
 
 /// A trait for the storage. Extends [ReadOnlyStorage] with write operations.
@@ -248,10 +252,6 @@ pub trait Storage: ReadOnlyStorage {
 
     /// If the storage is async, returns an instance of the async storage.
     fn get_async_self(&self) -> Option<impl AsyncStorage>;
-
-    /// If the storage supports concurrent task execution via [GatherableStorage::gather], returns
-    /// a mutable reference to it. Returns `None` otherwise.
-    fn as_gatherable_storage(&mut self) -> Option<&mut impl GatherableStorage>;
 }
 
 /// A trait wrapper for [Storage] that supports concurrency.
@@ -299,6 +299,10 @@ impl ReadOnlyStorage for NullStorage {
     async fn mget_mut(&mut self, keys: &[&DbKey]) -> PatriciaStorageResult<Vec<Option<DbValue>>> {
         ImmutableReadOnlyStorage::mget(self, keys).await
     }
+
+    fn as_gatherable_storage(&mut self) -> Option<&mut impl GatherableStorage> {
+        Some(self)
+    }
 }
 
 impl Storage for NullStorage {
@@ -331,10 +335,6 @@ impl Storage for NullStorage {
     fn get_async_self(&self) -> Option<impl AsyncStorage> {
         Some(self.clone())
     }
-
-    fn as_gatherable_storage(&mut self) -> Option<&mut impl GatherableStorage> {
-        Some(self)
-    }
 }
 
 impl GatherableStorage for NullStorage {}

crates/starknet_patricia_storage/src/two_layer_storage.rs

@@ -1,7 +1,9 @@
 use crate::storage_trait::{
     DbKey,
     DbValue,
+    GatherableStorage,
     ImmutableReadOnlyStorage,
+    NullStorage,
     PatriciaStorageResult,
     ReadOnlyStorage,
 };
@@ -67,4 +69,8 @@ where
         }
         Ok(out)
     }
+
+    fn as_gatherable_storage(&mut self) -> Option<&mut impl GatherableStorage> {
+        None::<&mut NullStorage>
+    }
 }