anvil_polkadot/substrate_node/
revert.rs

1use crate::substrate_node::service::{Backend, Client};
2use alloy_primitives::{B256, U256};
3use polkadot_sdk::{
4    polkadot_sdk_frame::runtime::types_common::OpaqueBlock,
5    sc_client_api::Backend as BackendT,
6    sp_blockchain::{HeaderBackend, Info, Result},
7};
8use std::{collections::BTreeMap, sync::Arc};
9
10// The snapshot contains the block number and the block hash
11type Snapshot = (u64, B256);
12
13pub struct RevertInfo {
14    pub info: Info<OpaqueBlock>,
15    pub reverted: u64,
16}
17
18pub struct RevertManager {
19    client: Arc<Client>,
20    backend: Arc<Backend>,
21    next_snapshot_id: U256,
22    snapshots: BTreeMap<U256, Snapshot>,
23}
24
25impl RevertManager {
26    pub fn new(client: Arc<Client>, backend: Arc<Backend>) -> Self {
27        Self { client, backend, next_snapshot_id: U256::ZERO, snapshots: BTreeMap::new() }
28    }
29}
30
31impl RevertManager {
32    /// Create a snapshot id corresponding to the best block number.
33    pub fn snapshot(&mut self) -> U256 {
34        let current_snapshot_id = self.next_snapshot_id;
35        self.next_snapshot_id += U256::ONE;
36        let block_number = self.client.info().best_number.into();
37        let block_hash = B256::from_slice(self.client.info().best_hash.as_ref());
38        self.snapshots.insert(current_snapshot_id, (block_number, block_hash));
39        current_snapshot_id
40    }
41
42    /// Revert the chain to the block number represented by the snapshot `id`.
43    pub fn revert(&mut self, snapshot_id: U256) -> Result<Option<RevertInfo>> {
44        let maybe_snapshot = self.snapshots.remove(&snapshot_id);
45        let Some((snapshot_block_number, _)) = maybe_snapshot else {
46            return Ok(None);
47        };
48
49        let current_best_number: u64 = self.client.info().best_number.into();
50        let number_of_blocks_to_revert = current_best_number - snapshot_block_number;
51
52        let (reverted, _) =
53            self.backend.revert(number_of_blocks_to_revert.try_into().unwrap_or(u32::MAX), true)?;
54
55        self.snapshots.retain(|_, (snap_to_remove, _)| *snap_to_remove < snapshot_block_number);
56
57        Ok(Some(RevertInfo { reverted: reverted.into(), info: self.client.info() }))
58    }
59
60    /// Revert from best block to a parent represented by current block height minus depth.
61    pub fn rollback(&self, depth: Option<u64>) -> Result<RevertInfo> {
62        let (reverted, _) =
63            self.backend.revert(depth.unwrap_or(1).try_into().unwrap_or(u32::MAX), true)?;
64        Ok(RevertInfo { reverted: reverted.into(), info: self.client.info() })
65    }
66
67    /// Will revert to genesis.
68    pub fn reset_to_genesis(&self) -> Result<RevertInfo> {
69        let current_block_number = self.client.info().best_number;
70        let (reverted, _) = self.backend.revert(current_block_number, true)?;
71
72        // The chain info can refer to a genesis block with a number different than 0, based on how
73        // the node was started, so we will query the state once more to return accurate info.
74        Ok(RevertInfo { reverted: reverted.into(), info: self.client.info() })
75    }
76
77    pub fn list_snapshots(&self) -> BTreeMap<U256, (u64, B256)> {
78        self.snapshots.clone()
79    }
80}
anvil_polkadot/substrate_node/revert.rs

anvil_polkadot/substrate_node/
revert.rs
