refactor(trin-execution): state import/export don't need TrinExecution (

#1505)

trin-execution/src/cli.rs

@@ -4,6 +4,8 @@ use clap::{Args, Parser, Subcommand};
 
 use crate::types::block_to_trace::BlockToTrace;
 
+pub const APP_NAME: &str = "trin-execution";
+
 #[derive(Parser, Debug, Clone)]
 #[command(name = "Trin Execution", about = "Executing blocks with no devp2p")]
 pub struct TrinExecutionConfig {

trin-execution/src/main.rs

@@ -1,14 +1,12 @@
 use clap::Parser;
 use tracing::info;
 use trin_execution::{
-    cli::{TrinExecutionConfig, TrinExecutionSubCommands},
-    era::manager::EraManager,
+    cli::{TrinExecutionConfig, TrinExecutionSubCommands, APP_NAME},
     execution::TrinExecution,
     subcommands::era2::{export::StateExporter, import::StateImporter},
 };
 use trin_utils::{dir::setup_data_dir, log::init_tracing_logger};
 
-const APP_NAME: &str = "trin-execution";
 const LATEST_BLOCK: u64 = 20_868_946;
 
 #[tokio::main]
@@ -28,34 +26,33 @@ async fn main() -> anyhow::Result<()> {
         trin_execution_config.ephemeral,
     )?;
 
-    let mut trin_execution =
-        TrinExecution::new(&data_dir, trin_execution_config.clone().into()).await?;
-
     if let Some(command) = trin_execution_config.command {
         match command {
-            TrinExecutionSubCommands::ImportState(import_state) => {
-                let mut state_importer = StateImporter::new(trin_execution, import_state);
-                state_importer.import_state()?;
-                state_importer.import_last_256_block_hashes().await?;
-
+            TrinExecutionSubCommands::ImportState(import_state_config) => {
+                let state_importer = StateImporter::new(import_state_config, &data_dir).await?;
+                let header = state_importer.import().await?;
                 info!(
                     "Imported state from era2: {} {}",
-                    state_importer.trin_execution.next_block_number() - 1,
-                    state_importer.trin_execution.get_root()?
+                    header.number, header.state_root,
                 );
                 return Ok(());
             }
-            TrinExecutionSubCommands::ExportState(export_state) => {
-                let mut era_manager =
-                    EraManager::new(trin_execution.next_block_number() - 1).await?;
-                let header = era_manager.get_next_block().await?.clone();
-                let mut state_exporter = StateExporter::new(trin_execution, export_state);
-                state_exporter.export_state(header.header)?;
+            TrinExecutionSubCommands::ExportState(export_state_config) => {
+                let state_exporter = StateExporter::new(export_state_config, &data_dir).await?;
+                state_exporter.export()?;
+                info!(
+                    "Exported state into era2: {} {}",
+                    state_exporter.header().number,
+                    state_exporter.header().state_root,
+                );
                 return Ok(());
             }
         }
     }
 
+    let mut trin_execution =
+        TrinExecution::new(&data_dir, trin_execution_config.clone().into()).await?;
+
     let (tx, rx) = tokio::sync::oneshot::channel();
     tokio::spawn(async move {
         tokio::signal::ctrl_c().await.unwrap();

trin-execution/src/storage/execution_position.rs

@@ -27,11 +27,7 @@ impl ExecutionPosition {
             Some(raw_execution_position) => {
                 Decodable::decode(&mut raw_execution_position.as_slice())?
             }
-            None => Self {
-                version: 0,
-                next_block_number: 0,
-                state_root: EMPTY_ROOT_HASH,
-            },
+            None => Self::default(),
         })
     }
 
@@ -54,3 +50,13 @@ impl ExecutionPosition {
         Ok(())
     }
 }
+
+impl Default for ExecutionPosition {
+    fn default() -> Self {
+        Self {
+            version: 0,
+            next_block_number: 0,
+            state_root: EMPTY_ROOT_HASH,
+        }
+    }
+}

trin-execution/src/subcommands/era2/export.rs

@@ -1,53 +1,85 @@
-use std::sync::Arc;
+use std::{
+    path::{Path, PathBuf},
+    sync::Arc,
+};
 
 use alloy_consensus::EMPTY_ROOT_HASH;
 use alloy_rlp::Decodable;
 use anyhow::ensure;
 use e2store::era2::{
     AccountEntry, AccountOrStorageEntry, Era2, StorageEntry, StorageItem, MAX_STORAGE_ITEMS,
 };
-use eth_trie::EthTrie;
+use eth_trie::{EthTrie, Trie};
 use ethportal_api::{types::state_trie::account_state::AccountState, Header};
 use parking_lot::Mutex;
 use revm_primitives::{B256, KECCAK_EMPTY, U256};
 use tracing::info;
 
-use crate::{cli::ExportStateConfig, execution::TrinExecution, storage::account_db::AccountDB};
+use crate::{
+    cli::ExportStateConfig,
+    config::StateConfig,
+    era::manager::EraManager,
+    storage::{
+        account_db::AccountDB, evm_db::EvmDB, execution_position::ExecutionPosition,
+        utils::setup_rocksdb,
+    },
+};
 
 pub struct StateExporter {
-    pub trin_execution: TrinExecution,
-    exporter_config: ExportStateConfig,
+    config: ExportStateConfig,
+    header: Header,
+    evm_db: EvmDB,
 }
 
 impl StateExporter {
-    pub fn new(trin_execution: TrinExecution, exporter_config: ExportStateConfig) -> Self {
-        Self {
-            trin_execution,
-            exporter_config,
-        }
-    }
+    pub async fn new(config: ExportStateConfig, data_dir: &Path) -> anyhow::Result<Self> {
+        let rocks_db = Arc::new(setup_rocksdb(data_dir)?);
+
+        let execution_position = ExecutionPosition::initialize_from_db(rocks_db.clone())?;
+        ensure!(
+            execution_position.next_block_number() > 0,
+            "Trin execution not initialized!"
+        );
 
-    pub fn export_state(&mut self, header: Header) -> anyhow::Result<()> {
+        let last_executed_block_number = execution_position.next_block_number() - 1;
+
+        let header = EraManager::new(last_executed_block_number)
+            .await?
+            .get_next_block()
+            .await?
+            .header
+            .clone();
+
+        let evm_db = EvmDB::new(StateConfig::default(), rocks_db, &execution_position)
+            .expect("Failed to create EVM database");
         ensure!(
-            header.state_root == self.trin_execution.get_root()?,
-            "State root mismatch fro block header we are trying to export"
+            evm_db.trie.lock().root_hash()? == header.state_root,
+            "State root mismatch from block header we are trying to export"
         );
+
+        Ok(Self {
+            config,
+            header,
+            evm_db,
+        })
+    }
+
+    pub fn export(&self) -> anyhow::Result<PathBuf> {
         info!(
             "Exporting state from block number: {} with state root: {}",
-            header.number, header.state_root
+            self.header.number, self.header.state_root
         );
-        let mut era2 = Era2::create(self.exporter_config.path_to_era2.clone(), header)?;
+        let mut era2 = Era2::create(self.config.path_to_era2.clone(), self.header.clone())?;
         info!("Era2 initiated");
         info!("Trie leaf iterator initiated");
         let mut accounts_exported = 0;
-        for key_hash_and_leaf_value in self.trin_execution.database.trie.lock().iter() {
+        for key_hash_and_leaf_value in self.evm_db.trie.lock().iter() {
             let (raw_account_hash, account_state) = key_hash_and_leaf_value?;
             let account_hash = B256::from_slice(&raw_account_hash);
 
             let account_state = AccountState::decode(&mut account_state.as_slice())?;
             let bytecode = if account_state.code_hash != KECCAK_EMPTY {
-                self.trin_execution
-                    .database
+                self.evm_db
                     .db
                     .get(account_state.code_hash)?
                     .expect("If code hash is not empty, code must be present")
@@ -57,8 +89,7 @@ impl StateExporter {
 
             let mut storage: Vec<StorageItem> = vec![];
             if account_state.storage_root != EMPTY_ROOT_HASH {
-                let account_db =
-                    AccountDB::new(account_hash, self.trin_execution.database.db.clone());
+                let account_db = AccountDB::new(account_hash, self.evm_db.db.clone());
                 let account_trie = Arc::new(Mutex::new(EthTrie::from(
                     Arc::new(account_db),
                     account_state.storage_root,
@@ -99,6 +130,10 @@ impl StateExporter {
 
         info!("Era2 snapshot exported");
 
-        Ok(())
+        Ok(era2.path().to_path_buf())
+    }
+
+    pub fn header(&self) -> &Header {
+        &self.header
     }
 }

trin-execution/src/subcommands/era2/import.rs

@@ -1,38 +1,62 @@
-use std::sync::Arc;
+use std::{path::Path, sync::Arc};
 
 use anyhow::{ensure, Error};
 use e2store::era2::{AccountEntry, AccountOrStorageEntry, Era2, StorageItem};
 use eth_trie::{EthTrie, Trie};
+use ethportal_api::Header;
 use revm_primitives::{keccak256, B256, U256};
 use tracing::info;
 
 use crate::{
-    cli::ImportStateConfig, era::manager::EraManager, evm::block_executor::BLOCKHASH_SERVE_WINDOW,
-    execution::TrinExecution, storage::account_db::AccountDB,
+    cli::ImportStateConfig,
+    config::StateConfig,
+    era::manager::EraManager,
+    evm::block_executor::BLOCKHASH_SERVE_WINDOW,
+    storage::{
+        account_db::AccountDB, evm_db::EvmDB, execution_position::ExecutionPosition,
+        utils::setup_rocksdb,
+    },
 };
 
 pub struct StateImporter {
-    pub trin_execution: TrinExecution,
-    importer_config: ImportStateConfig,
+    config: ImportStateConfig,
+    evm_db: EvmDB,
 }
 
 impl StateImporter {
-    pub fn new(trin_execution: TrinExecution, importer_config: ImportStateConfig) -> Self {
-        Self {
-            trin_execution,
-            importer_config,
-        }
+    pub async fn new(config: ImportStateConfig, data_dir: &Path) -> anyhow::Result<Self> {
+        let rocks_db = Arc::new(setup_rocksdb(data_dir)?);
+
+        let execution_position = ExecutionPosition::initialize_from_db(rocks_db.clone())?;
+        ensure!(
+            execution_position.next_block_number() == 0,
+            "Cannot import state from .era2, database is not empty",
+        );
+
+        let evm_db = EvmDB::new(StateConfig::default(), rocks_db, &execution_position)
+            .expect("Failed to create EVM database");
+
+        Ok(Self { config, evm_db })
+    }
+
+    pub async fn import(&self) -> anyhow::Result<Header> {
+        // Import state from era2 file
+        let header = self.import_state()?;
+
+        // Save execution position
+        let mut execution_position = ExecutionPosition::default();
+        execution_position.update_position(self.evm_db.db.clone(), &header)?;
+
+        // Import last 256 block hashes
+        self.import_last_256_block_hashes(header.number).await?;
+
+        Ok(header)
     }
 
-    pub fn import_state(&mut self) -> anyhow::Result<()> {
+    pub fn import_state(&self) -> anyhow::Result<Header> {
         info!("Importing state from .era2 file");
-        if self.trin_execution.next_block_number() != 0 {
-            return Err(Error::msg(
-                "Cannot import state from .era2, database is not empty",
-            ));
-        }
 
-        let mut era2 = Era2::open(self.importer_config.path_to_era2.clone())?;
+        let mut era2 = Era2::open(self.config.path_to_era2.clone())?;
         info!("Era2 reader initiated");
         let mut accounts_imported = 0;
         while let Some(account) = era2.next() {
@@ -47,7 +71,7 @@ impl StateImporter {
             } = account;
 
             // Build storage trie
-            let account_db = AccountDB::new(address_hash, self.trin_execution.database.db.clone());
+            let account_db = AccountDB::new(address_hash, self.evm_db.db.clone());
             let mut storage_trie = EthTrie::new(Arc::new(account_db));
             for _ in 0..storage_count {
                 let Some(AccountOrStorageEntry::Storage(storage_entry)) = era2.next() else {
@@ -76,21 +100,16 @@ impl StateImporter {
                 "Code hash mismatch, .era2 import failed"
             );
             if !bytecode.is_empty() {
-                self.trin_execution
-                    .database
-                    .db
-                    .put(keccak256(&bytecode), bytecode.clone())?;
+                self.evm_db.db.put(keccak256(&bytecode), bytecode.clone())?;
             }
 
             // Insert account into state trie
-            self.trin_execution
-                .database
+            self.evm_db
                 .trie
                 .lock()
                 .insert(address_hash.as_slice(), &alloy_rlp::encode(&account_state))?;
 
-            self.trin_execution
-                .database
+            self.evm_db
                 .db
                 .put(address_hash, alloy_rlp::encode(account_state))
                 .expect("Inserting account should never fail");
@@ -99,37 +118,29 @@ impl StateImporter {
             if accounts_imported % 1000 == 0 {
                 info!("Imported {} accounts", accounts_imported);
                 info!("Committing changes to database");
-                self.trin_execution.get_root()?;
+                self.evm_db.trie.lock().root_hash()?;
                 info!("Finished committing changes to database");
             }
         }
 
         // Check if the state root matches, if this fails it means either the .era2 is wrong or we
         // imported the state wrong
-        if era2.header.header.state_root != self.trin_execution.get_root()? {
+        if era2.header.header.state_root != self.evm_db.trie.lock().root_hash()? {
             return Err(Error::msg("State root mismatch, .era2 import failed"));
         }
 
-        // Save execution position
-        self.trin_execution
-            .execution_position
-            .update_position(self.trin_execution.database.db.clone(), &era2.header.header)?;
-
         info!("Done importing State from .era2 file");
 
-        Ok(())
+        Ok(era2.header.header)
     }
 
     /// insert the last 256 block hashes into the database
-    pub async fn import_last_256_block_hashes(&mut self) -> anyhow::Result<()> {
-        let first_block_hash_to_add = self
-            .trin_execution
-            .next_block_number()
-            .saturating_sub(BLOCKHASH_SERVE_WINDOW);
+    pub async fn import_last_256_block_hashes(&self, block_number: u64) -> anyhow::Result<()> {
+        let first_block_hash_to_add = block_number.saturating_sub(BLOCKHASH_SERVE_WINDOW);
         let mut era_manager = EraManager::new(first_block_hash_to_add).await?;
-        while era_manager.next_block_number() < self.trin_execution.next_block_number() {
+        while era_manager.next_block_number() < block_number {
             let block = era_manager.get_next_block().await?;
-            self.trin_execution.database.db.put(
+            self.evm_db.db.put(
                 keccak256(B256::from(U256::from(block.header.number))),
                 block.header.hash(),
             )?