From 7f68595704987172226c33c8fe7b9b381e63af40 Mon Sep 17 00:00:00 2001 From: Ariel Elperin Date: Thu, 7 May 2026 11:34:39 +0300 Subject: [PATCH] starknet_patricia_storage: two layer storage --- crates/starknet_patricia_storage/src/lib.rs | 1 + .../src/two_layer_storage.rs | 70 +++++++++++++++++++ .../src/two_layer_storage_test.rs | 56 +++++++++++++++ 3 files changed, 127 insertions(+) create mode 100644 crates/starknet_patricia_storage/src/two_layer_storage.rs create mode 100644 crates/starknet_patricia_storage/src/two_layer_storage_test.rs diff --git a/crates/starknet_patricia_storage/src/lib.rs b/crates/starknet_patricia_storage/src/lib.rs index 450de2ac5d2..04685fe4efd 100644 --- a/crates/starknet_patricia_storage/src/lib.rs +++ b/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; diff --git a/crates/starknet_patricia_storage/src/two_layer_storage.rs b/crates/starknet_patricia_storage/src/two_layer_storage.rs new file mode 100644 index 00000000000..add7fbad8f3 --- /dev/null +++ b/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 reads on top of a borrowed base [`ImmutableReadOnlyStorage`]: `overlay` is consulted +/// first via [`ImmutableReadOnlyStorage::get`] / [`ImmutableReadOnlyStorage::mget`]; on miss, +/// reads relay to `base`. +/// +/// [`ReadOnlyStorage::get_mut`] / [`ReadOnlyStorage::mget_mut`] use the same immutable overlay and +/// base paths on overlay misses so the composite implements [`ReadOnlyStorage`] while holding +/// `&'a Base`. Patricia paths reads do not mutate the underlying storage. +/// 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> { + 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>> { + 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) + } +} diff --git a/crates/starknet_patricia_storage/src/two_layer_storage_test.rs b/crates/starknet_patricia_storage/src/two_layer_storage_test.rs new file mode 100644 index 00000000000..38df3048b22 --- /dev/null +++ b/crates/starknet_patricia_storage/src/two_layer_storage_test.rs @@ -0,0 +1,56 @@ +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); + two.overlay.set(key.clone(), DbValue(vec![99])).await.unwrap(); + 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)]); +}