foundry_evm_core/
state_snapshot.rs1use alloy_primitives::{U256, map::HashMap};
4use std::ops::Add;
5
6#[derive(Clone, Debug)]
8pub struct StateSnapshots<T> {
9 id: U256,
10 state_snapshots: HashMap<U256, T>,
11}
12
13impl<T> StateSnapshots<T> {
14 fn next_id(&mut self) -> U256 {
15 let id = self.id;
16 self.id = id.saturating_add(U256::from(1));
17 id
18 }
19
20 pub fn get(&self, id: U256) -> Option<&T> {
22 self.state_snapshots.get(&id)
23 }
24
25 pub fn remove(&mut self, id: U256) -> Option<T> {
30 let snapshot_state = self.state_snapshots.remove(&id);
31
32 let mut to_revert = id.add(U256::from(1));
34 while to_revert < self.id {
35 self.state_snapshots.remove(&to_revert);
36 to_revert += U256::from(1);
37 }
38
39 snapshot_state
40 }
41
42 pub fn clear(&mut self) {
44 self.state_snapshots.clear();
45 }
46
47 pub fn remove_at(&mut self, id: U256) -> Option<T> {
51 self.state_snapshots.remove(&id)
52 }
53
54 pub fn insert(&mut self, state_snapshot: T) -> U256 {
56 let id = self.next_id();
57 self.state_snapshots.insert(id, state_snapshot);
58 id
59 }
60
61 pub fn insert_at(&mut self, state_snapshot: T, id: U256) -> U256 {
65 self.state_snapshots.insert(id, state_snapshot);
66 id
67 }
68}
69
70impl<T> Default for StateSnapshots<T> {
71 fn default() -> Self {
72 Self { id: U256::ZERO, state_snapshots: HashMap::default() }
73 }
74}