1use crate::{
2 executors::{
3 DURATION_BETWEEN_METRICS_REPORT, EarlyExit, EvmError, Executor, FuzzTestTimer,
4 RawCallResult, corpus::WorkerCorpus,
5 },
6 inspectors::Fuzzer,
7};
8use alloy_primitives::{
9 Address, Bytes, FixedBytes, Selector, U256,
10 map::{AddressMap, HashMap},
11};
12use alloy_sol_types::{SolCall, sol};
13use eyre::{ContextCompat, Result, eyre};
14use foundry_common::{
15 TestFunctionExt,
16 contracts::{ContractsByAddress, ContractsByArtifact},
17 sh_println,
18};
19use foundry_config::InvariantConfig;
20use foundry_evm_core::{
21 constants::{
22 CALLER, CHEATCODE_ADDRESS, DEFAULT_CREATE2_DEPLOYER, HARDHAT_CONSOLE_ADDRESS, MAGIC_ASSUME,
23 },
24 precompiles::PRECOMPILES,
25};
26use foundry_evm_fuzz::{
27 BasicTxDetails, FuzzCase, FuzzFixtures, FuzzedCases,
28 invariant::{
29 ArtifactFilters, FuzzRunIdentifiedContracts, InvariantContract, RandomCallGenerator,
30 SenderFilters, TargetedContract, TargetedContracts,
31 },
32 strategies::{EvmFuzzState, invariant_strat, override_call_strat},
33};
34use foundry_evm_traces::{CallTraceArena, SparsedTraceArena};
35use indicatif::ProgressBar;
36use parking_lot::RwLock;
37use proptest::{strategy::Strategy, test_runner::TestRunner};
38use result::{assert_after_invariant, assert_invariants, can_continue};
39use revm::state::Account;
40use serde::{Deserialize, Serialize};
41use serde_json::json;
42use std::{
43 collections::{HashMap as Map, btree_map::Entry},
44 sync::Arc,
45 time::{Instant, SystemTime, UNIX_EPOCH},
46};
47
48mod error;
49pub use error::{InvariantFailures, InvariantFuzzError};
50use foundry_evm_coverage::HitMaps;
51
52mod replay;
53pub use replay::{replay_error, replay_run};
54
55mod result;
56pub use result::InvariantFuzzTestResult;
57
58mod shrink;
59pub use shrink::check_sequence;
60
61sol! {
62 interface IInvariantTest {
63 #[derive(Default)]
64 struct FuzzSelector {
65 address addr;
66 bytes4[] selectors;
67 }
68
69 #[derive(Default)]
70 struct FuzzArtifactSelector {
71 string artifact;
72 bytes4[] selectors;
73 }
74
75 #[derive(Default)]
76 struct FuzzInterface {
77 address addr;
78 string[] artifacts;
79 }
80
81 function afterInvariant() external;
82
83 #[derive(Default)]
84 function excludeArtifacts() public view returns (string[] memory excludedArtifacts);
85
86 #[derive(Default)]
87 function excludeContracts() public view returns (address[] memory excludedContracts);
88
89 #[derive(Default)]
90 function excludeSelectors() public view returns (FuzzSelector[] memory excludedSelectors);
91
92 #[derive(Default)]
93 function excludeSenders() public view returns (address[] memory excludedSenders);
94
95 #[derive(Default)]
96 function targetArtifacts() public view returns (string[] memory targetedArtifacts);
97
98 #[derive(Default)]
99 function targetArtifactSelectors() public view returns (FuzzArtifactSelector[] memory targetedArtifactSelectors);
100
101 #[derive(Default)]
102 function targetContracts() public view returns (address[] memory targetedContracts);
103
104 #[derive(Default)]
105 function targetSelectors() public view returns (FuzzSelector[] memory targetedSelectors);
106
107 #[derive(Default)]
108 function targetSenders() public view returns (address[] memory targetedSenders);
109
110 #[derive(Default)]
111 function targetInterfaces() public view returns (FuzzInterface[] memory targetedInterfaces);
112 }
113}
114
115#[derive(Default, Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
117pub struct InvariantMetrics {
118 pub calls: usize,
120 pub reverts: usize,
122 pub discards: usize,
124}
125
126struct InvariantTestData {
128 fuzz_cases: Vec<FuzzedCases>,
130 failures: InvariantFailures,
132 last_run_inputs: Vec<BasicTxDetails>,
134 gas_report_traces: Vec<Vec<CallTraceArena>>,
136 last_call_results: Option<RawCallResult>,
138 line_coverage: Option<HitMaps>,
140 metrics: Map<String, InvariantMetrics>,
142
143 branch_runner: TestRunner,
148}
149
150struct InvariantTest {
152 fuzz_state: EvmFuzzState,
154 targeted_contracts: FuzzRunIdentifiedContracts,
156 test_data: InvariantTestData,
158}
159
160impl InvariantTest {
161 fn new(
163 fuzz_state: EvmFuzzState,
164 targeted_contracts: FuzzRunIdentifiedContracts,
165 failures: InvariantFailures,
166 last_call_results: Option<RawCallResult>,
167 branch_runner: TestRunner,
168 ) -> Self {
169 let mut fuzz_cases = vec![];
170 if last_call_results.is_none() {
171 fuzz_cases.push(FuzzedCases::new(vec![]));
172 }
173 let test_data = InvariantTestData {
174 fuzz_cases,
175 failures,
176 last_run_inputs: vec![],
177 gas_report_traces: vec![],
178 last_call_results,
179 line_coverage: None,
180 metrics: Map::default(),
181 branch_runner,
182 };
183 Self { fuzz_state, targeted_contracts, test_data }
184 }
185
186 fn reverts(&self) -> usize {
188 self.test_data.failures.reverts
189 }
190
191 fn has_errors(&self) -> bool {
193 self.test_data.failures.error.is_some()
194 }
195
196 fn set_error(&mut self, error: InvariantFuzzError) {
198 self.test_data.failures.error = Some(error);
199 }
200
201 fn set_last_call_results(&mut self, call_result: Option<RawCallResult>) {
203 self.test_data.last_call_results = call_result;
204 }
205
206 fn set_last_run_inputs(&mut self, inputs: &Vec<BasicTxDetails>) {
208 self.test_data.last_run_inputs.clone_from(inputs);
209 }
210
211 fn merge_line_coverage(&mut self, new_coverage: Option<HitMaps>) {
213 HitMaps::merge_opt(&mut self.test_data.line_coverage, new_coverage);
214 }
215
216 fn record_metrics(&mut self, tx_details: &BasicTxDetails, reverted: bool, discarded: bool) {
220 if let Some(metric_key) =
221 self.targeted_contracts.targets.lock().fuzzed_metric_key(tx_details)
222 {
223 let test_metrics = &mut self.test_data.metrics;
224 let invariant_metrics = test_metrics.entry(metric_key).or_default();
225 invariant_metrics.calls += 1;
226 if discarded {
227 invariant_metrics.discards += 1;
228 } else if reverted {
229 invariant_metrics.reverts += 1;
230 }
231 }
232 }
233
234 fn end_run(&mut self, run: InvariantTestRun, gas_samples: usize) {
237 self.targeted_contracts.clear_created_contracts(run.created_contracts);
239
240 if self.test_data.gas_report_traces.len() < gas_samples {
241 self.test_data
242 .gas_report_traces
243 .push(run.run_traces.into_iter().map(|arena| arena.arena).collect());
244 }
245 self.test_data.fuzz_cases.push(FuzzedCases::new(run.fuzz_runs));
246
247 self.fuzz_state.revert();
249 }
250}
251
252struct InvariantTestRun {
254 inputs: Vec<BasicTxDetails>,
256 executor: Executor,
258 fuzz_runs: Vec<FuzzCase>,
260 created_contracts: Vec<Address>,
262 run_traces: Vec<SparsedTraceArena>,
264 depth: u32,
266 rejects: u32,
268 new_coverage: bool,
270}
271
272impl InvariantTestRun {
273 fn new(first_input: BasicTxDetails, executor: Executor, depth: usize) -> Self {
275 Self {
276 inputs: vec![first_input],
277 executor,
278 fuzz_runs: Vec::with_capacity(depth),
279 created_contracts: vec![],
280 run_traces: vec![],
281 depth: 0,
282 rejects: 0,
283 new_coverage: false,
284 }
285 }
286}
287
288pub struct InvariantExecutor<'a> {
295 pub executor: Executor,
296 runner: TestRunner,
298 config: InvariantConfig,
300 setup_contracts: &'a ContractsByAddress,
302 project_contracts: &'a ContractsByArtifact,
305 artifact_filters: ArtifactFilters,
307}
308
309impl<'a> InvariantExecutor<'a> {
310 pub fn new(
312 executor: Executor,
313 runner: TestRunner,
314 config: InvariantConfig,
315 setup_contracts: &'a ContractsByAddress,
316 project_contracts: &'a ContractsByArtifact,
317 ) -> Self {
318 Self {
319 executor,
320 runner,
321 config,
322 setup_contracts,
323 project_contracts,
324 artifact_filters: ArtifactFilters::default(),
325 }
326 }
327
328 pub fn config(self) -> InvariantConfig {
329 self.config
330 }
331
332 pub fn invariant_fuzz(
334 &mut self,
335 invariant_contract: InvariantContract<'_>,
336 fuzz_fixtures: &FuzzFixtures,
337 fuzz_state: EvmFuzzState,
338 progress: Option<&ProgressBar>,
339 early_exit: &EarlyExit,
340 ) -> Result<InvariantFuzzTestResult> {
341 if !invariant_contract.invariant_function.inputs.is_empty() {
343 return Err(eyre!("Invariant test function should have no inputs"));
344 }
345
346 let (mut invariant_test, mut corpus_manager) =
347 self.prepare_test(&invariant_contract, fuzz_fixtures, fuzz_state)?;
348
349 let mut runs = 0;
351 let timer = FuzzTestTimer::new(self.config.timeout);
352 let mut last_metrics_report = Instant::now();
353 let continue_campaign = |runs: u32| {
354 if early_exit.should_stop() {
355 return false;
356 }
357
358 if timer.is_enabled() { !timer.is_timed_out() } else { runs < self.config.runs }
359 };
360
361 let edge_coverage_enabled = self.config.corpus.collect_edge_coverage();
363
364 'stop: while continue_campaign(runs) {
365 let initial_seq = corpus_manager.new_inputs(
366 &mut invariant_test.test_data.branch_runner,
367 &invariant_test.fuzz_state,
368 &invariant_test.targeted_contracts,
369 )?;
370
371 let mut current_run = InvariantTestRun::new(
373 initial_seq[0].clone(),
374 self.executor.clone(),
376 self.config.depth as usize,
377 );
378
379 if self.config.fail_on_revert && invariant_test.reverts() > 0 {
381 return Err(eyre!("call reverted"));
382 }
383
384 while current_run.depth < self.config.depth {
385 if timer.is_timed_out() {
387 break 'stop;
392 }
393
394 let tx = current_run
395 .inputs
396 .last()
397 .ok_or_else(|| eyre!("no input generated to call fuzzed target."))?;
398
399 let mut call_result = execute_tx(&mut current_run.executor, tx)?;
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 is_last_call = current_run.depth == self.config.depth - 1;
470 let should_check_invariant = if self.config.check_interval == 0 {
471 is_last_call
472 } else {
473 self.config.check_interval == 1
474 || (current_run.depth + 1).is_multiple_of(self.config.check_interval)
475 || is_last_call
476 };
477
478 let result = if should_check_invariant {
479 can_continue(
480 &invariant_contract,
481 &mut invariant_test,
482 &mut current_run,
483 &self.config,
484 call_result,
485 &state_changeset,
486 )
487 .map_err(|e| eyre!(e.to_string()))?
488 } else {
489 if call_result.reverted {
491 invariant_test.test_data.failures.reverts += 1;
492 if self.config.fail_on_revert {
493 let case_data = error::FailedInvariantCaseData::new(
494 &invariant_contract,
495 &self.config,
496 &invariant_test.targeted_contracts,
497 ¤t_run.inputs,
498 call_result,
499 &[],
500 );
501 invariant_test.test_data.failures.revert_reason =
502 Some(case_data.revert_reason.clone());
503 invariant_test.test_data.failures.error =
504 Some(InvariantFuzzError::Revert(case_data));
505 result::RichInvariantResults::new(false, None)
506 } else {
507 current_run.inputs.pop();
508 result::RichInvariantResults::new(true, None)
509 }
510 } else {
511 result::RichInvariantResults::new(true, None)
512 }
513 };
514
515 if !result.can_continue || current_run.depth == self.config.depth - 1 {
516 invariant_test.set_last_run_inputs(¤t_run.inputs);
517 }
518 if !result.can_continue {
520 break 'stop;
521 }
522
523 invariant_test.set_last_call_results(result.call_result);
524 current_run.depth += 1;
525 }
526
527 current_run.inputs.push(corpus_manager.generate_next_input(
528 &mut invariant_test.test_data.branch_runner,
529 &initial_seq,
530 discarded,
531 current_run.depth as usize,
532 )?);
533 }
534
535 corpus_manager.process_inputs(¤t_run.inputs, current_run.new_coverage);
537
538 if invariant_contract.call_after_invariant && !invariant_test.has_errors() {
540 assert_after_invariant(
541 &invariant_contract,
542 &mut invariant_test,
543 ¤t_run,
544 &self.config,
545 )
546 .map_err(|_| eyre!("Failed to call afterInvariant"))?;
547 }
548
549 invariant_test.end_run(current_run, self.config.gas_report_samples as usize);
551 if let Some(progress) = progress {
552 progress.inc(1);
554 if edge_coverage_enabled {
556 progress.set_message(format!("{}", &corpus_manager.metrics));
557 }
558 } else if edge_coverage_enabled
559 && last_metrics_report.elapsed() > DURATION_BETWEEN_METRICS_REPORT
560 {
561 let metrics = json!({
563 "timestamp": SystemTime::now()
564 .duration_since(UNIX_EPOCH)?
565 .as_secs(),
566 "invariant": invariant_contract.invariant_function.name,
567 "metrics": &corpus_manager.metrics,
568 });
569 let _ = sh_println!("{}", serde_json::to_string(&metrics)?);
570 last_metrics_report = Instant::now();
571 }
572
573 runs += 1;
574 }
575
576 trace!(?fuzz_fixtures);
577 invariant_test.fuzz_state.log_stats();
578
579 let result = invariant_test.test_data;
580 Ok(InvariantFuzzTestResult {
581 error: result.failures.error,
582 cases: result.fuzz_cases,
583 reverts: result.failures.reverts,
584 last_run_inputs: result.last_run_inputs,
585 gas_report_traces: result.gas_report_traces,
586 line_coverage: result.line_coverage,
587 metrics: result.metrics,
588 failed_corpus_replays: corpus_manager.failed_replays,
589 })
590 }
591
592 fn prepare_test(
596 &mut self,
597 invariant_contract: &InvariantContract<'_>,
598 fuzz_fixtures: &FuzzFixtures,
599 fuzz_state: EvmFuzzState,
600 ) -> Result<(InvariantTest, WorkerCorpus)> {
601 self.select_contract_artifacts(invariant_contract.address)?;
603 let (targeted_senders, targeted_contracts) =
604 self.select_contracts_and_senders(invariant_contract.address)?;
605
606 let strategy = invariant_strat(
608 fuzz_state.clone(),
609 targeted_senders,
610 targeted_contracts.clone(),
611 self.config.clone(),
612 fuzz_fixtures.clone(),
613 )
614 .no_shrink();
615
616 let fuzz_state =
619 if targeted_contracts.targets.lock().iter().any(|(_, t)| t.storage_layout.is_some()) {
620 fuzz_state.with_mapping_slots(AddressMap::default())
621 } else {
622 fuzz_state
623 };
624
625 self.executor.inspector_mut().set_fuzzer(Fuzzer {
629 call_generator: None,
630 fuzz_state: fuzz_state.clone(),
631 collect: true,
632 });
633
634 let mut failures = InvariantFailures::new();
639 let last_call_results = assert_invariants(
640 invariant_contract,
641 &self.config,
642 &targeted_contracts,
643 &self.executor,
644 &[],
645 &mut failures,
646 )?;
647 if let Some(error) = failures.error {
648 return Err(eyre!(error.revert_reason().unwrap_or_default()));
649 }
650
651 if self.config.call_override {
655 let target_contract_ref = Arc::new(RwLock::new(Address::ZERO));
656
657 let handler_addresses: std::collections::HashSet<Address> =
660 targeted_contracts.targets.lock().keys().copied().collect();
661
662 let call_generator = RandomCallGenerator::new(
663 invariant_contract.address,
664 handler_addresses,
665 self.runner.clone(),
666 override_call_strat(
667 fuzz_state.clone(),
668 targeted_contracts.clone(),
669 target_contract_ref.clone(),
670 fuzz_fixtures.clone(),
671 ),
672 target_contract_ref,
673 );
674
675 if let Some(fuzzer) = self.executor.inspector_mut().fuzzer.as_mut() {
676 fuzzer.call_generator = Some(call_generator);
677 }
678 }
679
680 let worker = WorkerCorpus::new(
681 0,
682 self.config.corpus.clone(),
683 strategy.boxed(),
684 Some(&self.executor),
685 None,
686 Some(&targeted_contracts),
687 )?;
688
689 let invariant_test = InvariantTest::new(
690 fuzz_state,
691 targeted_contracts,
692 failures,
693 last_call_results,
694 self.runner.clone(),
695 );
696
697 Ok((invariant_test, worker))
698 }
699
700 pub fn select_contract_artifacts(&mut self, invariant_address: Address) -> Result<()> {
710 let targeted_artifact_selectors = self
711 .executor
712 .call_sol_default(invariant_address, &IInvariantTest::targetArtifactSelectorsCall {});
713
714 for IInvariantTest::FuzzArtifactSelector { artifact, selectors } in
716 targeted_artifact_selectors
717 {
718 let identifier = self.validate_selected_contract(artifact, &selectors)?;
719 self.artifact_filters.targeted.entry(identifier).or_default().extend(selectors);
720 }
721
722 let targeted_artifacts = self
723 .executor
724 .call_sol_default(invariant_address, &IInvariantTest::targetArtifactsCall {});
725 let excluded_artifacts = self
726 .executor
727 .call_sol_default(invariant_address, &IInvariantTest::excludeArtifactsCall {});
728
729 for contract in excluded_artifacts {
731 let identifier = self.validate_selected_contract(contract, &[])?;
732
733 if !self.artifact_filters.excluded.contains(&identifier) {
734 self.artifact_filters.excluded.push(identifier);
735 }
736 }
737
738 for (artifact, contract) in self.project_contracts.iter() {
740 if contract
741 .abi
742 .functions()
743 .filter(|func| {
744 !matches!(
745 func.state_mutability,
746 alloy_json_abi::StateMutability::Pure
747 | alloy_json_abi::StateMutability::View
748 )
749 })
750 .count()
751 == 0
752 && !self.artifact_filters.excluded.contains(&artifact.identifier())
753 {
754 self.artifact_filters.excluded.push(artifact.identifier());
755 }
756 }
757
758 for contract in targeted_artifacts {
761 let identifier = self.validate_selected_contract(contract, &[])?;
762
763 if !self.artifact_filters.targeted.contains_key(&identifier)
764 && !self.artifact_filters.excluded.contains(&identifier)
765 {
766 self.artifact_filters.targeted.insert(identifier, vec![]);
767 }
768 }
769 Ok(())
770 }
771
772 fn validate_selected_contract(
775 &mut self,
776 contract: String,
777 selectors: &[FixedBytes<4>],
778 ) -> Result<String> {
779 if let Some((artifact, contract_data)) =
780 self.project_contracts.find_by_name_or_identifier(&contract)?
781 {
782 for selector in selectors {
784 contract_data
785 .abi
786 .functions()
787 .find(|func| func.selector().as_slice() == selector.as_slice())
788 .wrap_err(format!("{contract} does not have the selector {selector:?}"))?;
789 }
790
791 return Ok(artifact.identifier());
792 }
793 eyre::bail!(
794 "{contract} not found in the project. Allowed format: `contract_name` or `contract_path:contract_name`."
795 );
796 }
797
798 pub fn select_contracts_and_senders(
801 &self,
802 to: Address,
803 ) -> Result<(SenderFilters, FuzzRunIdentifiedContracts)> {
804 let targeted_senders =
805 self.executor.call_sol_default(to, &IInvariantTest::targetSendersCall {});
806 let mut excluded_senders =
807 self.executor.call_sol_default(to, &IInvariantTest::excludeSendersCall {});
808 excluded_senders.extend([
810 CHEATCODE_ADDRESS,
811 HARDHAT_CONSOLE_ADDRESS,
812 DEFAULT_CREATE2_DEPLOYER,
813 ]);
814 excluded_senders.extend(PRECOMPILES);
816 let sender_filters = SenderFilters::new(targeted_senders, excluded_senders);
817
818 let selected = self.executor.call_sol_default(to, &IInvariantTest::targetContractsCall {});
819 let excluded = self.executor.call_sol_default(to, &IInvariantTest::excludeContractsCall {});
820
821 let contracts = self
822 .setup_contracts
823 .iter()
824 .filter(|&(addr, (identifier, _))| {
825 if *addr == to && selected.contains(&to) {
827 return true;
828 }
829
830 *addr != to
831 && *addr != CHEATCODE_ADDRESS
832 && *addr != HARDHAT_CONSOLE_ADDRESS
833 && (selected.is_empty() || selected.contains(addr))
834 && (excluded.is_empty() || !excluded.contains(addr))
835 && self.artifact_filters.matches(identifier)
836 })
837 .map(|(addr, (identifier, abi))| {
838 (
839 *addr,
840 TargetedContract::new(identifier.clone(), abi.clone())
841 .with_project_contracts(self.project_contracts),
842 )
843 })
844 .collect();
845 let mut contracts = TargetedContracts { inner: contracts };
846
847 self.target_interfaces(to, &mut contracts)?;
848
849 self.select_selectors(to, &mut contracts)?;
850
851 if contracts.is_empty() {
853 eyre::bail!("No contracts to fuzz.");
854 }
855
856 Ok((sender_filters, FuzzRunIdentifiedContracts::new(contracts, selected.is_empty())))
857 }
858
859 pub fn target_interfaces(
864 &self,
865 invariant_address: Address,
866 targeted_contracts: &mut TargetedContracts,
867 ) -> Result<()> {
868 let interfaces = self
869 .executor
870 .call_sol_default(invariant_address, &IInvariantTest::targetInterfacesCall {});
871
872 let mut combined = TargetedContracts::new();
878
879 for IInvariantTest::FuzzInterface { addr, artifacts } in &interfaces {
882 for identifier in artifacts {
884 if let Some((_, contract_data)) =
886 self.project_contracts.iter().find(|(artifact, _)| {
887 &artifact.name == identifier || &artifact.identifier() == identifier
888 })
889 {
890 let abi = &contract_data.abi;
891 combined
892 .entry(*addr)
894 .and_modify(|entry| {
896 entry.abi.functions.extend(abi.functions.clone());
898 })
899 .or_insert_with(|| {
901 let mut contract =
902 TargetedContract::new(identifier.to_string(), abi.clone());
903 contract.storage_layout =
904 contract_data.storage_layout.as_ref().map(Arc::clone);
905 contract
906 });
907 }
908 }
909 }
910
911 targeted_contracts.extend(combined.inner);
912
913 Ok(())
914 }
915
916 pub fn select_selectors(
919 &self,
920 address: Address,
921 targeted_contracts: &mut TargetedContracts,
922 ) -> Result<()> {
923 for (address, (identifier, _)) in self.setup_contracts {
924 if let Some(selectors) = self.artifact_filters.targeted.get(identifier) {
925 self.add_address_with_functions(*address, selectors, false, targeted_contracts)?;
926 }
927 }
928
929 let mut target_test_selectors = vec![];
930 let mut excluded_test_selectors = vec![];
931
932 let selectors =
934 self.executor.call_sol_default(address, &IInvariantTest::targetSelectorsCall {});
935 for IInvariantTest::FuzzSelector { addr, selectors } in selectors {
936 if addr == address {
937 target_test_selectors = selectors.clone();
938 }
939 self.add_address_with_functions(addr, &selectors, false, targeted_contracts)?;
940 }
941
942 let excluded_selectors =
944 self.executor.call_sol_default(address, &IInvariantTest::excludeSelectorsCall {});
945 for IInvariantTest::FuzzSelector { addr, selectors } in excluded_selectors {
946 if addr == address {
947 excluded_test_selectors = selectors.clone();
950 }
951 self.add_address_with_functions(addr, &selectors, true, targeted_contracts)?;
952 }
953
954 if target_test_selectors.is_empty()
955 && let Some(target) = targeted_contracts.get(&address)
956 {
957 let selectors: Vec<_> = target
961 .abi
962 .functions()
963 .filter_map(|func| {
964 if matches!(
965 func.state_mutability,
966 alloy_json_abi::StateMutability::Pure
967 | alloy_json_abi::StateMutability::View
968 ) || func.is_reserved()
969 || excluded_test_selectors.contains(&func.selector())
970 {
971 None
972 } else {
973 Some(func.selector())
974 }
975 })
976 .collect();
977 self.add_address_with_functions(address, &selectors, false, targeted_contracts)?;
978 }
979
980 Ok(())
981 }
982
983 fn add_address_with_functions(
985 &self,
986 address: Address,
987 selectors: &[Selector],
988 should_exclude: bool,
989 targeted_contracts: &mut TargetedContracts,
990 ) -> eyre::Result<()> {
991 if selectors.is_empty() {
993 return Ok(());
994 }
995
996 let contract = match targeted_contracts.entry(address) {
997 Entry::Occupied(entry) => entry.into_mut(),
998 Entry::Vacant(entry) => {
999 let (identifier, abi) = self.setup_contracts.get(&address).ok_or_else(|| {
1000 eyre::eyre!(
1001 "[{}] address does not have an associated contract: {}",
1002 if should_exclude { "excludeSelectors" } else { "targetSelectors" },
1003 address
1004 )
1005 })?;
1006 entry.insert(
1007 TargetedContract::new(identifier.clone(), abi.clone())
1008 .with_project_contracts(self.project_contracts),
1009 )
1010 }
1011 };
1012 contract.add_selectors(selectors.iter().copied(), should_exclude)?;
1013 Ok(())
1014 }
1015}
1016
1017fn collect_data(
1021 invariant_test: &InvariantTest,
1022 state_changeset: &mut HashMap<Address, Account>,
1023 tx: &BasicTxDetails,
1024 call_result: &RawCallResult,
1025 run_depth: u32,
1026) {
1027 let mut has_code = false;
1029 if let Some(Some(code)) =
1030 state_changeset.get(&tx.sender).map(|account| account.info.code.as_ref())
1031 {
1032 has_code = !code.is_empty();
1033 }
1034
1035 let mut sender_changeset = None;
1037 if !has_code {
1038 sender_changeset = state_changeset.remove(&tx.sender);
1039 }
1040
1041 invariant_test.fuzz_state.collect_values_from_call(
1043 &invariant_test.targeted_contracts,
1044 tx,
1045 &call_result.result,
1046 &call_result.logs,
1047 &*state_changeset,
1048 run_depth,
1049 );
1050
1051 if let Some(changed) = sender_changeset {
1053 state_changeset.insert(tx.sender, changed);
1054 }
1055}
1056
1057pub(crate) fn call_after_invariant_function(
1061 executor: &Executor,
1062 to: Address,
1063) -> Result<(RawCallResult, bool), EvmError> {
1064 let calldata = Bytes::from_static(&IInvariantTest::afterInvariantCall::SELECTOR);
1065 let mut call_result = executor.call_raw(CALLER, to, calldata, U256::ZERO)?;
1066 let success = executor.is_raw_call_mut_success(to, &mut call_result, false);
1067 Ok((call_result, success))
1068}
1069
1070pub(crate) fn call_invariant_function(
1072 executor: &Executor,
1073 address: Address,
1074 calldata: Bytes,
1075) -> Result<(RawCallResult, bool)> {
1076 let mut call_result = executor.call_raw(CALLER, address, calldata, U256::ZERO)?;
1077 let success = executor.is_raw_call_mut_success(address, &mut call_result, false);
1078 Ok((call_result, success))
1079}
1080
1081pub(crate) fn execute_tx(executor: &mut Executor, tx: &BasicTxDetails) -> Result<RawCallResult> {
1084 if let Some(warp) = tx.warp {
1086 executor.env_mut().evm_env.block_env.timestamp += warp;
1087 }
1088 if let Some(roll) = tx.roll {
1089 executor.env_mut().evm_env.block_env.number += roll;
1090 }
1091
1092 let mut call_result = executor
1094 .call_raw(tx.sender, tx.call_details.target, tx.call_details.calldata.clone(), U256::ZERO)
1095 .map_err(|e| eyre!(format!("Could not make raw evm call: {e}")))?;
1096
1097 if let Some(warp) = tx.warp {
1099 call_result.env.evm_env.block_env.timestamp += warp;
1100 }
1101 if let Some(roll) = tx.roll {
1102 call_result.env.evm_env.block_env.number += roll;
1103 }
1104 Ok(call_result)
1105}