starknet_patricia_storage: two layer storage

crates/starknet_patricia_storage/src/lib.rs

@@ -15,3 +15,4 @@ pub mod short_key_storage;
 #[cfg(test)]
 pub mod storage_test;
 pub mod storage_trait;
+pub mod two_layer_storage;

crates/starknet_patricia_storage/src/two_layer_storage.rs

@@ -0,0 +1,70 @@
+use crate::storage_trait::{
+    DbKey,
+    DbValue,
+    ImmutableReadOnlyStorage,
+    PatriciaStorageResult,
+    ReadOnlyStorage,
+};
+
+#[cfg(test)]
+#[path = "two_layer_storage_test.rs"]
+mod two_layer_storage_test;
+
+/// Overlay storage on top of a borrowed base [`ImmutableReadOnlyStorage`]: `overlay` is consulted
+/// first via [`ImmutableReadOnlyStorage::get`] / [`ImmutableReadOnlyStorage::mget`]; on miss,
+/// reads relay to `base`.
+///
+/// On misses, [`ReadOnlyStorage::get_mut`] / [`ReadOnlyStorage::mget_mut`] use get/mget on the
+/// underlying storage to allow the composite to implement [`ReadOnlyStorage`] while holding `&'a
+/// Base`. This allows passing `TwoLayerStorage` to `fetch_all_patricia_paths`, which requires
+/// `&mut` [`ReadOnlyStorage`].
+pub struct TwoLayerStorage<'a, Overlay, Base>
+where
+    Overlay: ImmutableReadOnlyStorage + Sync,
+    Base: ImmutableReadOnlyStorage + Sync + ?Sized,
+{
+    overlay: Overlay,
+    base: &'a Base,
+}
+
+impl<'a, Overlay, Base> TwoLayerStorage<'a, Overlay, Base>
+where
+    Overlay: ImmutableReadOnlyStorage + Sync,
+    Base: ImmutableReadOnlyStorage + Sync + ?Sized,
+{
+    pub fn new(overlay: Overlay, base: &'a Base) -> Self {
+        Self { overlay, base }
+    }
+}
+
+impl<'a, Overlay, Base> ReadOnlyStorage for TwoLayerStorage<'a, Overlay, Base>
+where
+    Overlay: ImmutableReadOnlyStorage + Sync,
+    Base: ImmutableReadOnlyStorage + Sync + ?Sized,
+{
+    async fn get_mut(&mut self, key: &DbKey) -> PatriciaStorageResult<Option<DbValue>> {
+        Ok(match self.overlay.get(key).await? {
+            Some(v) => Some(v),
+            None => self.base.get(key).await?,
+        })
+    }
+
+    async fn mget_mut(&mut self, keys: &[&DbKey]) -> PatriciaStorageResult<Vec<Option<DbValue>>> {
+        let mut out = self.overlay.mget(keys).await?;
+        let mut miss_indices = Vec::new();
+        let mut miss_keys = Vec::new();
+        for (i, v) in out.iter().enumerate() {
+            if v.is_none() {
+                miss_indices.push(i);
+                miss_keys.push(keys[i]);
+            }
+        }
+        if !miss_keys.is_empty() {
+            let fetched = self.base.mget(&miss_keys).await?;
+            for (idx, val) in miss_indices.into_iter().zip(fetched) {
+                out[idx] = val;
+            }
+        }
+        Ok(out)
+    }
+}

crates/starknet_patricia_storage/src/two_layer_storage_test.rs

@@ -0,0 +1,58 @@
+use super::TwoLayerStorage;
+use crate::map_storage::MapStorage;
+use crate::storage_trait::{DbKey, DbValue, ReadOnlyStorage, Storage};
+
+#[tokio::test]
+async fn read_falls_through_to_base() {
+    let key = DbKey(vec![1, 2, 3]);
+    let val = DbValue(vec![9]);
+    let mut base = MapStorage::default();
+    base.0.insert(key.clone(), val.clone());
+
+    let mut two = TwoLayerStorage::new(MapStorage::default(), &base);
+    assert_eq!(two.get_mut(&key).await.unwrap(), Some(val));
+}
+
+#[tokio::test]
+async fn overlay_shadows_base() {
+    let key = DbKey(vec![1]);
+    let base_val = DbValue(vec![1]);
+    let over_val = DbValue(vec![2]);
+    let mut base = MapStorage::default();
+    base.0.insert(key.clone(), base_val);
+
+    let mut two = TwoLayerStorage::new(MapStorage::default(), &base);
+    two.overlay.set(key.clone(), over_val.clone()).await.unwrap();
+    assert_eq!(two.get_mut(&key).await.unwrap(), Some(over_val));
+}
+
+#[tokio::test]
+async fn delete_drops_overlay_entry_and_sees_base() {
+    let key = DbKey(vec![7]);
+    let base_val = DbValue(vec![42]);
+    let mut base = MapStorage::default();
+    base.0.insert(key.clone(), base_val.clone());
+
+    let mut two = TwoLayerStorage::new(MapStorage::default(), &base);
+    let overlay_val = DbValue(vec![99]);
+    two.overlay.set(key.clone(), overlay_val.clone()).await.unwrap();
+    assert_eq!(two.get_mut(&key).await.unwrap(), Some(overlay_val));
+    two.overlay.delete(&key).await.unwrap();
+    assert_eq!(two.get_mut(&key).await.unwrap(), Some(base_val));
+}
+
+#[tokio::test]
+async fn mget_mut_uses_immutable_base_mget_on_miss() {
+    let key_base_only = DbKey(vec![3]);
+    let key_overlay = DbKey(vec![4]);
+    let base_val = DbValue(vec![11]);
+    let over_val = DbValue(vec![22]);
+    let mut base = MapStorage::default();
+    base.0.insert(key_base_only.clone(), base_val.clone());
+
+    let mut layered = TwoLayerStorage::new(MapStorage::default(), &base);
+    layered.overlay.set(key_overlay.clone(), over_val.clone()).await.unwrap();
+
+    let keys = [&key_base_only, &key_overlay];
+    assert_eq!(layered.mget_mut(&keys).await.unwrap(), vec![Some(base_val), Some(over_val)]);
+}