feat(cli): adds load and state cli options (#502)

* feat: adds load and state cli options
* fix: batch number null issue

crates/cli/src/cli.rs

@@ -220,6 +220,21 @@ pub struct Cli {
     #[arg(long, value_name = "PATH", value_parser= parse_genesis_file)]
     pub init: Option<Genesis>,
 
+    /// This is an alias for both --load-state and --dump-state.
+    ///
+    /// It initializes the chain with the state and block environment stored at the file, if it
+    /// exists, and dumps the chain's state on exit.
+    #[arg(
+        long,
+        value_name = "PATH",
+        conflicts_with_all = &[
+            "init",
+            "dump_state",
+            "load_state"
+        ]
+    )]
+    pub state: Option<PathBuf>,
+
     /// Interval in seconds at which the state and block environment is to be dumped to disk.
     ///
     /// See --state and --dump-state
@@ -241,6 +256,10 @@ pub struct Cli {
     #[arg(long, conflicts_with = "init", default_value = "false")]
     pub preserve_historical_states: bool,
 
+    /// Initialize the chain from a previously saved state snapshot.
+    #[arg(long, value_name = "PATH", conflicts_with = "init")]
+    pub load_state: Option<PathBuf>,
+
     /// BIP39 mnemonic phrase used for generating accounts.
     /// Cannot be used if `mnemonic_random` or `mnemonic_seed` are used.
     #[arg(long, short, conflicts_with_all = &["mnemonic_seed", "mnemonic_random"], help_heading = "Account Configuration")]
@@ -436,9 +455,11 @@ impl Cli {
             .with_allow_origin(self.allow_origin)
             .with_no_cors(self.no_cors)
             .with_transaction_order(self.order)
+            .with_state(self.state)
             .with_state_interval(self.state_interval)
             .with_dump_state(self.dump_state)
-            .with_preserve_historical_states(self.preserve_historical_states);
+            .with_preserve_historical_states(self.preserve_historical_states)
+            .with_load_state(self.load_state);
 
         if self.emulate_evm && self.dev_system_contracts != Some(SystemContractsOptions::Local) {
             return Err(eyre::eyre!(
@@ -623,6 +644,7 @@ mod tests {
         net::{IpAddr, Ipv4Addr},
     };
     use tempdir::TempDir;
+    use zksync_types::{H160, U256};
 
     #[test]
     fn can_parse_host() {
@@ -689,8 +711,6 @@ mod tests {
             TxPool::new(ImpersonationManager::default(), config.transaction_order),
             BlockSealer::new(BlockSealerMode::noop()),
         );
-        let test_address = zksync_types::H160::random();
-        node.set_rich_account(test_address, 1000000u64.into());
 
         let mut state_dumper = PeriodicStateDumper::new(
             node.clone(),
@@ -718,4 +738,70 @@ mod tests {
 
         Ok(())
     }
+
+    #[tokio::test]
+    async fn test_load_state() -> anyhow::Result<()> {
+        let temp_dir = TempDir::new("state-load-test").expect("failed creating temporary dir");
+        let state_path = temp_dir.path().join("state.json");
+
+        let config = anvil_zksync_config::TestNodeConfig {
+            dump_state: Some(state_path.clone()),
+            state_interval: Some(1),
+            preserve_historical_states: true,
+            ..Default::default()
+        };
+
+        let node = InMemoryNode::new(
+            None,
+            None,
+            &config,
+            TimestampManager::default(),
+            ImpersonationManager::default(),
+            TxPool::new(ImpersonationManager::default(), config.transaction_order),
+            BlockSealer::new(BlockSealerMode::noop()),
+        );
+        let test_address = H160::from_low_u64_be(12345);
+        node.set_rich_account(test_address, U256::from(1000000u64));
+
+        let mut state_dumper = PeriodicStateDumper::new(
+            node.clone(),
+            config.dump_state.clone(),
+            std::time::Duration::from_secs(1),
+            config.preserve_historical_states,
+        );
+
+        let dumper_handle = tokio::spawn(async move {
+            tokio::select! {
+                _ = &mut state_dumper => {}
+            }
+            state_dumper.dump().await;
+        });
+
+        tokio::time::sleep(std::time::Duration::from_secs(2)).await;
+
+        dumper_handle.abort();
+        let _ = dumper_handle.await;
+
+        // assert the state json file was created
+        std::fs::read_to_string(&state_path).expect("Expected state file to be created");
+
+        let new_config = anvil_zksync_config::TestNodeConfig::default();
+        let new_node = InMemoryNode::new(
+            None,
+            None,
+            &new_config,
+            TimestampManager::default(),
+            ImpersonationManager::default(),
+            TxPool::new(ImpersonationManager::default(), config.transaction_order),
+            BlockSealer::new(BlockSealerMode::noop()),
+        );
+
+        new_node.load_state(zksync_types::web3::Bytes(std::fs::read(&state_path)?))?;
+
+        // assert the balance from the loaded state is correctly applied
+        let balance = new_node.get_balance_impl(test_address, None).await.unwrap();
+        assert_eq!(balance, U256::from(1000000u64));
+
+        Ok(())
+    }
 }

crates/cli/src/main.rs

@@ -278,28 +278,33 @@ async fn main() -> anyhow::Result<()> {
     let any_server_stopped =
         futures::future::select_all(server_handles.into_iter().map(|h| Box::pin(h.stopped())));
 
-    let dump_state = config.dump_state.clone();
+    // Load state from `--load-state` if provided
+    if let Some(ref load_state_path) = config.load_state {
+        let bytes = std::fs::read(load_state_path).expect("Failed to read load state file");
+        node.load_state(zksync_types::web3::Bytes(bytes))?;
+    }
+    if let Some(ref state_path) = config.state {
+        let bytes = std::fs::read(state_path).expect("Failed to read load state file");
+        node.load_state(zksync_types::web3::Bytes(bytes))?;
+    }
+
+    let state_path = config.dump_state.clone().or_else(|| config.state.clone());
     let dump_interval = config
         .state_interval
         .map(Duration::from_secs)
         .unwrap_or(Duration::from_secs(60)); // Default to 60 seconds
     let preserve_historical_states = config.preserve_historical_states;
     let node_for_dumper = node.clone();
-    let state_dumper = tokio::task::spawn(PeriodicStateDumper::new(
+    let state_dumper = PeriodicStateDumper::new(
         node_for_dumper,
-        dump_state,
+        state_path,
         dump_interval,
         preserve_historical_states,
-    ));
+    );
 
     let system_contracts =
         SystemContracts::from_options(&config.system_contracts_options, config.use_evm_emulator);
-    let block_producer_handle = tokio::task::spawn(BlockProducer::new(
-        node,
-        pool,
-        block_sealer,
-        system_contracts,
-    ));
+    let block_producer_handle = BlockProducer::new(node, pool, block_sealer, system_contracts);
 
     config.print(fork_print_info.as_ref());
 

crates/config/src/config.rs

@@ -121,13 +121,16 @@ pub struct TestNodeConfig {
     pub no_cors: bool,
     /// How transactions are sorted in the mempool
     pub transaction_order: TransactionOrder,
-    /// State configuration
+    /// Path to load/dump the state from
+    pub state: Option<PathBuf>,
     /// Path to dump the state to
     pub dump_state: Option<PathBuf>,
     /// Interval to dump the state
     pub state_interval: Option<u64>,
     /// Preserve historical states
     pub preserve_historical_states: bool,
+    /// State to load
+    pub load_state: Option<PathBuf>,
 }
 
 impl Default for TestNodeConfig {
@@ -195,9 +198,11 @@ impl Default for TestNodeConfig {
             no_cors: false,
 
             // state configuration
+            state: None,
             dump_state: None,
             state_interval: None,
             preserve_historical_states: false,
+            load_state: None,
         }
     }
 }
@@ -918,6 +923,13 @@ impl TestNodeConfig {
         self
     }
 
+    /// Set the state
+    #[must_use]
+    pub fn with_state(mut self, state: Option<PathBuf>) -> Self {
+        self.state = state;
+        self
+    }
+
     /// Set the state dump path
     #[must_use]
     pub fn with_dump_state(mut self, dump_state: Option<PathBuf>) -> Self {
@@ -938,4 +950,11 @@ impl TestNodeConfig {
         self.preserve_historical_states = preserve_historical_states;
         self
     }
+
+    /// Set the state to load
+    #[must_use]
+    pub fn with_load_state(mut self, load_state: Option<PathBuf>) -> Self {
+        self.load_state = load_state;
+        self
+    }
 }

crates/core/src/fork.rs

@@ -539,7 +539,7 @@ impl ForkDetails {
             .get_block_by_hash(root_hash, true)
             .await
             .map_err(|error| eyre!(error))?;
-        let block = opt_block.ok_or_else(|| {
+        let mut block = opt_block.ok_or_else(|| {
             eyre!(
                 "Could not find block #{:?} ({:#x}) in {:?}",
                 miniblock,
@@ -548,6 +548,7 @@ impl ForkDetails {
             )
         })?;
         let l1_batch_number = block_details.l1_batch_number;
+        block.l1_batch_number = Some(l1_batch_number.0.into());
 
         if !block_details
             .protocol_version

e2e-tests-rust/Cargo.toml

@@ -25,6 +25,8 @@ tower = "0.5"
 http = "1.1.0"
 anvil_zksync_core = { path = "../crates/core" }
 tempdir = "0.3.7"
+flate2 = "1.0"
+hex = "0.4"
 
 [dev-dependencies]