anvil-polkadot: add chain reversion RPCs

crates/anvil-polkadot/Cargo.toml

@@ -20,6 +20,7 @@ path = "bin/main.rs"
 [dependencies]
 # foundry internal
 substrate-runtime = { path = "substrate-runtime" }
+pallet-revive-eth-rpc = { git = "https://github.com/paritytech/polkadot-sdk.git", branch = "master" }
 secp256k1 = { version = "0.28.0", default-features = false }
 libsecp256k1 = { version = "0.7.0", default-features = false }
 sp-runtime-interface = { git = "https://github.com/paritytech/polkadot-sdk.git", branch = "master", default-features = false }
@@ -60,7 +61,6 @@ polkadot-sdk = { git = "https://github.com/paritytech/polkadot-sdk.git", branch
 	"substrate-frame-rpc-system",
 	"substrate-rpc-client",
 	"substrate-wasm-builder",
-	"pallet-revive-eth-rpc"
 ] }
 anvil.workspace = true
 anvil-core.workspace = true

crates/anvil-polkadot/src/api_server/error.rs

@@ -1,6 +1,7 @@
 use crate::substrate_node::mining_engine::MiningError;
 use anvil_rpc::{error::RpcError, response::ResponseResult};
-use polkadot_sdk::pallet_revive_eth_rpc::{EthRpcError, client::ClientError};
+use pallet_revive_eth_rpc::{EthRpcError, client::ClientError};
+use polkadot_sdk::sp_blockchain;
 use serde::Serialize;
 
 #[derive(Debug, thiserror::Error)]
@@ -13,6 +14,8 @@ pub enum Error {
     InvalidParams(String),
     #[error("Revive call failed: {0}")]
     ReviveRpc(#[from] EthRpcError),
+    #[error("Snapshot error: {0}")]
+    SnapshotRpc(sp_blockchain::Error),
 }
 impl From<subxt::Error> for Error {
     fn from(err: subxt::Error) -> Self {
@@ -69,6 +72,9 @@ impl<T: Serialize> ToRpcResponseResult for Result<T> {
                 Error::InvalidParams(error_message) => {
                     RpcError::invalid_params(error_message).into()
                 }
+                Error::SnapshotRpc(blockchain_err) => {
+                    RpcError::internal_error_with(format!("{blockchain_err}")).into()
+                }
                 Error::ReviveRpc(client_error) => {
                     RpcError::internal_error_with(format!("{client_error}")).into()
                 }

crates/anvil-polkadot/src/api_server/mod.rs

@@ -1,7 +1,11 @@
 use crate::{
-    AnvilNodeConfig,
     logging::LoggingManager,
-    substrate_node::{impersonation::ImpersonationManager, service::Service},
+    substrate_node::{
+        service::{FullClient, Service},
+        impersonation::ImpersonationManager,
+        snapshot::SnapshotManager,
+    },
+    AnvilNodeConfig,
 };
 use anvil_core::eth::EthRequest;
 use anvil_rpc::response::ResponseResult;
@@ -20,17 +24,18 @@ pub struct ApiRequest {
 }
 
 pub fn spawn(
+    config: &AnvilNodeConfig,
     substrate_service: &Service,
     logging_manager: LoggingManager,
-    config: &AnvilNodeConfig,
+    snapshot_manager: SnapshotManager<FullClient>,
 ) -> ApiHandle {
     let (api_handle, receiver) = mpsc::channel(100);
 
     let service = substrate_service.clone();
     let mut impersonation_manager = ImpersonationManager::default();
     impersonation_manager.set_auto_impersonate_account(config.enable_auto_impersonate);
     substrate_service.spawn_handle.spawn("anvil-api-server", "anvil", async move {
-        let api_server = ApiServer::new(service, receiver, logging_manager, impersonation_manager)
+        let api_server = ApiServer::new(service, receiver, logging_manager, snapshot_manager, impersonation_manager)
             .await
             .unwrap_or_else(|err| panic!("Failed to spawn the API server: {err}"));
         api_server.run().await;

crates/anvil-polkadot/src/api_server/server.rs

@@ -1,3 +1,4 @@
+
 use crate::{
     api_server::{
         ApiRequest,
@@ -9,8 +10,11 @@ use crate::{
     logging::LoggingManager,
     macros::node_info,
     substrate_node::{
-        impersonation::ImpersonationManager, in_mem_rpc::InMemoryRpcClient,
-        mining_engine::MiningEngine, service::Service,
+        in_mem_rpc::InMemoryRpcClient,
+        mining_engine::MiningEngine,
+        service::{FullClient, Service},
+        snapshot::SnapshotManager,
+        impersonation::ImpersonationManager,
     },
 };
 use alloy_eips::{BlockId, BlockNumberOrTag};
@@ -20,13 +24,13 @@ use alloy_serde::WithOtherFields;
 use anvil_core::eth::{EthRequest, Params as MineParams};
 use anvil_rpc::response::ResponseResult;
 use futures::{StreamExt, channel::mpsc};
+use pallet_revive_eth_rpc::{
+    EthRpcError, ReceiptExtractor, ReceiptProvider, SubxtBlockInfoProvider,
+    client::{Client as EthRpcClient, ClientError, SubscriptionType},
+    subxt_client::{self, SrcChainConfig},
+};
 use polkadot_sdk::{
     pallet_revive::evm::{Account, Block, Bytes, ReceiptInfo, TransactionSigned},
-    pallet_revive_eth_rpc::{
-        EthRpcError, ReceiptExtractor, ReceiptProvider, SubxtBlockInfoProvider,
-        client::{Client as EthRpcClient, ClientError, SubscriptionType},
-        subxt_client::{self, SrcChainConfig},
-    },
     sp_core::{self, keccak_256},
 };
 use sqlx::sqlite::SqlitePoolOptions;
@@ -44,6 +48,7 @@ pub struct ApiServer {
     req_receiver: mpsc::Receiver<ApiRequest>,
     logging_manager: LoggingManager,
     mining_engine: Arc<MiningEngine>,
+    snapshot_manager: SnapshotManager<FullClient>,
     eth_rpc_client: EthRpcClient,
     wallet: Wallet,
     impersonation_manager: ImpersonationManager,
@@ -54,6 +59,7 @@ impl ApiServer {
         substrate_service: Service,
         req_receiver: mpsc::Receiver<ApiRequest>,
         logging_manager: LoggingManager,
+        snapshot_manager: SnapshotManager<FullClient>,
         impersonation_manager: ImpersonationManager,
     ) -> Result<Self> {
         let eth_rpc_client = create_revive_rpc_client(&substrate_service).await?;
@@ -63,6 +69,7 @@ impl ApiServer {
             logging_manager,
             mining_engine: substrate_service.mining_engine.clone(),
             eth_rpc_client,
+            snapshot_manager,
             impersonation_manager,
             wallet: Wallet {
                 accounts: vec![
@@ -132,6 +139,10 @@ impl ApiServer {
             EthRequest::EthSendTransaction(request) => {
                 self.send_transaction(*request.clone()).await.to_rpc_result()
             }
+            // --- Snapshot ---
+            EthRequest::EvmSnapshot(_) => self.snapshot().await.to_rpc_result(),
+            // EthRequest::Rollback(depth) => {}
+            EthRequest::EvmRevert(id) => self.revert(id).await.to_rpc_result(),
             // -- Impersonation --
             EthRequest::ImpersonateAccount(addr) => {
                 self.impersonate_account(H160::from_slice(addr.0.as_ref())).to_rpc_result()
@@ -421,6 +432,16 @@ impl ApiServer {
         self.send_raw_transaction(Bytes(payload)).await
     }
 
+    pub(crate) async fn snapshot(&mut self) -> Result<U256> {
+        node_info!("evm_snapshot");
+        Ok(self.snapshot_manager.snapshot())
+    }
+
+    pub(crate) async fn revert(&mut self, id: U256) -> Result<bool> {
+        node_info!("evm_revert");
+        self.snapshot_manager.revert(id).map_err(Error::SnapshotRpc)
+    }
+
     // Helpers
     async fn get_block_hash_for_tag(&self, block_id: Option<BlockId>) -> Result<H256> {
         self.eth_rpc_client

crates/anvil-polkadot/src/lib.rs

@@ -6,7 +6,10 @@ use crate::{
     api_server::ApiHandle,
     config::AnvilNodeConfig,
     logging::{LoggingManager, NodeLogLayer},
-    substrate_node::service::Service,
+    substrate_node::{
+        service::{FullClient, Service},
+        snapshot::SnapshotManager,
+    },
 };
 use clap::{CommandFactory, Parser};
 use eyre::Result;
@@ -97,6 +100,7 @@ pub fn run_command(args: Anvil) -> Result<()> {
     } else {
         LoggingManager::default()
     };
+
     let runner: sc_cli::Runner<opts::SubstrateCli> =
         sc_cli::Runner::new(config, tokio_runtime, signals)?;
 
@@ -115,12 +119,22 @@ pub async fn spawn(
     let (substrate_service, task_manager) =
         substrate_node::service::new(&anvil_config, substrate_config)
             .map_err(sc_cli::Error::Service)?;
+    let snapshot_manager = SnapshotManager::new(
+        substrate_service.client.clone(),
+        substrate_service.backend.clone(),
+        anvil_config.get_genesis_number(),
+    );
 
     // Spawn the other tasks.
-    let api_handle =
-        spawn_anvil_tasks(anvil_config, &substrate_service, &task_manager, logging_manager)
-            .await
-            .map_err(|err| sc_cli::Error::Application(err.into()))?;
+    let api_handle = spawn_anvil_tasks(
+        anvil_config,
+        &substrate_service,
+        &task_manager,
+        logging_manager,
+        snapshot_manager,
+    )
+    .await
+    .map_err(|err| sc_cli::Error::Application(err.into()))?;
 
     Ok((substrate_service, task_manager, api_handle))
 }
@@ -130,9 +144,10 @@ pub async fn spawn_anvil_tasks(
     service: &Service,
     task_manager: &TaskManager,
     logging_manager: LoggingManager,
+    snapshot_manager: SnapshotManager<FullClient>,
 ) -> Result<ApiHandle> {
     // Spawn the api server.
-    let api_handle = api_server::spawn(service, logging_manager, &anvil_config);
+    let api_handle = api_server::spawn(&anvil_config, service, logging_manager, snapshot_manager);
 
     // Spawn the network servers.
     for addr in &anvil_config.host {

crates/anvil-polkadot/src/substrate_node/mining_engine.rs

@@ -346,7 +346,7 @@ async fn seal_now(
     let (sender, receiver) = oneshot::channel();
     let seal_command = EngineCommand::SealNewBlock {
         create_empty: true,
-        finalize: true,
+        finalize: false,
         parent_hash: None,
         sender: Some(sender),
     };

crates/anvil-polkadot/src/substrate_node/mod.rs

@@ -4,3 +4,4 @@ pub mod impersonation;
 pub mod in_mem_rpc;
 pub mod mining_engine;
 pub mod service;
+pub mod snapshot;

crates/anvil-polkadot/src/substrate_node/service.rs

@@ -164,7 +164,6 @@ fn spawn_rpc_server(
     let rpc_extensions_builder = {
         let client = client.clone();
         let pool = transaction_pool.clone();
-
         Box::new(move |_| {
             Ok(polkadot_sdk::substrate_frame_rpc_system::System::new(client.clone(), pool.clone())
                 .into_rpc())

crates/anvil-polkadot/src/substrate_node/snapshot.rs

@@ -0,0 +1,84 @@
+use crate::substrate_node::service::Backend;
+use alloy_primitives::U256;
+use polkadot_sdk::{
+    sc_client_api::Backend as BackendT,
+    sp_blockchain::{Error as BlockChainError, HeaderBackend, HeaderMetadata, Result},
+};
+use std::{collections::BTreeMap, sync::Arc};
+use substrate_runtime::OpaqueBlock;
+
+#[derive(Clone, Debug)]
+pub struct Snapshot {
+    pub best_number: u64,
+}
+
+pub struct SnapshotManager<C> {
+    client: Arc<C>,
+    backend: Arc<Backend>,
+    next_snapshot_id: U256,
+    snapshots: BTreeMap<U256, Snapshot>,
+}
+
+impl<C> SnapshotManager<C> {
+    pub fn new(client: Arc<C>, backend: Arc<Backend>, genesis_block_number: u64) -> Self {
+        let snapshot = Snapshot { best_number: genesis_block_number };
+        let mut map = BTreeMap::new();
+        map.insert(U256::ZERO, snapshot);
+
+        Self {
+            client,
+            backend,
+            // Start with 1 to mimic Ganache
+            next_snapshot_id: U256::ONE,
+            snapshots: map,
+        }
+    }
+}
+
+impl<C> SnapshotManager<C>
+where
+    C: HeaderBackend<OpaqueBlock>
+        + HeaderMetadata<OpaqueBlock, Error = BlockChainError>
+        + Send
+        + Sync
+        + 'static,
+{
+    /// Create a snapshot id corresponding to the best block number.
+    pub fn snapshot(&mut self) -> U256 {
+        let one = U256::ONE;
+        self.next_snapshot_id += one;
+        let snapshot = Snapshot { best_number: self.client.info().best_number.into() };
+        self.snapshots.insert(self.next_snapshot_id - one, snapshot);
+        // Safe since there is always a snapshot representing genesis.
+        self.next_snapshot_id - one
+    }
+
+    /// Revert the chain to the block number represented by the snapshot `id`.
+    pub fn revert(&mut self, id: U256) -> Result<bool> {
+        // Remove the snapshot when reverting. We do not want to keep it around
+        // since reverting to an existing snapshot could mean going back to future,
+        // which is not supported.
+        let maybe_snapshot = self.snapshots.remove(&id);
+        let Some(snap) = maybe_snapshot else { return Ok(false) };
+
+        let current_best_number: u64 = self.client.info().best_number.into();
+        let number_of_blocks_to_revert = current_best_number - snap.best_number;
+
+        self.backend.revert(
+            number_of_blocks_to_revert.try_into().expect("to not surpass u32 bounds"),
+            true,
+        )?;
+
+        self.snapshots
+            .retain(|&k, snap_to_remove| k < id || snap_to_remove.best_number >= snap.best_number);
+
+        Ok(true)
+    }
+
+    pub fn rollback(&self) -> Result<bool> {
+        let current_best_number = self.client.info().best_number;
+        self.backend.revert(current_best_number, true)?;
+
+        Ok(true)
+    }
+}

crates/anvil-polkadot/tests/it/main.rs

@@ -1,5 +1,6 @@
 mod impersonation;
 mod mining;
+mod snapshot;
 mod standard_rpc;
 mod time_machine;
 mod utils;