feat(anvil): implement anvil_rollback (#9783)

* implement anvil_rollback

* PR improvements

---------

Co-authored-by: grandizzy <[email protected]>

Author: jorgemmsilva

crates/anvil/core/src/eth/mod.rs

@@ -777,6 +777,10 @@ pub enum EthRequest {
     #[cfg_attr(feature = "serde", serde(rename = "anvil_reorg",))]
     Reorg(ReorgOptions),
 
+    /// Rollback the chain
+    #[cfg_attr(feature = "serde", serde(rename = "anvil_rollback",))]
+    Rollback(Option<u64>),
+
     /// Wallet
     #[cfg_attr(feature = "serde", serde(rename = "wallet_getCapabilities", with = "empty_params"))]
     WalletGetCapabilities(()),

crates/anvil/src/eth/api.rs

@@ -456,6 +456,7 @@ impl EthApi {
             EthRequest::Reorg(reorg_options) => {
                 self.anvil_reorg(reorg_options).await.to_rpc_result()
             }
+            EthRequest::Rollback(depth) => self.anvil_rollback(depth).await.to_rpc_result(),
             EthRequest::WalletGetCapabilities(()) => self.get_capabilities().to_rpc_result(),
             EthRequest::WalletSendTransaction(tx) => {
                 self.wallet_send_transaction(*tx).await.to_rpc_result()
@@ -2067,6 +2068,36 @@ impl EthApi {
         Ok(())
     }
 
+    /// Rollback the chain to a specific depth.
+    ///
+    /// e.g depth = 3
+    ///     A  -> B  -> C  -> D  -> E
+    ///     A  -> B
+    ///
+    /// Depth specifies the height to rollback the chain back to. Depth must not exceed the current
+    /// chain height, i.e. can't rollback past the genesis block.
+    ///
+    /// Handler for RPC call: `anvil_rollback`
+    pub async fn anvil_rollback(&self, depth: Option<u64>) -> Result<()> {
+        node_info!("anvil_rollback");
+        let depth = depth.unwrap_or(1);
+
+        // Check reorg depth doesn't exceed current chain height
+        let current_height = self.backend.best_number();
+        let common_height = current_height.checked_sub(depth).ok_or(BlockchainError::RpcError(
+            RpcError::invalid_params(format!(
+                "Rollback depth must not exceed current chain height: current height {current_height}, depth {depth}"
+            )),
+        ))?;
+
+        // Get the common ancestor block
+        let common_block =
+            self.backend.get_block(common_height).ok_or(BlockchainError::BlockNotFound)?;
+
+        self.backend.rollback(common_block).await?;
+        Ok(())
+    }
+
     /// Snapshot the state of the blockchain at the current block.
     ///
     /// Handler for RPC call: `evm_snapshot`

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

@@ -2631,6 +2631,27 @@ impl Backend {
         tx_pairs: HashMap<u64, Vec<Arc<PoolTransaction>>>,
         common_block: Block,
     ) -> Result<(), BlockchainError> {
+        self.rollback(common_block).await?;
+        // Create the new reorged chain, filling the blocks with transactions if supplied
+        for i in 0..depth {
+            let to_be_mined = tx_pairs.get(&i).cloned().unwrap_or_else(Vec::new);
+            let outcome = self.do_mine_block(to_be_mined).await;
+            node_info!(
+                "    Mined reorg block number {}. With {} valid txs and with invalid {} txs",
+                outcome.block_number,
+                outcome.included.len(),
+                outcome.invalid.len()
+            );
+        }
+
+        Ok(())
+    }
+
+    /// Rollback the chain to a common height.
+    ///
+    /// The state of the chain is rewound using `rewind` to the common block, including the db,
+    /// storage, and env.
+    pub async fn rollback(&self, common_block: Block) -> Result<(), BlockchainError> {
         // Get the database at the common block
         let common_state = {
             let mut state = self.states.write();
@@ -2661,31 +2682,14 @@ impl Backend {
 
             // Set environment back to common block
             let mut env = self.env.write();
-            env.block = BlockEnv {
-                number: U256::from(common_block.header.number),
-                timestamp: U256::from(common_block.header.timestamp),
-                gas_limit: U256::from(common_block.header.gas_limit),
-                difficulty: common_block.header.difficulty,
-                prevrandao: Some(common_block.header.mix_hash),
-                coinbase: env.block.coinbase,
-                basefee: env.block.basefee,
-                ..env.block.clone()
-            };
-            self.time.reset(env.block.timestamp.to::<u64>());
-        }
+            env.block.number = U256::from(common_block.header.number);
+            env.block.timestamp = U256::from(common_block.header.timestamp);
+            env.block.gas_limit = U256::from(common_block.header.gas_limit);
+            env.block.difficulty = common_block.header.difficulty;
+            env.block.prevrandao = Some(common_block.header.mix_hash);
 
-        // Create the new reorged chain, filling the blocks with transactions if supplied
-        for i in 0..depth {
-            let to_be_mined = tx_pairs.get(&i).cloned().unwrap_or_else(Vec::new);
-            let outcome = self.do_mine_block(to_be_mined).await;
-            node_info!(
-                "    Mined reorg block number {}. With {} valid txs and with invalid {} txs",
-                outcome.block_number,
-                outcome.included.len(),
-                outcome.invalid.len()
-            );
+            self.time.reset(env.block.timestamp.to::<u64>());
         }
-
         Ok(())
     }
 }

crates/anvil/tests/it/anvil_api.rs

@@ -805,6 +805,38 @@ async fn test_reorg() {
     assert!(res.is_err());
 }
 
+#[tokio::test(flavor = "multi_thread")]
+async fn test_rollback() {
+    let (api, handle) = spawn(NodeConfig::test()).await;
+    let provider = handle.http_provider();
+
+    // Mine 5 blocks
+    for _ in 0..5 {
+        api.mine_one().await;
+    }
+
+    // Get block 4 for later comparison
+    let block4 = provider.get_block(4.into(), false.into()).await.unwrap().unwrap();
+
+    // Rollback with None should rollback 1 block
+    api.anvil_rollback(None).await.unwrap();
+
+    // Assert we're at block 4 and the block contents are kept the same
+    let head = provider.get_block(BlockId::latest(), false.into()).await.unwrap().unwrap();
+    assert_eq!(head, block4);
+
+    // Get block 1 for comparison
+    let block1 = provider.get_block(1.into(), false.into()).await.unwrap().unwrap();
+
+    // Rollback to block 1
+    let depth = 3; // from block 4 to block 1
+    api.anvil_rollback(Some(depth)).await.unwrap();
+
+    // Assert we're at block 1 and the block contents are kept the same
+    let head = provider.get_block(BlockId::latest(), false.into()).await.unwrap().unwrap();
+    assert_eq!(head, block1);
+}
+
 // === wallet endpoints === //
 #[tokio::test(flavor = "multi_thread")]
 async fn can_get_wallet_capabilities() {