1use crate::{
2 executors::{Executor, RawCallResult},
3 inspectors::Fuzzer,
4};
5use alloy_primitives::{
6 Address, Bytes, FixedBytes, Selector, U256,
7 map::{AddressMap, HashMap},
8};
9use alloy_sol_types::{SolCall, sol};
10use eyre::{ContextCompat, Result, eyre};
11use foundry_common::contracts::{ContractsByAddress, ContractsByArtifact};
12use foundry_config::InvariantConfig;
13use foundry_evm_core::{
14 constants::{
15 CALLER, CHEATCODE_ADDRESS, DEFAULT_CREATE2_DEPLOYER, HARDHAT_CONSOLE_ADDRESS, MAGIC_ASSUME,
16 },
17 precompiles::PRECOMPILES,
18};
19use foundry_evm_fuzz::{
20 BasicTxDetails, FuzzCase, FuzzFixtures, FuzzedCases,
21 invariant::{
22 ArtifactFilters, FuzzRunIdentifiedContracts, InvariantContract, RandomCallGenerator,
23 SenderFilters, TargetedContract, TargetedContracts,
24 },
25 strategies::{EvmFuzzState, invariant_strat, override_call_strat},
26};
27use foundry_evm_traces::{CallTraceArena, SparsedTraceArena};
28use indicatif::ProgressBar;
29use parking_lot::RwLock;
30use proptest::{strategy::Strategy, test_runner::TestRunner};
31use result::{assert_after_invariant, assert_invariants, can_continue};
32use revm::state::Account;
33use std::{
34 collections::{HashMap as Map, btree_map::Entry},
35 sync::Arc,
36 time::{Instant, SystemTime, UNIX_EPOCH},
37};
38
39mod error;
40pub use error::{InvariantFailures, InvariantFuzzError};
41use foundry_evm_coverage::HitMaps;
42
43mod replay;
44pub use replay::{replay_error, replay_run};
45
46mod result;
47use foundry_common::{TestFunctionExt, sh_println};
48pub use result::InvariantFuzzTestResult;
49use serde::{Deserialize, Serialize};
50use serde_json::json;
51
52mod shrink;
53use crate::executors::{
54 DURATION_BETWEEN_METRICS_REPORT, EarlyExit, EvmError, FuzzTestTimer, corpus::CorpusManager,
55};
56pub use shrink::check_sequence;
57
58sol! {
59 interface IInvariantTest {
60 #[derive(Default)]
61 struct FuzzSelector {
62 address addr;
63 bytes4[] selectors;
64 }
65
66 #[derive(Default)]
67 struct FuzzArtifactSelector {
68 string artifact;
69 bytes4[] selectors;
70 }
71
72 #[derive(Default)]
73 struct FuzzInterface {
74 address addr;
75 string[] artifacts;
76 }
77
78 function afterInvariant() external;
79
80 #[derive(Default)]
81 function excludeArtifacts() public view returns (string[] memory excludedArtifacts);
82
83 #[derive(Default)]
84 function excludeContracts() public view returns (address[] memory excludedContracts);
85
86 #[derive(Default)]
87 function excludeSelectors() public view returns (FuzzSelector[] memory excludedSelectors);
88
89 #[derive(Default)]
90 function excludeSenders() public view returns (address[] memory excludedSenders);
91
92 #[derive(Default)]
93 function targetArtifacts() public view returns (string[] memory targetedArtifacts);
94
95 #[derive(Default)]
96 function targetArtifactSelectors() public view returns (FuzzArtifactSelector[] memory targetedArtifactSelectors);
97
98 #[derive(Default)]
99 function targetContracts() public view returns (address[] memory targetedContracts);
100
101 #[derive(Default)]
102 function targetSelectors() public view returns (FuzzSelector[] memory targetedSelectors);
103
104 #[derive(Default)]
105 function targetSenders() public view returns (address[] memory targetedSenders);
106
107 #[derive(Default)]
108 function targetInterfaces() public view returns (FuzzInterface[] memory targetedInterfaces);
109 }
110}
111
112#[derive(Default, Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
114pub struct InvariantMetrics {
115 pub calls: usize,
117 pub reverts: usize,
119 pub discards: usize,
121}
122
123struct InvariantTestData {
125 fuzz_cases: Vec<FuzzedCases>,
127 failures: InvariantFailures,
129 last_run_inputs: Vec<BasicTxDetails>,
131 gas_report_traces: Vec<Vec<CallTraceArena>>,
133 last_call_results: Option<RawCallResult>,
135 line_coverage: Option<HitMaps>,
137 metrics: Map<String, InvariantMetrics>,
139
140 branch_runner: TestRunner,
145}
146
147struct InvariantTest {
149 fuzz_state: EvmFuzzState,
151 targeted_contracts: FuzzRunIdentifiedContracts,
153 test_data: InvariantTestData,
155}
156
157impl InvariantTest {
158 fn new(
160 fuzz_state: EvmFuzzState,
161 targeted_contracts: FuzzRunIdentifiedContracts,
162 failures: InvariantFailures,
163 last_call_results: Option<RawCallResult>,
164 branch_runner: TestRunner,
165 ) -> Self {
166 let mut fuzz_cases = vec![];
167 if last_call_results.is_none() {
168 fuzz_cases.push(FuzzedCases::new(vec![]));
169 }
170 let test_data = InvariantTestData {
171 fuzz_cases,
172 failures,
173 last_run_inputs: vec![],
174 gas_report_traces: vec![],
175 last_call_results,
176 line_coverage: None,
177 metrics: Map::default(),
178 branch_runner,
179 };
180 Self { fuzz_state, targeted_contracts, test_data }
181 }
182
183 fn reverts(&self) -> usize {
185 self.test_data.failures.reverts
186 }
187
188 fn has_errors(&self) -> bool {
190 self.test_data.failures.error.is_some()
191 }
192
193 fn set_error(&mut self, error: InvariantFuzzError) {
195 self.test_data.failures.error = Some(error);
196 }
197
198 fn set_last_call_results(&mut self, call_result: Option<RawCallResult>) {
200 self.test_data.last_call_results = call_result;
201 }
202
203 fn set_last_run_inputs(&mut self, inputs: &Vec<BasicTxDetails>) {
205 self.test_data.last_run_inputs.clone_from(inputs);
206 }
207
208 fn merge_line_coverage(&mut self, new_coverage: Option<HitMaps>) {
210 HitMaps::merge_opt(&mut self.test_data.line_coverage, new_coverage);
211 }
212
213 fn record_metrics(&mut self, tx_details: &BasicTxDetails, reverted: bool, discarded: bool) {
217 if let Some(metric_key) =
218 self.targeted_contracts.targets.lock().fuzzed_metric_key(tx_details)
219 {
220 let test_metrics = &mut self.test_data.metrics;
221 let invariant_metrics = test_metrics.entry(metric_key).or_default();
222 invariant_metrics.calls += 1;
223 if discarded {
224 invariant_metrics.discards += 1;
225 } else if reverted {
226 invariant_metrics.reverts += 1;
227 }
228 }
229 }
230
231 fn end_run(&mut self, run: InvariantTestRun, gas_samples: usize) {
234 self.targeted_contracts.clear_created_contracts(run.created_contracts);
236
237 if self.test_data.gas_report_traces.len() < gas_samples {
238 self.test_data
239 .gas_report_traces
240 .push(run.run_traces.into_iter().map(|arena| arena.arena).collect());
241 }
242 self.test_data.fuzz_cases.push(FuzzedCases::new(run.fuzz_runs));
243
244 self.fuzz_state.revert();
246 }
247}
248
249struct InvariantTestRun {
251 inputs: Vec<BasicTxDetails>,
253 executor: Executor,
255 fuzz_runs: Vec<FuzzCase>,
257 created_contracts: Vec<Address>,
259 run_traces: Vec<SparsedTraceArena>,
261 depth: u32,
263 rejects: u32,
265 new_coverage: bool,
267}
268
269impl InvariantTestRun {
270 fn new(first_input: BasicTxDetails, executor: Executor, depth: usize) -> Self {
272 Self {
273 inputs: vec![first_input],
274 executor,
275 fuzz_runs: Vec::with_capacity(depth),
276 created_contracts: vec![],
277 run_traces: vec![],
278 depth: 0,
279 rejects: 0,
280 new_coverage: false,
281 }
282 }
283}
284
285pub struct InvariantExecutor<'a> {
292 pub executor: Executor,
293 runner: TestRunner,
295 config: InvariantConfig,
297 setup_contracts: &'a ContractsByAddress,
299 project_contracts: &'a ContractsByArtifact,
302 artifact_filters: ArtifactFilters,
304}
305
306impl<'a> InvariantExecutor<'a> {
307 pub fn new(
309 executor: Executor,
310 runner: TestRunner,
311 config: InvariantConfig,
312 setup_contracts: &'a ContractsByAddress,
313 project_contracts: &'a ContractsByArtifact,
314 ) -> Self {
315 Self {
316 executor,
317 runner,
318 config,
319 setup_contracts,
320 project_contracts,
321 artifact_filters: ArtifactFilters::default(),
322 }
323 }
324
325 pub fn config(self) -> InvariantConfig {
326 self.config
327 }
328
329 pub fn invariant_fuzz(
331 &mut self,
332 invariant_contract: InvariantContract<'_>,
333 fuzz_fixtures: &FuzzFixtures,
334 fuzz_state: EvmFuzzState,
335 progress: Option<&ProgressBar>,
336 early_exit: &EarlyExit,
337 ) -> Result<InvariantFuzzTestResult> {
338 if !invariant_contract.invariant_function.inputs.is_empty() {
340 return Err(eyre!("Invariant test function should have no inputs"));
341 }
342
343 let (mut invariant_test, mut corpus_manager) =
344 self.prepare_test(&invariant_contract, fuzz_fixtures, fuzz_state)?;
345
346 let mut runs = 0;
348 let timer = FuzzTestTimer::new(self.config.timeout);
349 let mut last_metrics_report = Instant::now();
350 let continue_campaign = |runs: u32| {
351 if early_exit.should_stop() {
352 return false;
353 }
354
355 if timer.is_enabled() { !timer.is_timed_out() } else { runs < self.config.runs }
356 };
357
358 let edge_coverage_enabled = self.config.corpus.collect_edge_coverage();
360
361 'stop: while continue_campaign(runs) {
362 let initial_seq = corpus_manager.new_inputs(
363 &mut invariant_test.test_data.branch_runner,
364 &invariant_test.fuzz_state,
365 &invariant_test.targeted_contracts,
366 )?;
367
368 let mut current_run = InvariantTestRun::new(
370 initial_seq[0].clone(),
371 self.executor.clone(),
373 self.config.depth as usize,
374 );
375
376 if self.config.fail_on_revert && invariant_test.reverts() > 0 {
378 return Err(eyre!("call reverted"));
379 }
380
381 while current_run.depth < self.config.depth {
382 if timer.is_timed_out() {
384 break 'stop;
389 }
390
391 let tx = current_run
392 .inputs
393 .last()
394 .ok_or_else(|| eyre!("no input generated to call fuzzed target."))?;
395
396 let mut call_result = execute_tx(&mut current_run.executor, tx)?;
399 let discarded = call_result.result.as_ref() == MAGIC_ASSUME;
400 if self.config.show_metrics {
401 invariant_test.record_metrics(tx, call_result.reverted, discarded);
402 }
403
404 invariant_test.merge_line_coverage(call_result.line_coverage.clone());
406 if corpus_manager.merge_edge_coverage(&mut call_result) {
408 current_run.new_coverage = true;
409 }
410
411 if discarded {
412 current_run.inputs.pop();
413 current_run.rejects += 1;
414 if current_run.rejects > self.config.max_assume_rejects {
415 invariant_test.set_error(InvariantFuzzError::MaxAssumeRejects(
416 self.config.max_assume_rejects,
417 ));
418 break 'stop;
419 }
420 } else {
421 current_run.executor.commit(&mut call_result);
423
424 let mut state_changeset = call_result.state_changeset.clone();
432 if !call_result.reverted {
433 collect_data(
434 &invariant_test,
435 &mut state_changeset,
436 tx,
437 &call_result,
438 self.config.depth,
439 );
440 }
441
442 if let Err(error) =
445 &invariant_test.targeted_contracts.collect_created_contracts(
446 &state_changeset,
447 self.project_contracts,
448 self.setup_contracts,
449 &self.artifact_filters,
450 &mut current_run.created_contracts,
451 )
452 {
453 warn!(target: "forge::test", "{error}");
454 }
455 current_run.fuzz_runs.push(FuzzCase {
456 calldata: tx.call_details.calldata.clone(),
457 gas: call_result.gas_used,
458 stipend: call_result.stipend,
459 });
460
461 let result = can_continue(
463 &invariant_contract,
464 &mut invariant_test,
465 &mut current_run,
466 &self.config,
467 call_result,
468 &state_changeset,
469 )
470 .map_err(|e| eyre!(e.to_string()))?;
471 if !result.can_continue || current_run.depth == self.config.depth - 1 {
472 invariant_test.set_last_run_inputs(¤t_run.inputs);
473 }
474 if !result.can_continue {
476 break 'stop;
477 }
478
479 invariant_test.set_last_call_results(result.call_result);
480 current_run.depth += 1;
481 }
482
483 current_run.inputs.push(corpus_manager.generate_next_input(
484 &mut invariant_test.test_data.branch_runner,
485 &initial_seq,
486 discarded,
487 current_run.depth as usize,
488 )?);
489 }
490
491 corpus_manager.process_inputs(¤t_run.inputs, current_run.new_coverage);
493
494 if invariant_contract.call_after_invariant && !invariant_test.has_errors() {
496 assert_after_invariant(
497 &invariant_contract,
498 &mut invariant_test,
499 ¤t_run,
500 &self.config,
501 )
502 .map_err(|_| eyre!("Failed to call afterInvariant"))?;
503 }
504
505 invariant_test.end_run(current_run, self.config.gas_report_samples as usize);
507 if let Some(progress) = progress {
508 progress.inc(1);
510 if edge_coverage_enabled {
512 progress.set_message(format!("{}", &corpus_manager.metrics));
513 }
514 } else if edge_coverage_enabled
515 && last_metrics_report.elapsed() > DURATION_BETWEEN_METRICS_REPORT
516 {
517 let metrics = json!({
519 "timestamp": SystemTime::now()
520 .duration_since(UNIX_EPOCH)?
521 .as_secs(),
522 "invariant": invariant_contract.invariant_function.name,
523 "metrics": &corpus_manager.metrics,
524 });
525 let _ = sh_println!("{}", serde_json::to_string(&metrics)?);
526 last_metrics_report = Instant::now();
527 }
528
529 runs += 1;
530 }
531
532 trace!(?fuzz_fixtures);
533 invariant_test.fuzz_state.log_stats();
534
535 let result = invariant_test.test_data;
536 Ok(InvariantFuzzTestResult {
537 error: result.failures.error,
538 cases: result.fuzz_cases,
539 reverts: result.failures.reverts,
540 last_run_inputs: result.last_run_inputs,
541 gas_report_traces: result.gas_report_traces,
542 line_coverage: result.line_coverage,
543 metrics: result.metrics,
544 failed_corpus_replays: corpus_manager.failed_replays(),
545 })
546 }
547
548 fn prepare_test(
552 &mut self,
553 invariant_contract: &InvariantContract<'_>,
554 fuzz_fixtures: &FuzzFixtures,
555 fuzz_state: EvmFuzzState,
556 ) -> Result<(InvariantTest, CorpusManager)> {
557 self.select_contract_artifacts(invariant_contract.address)?;
559 let (targeted_senders, targeted_contracts) =
560 self.select_contracts_and_senders(invariant_contract.address)?;
561
562 let strategy = invariant_strat(
564 fuzz_state.clone(),
565 targeted_senders,
566 targeted_contracts.clone(),
567 self.config.clone(),
568 fuzz_fixtures.clone(),
569 )
570 .no_shrink();
571
572 let mut call_generator = None;
575 if self.config.call_override {
576 let target_contract_ref = Arc::new(RwLock::new(Address::ZERO));
577
578 call_generator = Some(RandomCallGenerator::new(
579 invariant_contract.address,
580 self.runner.clone(),
581 override_call_strat(
582 fuzz_state.clone(),
583 targeted_contracts.clone(),
584 target_contract_ref.clone(),
585 fuzz_fixtures.clone(),
586 ),
587 target_contract_ref,
588 ));
589 }
590
591 let fuzz_state =
594 if targeted_contracts.targets.lock().iter().any(|(_, t)| t.storage_layout.is_some()) {
595 fuzz_state.with_mapping_slots(AddressMap::default())
596 } else {
597 fuzz_state
598 };
599
600 self.executor.inspector_mut().set_fuzzer(Fuzzer {
601 call_generator,
602 fuzz_state: fuzz_state.clone(),
603 collect: true,
604 });
605
606 let mut failures = InvariantFailures::new();
611 let last_call_results = assert_invariants(
612 invariant_contract,
613 &self.config,
614 &targeted_contracts,
615 &self.executor,
616 &[],
617 &mut failures,
618 )?;
619 if let Some(error) = failures.error {
620 return Err(eyre!(error.revert_reason().unwrap_or_default()));
621 }
622
623 let corpus_manager = CorpusManager::new(
624 self.config.corpus.clone(),
625 strategy.boxed(),
626 &self.executor,
627 None,
628 Some(&targeted_contracts),
629 )?;
630 let invariant_test = InvariantTest::new(
631 fuzz_state,
632 targeted_contracts,
633 failures,
634 last_call_results,
635 self.runner.clone(),
636 );
637
638 Ok((invariant_test, corpus_manager))
639 }
640
641 pub fn select_contract_artifacts(&mut self, invariant_address: Address) -> Result<()> {
651 let targeted_artifact_selectors = self
652 .executor
653 .call_sol_default(invariant_address, &IInvariantTest::targetArtifactSelectorsCall {});
654
655 for IInvariantTest::FuzzArtifactSelector { artifact, selectors } in
657 targeted_artifact_selectors
658 {
659 let identifier = self.validate_selected_contract(artifact, &selectors)?;
660 self.artifact_filters.targeted.entry(identifier).or_default().extend(selectors);
661 }
662
663 let targeted_artifacts = self
664 .executor
665 .call_sol_default(invariant_address, &IInvariantTest::targetArtifactsCall {});
666 let excluded_artifacts = self
667 .executor
668 .call_sol_default(invariant_address, &IInvariantTest::excludeArtifactsCall {});
669
670 for contract in excluded_artifacts {
672 let identifier = self.validate_selected_contract(contract, &[])?;
673
674 if !self.artifact_filters.excluded.contains(&identifier) {
675 self.artifact_filters.excluded.push(identifier);
676 }
677 }
678
679 for (artifact, contract) in self.project_contracts.iter() {
681 if contract
682 .abi
683 .functions()
684 .filter(|func| {
685 !matches!(
686 func.state_mutability,
687 alloy_json_abi::StateMutability::Pure
688 | alloy_json_abi::StateMutability::View
689 )
690 })
691 .count()
692 == 0
693 && !self.artifact_filters.excluded.contains(&artifact.identifier())
694 {
695 self.artifact_filters.excluded.push(artifact.identifier());
696 }
697 }
698
699 for contract in targeted_artifacts {
702 let identifier = self.validate_selected_contract(contract, &[])?;
703
704 if !self.artifact_filters.targeted.contains_key(&identifier)
705 && !self.artifact_filters.excluded.contains(&identifier)
706 {
707 self.artifact_filters.targeted.insert(identifier, vec![]);
708 }
709 }
710 Ok(())
711 }
712
713 fn validate_selected_contract(
716 &mut self,
717 contract: String,
718 selectors: &[FixedBytes<4>],
719 ) -> Result<String> {
720 if let Some((artifact, contract_data)) =
721 self.project_contracts.find_by_name_or_identifier(&contract)?
722 {
723 for selector in selectors {
725 contract_data
726 .abi
727 .functions()
728 .find(|func| func.selector().as_slice() == selector.as_slice())
729 .wrap_err(format!("{contract} does not have the selector {selector:?}"))?;
730 }
731
732 return Ok(artifact.identifier());
733 }
734 eyre::bail!(
735 "{contract} not found in the project. Allowed format: `contract_name` or `contract_path:contract_name`."
736 );
737 }
738
739 pub fn select_contracts_and_senders(
742 &self,
743 to: Address,
744 ) -> Result<(SenderFilters, FuzzRunIdentifiedContracts)> {
745 let targeted_senders =
746 self.executor.call_sol_default(to, &IInvariantTest::targetSendersCall {});
747 let mut excluded_senders =
748 self.executor.call_sol_default(to, &IInvariantTest::excludeSendersCall {});
749 excluded_senders.extend([
751 CHEATCODE_ADDRESS,
752 HARDHAT_CONSOLE_ADDRESS,
753 DEFAULT_CREATE2_DEPLOYER,
754 ]);
755 excluded_senders.extend(PRECOMPILES);
757 let sender_filters = SenderFilters::new(targeted_senders, excluded_senders);
758
759 let selected = self.executor.call_sol_default(to, &IInvariantTest::targetContractsCall {});
760 let excluded = self.executor.call_sol_default(to, &IInvariantTest::excludeContractsCall {});
761
762 let contracts = self
763 .setup_contracts
764 .iter()
765 .filter(|&(addr, (identifier, _))| {
766 if *addr == to && selected.contains(&to) {
768 return true;
769 }
770
771 *addr != to
772 && *addr != CHEATCODE_ADDRESS
773 && *addr != HARDHAT_CONSOLE_ADDRESS
774 && (selected.is_empty() || selected.contains(addr))
775 && (excluded.is_empty() || !excluded.contains(addr))
776 && self.artifact_filters.matches(identifier)
777 })
778 .map(|(addr, (identifier, abi))| {
779 (
780 *addr,
781 TargetedContract::new(identifier.clone(), abi.clone())
782 .with_project_contracts(self.project_contracts),
783 )
784 })
785 .collect();
786 let mut contracts = TargetedContracts { inner: contracts };
787
788 self.target_interfaces(to, &mut contracts)?;
789
790 self.select_selectors(to, &mut contracts)?;
791
792 if contracts.is_empty() {
794 eyre::bail!("No contracts to fuzz.");
795 }
796
797 Ok((sender_filters, FuzzRunIdentifiedContracts::new(contracts, selected.is_empty())))
798 }
799
800 pub fn target_interfaces(
805 &self,
806 invariant_address: Address,
807 targeted_contracts: &mut TargetedContracts,
808 ) -> Result<()> {
809 let interfaces = self
810 .executor
811 .call_sol_default(invariant_address, &IInvariantTest::targetInterfacesCall {});
812
813 let mut combined = TargetedContracts::new();
819
820 for IInvariantTest::FuzzInterface { addr, artifacts } in &interfaces {
823 for identifier in artifacts {
825 if let Some((_, contract_data)) =
827 self.project_contracts.iter().find(|(artifact, _)| {
828 &artifact.name == identifier || &artifact.identifier() == identifier
829 })
830 {
831 let abi = &contract_data.abi;
832 combined
833 .entry(*addr)
835 .and_modify(|entry| {
837 entry.abi.functions.extend(abi.functions.clone());
839 })
840 .or_insert_with(|| {
842 let mut contract =
843 TargetedContract::new(identifier.to_string(), abi.clone());
844 contract.storage_layout =
845 contract_data.storage_layout.as_ref().map(Arc::clone);
846 contract
847 });
848 }
849 }
850 }
851
852 targeted_contracts.extend(combined.inner);
853
854 Ok(())
855 }
856
857 pub fn select_selectors(
860 &self,
861 address: Address,
862 targeted_contracts: &mut TargetedContracts,
863 ) -> Result<()> {
864 for (address, (identifier, _)) in self.setup_contracts {
865 if let Some(selectors) = self.artifact_filters.targeted.get(identifier) {
866 self.add_address_with_functions(*address, selectors, false, targeted_contracts)?;
867 }
868 }
869
870 let mut target_test_selectors = vec![];
871 let mut excluded_test_selectors = vec![];
872
873 let selectors =
875 self.executor.call_sol_default(address, &IInvariantTest::targetSelectorsCall {});
876 for IInvariantTest::FuzzSelector { addr, selectors } in selectors {
877 if addr == address {
878 target_test_selectors = selectors.clone();
879 }
880 self.add_address_with_functions(addr, &selectors, false, targeted_contracts)?;
881 }
882
883 let excluded_selectors =
885 self.executor.call_sol_default(address, &IInvariantTest::excludeSelectorsCall {});
886 for IInvariantTest::FuzzSelector { addr, selectors } in excluded_selectors {
887 if addr == address {
888 excluded_test_selectors = selectors.clone();
891 }
892 self.add_address_with_functions(addr, &selectors, true, targeted_contracts)?;
893 }
894
895 if target_test_selectors.is_empty()
896 && let Some(target) = targeted_contracts.get(&address)
897 {
898 let selectors: Vec<_> = target
902 .abi
903 .functions()
904 .filter_map(|func| {
905 if matches!(
906 func.state_mutability,
907 alloy_json_abi::StateMutability::Pure
908 | alloy_json_abi::StateMutability::View
909 ) || func.is_reserved()
910 || excluded_test_selectors.contains(&func.selector())
911 {
912 None
913 } else {
914 Some(func.selector())
915 }
916 })
917 .collect();
918 self.add_address_with_functions(address, &selectors, false, targeted_contracts)?;
919 }
920
921 Ok(())
922 }
923
924 fn add_address_with_functions(
926 &self,
927 address: Address,
928 selectors: &[Selector],
929 should_exclude: bool,
930 targeted_contracts: &mut TargetedContracts,
931 ) -> eyre::Result<()> {
932 if selectors.is_empty() {
934 return Ok(());
935 }
936
937 let contract = match targeted_contracts.entry(address) {
938 Entry::Occupied(entry) => entry.into_mut(),
939 Entry::Vacant(entry) => {
940 let (identifier, abi) = self.setup_contracts.get(&address).ok_or_else(|| {
941 eyre::eyre!(
942 "[{}] address does not have an associated contract: {}",
943 if should_exclude { "excludeSelectors" } else { "targetSelectors" },
944 address
945 )
946 })?;
947 entry.insert(
948 TargetedContract::new(identifier.clone(), abi.clone())
949 .with_project_contracts(self.project_contracts),
950 )
951 }
952 };
953 contract.add_selectors(selectors.iter().copied(), should_exclude)?;
954 Ok(())
955 }
956}
957
958fn collect_data(
962 invariant_test: &InvariantTest,
963 state_changeset: &mut HashMap<Address, Account>,
964 tx: &BasicTxDetails,
965 call_result: &RawCallResult,
966 run_depth: u32,
967) {
968 let mut has_code = false;
970 if let Some(Some(code)) =
971 state_changeset.get(&tx.sender).map(|account| account.info.code.as_ref())
972 {
973 has_code = !code.is_empty();
974 }
975
976 let mut sender_changeset = None;
978 if !has_code {
979 sender_changeset = state_changeset.remove(&tx.sender);
980 }
981
982 invariant_test.fuzz_state.collect_values_from_call(
984 &invariant_test.targeted_contracts,
985 tx,
986 &call_result.result,
987 &call_result.logs,
988 &*state_changeset,
989 run_depth,
990 );
991
992 if let Some(changed) = sender_changeset {
994 state_changeset.insert(tx.sender, changed);
995 }
996}
997
998pub(crate) fn call_after_invariant_function(
1002 executor: &Executor,
1003 to: Address,
1004) -> Result<(RawCallResult, bool), EvmError> {
1005 let calldata = Bytes::from_static(&IInvariantTest::afterInvariantCall::SELECTOR);
1006 let mut call_result = executor.call_raw(CALLER, to, calldata, U256::ZERO)?;
1007 let success = executor.is_raw_call_mut_success(to, &mut call_result, false);
1008 Ok((call_result, success))
1009}
1010
1011pub(crate) fn call_invariant_function(
1013 executor: &Executor,
1014 address: Address,
1015 calldata: Bytes,
1016) -> Result<(RawCallResult, bool)> {
1017 let mut call_result = executor.call_raw(CALLER, address, calldata, U256::ZERO)?;
1018 let success = executor.is_raw_call_mut_success(address, &mut call_result, false);
1019 Ok((call_result, success))
1020}
1021
1022pub(crate) fn execute_tx(executor: &mut Executor, tx: &BasicTxDetails) -> Result<RawCallResult> {
1025 if let Some(warp) = tx.warp {
1027 executor.env_mut().evm_env.block_env.timestamp += warp;
1028 }
1029 if let Some(roll) = tx.roll {
1030 executor.env_mut().evm_env.block_env.number += roll;
1031 }
1032
1033 let mut call_result = executor
1035 .call_raw(tx.sender, tx.call_details.target, tx.call_details.calldata.clone(), U256::ZERO)
1036 .map_err(|e| eyre!(format!("Could not make raw evm call: {e}")))?;
1037
1038 if let Some(warp) = tx.warp {
1040 call_result.env.evm_env.block_env.timestamp += warp;
1041 }
1042 if let Some(roll) = tx.roll {
1043 call_result.env.evm_env.block_env.number += roll;
1044 }
1045 Ok(call_result)
1046}