Merge f5790e7 into 855a97f

Author: pablodeymo

crates/blockchain/Cargo.toml

@@ -25,7 +25,7 @@ ethrex-metrics = { path = "./metrics", default-features = false }
 [dev-dependencies]
 serde_json.workspace = true
 hex = "0.4.3"
-tokio.workspace = true
+tokio = { workspace = true, features = [] }
 
 [lib]
 path = "./blockchain.rs"

crates/blockchain/blockchain.rs

@@ -21,8 +21,7 @@ use ethrex_common::types::{
 };
 use ethrex_common::types::{BlobsBundle, ELASTICITY_MULTIPLIER};
 use ethrex_common::{Address, H256};
-use ethrex_storage::error::StoreError;
-use ethrex_storage::Store;
+use ethrex_storage::{error::StoreError, query_plan::QueryPlan, Store};
 use ethrex_vm::{BlockExecutionResult, Evm, EvmEngine};
 use mempool::Mempool;
 use std::collections::HashMap;
@@ -120,6 +119,33 @@ impl Blockchain {
         Ok(execution_result)
     }
 
+    async fn store_block_with_query_plan(
+        &self,
+        block: &Block,
+        execution_result: BlockExecutionResult,
+        account_updates: &[AccountUpdate],
+    ) -> Result<(), ChainError> {
+        // Apply the account updates over the last block's state and compute the new state root
+        let (new_state_root, state_query_plan, accounts_query_plan) = self
+            .storage
+            .apply_account_updates_batch(block.header.parent_hash, account_updates)
+            .await?
+            .ok_or(ChainError::ParentStateNotFound)?;
+
+        // Check state root matches the one in block header
+        validate_state_root(&block.header, new_state_root)?;
+
+        let query_plan = QueryPlan {
+            account_updates: (state_query_plan, accounts_query_plan),
+            block: block.clone(),
+            receipts: (block.hash(), execution_result.receipts),
+        };
+
+        query_plan.apply_to_store(self.storage.clone()).await?;
+
+        Ok(())
+    }
+
     pub async fn store_block(
         &self,
         block: &Block,
@@ -150,7 +176,7 @@ impl Blockchain {
         let since = Instant::now();
         let (res, updates) = self.execute_block(block).await?;
         let executed = Instant::now();
-        let result = self.store_block(block, res, &updates).await;
+        let result = self.store_block_with_query_plan(block, res, &updates).await;
         let stored = Instant::now();
         Self::print_add_block_logs(block, since, executed, stored);
         result

crates/common/trie/trie.rs

@@ -112,19 +112,17 @@ impl Trie {
     /// Remove a value from the trie given its RLP-encoded path.
     /// Returns the value if it was succesfully removed or None if it wasn't part of the trie
     pub fn remove(&mut self, path: PathRLP) -> Result<Option<ValueRLP>, TrieError> {
-        let value;
-        (self.root, value) = if self.root.is_valid() {
-            // If the trie is not empty, call the root node's removal logic.
-            let (node, value) = self
-                .root
-                .get_node(self.db.as_ref())?
-                .ok_or(TrieError::InconsistentTree)?
-                .remove(self.db.as_ref(), Nibbles::from_bytes(&path))?;
+        if !self.root.is_valid() {
+            return Ok(None);
+        }
 
-            (node.map(Into::into).unwrap_or_default(), value)
-        } else {
-            (NodeRef::default(), None)
-        };
+        // If the trie is not empty, call the root node's removal logic.
+        let (node, value) = self
+            .root
+            .get_node(self.db.as_ref())?
+            .ok_or(TrieError::InconsistentTree)?
+            .remove(self.db.as_ref(), Nibbles::from_bytes(&path))?;
+        self.root = node.map(Into::into).unwrap_or_default();
 
         Ok(value)
     }
@@ -147,6 +145,12 @@ impl Trie {
         }
     }
 
+    pub fn hash_prepare_batch(&mut self) -> (H256, Vec<(NodeHash, Vec<u8>)>) {
+        let query_plan = self.commit_vec_aka_state_diff();
+        let ret_hash = self.hash_no_commit();
+        (ret_hash, query_plan)
+    }
+
     /// Compute the hash of the root node and flush any changes into the database.
     ///
     /// This method will also compute the hash of all internal nodes indirectly. It will not clear
@@ -155,12 +159,21 @@ impl Trie {
         if self.root.is_valid() {
             let mut acc = Vec::new();
             self.root.commit(&mut acc);
-            self.db.put_batch(acc)?;
+            self.db.put_batch(acc)?; // we'll try to avoid calling this for every commit
         }
 
         Ok(())
     }
 
+    pub fn commit_vec_aka_state_diff(&mut self) -> Vec<(NodeHash, Vec<u8>)> {
+        let mut acc = Vec::new();
+        if self.root.is_valid() {
+            self.root.commit(&mut acc);
+        }
+
+        acc
+    }
+
     /// Obtain a merkle proof for the given path.
     /// The proof will contain all the encoded nodes traversed until reaching the node where the path is stored (including this last node).
     /// The proof will still be constructed even if the path is not stored in the trie, proving its absence.

crates/storage/api.rs

@@ -6,13 +6,15 @@ use ethrex_common::types::{
 };
 use std::{collections::HashMap, fmt::Debug, panic::RefUnwindSafe};
 
-use crate::{error::StoreError, store::STATE_TRIE_SEGMENTS};
+use crate::{error::StoreError, query_plan::QueryPlan, store::STATE_TRIE_SEGMENTS};
 use ethrex_trie::{Nibbles, Trie};
 
 // We need async_trait because the stabilized feature lacks support for object safety
 // (i.e. dyn StoreEngine)
 #[async_trait::async_trait]
 pub trait StoreEngine: Debug + Send + Sync + RefUnwindSafe {
+    /// Store changes in a batch
+    async fn store_changes(&self, query_plan: QueryPlan) -> Result<(), StoreError>;
     /// Add a batch of blocks in a single transaction.
     /// This will store -> BlockHeader, BlockBody, BlockTransactions, BlockNumber.
     async fn add_blocks(&self, blocks: Vec<Block>) -> Result<(), StoreError>;

crates/storage/lib.rs

@@ -1,5 +1,6 @@
 mod api;
 
+pub mod query_plan;
 #[cfg(any(feature = "libmdbx", feature = "redb"))]
 mod rlp;
 mod store;

crates/storage/query_plan.rs

@@ -0,0 +1,22 @@
+use crate::{error::StoreError, Store};
+use ethrex_common::{
+    types::{Block, Receipt},
+    H256,
+};
+use ethrex_trie::NodeHash;
+
+pub struct QueryPlan {
+    pub account_updates: (
+        Vec<(NodeHash, Vec<u8>)>,                 // vec<(node_hash, node_data)>
+        Vec<(Vec<u8>, Vec<(NodeHash, Vec<u8>)>)>, // hashed_address, vec<(node_hash, node_data)>
+    ),
+    pub block: Block,
+    pub receipts: (H256, Vec<Receipt>),
+}
+
+impl QueryPlan {
+    pub async fn apply_to_store(self, store: Store) -> Result<(), StoreError> {
+        store.store_changes(self).await?;
+        Ok(())
+    }
+}

crates/storage/store.rs

@@ -1,5 +1,6 @@
 use crate::api::StoreEngine;
 use crate::error::StoreError;
+use crate::query_plan::QueryPlan;
 use crate::store_db::in_memory::Store as InMemoryStore;
 #[cfg(feature = "libmdbx")]
 use crate::store_db::libmdbx::Store as LibmdbxStore;
@@ -15,7 +16,7 @@ use ethrex_common::types::{
 };
 use ethrex_rlp::decode::RLPDecode;
 use ethrex_rlp::encode::RLPEncode;
-use ethrex_trie::{Nibbles, Trie};
+use ethrex_trie::{Nibbles, NodeHash, Trie};
 use sha3::{Digest as _, Keccak256};
 use std::collections::{BTreeMap, HashMap};
 use std::fmt::Debug;
@@ -44,6 +45,10 @@ pub enum EngineType {
 }
 
 impl Store {
+    pub async fn store_changes(&self, query_plan: QueryPlan) -> Result<(), StoreError> {
+        self.engine.store_changes(query_plan).await
+    }
+    
     pub fn new(_path: &str, engine_type: EngineType) -> Result<Self, StoreError> {
         info!("Starting storage engine ({engine_type:?})");
         let store = match engine_type {
@@ -312,6 +317,91 @@ impl Store {
         Ok(Some(account_state.nonce))
     }
 
+    /// Applies account updates based on the block's latest storage state
+    /// and returns the new state root after the updates have been applied.
+    pub async fn apply_account_updates_batch(
+        &self,
+        block_hash: BlockHash,
+        account_updates: &[AccountUpdate],
+    ) -> Result<
+        Option<(
+            H256,
+            Vec<(NodeHash, Vec<u8>)>,
+            Vec<(Vec<u8>, Vec<(NodeHash, Vec<u8>)>)>,
+        )>,
+        StoreError,
+    > {
+        let Some(state_trie) = self.state_trie(block_hash)? else {
+            return Ok(None);
+        };
+
+        let (state_trie_hash, state_query_plan, query_plan) = self
+            .apply_account_updates_from_trie_batch(state_trie, account_updates)
+            .await?;
+
+        Ok(Some((state_trie_hash, state_query_plan, query_plan)))
+    }
+
+    pub async fn apply_account_updates_from_trie_batch(
+        &self,
+        mut state_trie: Trie,
+        account_updates: &[AccountUpdate],
+    ) -> Result<
+        (
+            H256,
+            Vec<(NodeHash, Vec<u8>)>,
+            Vec<(Vec<u8>, Vec<(NodeHash, Vec<u8>)>)>,
+        ),
+        StoreError,
+    > {
+        let mut ret_vec_account = Vec::new();
+        for update in account_updates.iter() {
+            let hashed_address = hash_address(&update.address);
+            if update.removed {
+                // Remove account from trie
+                state_trie.remove(hashed_address)?;
+                continue;
+            }
+            // Add or update AccountState in the trie
+            // Fetch current state or create a new state to be inserted
+            let mut account_state = match state_trie.get(&hashed_address)? {
+                Some(encoded_state) => AccountState::decode(&encoded_state)?,
+                None => AccountState::default(),
+            };
+            if let Some(info) = &update.info {
+                account_state.nonce = info.nonce;
+                account_state.balance = info.balance;
+                account_state.code_hash = info.code_hash;
+                // Store updated code in DB
+                if let Some(code) = &update.code {
+                    self.add_account_code(info.code_hash, code.clone()).await?;
+                }
+            }
+            // Store the added storage in the account's storage trie and compute its new root
+            if !update.added_storage.is_empty() {
+                let mut storage_trie = self.engine.open_storage_trie(
+                    H256::from_slice(&hashed_address),
+                    account_state.storage_root,
+                );
+                for (storage_key, storage_value) in &update.added_storage {
+                    let hashed_key = hash_key(storage_key);
+                    if storage_value.is_zero() {
+                        storage_trie.remove(hashed_key)?;
+                    } else {
+                        storage_trie.insert(hashed_key, storage_value.encode_to_vec())?;
+                    }
+                }
+                let (storage_hash, storage_query_plan) = storage_trie.hash_prepare_batch();
+                account_state.storage_root = storage_hash;
+                ret_vec_account.push((hashed_address.clone(), storage_query_plan));
+            }
+            state_trie.insert(hashed_address, account_state.encode_to_vec())?;
+        }
+        let (state_hash, state_query_plan) = state_trie.hash_prepare_batch();
+
+        Ok((state_hash, state_query_plan, ret_vec_account))
+    }
+
     /// Applies account updates based on the block's latest storage state
     /// and returns the new state root after the updates have been applied.
     pub async fn apply_account_updates(

crates/storage/store_db/in_memory.rs

@@ -1,7 +1,5 @@
 use crate::{
-    api::StoreEngine,
-    error::StoreError,
-    store::{MAX_SNAPSHOT_READS, STATE_TRIE_SEGMENTS},
+    api::StoreEngine, error::StoreError, query_plan::QueryPlan, store::{MAX_SNAPSHOT_READS, STATE_TRIE_SEGMENTS}
 };
 use bytes::Bytes;
 use ethereum_types::{H256, U256};
@@ -87,6 +85,10 @@ impl Store {
 
 #[async_trait::async_trait]
 impl StoreEngine for Store {
+    async fn store_changes(&self, _query_plan: QueryPlan) -> Result<(), StoreError> {
+       todo!()
+    }
+    
     fn get_block_header(&self, block_number: u64) -> Result<Option<BlockHeader>, StoreError> {
         let store = self.inner();
         if let Some(hash) = store.canonical_hashes.get(&block_number) {