feat: include block env in --dump state (#6763)

crates/anvil/src/cmd.rs

@@ -19,7 +19,7 @@ use rand::{rngs::StdRng, SeedableRng};
 use std::{
     future::Future,
     net::IpAddr,
-    path::PathBuf,
+    path::{Path, PathBuf},
     pin::Pin,
     str::FromStr,
     sync::{
@@ -119,8 +119,8 @@ pub struct NodeArgs {
 
     /// This is an alias for both --load-state and --dump-state.
     ///
-    /// It initializes the chain with the state stored at the file, if it exists, and dumps the
-    /// chain's state on exit.
+    /// 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.
     #[clap(
         long,
         value_name = "PATH",
@@ -133,13 +133,13 @@ pub struct NodeArgs {
     )]
     pub state: Option<StateFile>,
 
-    /// Interval in seconds at which the status is to be dumped to disk.
+    /// Interval in seconds at which the state and block environment is to be dumped to disk.
     ///
     /// See --state and --dump-state
     #[clap(short, long, value_name = "SECONDS")]
     pub state_interval: Option<u64>,
 
-    /// Dump the state of chain on exit to the given file.
+    /// Dump the state and block environment of chain on exit to the given file.
     ///
     /// If the value is a directory, the state will be written to `<VALUE>/state.json`.
     #[clap(long, value_name = "PATH", conflicts_with = "init")]
@@ -616,7 +616,12 @@ impl StateFile {
     /// This is used as the clap `value_parser` implementation to parse from file but only if it
     /// exists
     fn parse(path: &str) -> Result<Self, String> {
-        let mut path = PathBuf::from(path);
+        Self::parse_path(path)
+    }
+
+    /// Parse from file but only if it exists
+    pub fn parse_path(path: impl AsRef<Path>) -> Result<Self, String> {
+        let mut path = path.as_ref().to_path_buf();
         if path.is_dir() {
             path = path.join("state.json");
         }

crates/anvil/src/config.rs

@@ -1,4 +1,5 @@
 use crate::{
+    cmd::StateFile,
     eth::{
         backend::{
             db::{Db, SerializableState},
@@ -48,7 +49,7 @@ use std::{
     fmt::Write as FmtWrite,
     fs::File,
     net::{IpAddr, Ipv4Addr},
-    path::PathBuf,
+    path::{Path, PathBuf},
     sync::Arc,
     time::Duration,
 };
@@ -437,13 +438,20 @@ impl NodeConfig {
         self
     }
 
-    /// Sets a custom code size limit
+    /// Sets the init state if any
     #[must_use]
     pub fn with_init_state(mut self, init_state: Option<SerializableState>) -> Self {
         self.init_state = init_state;
         self
     }
 
+    /// Loads the init state from a file if it exists
+    #[must_use]
+    pub fn with_init_state_path(mut self, path: impl AsRef<Path>) -> Self {
+        self.init_state = StateFile::parse_path(path).ok().and_then(|file| file.state);
+        self
+    }
+
     /// Sets the chain ID
     #[must_use]
     pub fn with_chain_id<U: Into<u64>>(mut self, chain_id: Option<U>) -> Self {
@@ -866,13 +874,8 @@ impl NodeConfig {
                 .expect("Failed to create default create2 deployer");
         }
 
-        if let Some(ref state) = self.init_state {
-            backend
-                .get_db()
-                .write()
-                .await
-                .load_state(state.clone())
-                .expect("Failed to load init state");
+        if let Some(state) = self.init_state.clone() {
+            backend.load_state(state).await.expect("Failed to load init state");
         }
 
         backend

crates/anvil/src/eth/api.rs

@@ -1721,7 +1721,7 @@ impl EthApi {
     /// Handler for RPC call: `anvil_loadState`
     pub async fn anvil_load_state(&self, buf: Bytes) -> Result<bool> {
         node_info!("anvil_loadState");
-        self.backend.load_state(buf).await
+        self.backend.load_state_bytes(buf).await
     }
 
     /// Retrieves the Anvil node configuration params.

crates/anvil/src/eth/backend/db.rs

@@ -11,7 +11,7 @@ use foundry_evm::{
     hashbrown::HashMap,
     revm::{
         db::{CacheDB, DatabaseRef, DbAccount},
-        primitives::{Bytecode, KECCAK_EMPTY},
+        primitives::{BlockEnv, Bytecode, KECCAK_EMPTY},
         Database, DatabaseCommit,
     },
 };
@@ -126,7 +126,7 @@ pub trait Db:
     fn insert_block_hash(&mut self, number: U256, hash: B256);
 
     /// Write all chain data to serialized bytes buffer
-    fn dump_state(&self) -> DatabaseResult<Option<SerializableState>>;
+    fn dump_state(&self, at: BlockEnv) -> DatabaseResult<Option<SerializableState>>;
 
     /// Deserialize and add all chain data to the backend storage
     fn load_state(&mut self, state: SerializableState) -> DatabaseResult<bool> {
@@ -196,7 +196,7 @@ impl<T: DatabaseRef<Error = DatabaseError> + Send + Sync + Clone + fmt::Debug> D
         self.block_hashes.insert(number, hash);
     }
 
-    fn dump_state(&self) -> DatabaseResult<Option<SerializableState>> {
+    fn dump_state(&self, _at: BlockEnv) -> DatabaseResult<Option<SerializableState>> {
         Ok(None)
     }
 
@@ -324,6 +324,10 @@ impl MaybeHashDatabase for StateDb {
 
 #[derive(Clone, Debug, Default, Serialize, Deserialize)]
 pub struct SerializableState {
+    /// The block number of the state
+    ///
+    /// Note: This is an Option for backwards compatibility: <https://github.com/foundry-rs/foundry/issues/5460>
+    pub block: Option<BlockEnv>,
     pub accounts: BTreeMap<Address, SerializableAccountRecord>,
 }
 

crates/anvil/src/eth/backend/mem/fork_db.rs

@@ -14,6 +14,7 @@ use foundry_evm::{
 };
 
 pub use foundry_evm::fork::database::ForkedDatabase;
+use foundry_evm::revm::primitives::BlockEnv;
 
 /// Implement the helper for the fork database
 impl Db for ForkedDatabase {
@@ -31,7 +32,7 @@ impl Db for ForkedDatabase {
         self.inner().block_hashes().write().insert(number, hash);
     }
 
-    fn dump_state(&self) -> DatabaseResult<Option<SerializableState>> {
+    fn dump_state(&self, at: BlockEnv) -> DatabaseResult<Option<SerializableState>> {
         let mut db = self.database().clone();
         let accounts = self
             .database()
@@ -56,7 +57,7 @@ impl Db for ForkedDatabase {
                 ))
             })
             .collect::<Result<_, _>>()?;
-        Ok(Some(SerializableState { accounts }))
+        Ok(Some(SerializableState { block: Some(at), accounts }))
     }
 
     fn snapshot(&mut self) -> U256 {

crates/anvil/src/eth/backend/mem/in_memory_db.rs

@@ -16,8 +16,8 @@ use foundry_evm::{
 };
 
 // reexport for convenience
-use foundry_evm::backend::RevertSnapshotAction;
 pub use foundry_evm::{backend::MemDb, revm::db::DatabaseRef};
+use foundry_evm::{backend::RevertSnapshotAction, revm::primitives::BlockEnv};
 
 impl Db for MemDb {
     fn insert_account(&mut self, address: Address, account: AccountInfo) {
@@ -32,7 +32,7 @@ impl Db for MemDb {
         self.inner.block_hashes.insert(number, hash);
     }
 
-    fn dump_state(&self) -> DatabaseResult<Option<SerializableState>> {
+    fn dump_state(&self, at: BlockEnv) -> DatabaseResult<Option<SerializableState>> {
         let accounts = self
             .inner
             .accounts
@@ -57,7 +57,7 @@ impl Db for MemDb {
             })
             .collect::<Result<_, _>>()?;
 
-        Ok(Some(SerializableState { accounts }))
+        Ok(Some(SerializableState { block: Some(at), accounts }))
     }
 
     /// Creates a new snapshot
@@ -167,7 +167,7 @@ mod tests {
 
         dump_db.set_storage_at(test_addr, U256::from(1234567), U256::from(1)).unwrap();
 
-        let state = dump_db.dump_state().unwrap().unwrap();
+        let state = dump_db.dump_state(Default::default()).unwrap().unwrap();
 
         let mut load_db = MemDb::default();
 

crates/anvil/src/eth/backend/mem/mod.rs

@@ -718,7 +718,8 @@ impl Backend {
 
     /// Get the current state.
     pub async fn serialized_state(&self) -> Result<SerializableState, BlockchainError> {
-        let state = self.db.read().await.dump_state()?;
+        let at = self.env.read().block.clone();
+        let state = self.db.read().await.dump_state(at)?;
         state.ok_or_else(|| {
             RpcError::invalid_params("Dumping state not supported with the current configuration")
                 .into()
@@ -735,8 +736,25 @@ impl Backend {
         Ok(encoder.finish().unwrap_or_default().into())
     }
 
+    /// Apply [SerializableState] data to the backend storage.
+    pub async fn load_state(&self, state: SerializableState) -> Result<bool, BlockchainError> {
+        // reset the block env
+        if let Some(block) = state.block.clone() {
+            self.env.write().block = block;
+        }
+
+        if !self.db.write().await.load_state(state)? {
+            Err(RpcError::invalid_params(
+                "Loading state not supported with the current configuration",
+            )
+            .into())
+        } else {
+            Ok(true)
+        }
+    }
+
     /// Deserialize and add all chain data to the backend storage
-    pub async fn load_state(&self, buf: Bytes) -> Result<bool, BlockchainError> {
+    pub async fn load_state_bytes(&self, buf: Bytes) -> Result<bool, BlockchainError> {
         let orig_buf = &buf.0[..];
         let mut decoder = GzDecoder::new(orig_buf);
         let mut decoded_data = Vec::new();
@@ -751,14 +769,7 @@ impl Backend {
         })
         .map_err(|_| BlockchainError::FailedToDecodeStateDump)?;
 
-        if !self.db.write().await.load_state(state)? {
-            Err(RpcError::invalid_params(
-                "Loading state not supported with the current configuration",
-            )
-            .into())
-        } else {
-            Ok(true)
-        }
+        self.load_state(state).await
     }
 
     /// Returns the environment for the next block

crates/anvil/tests/it/main.rs

@@ -14,7 +14,9 @@ mod proof;
 mod pubsub;
 // mod revert; // TODO uncomment <https://github.com/gakonst/ethers-rs/issues/2186>
 mod otterscan;
+
 mod sign;
+mod state;
 mod traces;
 mod transaction;
 mod txpool;

crates/anvil/tests/it/state.rs

@@ -0,0 +1,23 @@
+//! general eth api tests
+
+use anvil::{spawn, NodeConfig};
+
+#[tokio::test(flavor = "multi_thread")]
+async fn can_load_state() {
+    let tmp = tempfile::tempdir().unwrap();
+    let state_file = tmp.path().join("state.json");
+
+    let (api, _handle) = spawn(NodeConfig::test()).await;
+
+    api.mine_one().await;
+
+    let num = api.block_number().unwrap();
+
+    let state = api.serialized_state().await.unwrap();
+    foundry_common::fs::write_json_file(&state_file, &state).unwrap();
+
+    let (api, _handle) = spawn(NodeConfig::test().with_init_state_path(state_file)).await;
+
+    let num2 = api.block_number().unwrap();
+    assert_eq!(num, num2);
+}