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