Skip to content

Commit ce46d9a

Browse files
pkhrypkhry
committed
snapshotting cheats
1 parent b153c11 commit ce46d9a

File tree

7 files changed

+266
-25
lines changed

7 files changed

+266
-25
lines changed

‎crates/forge/src/runner.rs‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -544,13 +544,17 @@ impl<'a> FunctionRunner<'a> {
544544
/// State modifications of before test txes and unit test function call are discarded after
545545
/// test ends, similar to `eth_call`.
546546
fn run_unit_test(mut self, func: &Function) -> TestResult {
547-
let binding = self.executor.clone().into_owned();
547+
let binding = self.executor.into_owned();
548548
self.executor = Cow::Owned(binding);
549+
self.executor.strategy.runner.start_transaction(self.executor.strategy.context.as_ref());
549550
// Prepare unit test execution.
550551
if self.prepare_test(func).is_err() {
552+
self.executor
553+
.strategy
554+
.runner
555+
.rollback_transaction(self.executor.strategy.context.as_ref());
551556
return self.result;
552557
}
553-
self.executor.strategy.runner.start_transaction(self.executor.strategy.context.as_ref());
554558

555559
// Run current unit test.
556560
let (mut raw_call_result, reason) = match self.executor.call(
@@ -598,7 +602,7 @@ impl<'a> FunctionRunner<'a> {
598602
/// - `bool[] public fixtureSwap = [true, false]` The `table_test` is then called with the pair
599603
/// of args `(2, true)` and `(5, false)`.
600604
fn run_table_test(mut self, func: &Function) -> TestResult {
601-
let binding = self.executor.clone().into_owned();
605+
let binding = self.executor.into_owned();
602606
self.executor = Cow::Owned(binding); // Prepare unit test execution.
603607
if self.prepare_test(func).is_err() {
604608
return self.result;
@@ -729,7 +733,7 @@ impl<'a> FunctionRunner<'a> {
729733
identified_contracts: &ContractsByAddress,
730734
test_bytecode: &Bytes,
731735
) -> TestResult {
732-
let binding = self.executor.clone().into_owned();
736+
let binding = self.executor.into_owned();
733737
self.executor = Cow::Owned(binding);
734738
// First, run the test normally to see if it needs to be skipped.
735739
if let Err(EvmError::Skip(reason)) = self.executor.call(
@@ -955,7 +959,7 @@ impl<'a> FunctionRunner<'a> {
955959
/// State modifications of before test txes and fuzz test are discarded after test ends,
956960
/// similar to `eth_call`.
957961
fn run_fuzz_test(mut self, func: &Function) -> TestResult {
958-
let binding = self.executor.clone().into_owned();
962+
let binding = self.executor.into_owned();
959963
self.executor = Cow::Owned(binding);
960964
// Prepare fuzz test execution.
961965
if self.prepare_test(func).is_err() {

‎crates/forge/tests/it/revive/cheat_snapshot.rs‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
use crate::{config::*, test_helpers::TEST_DATA_REVIVE};
2+
use foundry_test_utils::Filter;
3+
use revive_strategy::ReviveRuntimeMode;
4+
use revm::primitives::hardfork::SpecId;
5+
use rstest::rstest;
6+
7+
#[rstest]
8+
#[case::pvm(ReviveRuntimeMode::Pvm)]
9+
#[case::evm(ReviveRuntimeMode::Evm)]
10+
#[tokio::test(flavor = "multi_thread")]
11+
async fn test_snapshot_cheats(#[case] runtime_mode: ReviveRuntimeMode) {
12+
let runner: forge::MultiContractRunner = TEST_DATA_REVIVE.runner_revive(runtime_mode);
13+
let filter = Filter::new(".*", "StateSnapshotTest", ".*/revive/.*");
14+
15+
TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
16+
}

‎crates/forge/tests/it/revive/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ pub mod cheat_mock_call;
66
pub mod cheat_mock_calls;
77
pub mod cheat_mock_functions;
88
pub mod cheat_prank;
9+
mod cheat_snapshot;
910
pub mod cheat_store;
1011
pub mod migration;
1112
pub mod tx_gas_price;

‎crates/revive-strategy/src/cheatcodes/mod.rs‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,8 @@ use foundry_cheatcodes::{
99
CommonCreateInput, Ecx, EvmCheatcodeInspectorStrategyRunner, Result,
1010
Vm::{
1111
chainIdCall, dealCall, etchCall, getNonce_0Call, loadCall, pvmCall, resetNonceCall,
12-
rollCall, setNonceCall, setNonceUnsafeCall, storeCall, warpCall,
12+
revertToStateAndDeleteCall, revertToStateCall, rollCall, setNonceCall, setNonceUnsafeCall,
13+
snapshotStateCall, storeCall, warpCall,
1314
},
1415
journaled_account, precompile_error,
1516
};
@@ -318,6 +319,23 @@ impl CheatcodeInspectorStrategyRunner for PvmCheatcodeInspectorStrategyRunner {
318319

319320
cheatcode.dyn_apply(ccx, executor)
320321
}
322+
t if using_pvm && is::<snapshotStateCall>(t) => {
323+
ctx.externalities.start_snapshotting();
324+
cheatcode.dyn_apply(ccx, executor)
325+
}
326+
t if using_pvm && is::<revertToStateAndDeleteCall>(t) => {
327+
let &revertToStateAndDeleteCall { snapshotId } =
328+
cheatcode.as_any().downcast_ref().unwrap();
329+
330+
ctx.externalities.revert(snapshotId.try_into().unwrap());
331+
cheatcode.dyn_apply(ccx, executor)
332+
}
333+
t if using_pvm && is::<revertToStateCall>(t) => {
334+
let &revertToStateCall { snapshotId } = cheatcode.as_any().downcast_ref().unwrap();
335+
336+
ctx.externalities.revert(snapshotId.try_into().unwrap());
337+
cheatcode.dyn_apply(ccx, executor)
338+
}
321339
t if using_pvm && is::<warpCall>(t) => {
322340
let &warpCall { newTimestamp } = cheatcode.as_any().downcast_ref().unwrap();
323341

‎crates/revive-strategy/src/executor/runner.rs‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -157,12 +157,13 @@ impl ExecutorStrategyExt for ReviveExecutorStrategyRunner {
157157
fn start_transaction(&self, ctx: &dyn ExecutorStrategyContext) {
158158
let ctx = get_context_ref(ctx);
159159
let mut externalities = ctx.externalties.0.lock().unwrap();
160-
externalities.ext().storage_start_transaction();
160+
externalities.externalities.ext().storage_start_transaction();
161161
}
162162

163163
fn rollback_transaction(&self, ctx: &dyn ExecutorStrategyContext) {
164164
let ctx = get_context_ref(ctx);
165-
let mut externalities = ctx.externalties.0.lock().unwrap();
166-
externalities.ext().storage_rollback_transaction().unwrap();
165+
let mut state = ctx.externalties.0.lock().unwrap();
166+
state.depth = 0;
167+
while let Ok(_) = state.externalities.ext().storage_rollback_transaction() {}
167168
}
168169
}

‎crates/revive-strategy/src/state.rs‎

Lines changed: 45 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
use alloy_primitives::{Address, Bytes, FixedBytes, U256};
22
use foundry_cheatcodes::{Ecx, Error, Result};
3+
use polkadot_sdk::sp_externalities::Externalities;
34
use polkadot_sdk::{
45
pallet_revive::{
56
self, AccountInfo, AddressMapper, BalanceOf, BytecodeType, ContractInfo, ExecConfig,
@@ -8,20 +9,29 @@ use polkadot_sdk::{
89
sp_core::{self, H160},
910
sp_io::TestExternalities,
1011
};
12+
1113
use revive_env::{AccountId, ExtBuilder, Runtime, System, Timestamp};
1214
use std::{
1315
fmt::Debug,
1416
sync::{Arc, Mutex},
1517
};
16-
pub struct TestEnv(pub Arc<Mutex<TestExternalities>>);
1718

18-
impl Default for TestEnv {
19+
pub(crate) struct Inner {
20+
pub externalities: TestExternalities,
21+
pub depth: usize,
22+
}
23+
24+
#[derive(Default)]
25+
pub struct TestEnv(pub(crate) Arc<Mutex<Inner>>);
26+
27+
impl Default for Inner {
1928
fn default() -> Self {
20-
Self(Arc::new(Mutex::new(
21-
ExtBuilder::default()
29+
Self {
30+
externalities: ExtBuilder::default()
2231
.balance_genesis_config(vec![(H160::from_low_u64_be(1), 1000)])
2332
.build(),
24-
)))
33+
depth: 0,
34+
}
2535
}
2636
}
2737

@@ -33,9 +43,9 @@ impl Debug for TestEnv {
3343

3444
impl Clone for TestEnv {
3545
fn clone(&self) -> Self {
36-
let mut externalities = ExtBuilder::default().build();
37-
externalities.backend = self.0.lock().unwrap().as_backend();
38-
Self(Arc::new(Mutex::new(externalities)))
46+
let mut inner: Inner = Default::default();
47+
inner.externalities.backend = self.0.lock().unwrap().externalities.as_backend();
48+
Self(Arc::new(Mutex::new(inner)))
3949
}
4050
}
4151

@@ -44,20 +54,36 @@ impl TestEnv {
4454
Self(self.0.clone())
4555
}
4656

57+
pub fn start_snapshotting(&mut self) {
58+
let mut state = self.0.lock().unwrap();
59+
state.depth += 1;
60+
state.externalities.ext().storage_start_transaction();
61+
}
62+
63+
pub fn revert(&mut self, depth: usize) {
64+
let mut state = self.0.lock().unwrap();
65+
while state.depth > depth + 1 {
66+
state.externalities.ext().storage_rollback_transaction().unwrap();
67+
state.depth -= 1;
68+
}
69+
state.externalities.ext().storage_rollback_transaction().unwrap();
70+
state.externalities.ext().storage_start_transaction();
71+
}
72+
4773
pub fn execute_with<R, F: FnOnce() -> R>(&mut self, f: F) -> R {
48-
self.0.lock().unwrap().execute_with(f)
74+
self.0.lock().unwrap().externalities.execute_with(f)
4975
}
5076

5177
pub fn get_nonce(&mut self, account: Address) -> u32 {
52-
self.0.lock().unwrap().execute_with(|| {
78+
self.0.lock().unwrap().externalities.execute_with(|| {
5379
System::account_nonce(AccountId::to_fallback_account_id(&H160::from_slice(
5480
account.as_slice(),
5581
)))
5682
})
5783
}
5884

5985
pub fn set_nonce(&mut self, address: Address, nonce: u64) {
60-
self.0.lock().unwrap().execute_with(|| {
86+
self.0.lock().unwrap().externalities.execute_with(|| {
6187
let account_id =
6288
AccountId::to_fallback_account_id(&H160::from_slice(address.as_slice()));
6389

@@ -69,7 +95,7 @@ impl TestEnv {
6995

7096
pub fn set_chain_id(&mut self, new_chain_id: u64) {
7197
// Set chain id in pallet-revive runtime.
72-
self.0.lock().unwrap().execute_with(|| {
98+
self.0.lock().unwrap().externalities.execute_with(|| {
7399
<revive_env::Runtime as polkadot_sdk::pallet_revive::Config>::ChainId::set(
74100
&new_chain_id,
75101
);
@@ -78,14 +104,14 @@ impl TestEnv {
78104

79105
pub fn set_block_number(&mut self, new_height: U256) {
80106
// Set block number in pallet-revive runtime.
81-
self.0.lock().unwrap().execute_with(|| {
107+
self.0.lock().unwrap().externalities.execute_with(|| {
82108
System::set_block_number(new_height.try_into().expect("Block number exceeds u64"));
83109
});
84110
}
85111

86112
pub fn set_timestamp(&mut self, new_timestamp: U256) {
87113
// Set timestamp in pallet-revive runtime (milliseconds).
88-
self.0.lock().unwrap().execute_with(|| {
114+
self.0.lock().unwrap().externalities.execute_with(|| {
89115
let timestamp_ms = new_timestamp.saturating_to::<u64>().saturating_mul(1000);
90116
Timestamp::set_timestamp(timestamp_ms);
91117
});
@@ -97,7 +123,7 @@ impl TestEnv {
97123
new_runtime_code: &Bytes,
98124
ecx: Ecx<'_, '_, '_>,
99125
) -> Result {
100-
self.0.lock().unwrap().execute_with(|| {
126+
self.0.lock().unwrap().externalities.execute_with(|| {
101127
let origin_address = H160::from_slice(ecx.tx.caller.as_slice());
102128
let origin_account = AccountId::to_fallback_account_id(&origin_address);
103129

@@ -153,6 +179,7 @@ impl TestEnv {
153179
self.0
154180
.lock()
155181
.unwrap()
182+
.externalities
156183
.execute_with(|| {
157184
pallet_revive::Pallet::<Runtime>::get_storage(target_address_h160, slot.into())
158185
})
@@ -169,6 +196,7 @@ impl TestEnv {
169196
self.0
170197
.lock()
171198
.unwrap()
199+
.externalities
172200
.execute_with(|| {
173201
pallet_revive::Pallet::<Runtime>::set_storage(
174202
target_address_h160,
@@ -184,7 +212,7 @@ impl TestEnv {
184212
let amount_pvm =
185213
sp_core::U256::from_little_endian(&amount.as_le_bytes()).min(u128::MAX.into());
186214

187-
self.0.lock().unwrap().execute_with(|| {
215+
self.0.lock().unwrap().externalities.execute_with(|| {
188216
let h160_addr = H160::from_slice(address.as_slice());
189217
pallet_revive::Pallet::<Runtime>::set_evm_balance(&h160_addr, amount_pvm)
190218
.expect("failed to set evm balance");
@@ -195,6 +223,7 @@ impl TestEnv {
195223
self.0
196224
.lock()
197225
.unwrap()
226+
.externalities
198227
.execute_with(|| {
199228
let h160_addr = H160::from_slice(address.as_slice());
200229
pallet_revive::Pallet::<Runtime>::evm_balance(&h160_addr)

0 commit comments

Comments
 (0)