diff --git a/crates/cheatcodes/assets/cheatcodes.json b/crates/cheatcodes/assets/cheatcodes.json index d3a47288ff614..f612dda4a84c7 100644 --- a/crates/cheatcodes/assets/cheatcodes.json +++ b/crates/cheatcodes/assets/cheatcodes.json @@ -604,6 +604,27 @@ "description": "Address of the contract implementation that will be delegated to.\n Gets encoded into delegation code: 0xef0100 || implementation." } ] + }, + { + "name": "PotentialRevert", + "description": "Represents a \"potential\" revert reason from a single subsequent call when using `vm.assumeNoReverts`.\n Reverts that match will result in a FOUNDRY::ASSUME rejection, whereas unmatched reverts will be surfaced\n as normal.", + "fields": [ + { + "name": "reverter", + "ty": "address", + "description": "The allowed origin of the revert opcode; address(0) allows reverts from any address" + }, + { + "name": "partialMatch", + "ty": "bool", + "description": "When true, only matches on the beginning of the revert data, otherwise, matches on entire revert data" + }, + { + "name": "revertData", + "ty": "bytes", + "description": "The data to use to match encountered reverts" + } + ] } ], "cheatcodes": [ @@ -3089,7 +3110,7 @@ }, { "func": { - "id": "assumeNoRevert", + "id": "assumeNoRevert_0", "description": "Discard this run's fuzz inputs and generate new ones if next call reverted.", "declaration": "function assumeNoRevert() external pure;", "visibility": "external", @@ -3107,6 +3128,46 @@ "status": "stable", "safety": "safe" }, + { + "func": { + "id": "assumeNoRevert_1", + "description": "Discard this run's fuzz inputs and generate new ones if next call reverts with the potential revert parameters.", + "declaration": "function assumeNoRevert(PotentialRevert calldata potentialRevert) external pure;", + "visibility": "external", + "mutability": "pure", + "signature": "assumeNoRevert((address,bool,bytes))", + "selector": "0xd8591eeb", + "selectorBytes": [ + 216, + 89, + 30, + 235 + ] + }, + "group": "testing", + "status": "stable", + "safety": "safe" + }, + { + "func": { + "id": "assumeNoRevert_2", + "description": "Discard this run's fuzz inputs and generate new ones if next call reverts with the any of the potential revert parameters.", + "declaration": "function assumeNoRevert(PotentialRevert[] calldata potentialReverts) external pure;", + "visibility": "external", + "mutability": "pure", + "signature": "assumeNoRevert((address,bool,bytes)[])", + "selector": "0x8a4592cc", + "selectorBytes": [ + 138, + 69, + 146, + 204 + ] + }, + "group": "testing", + "status": "stable", + "safety": "safe" + }, { "func": { "id": "attachDelegation", diff --git a/crates/cheatcodes/spec/src/lib.rs b/crates/cheatcodes/spec/src/lib.rs index c4d7e9868fe13..e39cfc2d7ba91 100644 --- a/crates/cheatcodes/spec/src/lib.rs +++ b/crates/cheatcodes/spec/src/lib.rs @@ -88,6 +88,7 @@ impl Cheatcodes<'static> { Vm::DebugStep::STRUCT.clone(), Vm::BroadcastTxSummary::STRUCT.clone(), Vm::SignedDelegation::STRUCT.clone(), + Vm::PotentialRevert::STRUCT.clone(), ]), enums: Cow::Owned(vec![ Vm::CallerMode::ENUM.clone(), diff --git a/crates/cheatcodes/spec/src/vm.rs b/crates/cheatcodes/spec/src/vm.rs index 47e0b625b4de0..f1bb68f5ee8e0 100644 --- a/crates/cheatcodes/spec/src/vm.rs +++ b/crates/cheatcodes/spec/src/vm.rs @@ -324,6 +324,18 @@ interface Vm { address implementation; } + /// Represents a "potential" revert reason from a single subsequent call when using `vm.assumeNoReverts`. + /// Reverts that match will result in a FOUNDRY::ASSUME rejection, whereas unmatched reverts will be surfaced + /// as normal. + struct PotentialRevert { + /// The allowed origin of the revert opcode; address(0) allows reverts from any address + address reverter; + /// When true, only matches on the beginning of the revert data, otherwise, matches on entire revert data + bool partialMatch; + /// The data to use to match encountered reverts + bytes revertData; + } + // ======== EVM ======== /// Gets the address for a given private key. @@ -894,6 +906,14 @@ interface Vm { #[cheatcode(group = Testing, safety = Safe)] function assumeNoRevert() external pure; + /// Discard this run's fuzz inputs and generate new ones if next call reverts with the potential revert parameters. + #[cheatcode(group = Testing, safety = Safe)] + function assumeNoRevert(PotentialRevert calldata potentialRevert) external pure; + + /// Discard this run's fuzz inputs and generate new ones if next call reverts with the any of the potential revert parameters. + #[cheatcode(group = Testing, safety = Safe)] + function assumeNoRevert(PotentialRevert[] calldata potentialReverts) external pure; + /// Writes a breakpoint to jump to in the debugger. #[cheatcode(group = Testing, safety = Safe)] function breakpoint(string calldata char) external pure; diff --git a/crates/cheatcodes/src/inspector.rs b/crates/cheatcodes/src/inspector.rs index a0695d0d254f3..3cb74c6ac8550 100644 --- a/crates/cheatcodes/src/inspector.rs +++ b/crates/cheatcodes/src/inspector.rs @@ -15,6 +15,7 @@ use crate::{ self, ExpectedCallData, ExpectedCallTracker, ExpectedCallType, ExpectedEmitTracker, ExpectedRevert, ExpectedRevertKind, }, + revert_handlers, }, utils::IgnoredTraces, CheatsConfig, CheatsCtxt, DynCheatcode, Error, Result, @@ -755,16 +756,14 @@ where { matches!(expected_revert.kind, ExpectedRevertKind::Default) { let mut expected_revert = std::mem::take(&mut self.expected_revert).unwrap(); - let handler_result = expect::handle_expect_revert( + return match revert_handlers::handle_expect_revert( false, true, - &mut expected_revert, + &expected_revert, outcome.result.result, outcome.result.output.clone(), &self.config.available_artifacts, - ); - - return match handler_result { + ) { Ok((address, retdata)) => { expected_revert.actual_count += 1; if expected_revert.actual_count < expected_revert.count { @@ -1287,16 +1286,45 @@ impl Inspector<&mut dyn DatabaseExt> for Cheatcodes { } } - // Handle assume not revert cheatcode. - if let Some(assume_no_revert) = &self.assume_no_revert { - if ecx.journaled_state.depth() == assume_no_revert.depth && !cheatcode_call { - // Discard run if we're at the same depth as cheatcode and call reverted. + // Handle assume no revert cheatcode. + if let Some(assume_no_revert) = &mut self.assume_no_revert { + // Record current reverter address before processing the expect revert if call reverted, + // expect revert is set with expected reverter address and no actual reverter set yet. + if outcome.result.is_revert() && assume_no_revert.reverted_by.is_none() { + assume_no_revert.reverted_by = Some(call.target_address); + } + // allow multiple cheatcode calls at the same depth + if ecx.journaled_state.depth() <= assume_no_revert.depth && !cheatcode_call { + // Discard run if we're at the same depth as cheatcode, call reverted, and no + // specific reason was supplied if outcome.result.is_revert() { - outcome.result.output = Error::from(MAGIC_ASSUME).abi_encode().into(); + let assume_no_revert = std::mem::take(&mut self.assume_no_revert).unwrap(); + return match revert_handlers::handle_assume_no_revert( + &assume_no_revert, + outcome.result.result, + &outcome.result.output, + &self.config.available_artifacts, + ) { + // if result is Ok, it was an anticipated revert; return an "assume" error + // to reject this run + Ok(_) => { + outcome.result.output = Error::from(MAGIC_ASSUME).abi_encode().into(); + outcome + } + // if result is Error, it was an unanticipated revert; should revert + // normally + Err(error) => { + trace!(expected=?assume_no_revert, ?error, status=?outcome.result.result, "Expected revert mismatch"); + outcome.result.result = InstructionResult::Revert; + outcome.result.output = error.abi_encode().into(); + outcome + } + } + } else { + // Call didn't revert, reset `assume_no_revert` state. + self.assume_no_revert = None; return outcome; } - // Call didn't revert, reset `assume_no_revert` state. - self.assume_no_revert = None; } } @@ -1330,20 +1358,15 @@ impl Inspector<&mut dyn DatabaseExt> for Cheatcodes { }; if needs_processing { - // Only `remove` the expected revert from state if `expected_revert.count` == - // `expected_revert.actual_count` let mut expected_revert = std::mem::take(&mut self.expected_revert).unwrap(); - - let handler_result = expect::handle_expect_revert( + return match revert_handlers::handle_expect_revert( cheatcode_call, false, - &mut expected_revert, + &expected_revert, outcome.result.result, outcome.result.output.clone(), &self.config.available_artifacts, - ); - - return match handler_result { + ) { Err(error) => { trace!(expected=?expected_revert, ?error, status=?outcome.result.result, "Expected revert mismatch"); outcome.result.result = InstructionResult::Revert; diff --git a/crates/cheatcodes/src/test.rs b/crates/cheatcodes/src/test.rs index 417ae69833f73..f1876fc849cf1 100644 --- a/crates/cheatcodes/src/test.rs +++ b/crates/cheatcodes/src/test.rs @@ -10,6 +10,7 @@ use std::env; pub(crate) mod assert; pub(crate) mod assume; pub(crate) mod expect; +pub(crate) mod revert_handlers; impl Cheatcode for breakpoint_0Call { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { diff --git a/crates/cheatcodes/src/test/assume.rs b/crates/cheatcodes/src/test/assume.rs index a0321b5a1cd38..74bd79e0964b6 100644 --- a/crates/cheatcodes/src/test/assume.rs +++ b/crates/cheatcodes/src/test/assume.rs @@ -1,12 +1,46 @@ use crate::{Cheatcode, Cheatcodes, CheatsCtxt, Error, Result}; +use alloy_primitives::Address; use foundry_evm_core::constants::MAGIC_ASSUME; -use spec::Vm::{assumeCall, assumeNoRevertCall}; +use spec::Vm::{ + assumeCall, assumeNoRevert_0Call, assumeNoRevert_1Call, assumeNoRevert_2Call, PotentialRevert, +}; use std::fmt::Debug; #[derive(Clone, Debug)] pub struct AssumeNoRevert { /// The call depth at which the cheatcode was added. pub depth: u64, + /// Acceptable revert parameters for the next call, to be thrown out if they are encountered; + /// reverts with parameters not specified here will count as normal reverts and not rejects + /// towards the counter. + pub reasons: Vec, + /// Address that reverted the call. + pub reverted_by: Option
, +} + +/// Parameters for a single anticipated revert, to be thrown out if encountered. +#[derive(Clone, Debug)] +pub struct AcceptableRevertParameters { + /// The expected revert data returned by the revert + pub reason: Vec, + /// If true then only the first 4 bytes of expected data returned by the revert are checked. + pub partial_match: bool, + /// Contract expected to revert next call. + pub reverter: Option
, +} + +impl AcceptableRevertParameters { + fn from(potential_revert: &PotentialRevert) -> Self { + Self { + reason: potential_revert.revertData.to_vec(), + partial_match: potential_revert.partialMatch, + reverter: if potential_revert.reverter == Address::ZERO { + None + } else { + Some(potential_revert.reverter) + }, + } + } } impl Cheatcode for assumeCall { @@ -20,10 +54,45 @@ impl Cheatcode for assumeCall { } } -impl Cheatcode for assumeNoRevertCall { +impl Cheatcode for assumeNoRevert_0Call { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { - ccx.state.assume_no_revert = - Some(AssumeNoRevert { depth: ccx.ecx.journaled_state.depth() }); - Ok(Default::default()) + assume_no_revert(ccx.state, ccx.ecx.journaled_state.depth(), vec![]) } } + +impl Cheatcode for assumeNoRevert_1Call { + fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { + let Self { potentialRevert } = self; + assume_no_revert( + ccx.state, + ccx.ecx.journaled_state.depth(), + vec![AcceptableRevertParameters::from(potentialRevert)], + ) + } +} + +impl Cheatcode for assumeNoRevert_2Call { + fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { + let Self { potentialReverts } = self; + assume_no_revert( + ccx.state, + ccx.ecx.journaled_state.depth(), + potentialReverts.iter().map(AcceptableRevertParameters::from).collect(), + ) + } +} + +fn assume_no_revert( + state: &mut Cheatcodes, + depth: u64, + parameters: Vec, +) -> Result { + ensure!( + state.assume_no_revert.is_none(), + "you must make another external call prior to calling assumeNoRevert again" + ); + + state.assume_no_revert = Some(AssumeNoRevert { depth, reasons: parameters, reverted_by: None }); + + Ok(Default::default()) +} diff --git a/crates/cheatcodes/src/test/expect.rs b/crates/cheatcodes/src/test/expect.rs index e45298923cd69..f77eb4d402cff 100644 --- a/crates/cheatcodes/src/test/expect.rs +++ b/crates/cheatcodes/src/test/expect.rs @@ -2,29 +2,12 @@ use std::collections::VecDeque; use crate::{Cheatcode, Cheatcodes, CheatsCtxt, Error, Result, Vm::*}; use alloy_primitives::{ - address, hex, map::{hash_map::Entry, AddressHashMap, HashMap}, Address, Bytes, LogData as RawLog, U256, }; -use alloy_sol_types::{SolError, SolValue}; -use foundry_common::ContractsByArtifact; -use foundry_evm_core::decode::RevertDecoder; -use revm::interpreter::{ - return_ok, InstructionResult, Interpreter, InterpreterAction, InterpreterResult, -}; -use spec::Vm; - -/// For some cheatcodes we may internally change the status of the call, i.e. in `expectRevert`. -/// Solidity will see a successful call and attempt to decode the return data. Therefore, we need -/// to populate the return with dummy bytes so the decode doesn't fail. -/// -/// 8192 bytes was arbitrarily chosen because it is long enough for return values up to 256 words in -/// size. -static DUMMY_CALL_OUTPUT: Bytes = Bytes::from_static(&[0u8; 8192]); - -/// Same reasoning as [DUMMY_CALL_OUTPUT], but for creates. -const DUMMY_CREATE_ADDRESS: Address = address!("0000000000000000000000000000000000000001"); +use revm::interpreter::{InstructionResult, Interpreter, InterpreterAction, InterpreterResult}; +use super::revert_handlers::RevertParameters; /// Tracks the expected calls per address. /// /// For each address, we track the expected calls per call data. We track it in such manner @@ -87,7 +70,7 @@ pub struct ExpectedRevert { pub partial_match: bool, /// Contract expected to revert next call. pub reverter: Option
, - /// Actual reverter of the call. + /// Address that reverted the call. pub reverted_by: Option
, /// Number of times this revert is expected. pub count: u64, @@ -605,6 +588,20 @@ impl Cheatcode for expectSafeMemoryCallCall { } } +impl RevertParameters for ExpectedRevert { + fn reverter(&self) -> Option
{ + self.reverter + } + + fn reason(&self) -> Option<&[u8]> { + self.reason.as_deref() + } + + fn partial_match(&self) -> bool { + self.partial_match + } +} + /// Handles expected calls specified by the `expectCall` cheatcodes. /// /// It can handle calls in two ways: @@ -920,136 +917,6 @@ fn expect_revert( Ok(Default::default()) } -pub(crate) fn handle_expect_revert( - is_cheatcode: bool, - is_create: bool, - expected_revert: &mut ExpectedRevert, - status: InstructionResult, - retdata: Bytes, - known_contracts: &Option, -) -> Result<(Option
, Bytes)> { - let success_return = || { - if is_create { - (Some(DUMMY_CREATE_ADDRESS), Bytes::new()) - } else { - (None, DUMMY_CALL_OUTPUT.clone()) - } - }; - - let stringify = |data: &[u8]| { - if let Ok(s) = String::abi_decode(data, true) { - return s; - } - if data.is_ascii() { - return std::str::from_utf8(data).unwrap().to_owned(); - } - hex::encode_prefixed(data) - }; - - if expected_revert.count == 0 { - if expected_revert.reverter.is_none() && expected_revert.reason.is_none() { - ensure!( - matches!(status, return_ok!()), - "call reverted when it was expected not to revert" - ); - return Ok(success_return()); - } - - // Flags to track if the reason and reverter match. - let mut reason_match = expected_revert.reason.as_ref().map(|_| false); - let mut reverter_match = expected_revert.reverter.as_ref().map(|_| false); - - // Reverter check - if let (Some(expected_reverter), Some(actual_reverter)) = - (expected_revert.reverter, expected_revert.reverted_by) - { - if expected_reverter == actual_reverter { - reverter_match = Some(true); - } - } - - // Reason check - let expected_reason = expected_revert.reason.as_deref(); - if let Some(expected_reason) = expected_reason { - let mut actual_revert: Vec = retdata.into(); - actual_revert = decode_revert(actual_revert); - - if actual_revert == expected_reason { - reason_match = Some(true); - } - }; - - match (reason_match, reverter_match) { - (Some(true), Some(true)) => Err(fmt_err!( - "expected 0 reverts with reason: {}, from address: {}, but got one", - &stringify(expected_reason.unwrap_or_default()), - expected_revert.reverter.unwrap() - )), - (Some(true), None) => Err(fmt_err!( - "expected 0 reverts with reason: {}, but got one", - &stringify(expected_reason.unwrap_or_default()) - )), - (None, Some(true)) => Err(fmt_err!( - "expected 0 reverts from address: {}, but got one", - expected_revert.reverter.unwrap() - )), - _ => Ok(success_return()), - } - } else { - ensure!(!matches!(status, return_ok!()), "next call did not revert as expected"); - - // If expected reverter address is set then check it matches the actual reverter. - if let (Some(expected_reverter), Some(actual_reverter)) = - (expected_revert.reverter, expected_revert.reverted_by) - { - if expected_reverter != actual_reverter { - return Err(fmt_err!( - "Reverter != expected reverter: {} != {}", - actual_reverter, - expected_reverter - )); - } - } - - let expected_reason = expected_revert.reason.as_deref(); - // If None, accept any revert. - let Some(expected_reason) = expected_reason else { - return Ok(success_return()); - }; - - if !expected_reason.is_empty() && retdata.is_empty() { - bail!("call reverted as expected, but without data"); - } - - let mut actual_revert: Vec = retdata.into(); - - // Compare only the first 4 bytes if partial match. - if expected_revert.partial_match && actual_revert.get(..4) == expected_reason.get(..4) { - return Ok(success_return()) - } - - // Try decoding as known errors. - actual_revert = decode_revert(actual_revert); - - if actual_revert == expected_reason || - (is_cheatcode && memchr::memmem::find(&actual_revert, expected_reason).is_some()) - { - Ok(success_return()) - } else { - let (actual, expected) = if let Some(contracts) = known_contracts { - let decoder = RevertDecoder::new().with_abis(contracts.iter().map(|(_, c)| &c.abi)); - ( - &decoder.decode(actual_revert.as_slice(), Some(status)), - &decoder.decode(expected_reason, Some(status)), - ) - } else { - (&stringify(&actual_revert), &stringify(expected_reason)) - }; - Err(fmt_err!("Error != expected error: {} != {}", actual, expected,)) - } - } -} - fn checks_topics_and_data(checks: [bool; 5], expected: &RawLog, log: &RawLog) -> bool { if log.topics().len() != expected.topics().len() { return false @@ -1081,15 +948,3 @@ fn expect_safe_memory(state: &mut Cheatcodes, start: u64, end: u64, depth: u64) offsets.push(start..end); Ok(Default::default()) } - -fn decode_revert(revert: Vec) -> Vec { - if matches!( - revert.get(..4).map(|s| s.try_into().unwrap()), - Some(Vm::CheatcodeError::SELECTOR | alloy_sol_types::Revert::SELECTOR) - ) { - if let Ok(decoded) = Vec::::abi_decode(&revert[4..], false) { - return decoded; - } - } - revert -} diff --git a/crates/cheatcodes/src/test/revert_handlers.rs b/crates/cheatcodes/src/test/revert_handlers.rs new file mode 100644 index 0000000000000..78297b38872aa --- /dev/null +++ b/crates/cheatcodes/src/test/revert_handlers.rs @@ -0,0 +1,234 @@ +use crate::{Error, Result}; +use alloy_primitives::{address, hex, Address, Bytes}; +use alloy_sol_types::{SolError, SolValue}; +use foundry_common::ContractsByArtifact; +use foundry_evm_core::decode::RevertDecoder; +use revm::interpreter::{return_ok, InstructionResult}; +use spec::Vm; + +use super::{ + assume::{AcceptableRevertParameters, AssumeNoRevert}, + expect::ExpectedRevert, +}; + +/// For some cheatcodes we may internally change the status of the call, i.e. in `expectRevert`. +/// Solidity will see a successful call and attempt to decode the return data. Therefore, we need +/// to populate the return with dummy bytes so the decode doesn't fail. +/// +/// 8192 bytes was arbitrarily chosen because it is long enough for return values up to 256 words in +/// size. +static DUMMY_CALL_OUTPUT: Bytes = Bytes::from_static(&[0u8; 8192]); + +/// Same reasoning as [DUMMY_CALL_OUTPUT], but for creates. +const DUMMY_CREATE_ADDRESS: Address = address!("0000000000000000000000000000000000000001"); + +fn stringify(data: &[u8]) -> String { + if let Ok(s) = String::abi_decode(data, true) { + return s; + } + if data.is_ascii() { + return std::str::from_utf8(data).unwrap().to_owned(); + } + hex::encode_prefixed(data) +} + +/// Common parameters for expected or assumed reverts. Allows for code reuse. +pub(crate) trait RevertParameters { + fn reverter(&self) -> Option
; + fn reason(&self) -> Option<&[u8]>; + fn partial_match(&self) -> bool; +} + +impl RevertParameters for AcceptableRevertParameters { + fn reverter(&self) -> Option
{ + self.reverter + } + + fn reason(&self) -> Option<&[u8]> { + Some(&self.reason) + } + + fn partial_match(&self) -> bool { + self.partial_match + } +} + +/// Core logic for handling reverts that may or may not be expected (or assumed). +fn handle_revert( + is_cheatcode: bool, + revert_params: &impl RevertParameters, + status: InstructionResult, + retdata: &Bytes, + known_contracts: &Option, + reverter: Option<&Address>, +) -> Result<(), Error> { + // If expected reverter address is set then check it matches the actual reverter. + if let (Some(expected_reverter), Some(&actual_reverter)) = (revert_params.reverter(), reverter) + { + if expected_reverter != actual_reverter { + return Err(fmt_err!( + "Reverter != expected reverter: {} != {}", + actual_reverter, + expected_reverter + )); + } + } + + let expected_reason = revert_params.reason(); + // If None, accept any revert. + let Some(expected_reason) = expected_reason else { + return Ok(()); + }; + + if !expected_reason.is_empty() && retdata.is_empty() { + bail!("call reverted as expected, but without data"); + } + + let mut actual_revert: Vec = retdata.to_vec(); + + // Compare only the first 4 bytes if partial match. + if revert_params.partial_match() && actual_revert.get(..4) == expected_reason.get(..4) { + return Ok(()); + } + + // Try decoding as known errors. + actual_revert = decode_revert(actual_revert); + + if actual_revert == expected_reason || + (is_cheatcode && memchr::memmem::find(&actual_revert, expected_reason).is_some()) + { + Ok(()) + } else { + let (actual, expected) = if let Some(contracts) = known_contracts { + let decoder = RevertDecoder::new().with_abis(contracts.iter().map(|(_, c)| &c.abi)); + ( + &decoder.decode(actual_revert.as_slice(), Some(status)), + &decoder.decode(expected_reason, Some(status)), + ) + } else { + (&stringify(&actual_revert), &stringify(expected_reason)) + }; + Err(fmt_err!("Error != expected error: {} != {}", actual, expected,)) + } +} + +pub(crate) fn handle_assume_no_revert( + assume_no_revert: &AssumeNoRevert, + status: InstructionResult, + retdata: &Bytes, + known_contracts: &Option, +) -> Result<()> { + // if a generic AssumeNoRevert, return Ok(). Otherwise, iterate over acceptable reasons and try + // to match against any, otherwise, return an Error with the revert data + if assume_no_revert.reasons.is_empty() { + Ok(()) + } else { + assume_no_revert + .reasons + .iter() + .find_map(|reason| { + handle_revert( + false, + reason, + status, + retdata, + known_contracts, + assume_no_revert.reverted_by.as_ref(), + ) + .ok() + }) + .ok_or_else(|| retdata.clone().into()) + } +} + +pub(crate) fn handle_expect_revert( + is_cheatcode: bool, + is_create: bool, + expected_revert: &ExpectedRevert, + status: InstructionResult, + retdata: Bytes, + known_contracts: &Option, +) -> Result<(Option
, Bytes)> { + let success_return = || { + if is_create { + (Some(DUMMY_CREATE_ADDRESS), Bytes::new()) + } else { + (None, DUMMY_CALL_OUTPUT.clone()) + } + }; + + if expected_revert.count == 0 { + if expected_revert.reverter.is_none() && expected_revert.reason.is_none() { + ensure!( + matches!(status, return_ok!()), + "call reverted when it was expected not to revert" + ); + return Ok(success_return()); + } + + // Flags to track if the reason and reverter match. + let mut reason_match = expected_revert.reason.as_ref().map(|_| false); + let mut reverter_match = expected_revert.reverter.as_ref().map(|_| false); + + // Reverter check + if let (Some(expected_reverter), Some(actual_reverter)) = + (expected_revert.reverter, expected_revert.reverted_by) + { + if expected_reverter == actual_reverter { + reverter_match = Some(true); + } + } + + // Reason check + let expected_reason = expected_revert.reason.as_deref(); + if let Some(expected_reason) = expected_reason { + let mut actual_revert: Vec = retdata.into(); + actual_revert = decode_revert(actual_revert); + + if actual_revert == expected_reason { + reason_match = Some(true); + } + }; + + match (reason_match, reverter_match) { + (Some(true), Some(true)) => Err(fmt_err!( + "expected 0 reverts with reason: {}, from address: {}, but got one", + &stringify(expected_reason.unwrap_or_default()), + expected_revert.reverter.unwrap() + )), + (Some(true), None) => Err(fmt_err!( + "expected 0 reverts with reason: {}, but got one", + &stringify(expected_reason.unwrap_or_default()) + )), + (None, Some(true)) => Err(fmt_err!( + "expected 0 reverts from address: {}, but got one", + expected_revert.reverter.unwrap() + )), + _ => Ok(success_return()), + } + } else { + ensure!(!matches!(status, return_ok!()), "next call did not revert as expected"); + + handle_revert( + is_cheatcode, + expected_revert, + status, + &retdata, + known_contracts, + expected_revert.reverted_by.as_ref(), + )?; + Ok(success_return()) + } +} + +fn decode_revert(revert: Vec) -> Vec { + if matches!( + revert.get(..4).map(|s| s.try_into().unwrap()), + Some(Vm::CheatcodeError::SELECTOR | alloy_sol_types::Revert::SELECTOR) + ) { + if let Ok(decoded) = Vec::::abi_decode(&revert[4..], false) { + return decoded; + } + } + revert +} diff --git a/crates/forge/tests/cli/test.sol b/crates/forge/tests/cli/test.sol new file mode 100644 index 0000000000000..0b44c7aa75b4c --- /dev/null +++ b/crates/forge/tests/cli/test.sol @@ -0,0 +1,129 @@ +import {Test} from "forge-std/Test.sol"; + +interface Vm { + struct PotentialRevert { + bytes revertData; + bool partialMatch; + address reverter; + } + function expectRevert() external; + function assumeNoRevert() external pure; + function assumeNoRevert(bytes4 revertData) external pure; + function assumeNoRevert(bytes calldata revertData) external pure; + function assumeNoRevert(bytes4 revertData, address reverter) external pure; + function assumeNoRevert(bytes calldata revertData, address reverter) external pure; +} + +contract ReverterB { + /// @notice has same error selectors as contract below to test the `reverter` param + error MyRevert(); + error SpecialRevertWithData(uint256 x); + + function revertIf2(uint256 x) public pure returns (bool) { + if (x == 2) { + revert MyRevert(); + } + return true; + } + + function revertWithData() public pure returns (bool) { + revert SpecialRevertWithData(2); + } +} + +contract Reverter { + error MyRevert(); + error RevertWithData(uint256 x); + error UnusedError(); + + ReverterB public immutable subReverter; + + constructor() { + subReverter = new ReverterB(); + } + + function myFunction() public pure returns (bool) { + revert MyRevert(); + } + + function revertIf2(uint256 value) public pure returns (bool) { + if (value == 2) { + revert MyRevert(); + } + return true; + } + + function revertWithDataIf2(uint256 value) public pure returns (bool) { + if (value == 2) { + revert RevertWithData(2); + } + return true; + } + + function twoPossibleReverts(uint256 x) public pure returns (bool) { + if (x == 2) { + revert MyRevert(); + } else if (x == 3) { + revert RevertWithData(3); + } + return true; + } +} + +contract ReverterTest is Test { + Reverter reverter; + Vm _vm = Vm(VM_ADDRESS); + + function setUp() public { + reverter = new Reverter(); + } + + /// @dev Test that `assumeNoRevert` does not reject an unanticipated error selector + function testAssume_wrongSelector_fails(uint256 x) public view { + _vm.assumeNoRevert(PotentialRevert({revertData: abi.encodeWithSelector(Reverter.UnusedError.selector), partialMatch: false, reverter: address(0)})); + reverter.revertIf2(x); + } + + /// @dev Test that `assumeNoRevert` does not reject an unanticipated error with extra data + function testAssume_wrongData_fails(uint256 x) public view { + _vm.assumeNoRevert(PotentialRevert({revertData: abi.encodeWithSelector(Reverter.RevertWithData.selector, 3), partialMatch: false, reverter: address(0)})); + reverter.revertWithDataIf2(x); + } + + /// @dev Test that `assumeNoRevert` correctly rejects an error selector from a different contract + function testAssumeWithReverter_fails(uint256 x) public view { + ReverterB subReverter = (reverter.subReverter()); + _vm.assumeNoRevert(PotentialRevert({revertData: abi.encodeWithSelector(Reverter.MyRevert.selector), partialMatch: false, reverter: address(reverter)})); + subReverter.revertIf2(x); + } + + /// @dev Test that `assumeNoRevert` correctly rejects one of two different error selectors when supplying a specific reverter + function testMultipleAssumes_OneWrong_fails(uint256 x) public view { + _vm.assumeNoRevert(PotentialRevert({revertData: abi.encodeWithSelector(Reverter.MyRevert.selector), partialMatch: false, reverter: address(reverter)})); + _vm.assumeNoRevert(PotentialRevert({revertData: abi.encodeWithSelector(Reverter.RevertWithData.selector, 4), partialMatch: false, reverter: address(reverter)})); + reverter.twoPossibleReverts(x); + } + + /// @dev Test that `assumeNoRevert` assumptions are cleared after the first non-cheatcode external call + function testMultipleAssumesClearAfterCall_fails(uint256 x) public view { + _vm.assumeNoRevert(PotentialRevert({revertData: abi.encodeWithSelector(Reverter.MyRevert.selector), partialMatch: false, reverter: address(0)})); + _vm.assumeNoRevert(PotentialRevert({revertData: abi.encodeWithSelector(Reverter.RevertWithData.selector, 4), partialMatch: false, reverter: address(reverter)})); + reverter.twoPossibleReverts(x); + + reverter.twoPossibleReverts(2); + } + + /// @dev Test that `assumeNoRevert` correctly rejects a generic assumeNoRevert call after any specific reason is provided + function testMultipleAssumes_ThrowOnGenericNoRevert_AfterSpecific_fails(bytes4 selector) public view { + _vm.assumeNoRevert(PotentialRevert({revertData: selector, partialMatch: false, reverter: address(0)})); + _vm.assumeNoRevert(); + reverter.twoPossibleReverts(2); + } + + /// @dev Test that calling `expectRevert` after `assumeNoRevert` results in an error + function testAssumeThenExpect_fails(uint256) public { + _vm.assumeNoRevert(PotentialRevert({revertData: abi.encodeWithSelector(Reverter.MyRevert.selector), partialMatch: false, reverter: address(0)})); + _vm.expectRevert(); + reverter.revertIf2(1); + } +} diff --git a/crates/forge/tests/cli/test_cmd.rs b/crates/forge/tests/cli/test_cmd.rs index de066297fd936..64c12cda308f5 100644 --- a/crates/forge/tests/cli/test_cmd.rs +++ b/crates/forge/tests/cli/test_cmd.rs @@ -2449,6 +2449,197 @@ contract Dummy { assert!(dump_path.exists()); }); +forgetest_init!(test_assume_no_revert_with_data, |prj, cmd| { + let config = Config { + fuzz: { FuzzConfig { runs: 60, seed: Some(U256::from(100)), ..Default::default() } }, + ..Default::default() + }; + prj.write_config(config); + + prj.add_source( + "AssumeNoRevertTest.t.sol", + r#" +import {Test} from "forge-std/Test.sol"; + +interface Vm { + struct PotentialRevert { + address reverter; + bool partialMatch; + bytes revertData; + } + function expectRevert() external; + function assumeNoRevert() external pure; + function assumeNoRevert(PotentialRevert calldata revertData) external pure; + function assumeNoRevert(PotentialRevert[] calldata revertData) external pure; + function expectRevert(bytes4 revertData, uint64 count) external; +} + +contract ReverterB { + /// @notice has same error selectors as contract below to test the `reverter` param + error MyRevert(); + error SpecialRevertWithData(uint256 x); + + function revertIf2(uint256 x) public pure returns (bool) { + if (x == 2) { + revert MyRevert(); + } + return true; + } + + function revertWithData() public pure returns (bool) { + revert SpecialRevertWithData(2); + } +} + +contract Reverter { + error MyRevert(); + error RevertWithData(uint256 x); + error UnusedError(); + error ExpectedRevertCountZero(); + + ReverterB public immutable subReverter; + + constructor() { + subReverter = new ReverterB(); + } + + function myFunction() public pure returns (bool) { + revert MyRevert(); + } + + function revertIf2(uint256 value) public pure returns (bool) { + if (value == 2) { + revert MyRevert(); + } + return true; + } + + function revertWithDataIf2(uint256 value) public pure returns (bool) { + if (value == 2) { + revert RevertWithData(2); + } + return true; + } + + function twoPossibleReverts(uint256 x) public pure returns (bool) { + if (x == 2) { + revert MyRevert(); + } else if (x == 3) { + revert RevertWithData(3); + } + return true; + } + + function revertIf2Or3ExpectedRevertZero(uint256 x) public pure returns (bool) { + if (x == 2) { + revert ExpectedRevertCountZero(); + } else if (x == 3) { + revert MyRevert(); + } + return true; + } +} + +contract ReverterTest is Test { + Reverter reverter; + Vm _vm = Vm(VM_ADDRESS); + + function setUp() public { + reverter = new Reverter(); + } + + /// @dev Test that `assumeNoRevert` does not reject an unanticipated error selector + function testAssume_wrongSelector_fails(uint256 x) public view { + _vm.assumeNoRevert(Vm.PotentialRevert({revertData: abi.encodeWithSelector(Reverter.UnusedError.selector), partialMatch: false, reverter: address(0)})); + reverter.revertIf2(x); + } + + /// @dev Test that `assumeNoRevert` does not reject an unanticipated error with extra data + function testAssume_wrongData_fails(uint256 x) public view { + _vm.assumeNoRevert(Vm.PotentialRevert({revertData: abi.encodeWithSelector(Reverter.RevertWithData.selector, 3), partialMatch: false, reverter: address(0)})); + reverter.revertWithDataIf2(x); + } + + /// @dev Test that `assumeNoRevert` correctly rejects an error selector from a different contract + function testAssumeWithReverter_fails(uint256 x) public view { + ReverterB subReverter = (reverter.subReverter()); + _vm.assumeNoRevert(Vm.PotentialRevert({revertData: abi.encodeWithSelector(Reverter.MyRevert.selector), partialMatch: false, reverter: address(reverter)})); + subReverter.revertIf2(x); + } + + /// @dev Test that `assumeNoRevert` correctly rejects one of two different error selectors when supplying a specific reverter + function testMultipleAssumes_OneWrong_fails(uint256 x) public view { + Vm.PotentialRevert[] memory revertData = new Vm.PotentialRevert[](2); + revertData[0] = Vm.PotentialRevert({revertData: abi.encodeWithSelector(Reverter.MyRevert.selector), partialMatch: false, reverter: address(reverter)}); + revertData[1] = Vm.PotentialRevert({revertData: abi.encodeWithSelector(Reverter.RevertWithData.selector, 4), partialMatch: false, reverter: address(reverter)}); + _vm.assumeNoRevert(revertData); + reverter.twoPossibleReverts(x); + } + + /// @dev Test that `assumeNoRevert` assumptions are cleared after the first non-cheatcode external call + function testMultipleAssumesClearAfterCall_fails(uint256 x) public view { + Vm.PotentialRevert[] memory revertData = new Vm.PotentialRevert[](2); + revertData[0] = Vm.PotentialRevert({revertData: abi.encodeWithSelector(Reverter.MyRevert.selector), partialMatch: false, reverter: address(0)}); + revertData[1] = Vm.PotentialRevert({revertData: abi.encodeWithSelector(Reverter.RevertWithData.selector, 4), partialMatch: false, reverter: address(reverter)}); + _vm.assumeNoRevert(revertData); + reverter.twoPossibleReverts(x); + + reverter.twoPossibleReverts(2); + } + + /// @dev Test that `assumeNoRevert` correctly rejects a generic assumeNoRevert call after any specific reason is provided + function testMultipleAssumes_ThrowOnGenericNoRevert_AfterSpecific_fails(bytes4 selector) public view { + _vm.assumeNoRevert(Vm.PotentialRevert({revertData: abi.encode(selector), partialMatch: false, reverter: address(0)})); + _vm.assumeNoRevert(); + reverter.twoPossibleReverts(2); + } + + function testAssumeThenExpectCountZeroFails(uint256 x) public { + _vm.assumeNoRevert( + Vm.PotentialRevert({ + revertData: abi.encodeWithSelector(Reverter.MyRevert.selector), + partialMatch: false, + reverter: address(0) + }) + ); + _vm.expectRevert(Reverter.ExpectedRevertCountZero.selector, 0); + reverter.revertIf2Or3ExpectedRevertZero(x); + } + + function testExpectCountZeroThenAssumeFails(uint256 x) public { + _vm.expectRevert(Reverter.ExpectedRevertCountZero.selector, 0); + _vm.assumeNoRevert( + Vm.PotentialRevert({ + revertData: abi.encodeWithSelector(Reverter.MyRevert.selector), + partialMatch: false, + reverter: address(0) + }) + ); + reverter.revertIf2Or3ExpectedRevertZero(x); + } + +}"#, + ) + .unwrap(); + cmd.args(["test", "--mc", "ReverterTest"]).assert_failure().stdout_eq(str![[r#" +[COMPILING_FILES] with [SOLC_VERSION] +[SOLC_VERSION] [ELAPSED] +Compiler run successful! + +Ran 8 tests for src/AssumeNoRevertTest.t.sol:ReverterTest +[FAIL: expected 0 reverts with reason: 0x92fa317b, but got one; counterexample: [..]] testAssumeThenExpectCountZeroFails(uint256) (runs: [..], [AVG_GAS]) +[FAIL: MyRevert(); counterexample: calldata=[..]] testAssumeWithReverter_fails(uint256) (runs: [..], [AVG_GAS]) +[FAIL: RevertWithData(2); counterexample: [..]] testAssume_wrongData_fails(uint256) (runs: [..], [AVG_GAS]) +[FAIL: MyRevert(); counterexample: [..]] testAssume_wrongSelector_fails(uint256) (runs: [..], [AVG_GAS]) +[FAIL: expected 0 reverts with reason: 0x92fa317b, but got one; counterexample: [..]] testExpectCountZeroThenAssumeFails(uint256) (runs: [..], [AVG_GAS]) +[FAIL: MyRevert(); counterexample: [..]] testMultipleAssumesClearAfterCall_fails(uint256) (runs: 0, [AVG_GAS]) +[FAIL: RevertWithData(3); counterexample: [..]] testMultipleAssumes_OneWrong_fails(uint256) (runs: [..], [AVG_GAS]) +[FAIL: vm.assumeNoRevert: you must make another external call prior to calling assumeNoRevert again; counterexample: [..]] testMultipleAssumes_ThrowOnGenericNoRevert_AfterSpecific_fails(bytes4) (runs: [..], [AVG_GAS]) +... + +"#]]); +}); + forgetest_async!(can_get_broadcast_txs, |prj, cmd| { foundry_test_utils::util::initialize(prj.root()); diff --git a/crates/forge/tests/it/invariant.rs b/crates/forge/tests/it/invariant.rs index df0cba0111486..6cd8482938181 100644 --- a/crates/forge/tests/it/invariant.rs +++ b/crates/forge/tests/it/invariant.rs @@ -687,24 +687,6 @@ async fn test_invariant_after_invariant() { ); } -#[tokio::test(flavor = "multi_thread")] -async fn test_invariant_selectors_weight() { - let filter = Filter::new(".*", ".*", ".*fuzz/invariant/common/InvariantSelectorsWeight.t.sol"); - let mut runner = TEST_DATA_DEFAULT.runner_with(|config| { - config.fuzz.seed = Some(U256::from(119u32)); - config.invariant.runs = 1; - config.invariant.depth = 10; - }); - let results = runner.test_collect(&filter); - assert_multiple( - &results, - BTreeMap::from([( - "default/fuzz/invariant/common/InvariantSelectorsWeight.t.sol:InvariantSelectorsWeightTest", - vec![("invariant_selectors_weight()", true, None, None, None)], - )]), - ) -} - #[tokio::test(flavor = "multi_thread")] async fn test_no_reverts_in_counterexample() { let filter = @@ -1015,3 +997,80 @@ Ran 1 test suite [ELAPSED]: 1 tests passed, 0 failed, 0 skipped (1 total tests) "#]]); }); + +// Tests that selector hits are uniformly distributed +// +forgetest_init!(invariant_selectors_weight, |prj, cmd| { + prj.write_config(Config { + optimizer: Some(true), + invariant: { InvariantConfig { runs: 1, depth: 10, ..Default::default() } }, + ..Default::default() + }); + prj.add_source( + "InvariantHandlers.sol", + r#" +contract HandlerOne { + uint256 public hit1; + + function selector1() external { + hit1 += 1; + } +} + +contract HandlerTwo { + uint256 public hit2; + uint256 public hit3; + uint256 public hit4; + uint256 public hit5; + + function selector2() external { + hit2 += 1; + } + + function selector3() external { + hit3 += 1; + } + + function selector4() external { + hit4 += 1; + } + + function selector5() external { + hit5 += 1; + } +} + "#, + ) + .unwrap(); + + prj.add_test( + "InvariantSelectorsWeightTest.t.sol", + r#" +import {Test} from "forge-std/Test.sol"; +import "src/InvariantHandlers.sol"; + +contract InvariantSelectorsWeightTest is Test { + HandlerOne handlerOne; + HandlerTwo handlerTwo; + + function setUp() public { + handlerOne = new HandlerOne(); + handlerTwo = new HandlerTwo(); + } + + function afterInvariant() public { + assertEq(handlerOne.hit1(), 2); + assertEq(handlerTwo.hit2(), 2); + assertEq(handlerTwo.hit3(), 3); + assertEq(handlerTwo.hit4(), 1); + assertEq(handlerTwo.hit5(), 2); + } + + function invariant_selectors_weight() public view {} +} + "#, + ) + .unwrap(); + + cmd.args(["test", "--fuzz-seed", "119", "--mt", "invariant_selectors_weight"]).assert_success(); +}); diff --git a/testdata/cheats/Vm.sol b/testdata/cheats/Vm.sol index 260b5bb385600..d1a301e0f3df5 100644 --- a/testdata/cheats/Vm.sol +++ b/testdata/cheats/Vm.sol @@ -24,6 +24,7 @@ interface Vm { struct DebugStep { uint256[] stack; bytes memoryInput; uint8 opcode; uint64 depth; bool isOutOfGas; address contractAddr; } struct BroadcastTxSummary { bytes32 txHash; BroadcastTxType txType; address contractAddress; uint64 blockNumber; bool success; } struct SignedDelegation { uint8 v; bytes32 r; bytes32 s; uint64 nonce; address implementation; } + struct PotentialRevert { address reverter; bool partialMatch; bytes revertData; } function _expectCheatcodeRevert() external; function _expectCheatcodeRevert(bytes4 revertData) external; function _expectCheatcodeRevert(bytes calldata revertData) external; @@ -149,6 +150,8 @@ interface Vm { function assertTrue(bool condition, string calldata error) external pure; function assume(bool condition) external pure; function assumeNoRevert() external pure; + function assumeNoRevert(PotentialRevert calldata potentialRevert) external pure; + function assumeNoRevert(PotentialRevert[] calldata potentialReverts) external pure; function attachDelegation(SignedDelegation calldata signedDelegation) external; function blobBaseFee(uint256 newBlobBaseFee) external; function blobhashes(bytes32[] calldata hashes) external; diff --git a/testdata/default/cheats/AssumeNoRevert.t.sol b/testdata/default/cheats/AssumeNoRevert.t.sol new file mode 100644 index 0000000000000..ea6d2d9747bdd --- /dev/null +++ b/testdata/default/cheats/AssumeNoRevert.t.sol @@ -0,0 +1,153 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +pragma solidity ^0.8.18; + +import {DSTest as Test} from "ds-test/test.sol"; +import {Vm} from "cheats/Vm.sol"; + +contract ReverterB { + /// @notice has same error selectors as contract below to test the `reverter` param + error MyRevert(); + error SpecialRevertWithData(uint256 x); + + function revertIf2(uint256 x) public pure returns (bool) { + if (x == 2) { + revert MyRevert(); + } + return true; + } + + function revertWithData() public pure returns (bool) { + revert SpecialRevertWithData(2); + } +} + +contract Reverter { + error MyRevert(); + error RevertWithData(uint256 x); + error UnusedError(); + + ReverterB public immutable subReverter; + + constructor() { + subReverter = new ReverterB(); + } + + function myFunction() public pure returns (bool) { + revert MyRevert(); + } + + function revertIf2(uint256 value) public pure returns (bool) { + if (value == 2) { + revert MyRevert(); + } + return true; + } + + function revertWithDataIf2(uint256 value) public pure returns (bool) { + if (value == 2) { + revert RevertWithData(2); + } + return true; + } + + function twoPossibleReverts(uint256 x) public pure returns (bool) { + if (x == 2) { + revert MyRevert(); + } else if (x == 3) { + revert RevertWithData(3); + } + return true; + } +} + +contract ReverterTest is Test { + Reverter reverter; + Vm _vm = Vm(HEVM_ADDRESS); + + function setUp() public { + reverter = new Reverter(); + } + + /// @dev Test that `assumeNoRevert` anticipates and correctly rejects a specific error selector + function testAssumeSelector(uint256 x) public view { + _vm.assumeNoRevert( + Vm.PotentialRevert({ + revertData: abi.encodeWithSelector(Reverter.MyRevert.selector), + partialMatch: false, + reverter: address(0) + }) + ); + reverter.revertIf2(x); + } + + /// @dev Test that `assumeNoRevert` anticipates and correctly rejects a specific error selector and data + function testAssumeWithDataSingle(uint256 x) public view { + _vm.assumeNoRevert( + Vm.PotentialRevert({ + revertData: abi.encodeWithSelector(Reverter.RevertWithData.selector, 2), + partialMatch: false, + reverter: address(0) + }) + ); + reverter.revertWithDataIf2(x); + } + + /// @dev Test that `assumeNoRevert` anticipates and correctly rejects a specific error selector with any extra data (ie providing selector allows for arbitrary extra data) + function testAssumeWithDataPartial(uint256 x) public view { + _vm.assumeNoRevert( + Vm.PotentialRevert({ + revertData: abi.encodeWithSelector(Reverter.RevertWithData.selector), + partialMatch: true, + reverter: address(0) + }) + ); + reverter.revertWithDataIf2(x); + } + + /// @dev Test that `assumeNoRevert` assumptions are not cleared after a cheatcode call + function testAssumeNotClearedAfterCheatcodeCall(uint256 x) public { + _vm.assumeNoRevert( + Vm.PotentialRevert({ + revertData: abi.encodeWithSelector(Reverter.MyRevert.selector), + partialMatch: false, + reverter: address(0) + }) + ); + _vm.warp(block.timestamp + 1000); + reverter.revertIf2(x); + } + + /// @dev Test that `assumeNoRevert` correctly rejects two different error selectors + function testMultipleAssumesPasses(uint256 x) public view { + Vm.PotentialRevert[] memory revertData = new Vm.PotentialRevert[](2); + revertData[0] = Vm.PotentialRevert({ + revertData: abi.encodeWithSelector(Reverter.MyRevert.selector), + partialMatch: false, + reverter: address(reverter) + }); + revertData[1] = Vm.PotentialRevert({ + revertData: abi.encodeWithSelector(Reverter.RevertWithData.selector, 3), + partialMatch: false, + reverter: address(reverter) + }); + _vm.assumeNoRevert(revertData); + reverter.twoPossibleReverts(x); + } + + /// @dev Test that `assumeNoRevert` correctly interacts with itself when partially matching on the error selector + function testMultipleAssumes_Partial(uint256 x) public view { + Vm.PotentialRevert[] memory revertData = new Vm.PotentialRevert[](2); + revertData[0] = Vm.PotentialRevert({ + revertData: abi.encodeWithSelector(Reverter.RevertWithData.selector), + partialMatch: true, + reverter: address(reverter) + }); + revertData[1] = Vm.PotentialRevert({ + revertData: abi.encodeWithSelector(Reverter.MyRevert.selector), + partialMatch: false, + reverter: address(reverter) + }); + _vm.assumeNoRevert(revertData); + reverter.twoPossibleReverts(x); + } +} diff --git a/testdata/default/fuzz/invariant/common/InvariantSelectorsWeight.t.sol b/testdata/default/fuzz/invariant/common/InvariantSelectorsWeight.t.sol deleted file mode 100644 index aea46f41859b0..0000000000000 --- a/testdata/default/fuzz/invariant/common/InvariantSelectorsWeight.t.sol +++ /dev/null @@ -1,56 +0,0 @@ -// SPDX-License-Identifier: UNLICENSED -pragma solidity ^0.8.13; - -import "ds-test/test.sol"; - -contract HandlerOne { - uint256 public hit1; - - function selector1() external { - hit1 += 1; - } -} - -contract HandlerTwo { - uint256 public hit2; - uint256 public hit3; - uint256 public hit4; - uint256 public hit5; - - function selector2() external { - hit2 += 1; - } - - function selector3() external { - hit3 += 1; - } - - function selector4() external { - hit4 += 1; - } - - function selector5() external { - hit5 += 1; - } -} - -contract InvariantSelectorsWeightTest is DSTest { - HandlerOne handlerOne; - HandlerTwo handlerTwo; - - function setUp() public { - handlerOne = new HandlerOne(); - handlerTwo = new HandlerTwo(); - } - - function afterInvariant() public { - // selector hits uniformly distributed, see https://github.com/foundry-rs/foundry/issues/2986 - assertEq(handlerOne.hit1(), 2); - assertEq(handlerTwo.hit2(), 2); - assertEq(handlerTwo.hit3(), 3); - assertEq(handlerTwo.hit4(), 1); - assertEq(handlerTwo.hit5(), 2); - } - - function invariant_selectors_weight() public view {} -}