Fix ExpectRevert family of cheats (#458)

* fix expect_revert

Author: pkhry

crates/forge/tests/it/revive/cheats_individual.rs

@@ -41,6 +41,7 @@ macro_rules! revive_cheat_test_original {
 
 revive_cheat_test!(test_custom_nonce, "Nonce");
 revive_cheat_test_original!(test_nonce, "Nonce");
+revive_cheat_test_original!(test_expect_revert, "ExpectRevert");
 revive_cheat_test!(test_coinbase, "CoinBase");
 revive_cheat_test!(test_set_custom_blockhash, "SetBlockhash");
 revive_cheat_test_original!(test_set_blockhash, "SetBlockhash");

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

@@ -875,7 +875,7 @@ impl foundry_cheatcodes::CheatcodeInspectorStrategyExt for PvmCheatcodeInspector
 
         let gas_price_pvm =
             sp_core::U256::from_little_endian(&U256::from(ecx.tx.gas_price).as_le_bytes());
-        let mut tracer = Tracer::new(true);
+        let mut tracer = Tracer::new();
         let caller_h160 = H160::from_slice(input.caller().as_slice());
 
         let res = ctx.externalities.execute_with(|| {
@@ -1042,7 +1042,7 @@ impl foundry_cheatcodes::CheatcodeInspectorStrategyExt for PvmCheatcodeInspector
         let should_bump_nonce = !call.is_static;
         let caller_h160 = H160::from_slice(call.caller.as_slice());
 
-        let mut tracer = Tracer::new(true);
+        let mut tracer = Tracer::new();
         let res = ctx.externalities.execute_with(|| {
             // Watch the caller's address so its nonce changes get tracked in prestate trace
             tracer.watch_address(&caller_h160);
@@ -1238,8 +1238,11 @@ fn post_exec(
     }
 
     if let Some(expected_revert) = &mut state.expected_revert {
-        expected_revert.max_depth =
-            std::cmp::max(ecx.journaled_state.depth() + 1, expected_revert.max_depth);
+        expected_revert.max_depth = std::cmp::max(
+            ecx.journaled_state.depth() + tracer.revert_tracer.max_depth,
+            expected_revert.max_depth,
+        );
+        expected_revert.reverted_by = tracer.revert_tracer.has_reverted.map(|x| Address::from(x.0));
     }
 }
 

crates/revive-strategy/src/tracing/mod.rs

@@ -11,19 +11,22 @@ use polkadot_sdk::pallet_revive::{
     },
     tracing::{Tracing, trace as trace_revive},
 };
+use revert_tracer::RevertTracer;
 use revive_env::Runtime;
 use revm::{context::JournalTr, database::states::StorageSlot, state::Bytecode};
 use storage_tracer::{AccountAccess, StorageTracer};
+mod revert_tracer;
 pub mod storage_tracer;
 
 pub struct Tracer {
     pub call_tracer: CallTracer<U256, fn(Weight) -> U256>,
     pub prestate_tracer: PrestateTracer<Runtime>,
-    pub storage_accesses: Option<StorageTracer>,
+    pub storage_accesses: StorageTracer,
+    pub revert_tracer: RevertTracer,
 }
 
 impl Tracer {
-    pub fn new(is_recording: bool) -> Self {
+    pub fn new() -> Self {
         let call_tracer =
             match Pallet::<revive_env::Runtime>::evm_tracer(TracerType::CallTracer(None)) {
                 ReviveTracer::CallTracer(tracer) => tracer,
@@ -37,9 +40,12 @@ impl Tracer {
                 disable_code: false,
             });
 
-        let storage_tracer = if is_recording { Some(Default::default()) } else { None };
-
-        Self { call_tracer, prestate_tracer, storage_accesses: storage_tracer }
+        Self {
+            call_tracer,
+            prestate_tracer,
+            storage_accesses: Default::default(),
+            revert_tracer: RevertTracer::new(),
+        }
     }
 
     pub fn trace<R, F: FnOnce() -> R>(&mut self, f: F) -> R {
@@ -58,7 +64,7 @@ impl Tracer {
 
     /// Collects recorded accesses
     pub fn get_recorded_accesses(&mut self) -> Vec<AccountAccess> {
-        self.storage_accesses.take().unwrap_or_default().get_records()
+        self.storage_accesses.get_records()
     }
 
     /// Applies `PrestateTrace` diffs to the revm state
@@ -146,9 +152,21 @@ impl Tracing for Tracer {
     fn watch_address(&mut self, addr: &polkadot_sdk::sp_core::H160) {
         self.prestate_tracer.watch_address(addr);
         self.call_tracer.watch_address(addr);
-        if let Some(storage_tracer) = &mut self.storage_accesses {
-            storage_tracer.watch_address(addr);
-        }
+        self.storage_accesses.watch_address(addr);
+        self.revert_tracer.watch_address(addr);
+    }
+
+    fn terminate(
+        &mut self,
+        contract_address: polkadot_sdk::sp_core::H160,
+        beneficiary_address: polkadot_sdk::sp_core::H160,
+        gas_left: Weight,
+        value: U256,
+    ) {
+        self.prestate_tracer.terminate(contract_address, beneficiary_address, gas_left, value);
+        self.call_tracer.terminate(contract_address, beneficiary_address, gas_left, value);
+        self.storage_accesses.terminate(contract_address, beneficiary_address, gas_left, value);
+        self.revert_tracer.terminate(contract_address, beneficiary_address, gas_left, value);
     }
 
     fn enter_child_span(
@@ -179,17 +197,24 @@ impl Tracing for Tracer {
             input,
             gas,
         );
-        if let Some(storage_tracer) = &mut self.storage_accesses {
-            storage_tracer.enter_child_span(
-                from,
-                to,
-                is_delegate_call,
-                is_read_only,
-                value,
-                input,
-                gas,
-            )
-        }
+        self.storage_accesses.enter_child_span(
+            from,
+            to,
+            is_delegate_call,
+            is_read_only,
+            value,
+            input,
+            gas,
+        );
+        self.revert_tracer.enter_child_span(
+            from,
+            to,
+            is_delegate_call,
+            is_read_only,
+            value,
+            input,
+            gas,
+        )
     }
 
     fn instantiate_code(
@@ -199,25 +224,22 @@ impl Tracing for Tracer {
     ) {
         self.prestate_tracer.instantiate_code(code, salt);
         self.call_tracer.instantiate_code(code, salt);
-        if let Some(storage_tracer) = &mut self.storage_accesses {
-            storage_tracer.instantiate_code(code, salt);
-        }
+        self.storage_accesses.instantiate_code(code, salt);
+        self.revert_tracer.instantiate_code(code, salt);
     }
 
     fn balance_read(&mut self, addr: &polkadot_sdk::sp_core::H160, value: U256) {
         self.prestate_tracer.balance_read(addr, value);
         self.call_tracer.balance_read(addr, value);
-        if let Some(storage_tracer) = &mut self.storage_accesses {
-            storage_tracer.balance_read(addr, value);
-        }
+        self.storage_accesses.balance_read(addr, value);
+        self.revert_tracer.balance_read(addr, value);
     }
 
     fn storage_read(&mut self, key: &polkadot_sdk::pallet_revive::Key, value: Option<&[u8]>) {
         self.prestate_tracer.storage_read(key, value);
         self.call_tracer.storage_read(key, value);
-        if let Some(storage_tracer) = &mut self.storage_accesses {
-            storage_tracer.storage_read(key, value);
-        }
+        self.storage_accesses.storage_read(key, value);
+        self.revert_tracer.storage_read(key, value);
     }
 
     fn storage_write(
@@ -228,9 +250,8 @@ impl Tracing for Tracer {
     ) {
         self.prestate_tracer.storage_write(key, old_value.clone(), new_value);
         self.call_tracer.storage_write(key, old_value.clone(), new_value);
-        if let Some(storage_tracer) = &mut self.storage_accesses {
-            storage_tracer.storage_write(key, old_value, new_value);
-        }
+        self.storage_accesses.storage_write(key, old_value.clone(), new_value);
+        self.revert_tracer.storage_write(key, old_value, new_value);
     }
 
     fn log_event(
@@ -241,9 +262,8 @@ impl Tracing for Tracer {
     ) {
         self.prestate_tracer.log_event(event, topics, data);
         self.call_tracer.log_event(event, topics, data);
-        if let Some(storage_tracer) = &mut self.storage_accesses {
-            storage_tracer.log_event(event, topics, data);
-        }
+        self.storage_accesses.log_event(event, topics, data);
+        self.revert_tracer.log_event(event, topics, data);
     }
 
     fn exit_child_span(
@@ -253,9 +273,8 @@ impl Tracing for Tracer {
     ) {
         self.prestate_tracer.exit_child_span(output, gas_left);
         self.call_tracer.exit_child_span(output, gas_left);
-        if let Some(storage_tracer) = &mut self.storage_accesses {
-            storage_tracer.exit_child_span(output, gas_left);
-        }
+        self.storage_accesses.exit_child_span(output, gas_left);
+        self.revert_tracer.exit_child_span(output, gas_left);
     }
 
     fn exit_child_span_with_error(
@@ -265,8 +284,7 @@ impl Tracing for Tracer {
     ) {
         self.prestate_tracer.exit_child_span_with_error(error, gas_left);
         self.call_tracer.exit_child_span_with_error(error, gas_left);
-        if let Some(storage_tracer) = &mut self.storage_accesses {
-            storage_tracer.exit_child_span_with_error(error, gas_left);
-        }
+        self.storage_accesses.exit_child_span_with_error(error, gas_left);
+        self.revert_tracer.exit_child_span_with_error(error, gas_left);
     }
 }

crates/revive-strategy/src/tracing/revert_tracer.rs

@@ -0,0 +1,83 @@
+use polkadot_sdk::{
+    pallet_revive::{Code, tracing::Tracing},
+    sp_core::{H160, U256},
+    sp_weights::Weight,
+};
+
+#[derive(Debug)]
+pub(crate) struct RevertTracer {
+    pub max_depth: usize,
+    pub has_reverted: Option<polkadot_sdk::sp_core::H160>,
+    pub calls: Vec<polkadot_sdk::sp_core::H160>,
+    pub is_create: bool,
+    pub call_types: Vec<Type>,
+}
+
+#[derive(Debug)]
+pub enum Type {
+    Create,
+    Rest,
+}
+
+impl RevertTracer {
+    fn current_addr(&self) -> H160 {
+        self.calls.last().copied().unwrap_or_default()
+    }
+    pub fn new() -> Self {
+        Self {
+            max_depth: 1,
+            has_reverted: None,
+            calls: vec![],
+            is_create: false,
+            call_types: vec![],
+        }
+    }
+}
+
+impl Tracing for RevertTracer {
+    fn instantiate_code(&mut self, _code: &Code, _salt: Option<&[u8; 32]>) {
+        self.is_create = true;
+    }
+
+    fn enter_child_span(
+        &mut self,
+        _from: H160,
+        to: H160,
+        _is_delegate_call: bool,
+        _is_read_only: bool,
+        _value: U256,
+        _input: &[u8],
+        _gas: Weight,
+    ) {
+        self.call_types.push(if self.is_create { Type::Create } else { Type::Rest });
+
+        self.calls.push(if self.call_types.last().is_some_and(|x| matches!(x, Type::Create)) {
+            self.current_addr()
+        } else {
+            to
+        });
+
+        if self.has_reverted.is_none() {
+            self.max_depth += 1;
+        }
+    }
+
+    fn exit_child_span(
+        &mut self,
+        output: &polkadot_sdk::pallet_revive::ExecReturnValue,
+        _gas_left: Weight,
+    ) {
+        let addr = self.calls.pop().unwrap_or_default();
+
+        if output.did_revert() && self.has_reverted.is_none() {
+            self.has_reverted = Some(addr);
+        }
+        let typ = self.call_types.pop();
+        if typ.is_some_and(|x| matches!(x, Type::Create)) {
+            self.is_create = false;
+        }
+        if self.has_reverted.is_none() {
+            self.max_depth -= 1;
+        }
+    }
+}