opt(state root) compact trie node serialization and remove comptiable trie updates

Author: AshinGau

Cargo.lock

@@ -8,6 +8,7 @@ use reth_primitives_traits::NodePrimitives;
 use reth_provider::{
     providers::ProviderNodeTypes, writer::UnifiedStorageWriter, BlockHashReader,
     ChainStateBlockWriter, DatabaseProviderFactory, ProviderFactory, StaticFileProviderFactory,
+    PERSIST_BLOCK_CACHE,
 };
 use reth_prune::{PrunerError, PrunerOutput, PrunerWithFactory};
 use reth_stages_api::{MetricEvent, MetricEventsSender};
@@ -151,12 +152,13 @@ where
         });
 
         let num_blocks = blocks.len();
-        if last_block_hash_num.is_some() {
+        if let Some(last_block_hn) = &last_block_hash_num {
             let provider_rw = self.provider.database_provider_rw()?;
             let static_file_provider = self.provider.static_file_provider();
 
             UnifiedStorageWriter::from(&provider_rw, &static_file_provider).save_blocks(blocks)?;
             UnifiedStorageWriter::commit(provider_rw)?;
+            PERSIST_BLOCK_CACHE.persist_tip(last_block_hn.number);
         }
         let elapsed = start_time.elapsed();
         self.metrics.save_blocks_duration_seconds.record(elapsed);

crates/engine/tree/src/persistence.rs

@@ -9,19 +9,13 @@ pub struct Config {
     pub disable_pipe_execution: bool,
     /// Whether to disable the Grevm executor. default true.
     pub disable_grevm: bool,
-    /// Whether compatible with trie output. default false.
-    pub compatible_trie_output: bool,
 }
 
 /// Global configuration instance, initialized lazily.
 pub static CONFIG: LazyLock<Config> = LazyLock::new(|| {
     let config = Config {
         disable_pipe_execution: std::env::var("GRETH_DISABLE_PIPE_EXECUTION").is_ok(),
         disable_grevm: std::env::var("GRETH_DISABLE_GREVM").is_ok(),
-        compatible_trie_output: std::env::var("COMPATIBLE_TRIE_OUTPUT")
-            .ok()
-            .and_then(|v| v.parse().ok())
-            .unwrap_or(false),
     };
     println!("Gravity Reth config: {config:#?}");
     config

crates/gravity-primitives/src/config.rs

@@ -10,10 +10,7 @@ use reth_revm::{
     state::AccountInfo, DatabaseRef,
 };
 use reth_storage_api::StateProviderFactory;
-use reth_trie::{
-    updates::{TrieUpdates, TrieUpdatesV2},
-    HashedPostState,
-};
+use reth_trie::{updates::TrieUpdatesV2, HashedPostState};
 use reth_trie_parallel::nested_hash::NestedStateRoot;
 use std::{collections::BTreeMap, sync::Mutex};
 
@@ -58,14 +55,10 @@ where
         Ok(BlockViewProvider::new(StateProviderDatabase::new(state), Some(self.cache.clone())))
     }
 
-    fn state_root(
-        &self,
-        hashed_state: &HashedPostState,
-        compatible: bool,
-    ) -> ProviderResult<(B256, TrieUpdatesV2, Option<TrieUpdates>)> {
+    fn state_root(&self, hashed_state: &HashedPostState) -> ProviderResult<(B256, TrieUpdatesV2)> {
         let provider = || self.client.database_provider_ro().map(|db| db.into_tx());
         let nested_hash = NestedStateRoot::new(provider, Some(self.cache.clone()));
-        nested_hash.calculate(hashed_state, compatible)
+        nested_hash.calculate(hashed_state)
     }
 
     fn insert_block_id(&self, block_number: u64, block_id: B256) {

crates/gravity-storage/src/block_view_storage/mod.rs

@@ -6,10 +6,7 @@ pub mod block_view_storage;
 use alloy_primitives::B256;
 use reth_evm::ParallelDatabase;
 use reth_provider::ProviderResult;
-use reth_trie::{
-    updates::{TrieUpdates, TrieUpdatesV2},
-    HashedPostState,
-};
+use reth_trie::{updates::TrieUpdatesV2, HashedPostState};
 
 /// Gravity storage for pipeline execution
 pub trait GravityStorage: Send + Sync + 'static {
@@ -20,11 +17,7 @@ pub trait GravityStorage: Send + Sync + 'static {
     fn get_state_view(&self) -> ProviderResult<Self::StateView>;
 
     /// calculate state root
-    fn state_root(
-        &self,
-        hashed_state: &HashedPostState,
-        compatible: bool,
-    ) -> ProviderResult<(B256, TrieUpdatesV2, Option<TrieUpdates>)>;
+    fn state_root(&self, hashed_state: &HashedPostState) -> ProviderResult<(B256, TrieUpdatesV2)>;
 
     /// Insert the mapping from `block_number` to `block_id`
     fn insert_block_id(&self, block_number: u64, block_id: B256);

crates/gravity-storage/src/lib.rs

@@ -53,7 +53,6 @@ use tokio::sync::{
     oneshot, Mutex,
 };
 
-use gravity_primitives::CONFIG;
 use reth_trie::{HashedPostState, KeccakKeyHasher};
 use tracing::*;
 
@@ -333,8 +332,7 @@ impl<Storage: GravityStorage> Core<Storage> {
         // Merkling the state trie
         self.merklize_barrier.wait(block_number - 1).await.unwrap();
         let start_time = Instant::now();
-        let (state_root, trie_updates, compatible_trie_pudates) =
-            self.storage.state_root(&hashed_state, CONFIG.compatible_trie_output).unwrap();
+        let (state_root, trie_updates) = self.storage.state_root(&hashed_state).unwrap();
         let write_start = Instant::now();
         self.cache.write_trie_updates(&trie_updates, block_number);
         self.metrics.cache_trie_state.record(write_start.elapsed());
@@ -367,23 +365,13 @@ impl<Storage: GravityStorage> Core<Storage> {
         );
 
         let gas_used = sealed_block.gas_used;
-        let executed_block = if let Some(compatible_trie_updates) = compatible_trie_pudates {
-            ExecutedBlockWithTrieUpdates::new(
-                Arc::new(RecoveredBlock::new_sealed(sealed_block, senders)),
-                Arc::new(execution_outcome),
-                Arc::new(hashed_state),
-                ExecutedTrieUpdates::Present(Arc::new(compatible_trie_updates)),
-                Arc::new(trie_updates),
-            )
-        } else {
-            ExecutedBlockWithTrieUpdates::new(
-                Arc::new(RecoveredBlock::new_sealed(sealed_block, senders)),
-                Arc::new(execution_outcome),
-                Arc::new(hashed_state),
-                ExecutedTrieUpdates::empty(),
-                Arc::new(trie_updates),
-            )
-        };
+        let executed_block = ExecutedBlockWithTrieUpdates::new(
+            Arc::new(RecoveredBlock::new_sealed(sealed_block, senders)),
+            Arc::new(execution_outcome),
+            Arc::new(hashed_state),
+            ExecutedTrieUpdates::empty(),
+            Arc::new(trie_updates),
+        );
         let execution_result =
             ExecutionResult { block_id, block_number, block_hash, txs_info, gravity_events };
 

crates/pipe-exec-layer-ext-v2/execute/src/lib.rs

@@ -145,7 +145,7 @@ async fn run_pipe(
     let nested_provider = || provider.database_provider_ro().map(|db| db.into_tx());
     let nested_hash = NestedStateRoot::new(nested_provider, None);
     let hashed_state = nested_hash.read_hashed_state(None).unwrap();
-    let (root_hash, trie_updates, _) = nested_hash.calculate(&hashed_state, false).unwrap();
+    let (root_hash, trie_updates) = nested_hash.calculate(&hashed_state).unwrap();
     let trie_write = provider.database_provider_rw().unwrap();
     trie_write.write_trie_updatesv2(&trie_updates).unwrap();
     UnifiedStorageWriter::commit_unwind(trie_write)?;

crates/pipe-exec-layer-ext-v2/execute/tests/pipe_test.rs

@@ -8,8 +8,9 @@ use reth_db_api::{
 };
 use reth_primitives_traits::{GotExpected, SealedHeader};
 use reth_provider::{
-    DBProvider, HeaderProvider, ProviderError, StageCheckpointReader, StageCheckpointWriter,
-    StatsReader, TrieWriter, TrieWriterV2, EXECUTION_MERKLE_CHANNEL,
+    DBProvider, HeaderProvider, PersistBlockCache, ProviderError, StageCheckpointReader,
+    StageCheckpointWriter, StatsReader, TrieWriter, TrieWriterV2, EXECUTION_MERKLE_CHANNEL,
+    PERSIST_BLOCK_CACHE,
 };
 use reth_stages_api::{
     BlockErrorKind, BoxedConcurrentProvider, EntitiesCheckpoint, ExecInput, ExecOutput,
@@ -197,16 +198,21 @@ where
             (target_block_root, input.checkpoint().entities_stage_checkpoint().unwrap_or_default())
         } else {
             debug!(target: "sync::stages::merkle::exec", current = ?current_block_number, target = ?to_block, "Updating trie in chunks");
+            let cache: PersistBlockCache = PERSIST_BLOCK_CACHE.clone();
             // Use optimized nested hash algorithm for state root calculation
-            let nested_state_root = NestedStateRoot::new(
-                || provider_ro().map(|db| db.into_tx()),
-                None, // No cache for history sync
-            );
+            let nested_state_root =
+                NestedStateRoot::new(|| provider_ro().map(|db| db.into_tx()), Some(cache.clone()));
             // Read the hashed state from database for the specified range
-            let hashed_state = EXECUTION_MERKLE_CHANNEL.consume(from_block);
-            let (final_root, trie_updates_v2, _compatible_updates) =
-                nested_state_root.calculate(&hashed_state, false)?;
+            let hashed_state =
+                if let Some(hashed_state) = EXECUTION_MERKLE_CHANNEL.consume(from_block) {
+                    hashed_state
+                } else {
+                    nested_state_root.read_hashed_state(Some(range))?
+                };
+            let (final_root, trie_updates_v2) = nested_state_root.calculate(&hashed_state)?;
+            cache.write_trie_updates(&trie_updates_v2, to_block);
             provider.write_trie_updatesv2(&trie_updates_v2)?;
+            cache.persist_tip(to_block);
 
             let total_hashed_entries = (provider.count_entries::<tables::HashedAccounts>()? +
                 provider.count_entries::<tables::HashedStorages>()?)
@@ -278,8 +284,7 @@ where
             let nested_state_root =
                 NestedStateRoot::new(|| provider_ro().map(|db| db.into_tx()), None);
             let hashed_state = nested_state_root.read_hashed_state(Some(range))?;
-            let (block_root, trie_updates_v2, _compatible_updates) =
-                nested_state_root.calculate(&hashed_state, false)?;
+            let (block_root, trie_updates_v2) = nested_state_root.calculate(&hashed_state)?;
 
             // Validate the calculated state root
             let target = provider

crates/stages/stages/src/stages/merkle.rs

@@ -329,7 +329,7 @@ where
     }
     let hashed_state = HashedPostState { accounts, storages };
     let nested_hash = NestedStateRoot::new(provider, None);
-    let (root_hash, trie_updates, _) = nested_hash.calculate(&hashed_state, false)?;
+    let (root_hash, trie_updates) = nested_hash.calculate(&hashed_state)?;
 
     writer.write_trie_updatesv2(&trie_updates)?;
     info!(target: "reth::cli",

crates/storage/db-common/src/init.rs

@@ -1,7 +1,7 @@
 use crate::{
     providers::{StaticFileProvider, StaticFileWriter as SfWriter},
     BlockExecutionWriter, BlockWriter, HistoryWriter, StateWriter, StaticFileProviderFactory,
-    StorageLocation, TrieWriter, TrieWriterV2, PERSIST_BLOCK_CACHE,
+    StorageLocation, TrieWriter, TrieWriterV2,
 };
 use alloy_consensus::BlockHeader;
 use reth_chain_state::{ExecutedBlock, ExecutedBlockWithTrieUpdates};
@@ -167,7 +167,6 @@ where
             triev2,
         } in blocks
         {
-            let block_number = recovered_block.number();
             let block_hash = recovered_block.hash();
 
             #[cfg(not(feature = "pipe_test"))]
@@ -188,7 +187,6 @@ where
                 trie.as_ref().ok_or(ProviderError::MissingTrieUpdates(block_hash))?,
             )?;
             let _ = self.database().write_trie_updatesv2(triev2.as_ref())?;
-            PERSIST_BLOCK_CACHE.persist_tip(block_number);
         }
 
         // update history indices