feat(trin-execution): implement import/export for .era2

e2store/src/era2.rs

@@ -18,6 +18,8 @@ use crate::{
     utils::underlying_io_error_kind,
 };
 
+pub const MAX_STORAGE_ITEMS: usize = 10_000_000;
+
 // <network-name>-<block-number>-<short-state-root>.era2
 //
 // era2 := Version | CompressedHeader | account*
@@ -108,11 +110,11 @@ impl Era2 {
                 match self.pending_storage_entries {
                     0 => bail!("Invalid append entry state: expected an account entry, got a storage entry. No storage entries left to append for the account"),
                     1 => ensure!(
-                        storage.len() <= 10_000_000,
+                        storage.len() <= MAX_STORAGE_ITEMS,
                         "Storage entry can't have more than 10 million items",
                     ),
                     _ => ensure!(
-                        storage.len() == 10_000_000,
+                        storage.len() == MAX_STORAGE_ITEMS,
                         "Only last storage entry can have less than 10 million items",
                     ),
                 }
@@ -221,7 +223,7 @@ impl TryFrom<AccountEntry> for Entry {
 }
 
 #[derive(Clone, Eq, PartialEq, Debug)]
-pub struct StorageEntry(Vec<StorageItem>);
+pub struct StorageEntry(pub Vec<StorageItem>);
 
 impl Deref for StorageEntry {
     type Target = Vec<StorageItem>;

trin-execution/src/cli.rs

@@ -1,6 +1,6 @@
-use std::net::SocketAddr;
+use std::{net::SocketAddr, path::PathBuf};
 
-use clap::Parser;
+use clap::{Args, Parser, Subcommand};
 
 use crate::types::block_to_trace::BlockToTrace;
 
@@ -26,4 +26,34 @@ pub struct TrinExecutionConfig {
         help = "Enable prometheus metrics reporting (provide local IP/Port from which your Prometheus server is configured to fetch metrics)"
     )]
     pub enable_metrics_with_url: Option<SocketAddr>,
+
+    #[command(subcommand)]
+    pub command: Option<TrinExecutionSubCommands>,
+}
+
+#[derive(Subcommand, Debug, Clone, PartialEq)]
+#[allow(clippy::enum_variant_names)]
+pub enum TrinExecutionSubCommands {
+    /// Import a era2 state snapshot from a file, useful for bootstrapping a new node quickly
+    ImportState(ImportState),
+    /// Export the current state of the node to a era2 file
+    ExportState(ExportState),
+}
+
+#[derive(Args, Debug, Default, Clone, PartialEq)]
+pub struct ImportState {
+    #[arg(
+        long = "path-to-era2",
+        help = "path to where the era2 state snapshot is located"
+    )]
+    pub path_to_era2: PathBuf,
+}
+
+#[derive(Args, Debug, Default, Clone, PartialEq)]
+pub struct ExportState {
+    #[arg(
+        long = "path-to-era2",
+        help = "path to where the era2 state snapshot is located"
+    )]
+    pub path_to_era2: PathBuf,
 }

trin-execution/src/execution.rs

@@ -63,7 +63,7 @@ struct GenesisConfig {
 pub struct State {
     pub database: EvmDB,
     pub config: StateConfig,
-    execution_position: ExecutionPosition,
+    pub execution_position: ExecutionPosition,
     pub era_manager: Arc<Mutex<EraManager>>,
     pub node_data_directory: PathBuf,
 }

trin-execution/src/lib.rs

@@ -8,6 +8,7 @@ pub mod execution;
 pub mod metrics;
 pub mod spec_id;
 pub mod storage;
+pub mod subcommands;
 pub mod transaction;
 pub mod trie_walker;
 pub mod types;

trin-execution/src/main.rs

@@ -4,8 +4,12 @@ use e2store::era1::BLOCK_TUPLE_COUNT;
 use revm_primitives::SpecId;
 use tracing::info;
 use trin_execution::{
-    cli::TrinExecutionConfig, execution::State, spec_id::get_spec_block_number,
+    cli::{TrinExecutionConfig, TrinExecutionSubCommands},
+    era::manager::EraManager,
+    execution::State,
+    spec_id::get_spec_block_number,
     storage::utils::setup_temp_dir,
+    subcommands::era2::{StateExporter, StateImporter},
 };
 use trin_utils::log::init_tracing_logger;
 
@@ -27,7 +31,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     let mut state = State::new(
         directory.map(|temp_directory| temp_directory.path().to_path_buf()),
-        trin_execution_config.into(),
+        trin_execution_config.clone().into(),
     )
     .await?;
 
@@ -39,6 +43,28 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
             .expect("signal ctrl_c should never fail");
     });
 
+    if let Some(command) = trin_execution_config.command {
+        match command {
+            TrinExecutionSubCommands::ImportState(import_state) => {
+                let mut state_importer = StateImporter::new(state, import_state);
+                state_importer.import_state()?;
+                info!(
+                    "Imported state from era2: {} {}",
+                    state_importer.state.next_block_number() - 1,
+                    state_importer.state.get_root()?
+                );
+                return Ok(());
+            }
+            TrinExecutionSubCommands::ExportState(export_state) => {
+                let mut era_manager = EraManager::new(state.next_block_number() - 1).await?;
+                let header = era_manager.get_next_block().await?.clone();
+                let mut state_exporter = StateExporter::new(state, export_state);
+                state_exporter.export_state(header.header)?;
+                return Ok(());
+            }
+        }
+    }
+
     let mut block_number = state.next_block_number();
 
     let end_block = get_spec_block_number(SpecId::MERGE);

trin-execution/src/storage/account.rs

@@ -23,6 +23,17 @@ impl Default for Account {
     }
 }
 
+impl From<AccountStateInfo> for Account {
+    fn from(account: AccountStateInfo) -> Self {
+        Self {
+            nonce: account.nonce,
+            balance: account.balance,
+            storage_root: account.storage_root,
+            code_hash: account.code_hash,
+        }
+    }
+}
+
 impl From<&Account> for AccountStateInfo {
     fn from(account: &Account) -> Self {
         Self {

trin-execution/src/storage/evm_db.rs

@@ -381,7 +381,7 @@ impl DatabaseRef for EvmDB {
     fn storage_ref(&self, address: Address, index: U256) -> Result<U256, Self::Error> {
         let _timer = start_timer_vec(&TRANSACTION_PROCESSING_TIMES, &["database_get_storage"]);
         let address_hash = keccak256(address);
-        let account: RocksAccount = match self.db.get(keccak256(address))? {
+        let account: RocksAccount = match self.db.get(address_hash)? {
             Some(raw_account) => Decodable::decode(&mut raw_account.as_slice())?,
             None => return Err(Self::Error::NotFound("storage".to_string())),
         };