feat: adds load and state cli options

crates/cli/Cargo.toml

@@ -37,3 +37,4 @@ flate2.workspace = true
 
 [dev-dependencies]
 tempdir.workspace = true
+anvil_zksync_core.workspace = true

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,7 +278,19 @@ 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))?;
+    }
+
+    // Load state from `--state` if provided and contains a saved state
+    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)
@@ -287,7 +299,7 @@ async fn main() -> anyhow::Result<()> {
     let node_for_dumper = node.clone();
     let state_dumper = tokio::task::spawn(PeriodicStateDumper::new(
         node_for_dumper,
-        dump_state,
+        state_path,
         dump_interval,
         preserve_historical_states,
     ));

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
+    }
 }

e2e-tests-rust/Cargo.toml

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

e2e-tests-rust/tests/lib.rs

@@ -10,6 +10,7 @@ use anvil_zksync_e2e_tests::{
     init_testing_provider, init_testing_provider_with_client, AnvilZKsyncApi, ReceiptExt,
     ZksyncWalletProviderExt, DEFAULT_TX_VALUE, get_node_binary_path
 };
+use anvil_zksync_core::utils::write_json_file;
 use http::header::{
     HeaderMap, HeaderValue, ACCESS_CONTROL_ALLOW_HEADERS, ACCESS_CONTROL_ALLOW_METHODS,
     ACCESS_CONTROL_ALLOW_ORIGIN, ORIGIN,
@@ -22,6 +23,8 @@ const ANY_ORIGIN: HeaderValue = HeaderValue::from_static("*");
 use anvil_zksync_core::node::VersionedState;
 use std::{ fs, convert::identity, thread::sleep, time::Duration };
 use tempdir::TempDir;
+use flate2::read::GzDecoder;
+use std::io::Read;
 
 #[tokio::test]
 async fn interval_sealing_finalization() -> anyhow::Result<()> {
@@ -619,3 +622,109 @@ async fn dump_state_on_fork() -> anyhow::Result<()>  {
 
     Ok(())
 }
+
+#[tokio::test]
+async fn load_state_on_run() -> anyhow::Result<()> {
+    let temp_dir = TempDir::new("load-state-test").expect("failed creating temporary dir");
+    let dump_path = temp_dir.path().join("load_state_run.json");
+    let provider = init_testing_provider(identity).await?;
+    let receipts = [
+        provider.tx().finalize().await?,
+        provider.tx().finalize().await?,
+    ];
+    let blocks = provider.get_blocks_by_receipts(&receipts).await?;
+    let state_bytes = provider.anvil_dump_state().await?;
+    drop(provider);
+
+    let mut decoder = GzDecoder::new(&state_bytes.0[..]);
+    let mut json_str = String::new();
+    decoder.read_to_string(&mut json_str).unwrap();
+    let state: VersionedState = serde_json::from_str(&json_str).unwrap();
+    write_json_file(&dump_path, &state)?;
+
+    let dump_path_clone = dump_path.clone();
+    let new_provider = init_testing_provider(move |node| {
+        node
+            .path(get_node_binary_path())
+            .arg("--state-interval")
+            .arg("1")
+            .arg("--load-state")
+            .arg(dump_path_clone.to_str().unwrap())
+    })
+    .await?;
+
+    sleep(Duration::from_secs(2));
+
+    new_provider.assert_has_receipts(&receipts).await?;
+    new_provider.assert_has_blocks(&blocks).await?;
+    new_provider
+        .assert_balance(receipts[0].sender()?, DEFAULT_TX_VALUE)
+        .await?;
+    new_provider
+        .assert_balance(receipts[1].sender()?, DEFAULT_TX_VALUE)
+        .await?;
+
+    drop(new_provider);
+
+    assert!(
+        dump_path.exists(),
+        "State dump file should still exist at {:?}",
+        dump_path
+    );
+
+    Ok(())
+}
+
+#[tokio::test]
+#[ignore]
+async fn load_state_on_fork() -> anyhow::Result<()> {
+    let temp_dir = TempDir::new("load-state-fork-test").expect("failed creating temporary dir");
+    let dump_path = temp_dir.path().join("load_state_fork.json");
+    let provider = init_testing_provider(identity).await?;
+    let receipts = [
+        provider.tx().finalize().await?,
+        provider.tx().finalize().await?,
+    ];
+    let blocks = provider.get_blocks_by_receipts(&receipts).await?;
+    let state_bytes = provider.anvil_dump_state().await?;
+    drop(provider);
+
+    let mut decoder = GzDecoder::new(&state_bytes.0[..]);
+    let mut json_str = String::new();
+    decoder.read_to_string(&mut json_str).unwrap();
+    let state: VersionedState = serde_json::from_str(&json_str).unwrap();
+    write_json_file(&dump_path, &state)?;
+
+    let dump_path_clone = dump_path.clone();
+    let new_provider = init_testing_provider(move |node| {
+        node
+            .path(get_node_binary_path())
+            .arg("--state-interval")
+            .arg("1")
+            .arg("--load-state")
+            .arg(dump_path_clone.to_str().unwrap())
+            .fork("mainnet")
+    })
+    .await?;
+
+    sleep(Duration::from_secs(2));
+
+    new_provider.assert_has_receipts(&receipts).await?;
+    new_provider.assert_has_blocks(&blocks).await?;
+    new_provider
+        .assert_balance(receipts[0].sender()?, DEFAULT_TX_VALUE)
+        .await?;
+    new_provider
+        .assert_balance(receipts[1].sender()?, DEFAULT_TX_VALUE)
+        .await?;
+
+    drop(new_provider);
+
+    assert!(
+        dump_path.exists(),
+        "State dump file should still exist at {:?}",
+        dump_path
+    );
+
+    Ok(())
+}