perf(core): make disklayer in snapshots use the database

CHANGELOG.md

@@ -5,6 +5,7 @@
 ### 2025-05-22
 
 - Add block snapshots with diff layers. [2664](https://github.com/lambdaclass/ethrex/pull/2664)
+- Make disk layer in snapshots use the database. [2848](https://github.com/lambdaclass/ethrex/pull/2848)
 
 ### 2025-05-20
 

crates/l2/prover/bench/src/rpc/mod.rs

@@ -208,7 +208,7 @@ pub async fn get_account(
     let trie = Trie::from_nodes(Some(root), &other)
         .map_err(|err| format!("failed to build account proof trie: {err}"))?;
     if trie
-        .get(&hash_address(address))
+        .get(hash_address(address))
         .map_err(|err| format!("failed get account from proof trie: {err}"))?
         .is_none()
     {

crates/storage/api.rs

@@ -316,6 +316,13 @@ pub trait StoreEngine: Debug + Send + Sync + RefUnwindSafe {
     /// Clears all checkpoint data created during the last snap sync
     async fn clear_snap_state(&self) -> Result<(), StoreError>;
 
+    /// Write an account batch into the current state snapshot. Blocking non-async version.
+    fn write_snapshot_account_batch_blocking(
+        &self,
+        account_hashes: Vec<H256>,
+        account_states: Vec<AccountState>,
+    ) -> Result<(), StoreError>;
+
     /// Write an account batch into the current state snapshot
     async fn write_snapshot_account_batch(
         &self,
@@ -339,6 +346,14 @@ pub trait StoreEngine: Debug + Send + Sync + RefUnwindSafe {
         storage_values: Vec<Vec<U256>>,
     ) -> Result<(), StoreError>;
 
+    /// Write multiple storage batches belonging to different accounts into the current storage snapshot. Blocking non-async version.
+    fn write_snapshot_storage_batches_blocking(
+        &self,
+        account_hashes: Vec<H256>,
+        storage_keys: Vec<Vec<H256>>,
+        storage_values: Vec<Vec<U256>>,
+    ) -> Result<(), StoreError>;
+
     /// Set the latest root of the rebuilt state trie and the last downloaded hashes from each segment
     async fn set_state_trie_rebuild_checkpoint(
         &self,
@@ -374,6 +389,16 @@ pub trait StoreEngine: Debug + Send + Sync + RefUnwindSafe {
         account_hash: H256,
     ) -> Result<Vec<(H256, U256)>, StoreError>;
 
+    /// Gets a single account from the snapshot state.
+    fn get_account_snapshot(&self, account_hash: H256) -> Result<Option<AccountState>, StoreError>;
+
+    /// Gets a single storage value from the snapshot state.
+    fn get_storage_snapshot(
+        &self,
+        account_hash: H256,
+        storage_hash: H256,
+    ) -> Result<Option<U256>, StoreError>;
+
     /// The `forkchoice_update` and `new_payload` methods require the `latest_valid_hash`
     /// when processing an invalid payload. To provide this, we must track invalid chains.
     ///

crates/storage/snapshot/difflayer.rs

@@ -20,7 +20,7 @@ pub struct DiffLayer {
     state_root: H256,
     stale: bool,
     accounts: HashMap<H256, Option<AccountState>>, // None if deleted
-    storage: HashMap<H256, HashMap<H256, U256>>,
+    storage: HashMap<H256, HashMap<H256, Option<U256>>>,
     /// tracks all diffed items up to disk layer
     pub(crate) diffed: Bloom,
 }
@@ -32,7 +32,7 @@ impl DiffLayer {
         block_hash: BlockHash,
         state_root: H256,
         accounts: HashMap<H256, Option<AccountState>>,
-        storage: HashMap<H256, HashMap<H256, U256>>,
+        storage: HashMap<H256, HashMap<H256, Option<U256>>>,
     ) -> Self {
         DiffLayer {
             origin: origin.clone(),
@@ -97,7 +97,7 @@ impl DiffLayer {
 
         // If bloom misses we can skip diff layers
         if !hit {
-            return self.origin.get_account(hash, layers);
+            return self.origin.get_account(hash);
         }
 
         // Start traversing layers.
@@ -119,7 +119,7 @@ impl DiffLayer {
             .get(&account_hash)
             .and_then(|x| x.get(&storage_hash))
         {
-            return Ok(Some(*value));
+            return Ok(*value);
         }
 
         let bloom_hash = account_hash ^ storage_hash;
@@ -129,7 +129,7 @@ impl DiffLayer {
 
         // If bloom misses we can skip diff layers
         if !hit {
-            return self.origin.get_storage(account_hash, storage_hash, layers);
+            return self.origin.get_storage(account_hash, storage_hash);
         }
 
         // Start traversing layers.
@@ -155,7 +155,7 @@ impl DiffLayer {
         block: BlockHash,
         state_root: H256,
         accounts: HashMap<H256, Option<AccountState>>,
-        storage: HashMap<H256, HashMap<H256, U256>>,
+        storage: HashMap<H256, HashMap<H256, Option<U256>>>,
     ) -> DiffLayer {
         let mut layer = DiffLayer::new(
             self.block_hash,
@@ -196,7 +196,7 @@ impl DiffLayer {
 
         // delegate to parent
         match &layers[&self.parent] {
-            Layer::DiskLayer(disk_layer) => disk_layer.get_account(hash, layers),
+            Layer::DiskLayer(disk_layer) => disk_layer.get_account(hash),
             Layer::DiffLayer(diff_layer) => diff_layer
                 .read()
                 .map_err(|error| SnapshotError::LockError(error.to_string()))?
@@ -220,14 +220,12 @@ impl DiffLayer {
             .get(&account_hash)
             .and_then(|x| x.get(&storage_hash))
         {
-            return Ok(Some(*value));
+            return Ok(*value);
         }
 
         // delegate to parent
         match &layers[&self.parent] {
-            Layer::DiskLayer(disk_layer) => {
-                disk_layer.get_storage(account_hash, storage_hash, layers)
-            }
+            Layer::DiskLayer(disk_layer) => disk_layer.get_storage(account_hash, storage_hash),
             Layer::DiffLayer(diff_layer) => diff_layer
                 .read()
                 .map_err(|error| SnapshotError::LockError(error.to_string()))?
@@ -239,7 +237,7 @@ impl DiffLayer {
         self.accounts.extend(accounts);
     }
 
-    pub fn add_storage(&mut self, storage: HashMap<H256, HashMap<H256, U256>>) {
+    pub fn add_storage(&mut self, storage: HashMap<H256, HashMap<H256, Option<U256>>>) {
         for (address, st) in storage.iter() {
             let entry = self.storage.entry(*address).or_default();
             entry.extend(st);
@@ -250,7 +248,7 @@ impl DiffLayer {
         self.accounts.clone()
     }
 
-    pub fn storage(&self) -> HashMap<H256, HashMap<H256, U256>> {
+    pub fn storage(&self) -> HashMap<H256, HashMap<H256, Option<U256>>> {
         self.storage.clone()
     }