Skip to main content

foundry_evm/executors/invariant/
mod.rs

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    FuzzCase, FuzzFixtures, FuzzedCases,
18    invariant::{
19        ArtifactFilters, BasicTxDetails, FuzzRunIdentifiedContracts, InvariantContract,
20        RandomCallGenerator, 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    cell::RefCell,
33    collections::{HashMap as Map, btree_map::Entry},
34    sync::Arc,
35    time::{Duration, Instant, SystemTime, UNIX_EPOCH},
36};
37
38mod error;
39pub use error::{InvariantFailures, InvariantFuzzError};
40use foundry_evm_coverage::HitMaps;
41
42mod replay;
43pub use replay::{replay_error, replay_run};
44
45mod result;
46use foundry_common::{TestFunctionExt, sh_println};
47pub use result::InvariantFuzzTestResult;
48use serde::{Deserialize, Serialize};
49use serde_json::json;
50
51mod corpus;
52
53mod shrink;
54use crate::executors::{EvmError, FuzzTestTimer, invariant::corpus::TxCorpusManager};
55pub use shrink::check_sequence;
56
57sol! {
58    interface IInvariantTest {
59        #[derive(Default)]
60        struct FuzzSelector {
61            address addr;
62            bytes4[] selectors;
63        }
64
65        #[derive(Default)]
66        struct FuzzArtifactSelector {
67            string artifact;
68            bytes4[] selectors;
69        }
70
71        #[derive(Default)]
72        struct FuzzInterface {
73            address addr;
74            string[] artifacts;
75        }
76
77        function afterInvariant() external;
78
79        #[derive(Default)]
80        function excludeArtifacts() public view returns (string[] memory excludedArtifacts);
81
82        #[derive(Default)]
83        function excludeContracts() public view returns (address[] memory excludedContracts);
84
85        #[derive(Default)]
86        function excludeSelectors() public view returns (FuzzSelector[] memory excludedSelectors);
87
88        #[derive(Default)]
89        function excludeSenders() public view returns (address[] memory excludedSenders);
90
91        #[derive(Default)]
92        function targetArtifacts() public view returns (string[] memory targetedArtifacts);
93
94        #[derive(Default)]
95        function targetArtifactSelectors() public view returns (FuzzArtifactSelector[] memory targetedArtifactSelectors);
96
97        #[derive(Default)]
98        function targetContracts() public view returns (address[] memory targetedContracts);
99
100        #[derive(Default)]
101        function targetSelectors() public view returns (FuzzSelector[] memory targetedSelectors);
102
103        #[derive(Default)]
104        function targetSenders() public view returns (address[] memory targetedSenders);
105
106        #[derive(Default)]
107        function targetInterfaces() public view returns (FuzzInterface[] memory targetedInterfaces);
108    }
109}
110
111const DURATION_BETWEEN_METRICS_REPORT: Duration = Duration::from_secs(5);
112
113/// Contains invariant metrics for a single fuzzed selector.
114#[derive(Default, Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
115pub struct InvariantMetrics {
116    // Count of fuzzed selector calls.
117    pub calls: usize,
118    // Count of fuzzed selector reverts.
119    pub reverts: usize,
120    // Count of fuzzed selector discards (through assume cheatcodes).
121    pub discards: usize,
122}
123
124/// Contains data collected during invariant test runs.
125pub struct InvariantTestData {
126    // Consumed gas and calldata of every successful fuzz call.
127    pub fuzz_cases: Vec<FuzzedCases>,
128    // Data related to reverts or failed assertions of the test.
129    pub failures: InvariantFailures,
130    // Calldata in the last invariant run.
131    pub last_run_inputs: Vec<BasicTxDetails>,
132    // Additional traces for gas report.
133    pub gas_report_traces: Vec<Vec<CallTraceArena>>,
134    // Last call results of the invariant test.
135    pub last_call_results: Option<RawCallResult>,
136    // Line coverage information collected from all fuzzed calls.
137    pub line_coverage: Option<HitMaps>,
138    // Metrics for each fuzzed selector.
139    pub metrics: Map<String, InvariantMetrics>,
140
141    // Proptest runner to query for random values.
142    // The strategy only comes with the first `input`. We fill the rest of the `inputs`
143    // until the desired `depth` so we can use the evolving fuzz dictionary
144    // during the run.
145    pub branch_runner: TestRunner,
146}
147
148/// Contains invariant test data.
149pub struct InvariantTest {
150    // Fuzz state of invariant test.
151    pub fuzz_state: EvmFuzzState,
152    // Contracts fuzzed by the invariant test.
153    pub targeted_contracts: FuzzRunIdentifiedContracts,
154    // Data collected during invariant runs.
155    pub execution_data: RefCell<InvariantTestData>,
156}
157
158impl InvariantTest {
159    /// Instantiates an invariant test.
160    pub fn new(
161        fuzz_state: EvmFuzzState,
162        targeted_contracts: FuzzRunIdentifiedContracts,
163        failures: InvariantFailures,
164        last_call_results: Option<RawCallResult>,
165        branch_runner: TestRunner,
166    ) -> Self {
167        let mut fuzz_cases = vec![];
168        if last_call_results.is_none() {
169            fuzz_cases.push(FuzzedCases::new(vec![]));
170        }
171        let execution_data = RefCell::new(InvariantTestData {
172            fuzz_cases,
173            failures,
174            last_run_inputs: vec![],
175            gas_report_traces: vec![],
176            last_call_results,
177            line_coverage: None,
178            metrics: Map::default(),
179            branch_runner,
180        });
181        Self { fuzz_state, targeted_contracts, execution_data }
182    }
183
184    /// Returns number of invariant test reverts.
185    pub fn reverts(&self) -> usize {
186        self.execution_data.borrow().failures.reverts
187    }
188
189    /// Whether invariant test has errors or not.
190    pub fn has_errors(&self) -> bool {
191        self.execution_data.borrow().failures.error.is_some()
192    }
193
194    /// Set invariant test error.
195    pub fn set_error(&self, error: InvariantFuzzError) {
196        self.execution_data.borrow_mut().failures.error = Some(error);
197    }
198
199    /// Set last invariant test call results.
200    pub fn set_last_call_results(&self, call_result: Option<RawCallResult>) {
201        self.execution_data.borrow_mut().last_call_results = call_result;
202    }
203
204    /// Set last invariant run call sequence.
205    pub fn set_last_run_inputs(&self, inputs: &Vec<BasicTxDetails>) {
206        self.execution_data.borrow_mut().last_run_inputs.clone_from(inputs);
207    }
208
209    /// Merge current collected line coverage with the new coverage from last fuzzed call.
210    pub fn merge_coverage(&self, new_coverage: Option<HitMaps>) {
211        HitMaps::merge_opt(&mut self.execution_data.borrow_mut().line_coverage, new_coverage);
212    }
213
214    /// Update metrics for a fuzzed selector, extracted from tx details.
215    /// Always increments number of calls; discarded runs (through assume cheatcodes) are tracked
216    /// separated from reverts.
217    pub fn record_metrics(&self, tx_details: &BasicTxDetails, reverted: bool, discarded: bool) {
218        if let Some(metric_key) =
219            self.targeted_contracts.targets.lock().fuzzed_metric_key(tx_details)
220        {
221            let test_metrics = &mut self.execution_data.borrow_mut().metrics;
222            let invariant_metrics = test_metrics.entry(metric_key).or_default();
223            invariant_metrics.calls += 1;
224            if discarded {
225                invariant_metrics.discards += 1;
226            } else if reverted {
227                invariant_metrics.reverts += 1;
228            }
229        }
230    }
231
232    /// End invariant test run by collecting results, cleaning collected artifacts and reverting
233    /// created fuzz state.
234    pub fn end_run(&self, run: InvariantTestRun, gas_samples: usize) {
235        // We clear all the targeted contracts created during this run.
236        self.targeted_contracts.clear_created_contracts(run.created_contracts);
237
238        let mut invariant_data = self.execution_data.borrow_mut();
239        if invariant_data.gas_report_traces.len() < gas_samples {
240            invariant_data
241                .gas_report_traces
242                .push(run.run_traces.into_iter().map(|arena| arena.arena).collect());
243        }
244        invariant_data.fuzz_cases.push(FuzzedCases::new(run.fuzz_runs));
245
246        // Revert state to not persist values between runs.
247        self.fuzz_state.revert();
248    }
249}
250
251/// Contains data for an invariant test run.
252pub struct InvariantTestRun {
253    // Invariant run call sequence.
254    pub inputs: Vec<BasicTxDetails>,
255    // Current invariant run executor.
256    pub executor: Executor,
257    // Invariant run stat reports (eg. gas usage).
258    pub fuzz_runs: Vec<FuzzCase>,
259    // Contracts created during current invariant run.
260    pub created_contracts: Vec<Address>,
261    // Traces of each call of the invariant run call sequence.
262    pub run_traces: Vec<SparsedTraceArena>,
263    // Current depth of invariant run.
264    pub depth: u32,
265    // Current assume rejects of the invariant run.
266    pub assume_rejects_counter: u32,
267    // Whether new coverage was discovered during this run.
268    pub new_coverage: bool,
269}
270
271impl InvariantTestRun {
272    /// Instantiates an invariant test run.
273    pub fn new(first_input: BasicTxDetails, executor: Executor, depth: usize) -> Self {
274        Self {
275            inputs: vec![first_input],
276            executor,
277            fuzz_runs: Vec::with_capacity(depth),
278            created_contracts: vec![],
279            run_traces: vec![],
280            depth: 0,
281            assume_rejects_counter: 0,
282            new_coverage: false,
283        }
284    }
285}
286
287/// Wrapper around any [`Executor`] implementer which provides fuzzing support using [`proptest`].
288///
289/// After instantiation, calling `invariant_fuzz` will proceed to hammer the deployed smart
290/// contracts with inputs, until it finds a counterexample sequence. The provided [`TestRunner`]
291/// contains all the configuration which can be overridden via [environment
292/// variables](proptest::test_runner::Config)
293pub struct InvariantExecutor<'a> {
294    pub executor: Executor,
295    /// Proptest runner.
296    runner: TestRunner,
297    /// The invariant configuration
298    config: InvariantConfig,
299    /// Contracts deployed with `setUp()`
300    setup_contracts: &'a ContractsByAddress,
301    /// Contracts that are part of the project but have not been deployed yet. We need the bytecode
302    /// to identify them from the stateset changes.
303    project_contracts: &'a ContractsByArtifact,
304    /// Filters contracts to be fuzzed through their artifact identifiers.
305    artifact_filters: ArtifactFilters,
306    /// History of binned hitcount of edges seen during fuzzing.
307    history_map: Vec<u8>,
308}
309const COVERAGE_MAP_SIZE: usize = 65536;
310
311impl<'a> InvariantExecutor<'a> {
312    /// Instantiates a fuzzed executor EVM given a testrunner
313    pub fn new(
314        executor: Executor,
315        runner: TestRunner,
316        config: InvariantConfig,
317        setup_contracts: &'a ContractsByAddress,
318        project_contracts: &'a ContractsByArtifact,
319    ) -> Self {
320        Self {
321            executor,
322            runner,
323            config,
324            setup_contracts,
325            project_contracts,
326            artifact_filters: ArtifactFilters::default(),
327            history_map: vec![0u8; COVERAGE_MAP_SIZE],
328        }
329    }
330
331    /// Fuzzes any deployed contract and checks any broken invariant at `invariant_address`.
332    pub fn invariant_fuzz(
333        &mut self,
334        invariant_contract: InvariantContract<'_>,
335        fuzz_fixtures: &FuzzFixtures,
336        deployed_libs: &[Address],
337        progress: Option<&ProgressBar>,
338    ) -> Result<InvariantFuzzTestResult> {
339        // Throw an error to abort test run if the invariant function accepts input params
340        if !invariant_contract.invariant_function.inputs.is_empty() {
341            return Err(eyre!("Invariant test function should have no inputs"));
342        }
343
344        let (invariant_test, mut corpus_manager) =
345            self.prepare_test(&invariant_contract, fuzz_fixtures, deployed_libs)?;
346
347        // Start timer for this invariant test.
348        let mut runs = 0;
349        let timer = FuzzTestTimer::new(self.config.timeout);
350        let mut last_metrics_report = Instant::now();
351        let continue_campaign = |runs: u32| {
352            // If timeout is configured, then perform invariant runs until expires.
353            if self.config.timeout.is_some() {
354                return !timer.is_timed_out();
355            }
356            // If no timeout configured then loop until configured runs.
357            runs < self.config.runs
358        };
359
360        // Invariant runs with edge coverage if corpus dir is set or showing edge coverage.
361        let edge_coverage_enabled =
362            self.config.corpus_dir.is_some() || self.config.show_edge_coverage;
363
364        'stop: while continue_campaign(runs) {
365            let initial_seq = corpus_manager.new_sequence(&invariant_test)?;
366
367            // Create current invariant run data.
368            let mut current_run = InvariantTestRun::new(
369                initial_seq[0].clone(),
370                // Before each run, we must reset the backend state.
371                self.executor.clone(),
372                self.config.depth as usize,
373            );
374
375            // We stop the run immediately if we have reverted, and `fail_on_revert` is set.
376            if self.config.fail_on_revert && invariant_test.reverts() > 0 {
377                return Err(eyre!("call reverted"));
378            }
379
380            while current_run.depth < self.config.depth {
381                // Check if the timeout has been reached.
382                if timer.is_timed_out() {
383                    // Since we never record a revert here the test is still considered
384                    // successful even though it timed out. We *want*
385                    // this behavior for now, so that's ok, but
386                    // future developers should be aware of this.
387                    break 'stop;
388                }
389
390                let tx = current_run
391                    .inputs
392                    .last()
393                    .ok_or_else(|| eyre!("no input generated to call fuzzed target."))?;
394                // Execute call from the randomly generated sequence without committing state.
395                // State is committed only if call is not a magic assume.
396                current_run
397                    .executor
398                    .strategy
399                    .runner
400                    .start_transaction(current_run.executor.strategy.context.as_ref());
401                let call_result = current_run
402                    .executor
403                    .call_raw(
404                        tx.sender,
405                        tx.call_details.target,
406                        tx.call_details.calldata.clone(),
407                        U256::ZERO,
408                    )
409                    .map_err(|e| eyre!(format!("Could not make raw evm call: {e}")));
410                current_run
411                    .executor
412                    .strategy
413                    .runner
414                    .rollback_transaction(current_run.executor.strategy.context.as_ref());
415
416                let mut call_result = call_result?;
417                let discarded = call_result.result.as_ref() == MAGIC_ASSUME;
418                if self.config.show_metrics {
419                    invariant_test.record_metrics(tx, call_result.reverted, discarded);
420                }
421
422                // Collect line coverage from last fuzzed call.
423                invariant_test.merge_coverage(call_result.line_coverage.clone());
424                // If running with edge coverage then merge edge count with the current history
425                // map and set new coverage in current run.
426                if edge_coverage_enabled {
427                    let (new_coverage, is_edge) =
428                        call_result.merge_edge_coverage(&mut self.history_map);
429                    if new_coverage {
430                        current_run.new_coverage = true;
431                        corpus_manager.update_seen_metrics(is_edge);
432                    }
433                }
434
435                if discarded {
436                    current_run.inputs.pop();
437                    current_run.assume_rejects_counter += 1;
438                    if current_run.assume_rejects_counter > self.config.max_assume_rejects {
439                        invariant_test.set_error(InvariantFuzzError::MaxAssumeRejects(
440                            self.config.max_assume_rejects,
441                        ));
442                        break 'stop;
443                    }
444                } else {
445                    // Commit executed call result.
446                    current_run.executor.commit(&mut call_result);
447
448                    // Collect data for fuzzing from the state changeset.
449                    // This step updates the state dictionary and therefore invalidates the
450                    // ValueTree in use by the current run. This manifestsitself in proptest
451                    // observing a different input case than what it was called with, and creates
452                    // inconsistencies whenever proptest tries to use the input case after test
453                    // execution.
454                    // See <https://github.com/foundry-rs/foundry/issues/9764>.
455                    let mut state_changeset = call_result.state_changeset.clone();
456                    if !call_result.reverted {
457                        collect_data(
458                            &invariant_test,
459                            &mut state_changeset,
460                            tx,
461                            &call_result,
462                            self.config.depth,
463                        );
464                    }
465
466                    // Collect created contracts and add to fuzz targets only if targeted contracts
467                    // are updatable.
468                    if let Err(error) =
469                        &invariant_test.targeted_contracts.collect_created_contracts(
470                            &state_changeset,
471                            self.project_contracts,
472                            self.setup_contracts,
473                            &self.artifact_filters,
474                            &mut current_run.created_contracts,
475                        )
476                    {
477                        warn!(target: "forge::test", "{error}");
478                    }
479                    current_run.fuzz_runs.push(FuzzCase {
480                        calldata: tx.call_details.calldata.clone(),
481                        gas: call_result.gas_used,
482                        stipend: call_result.stipend,
483                    });
484
485                    // Determine if test can continue or should exit.
486                    let result = can_continue(
487                        &invariant_contract,
488                        &invariant_test,
489                        &mut current_run,
490                        &self.config,
491                        call_result,
492                        &state_changeset,
493                    )
494                    .map_err(|e| eyre!(e.to_string()))?;
495                    if !result.can_continue || current_run.depth == self.config.depth - 1 {
496                        invariant_test.set_last_run_inputs(&current_run.inputs);
497                    }
498                    // If test cannot continue then stop current run and exit test suite.
499                    if !result.can_continue {
500                        break 'stop;
501                    }
502
503                    invariant_test.set_last_call_results(result.call_result);
504                    current_run.depth += 1;
505                }
506
507                current_run.inputs.push(corpus_manager.generate_next_input(
508                    &invariant_test,
509                    &initial_seq,
510                    discarded,
511                    current_run.depth as usize,
512                )?);
513            }
514
515            // Extend corpus with current run data.
516            corpus_manager.collect_inputs(&current_run);
517
518            // Call `afterInvariant` only if it is declared and test didn't fail already.
519            if invariant_contract.call_after_invariant && !invariant_test.has_errors() {
520                assert_after_invariant(
521                    &invariant_contract,
522                    &invariant_test,
523                    &current_run,
524                    &self.config,
525                )
526                .map_err(|_| eyre!("Failed to call afterInvariant"))?;
527            }
528
529            // End current invariant test run.
530            invariant_test.end_run(current_run, self.config.gas_report_samples as usize);
531            if let Some(progress) = progress {
532                // If running with progress then increment completed runs.
533                progress.inc(1);
534                // Display metrics in progress bar.
535                if edge_coverage_enabled {
536                    progress.set_message(format!("{}", &corpus_manager.metrics));
537                }
538            } else if edge_coverage_enabled
539                && last_metrics_report.elapsed() > DURATION_BETWEEN_METRICS_REPORT
540            {
541                // Display metrics inline if corpus dir set.
542                let metrics = json!({
543                    "timestamp": SystemTime::now()
544                        .duration_since(UNIX_EPOCH)?
545                        .as_secs(),
546                    "invariant": invariant_contract.invariant_function.name,
547                    "metrics": &corpus_manager.metrics,
548                });
549                let _ = sh_println!("{}", serde_json::to_string(&metrics)?);
550                last_metrics_report = Instant::now();
551            }
552
553            runs += 1;
554        }
555
556        trace!(?fuzz_fixtures);
557        invariant_test.fuzz_state.log_stats();
558
559        let result = invariant_test.execution_data.into_inner();
560        Ok(InvariantFuzzTestResult {
561            error: result.failures.error,
562            cases: result.fuzz_cases,
563            reverts: result.failures.reverts,
564            last_run_inputs: result.last_run_inputs,
565            gas_report_traces: result.gas_report_traces,
566            line_coverage: result.line_coverage,
567            metrics: result.metrics,
568            failed_corpus_replays: corpus_manager.failed_replays(),
569        })
570    }
571
572    /// Prepares certain structures to execute the invariant tests:
573    /// * Invariant Fuzz Test.
574    /// * Invariant Corpus Manager.
575    fn prepare_test(
576        &mut self,
577        invariant_contract: &InvariantContract<'_>,
578        fuzz_fixtures: &FuzzFixtures,
579        deployed_libs: &[Address],
580    ) -> Result<(InvariantTest, TxCorpusManager)> {
581        // Finds out the chosen deployed contracts and/or senders.
582        self.select_contract_artifacts(invariant_contract.address)?;
583        let (targeted_senders, targeted_contracts) =
584            self.select_contracts_and_senders(invariant_contract.address)?;
585
586        // Stores fuzz state for use with [fuzz_calldata_from_state].
587        let fuzz_state = EvmFuzzState::new(
588            self.executor.backend().mem_db(),
589            self.config.dictionary,
590            deployed_libs,
591        );
592
593        // Creates the invariant strategy.
594        let strategy = invariant_strat(
595            fuzz_state.clone(),
596            targeted_senders,
597            targeted_contracts.clone(),
598            self.config.dictionary.dictionary_weight,
599            fuzz_fixtures.clone(),
600            self.config.max_fuzz_int,
601        )
602        .no_shrink();
603
604        // Allows `override_call_strat` to use the address given by the Fuzzer inspector during
605        // EVM execution.
606        let mut call_generator = None;
607        if self.config.call_override {
608            let target_contract_ref = Arc::new(RwLock::new(Address::ZERO));
609
610            call_generator = Some(RandomCallGenerator::new(
611                invariant_contract.address,
612                self.runner.clone(),
613                override_call_strat(
614                    fuzz_state.clone(),
615                    targeted_contracts.clone(),
616                    target_contract_ref.clone(),
617                    fuzz_fixtures.clone(),
618                    self.config.max_fuzz_int,
619                ),
620                target_contract_ref,
621            ));
622        }
623
624        self.executor.inspector_mut().fuzzer =
625            Some(Fuzzer { call_generator, fuzz_state: fuzz_state.clone(), collect: true });
626
627        // Let's make sure the invariant is sound before actually starting the run:
628        // We'll assert the invariant in its initial state, and if it fails, we'll
629        // already know if we can early exit the invariant run.
630        // This does not count as a fuzz run. It will just register the revert.
631        let mut failures = InvariantFailures::new();
632        let last_call_results = assert_invariants(
633            invariant_contract,
634            &self.config,
635            &targeted_contracts,
636            &self.executor,
637            &[],
638            &mut failures,
639        )?;
640        if let Some(error) = failures.error {
641            return Err(eyre!(error.revert_reason().unwrap_or_default()));
642        }
643
644        let corpus_manager = TxCorpusManager::new(
645            &self.config,
646            &invariant_contract.invariant_function.name,
647            &targeted_contracts,
648            strategy.boxed(),
649            &self.executor,
650            &mut self.history_map,
651        )?;
652
653        let invariant_test = InvariantTest::new(
654            fuzz_state,
655            targeted_contracts,
656            failures,
657            last_call_results,
658            self.runner.clone(),
659        );
660
661        Ok((invariant_test, corpus_manager))
662    }
663
664    /// Fills the `InvariantExecutor` with the artifact identifier filters (in `path:name` string
665    /// format). They will be used to filter contracts after the `setUp`, and more importantly,
666    /// during the runs.
667    ///
668    /// Also excludes any contract without any mutable functions.
669    ///
670    /// Priority:
671    ///
672    /// targetArtifactSelectors > excludeArtifacts > targetArtifacts
673    pub fn select_contract_artifacts(&mut self, invariant_address: Address) -> Result<()> {
674        let targeted_artifact_selectors = self
675            .executor
676            .call_sol_default(invariant_address, &IInvariantTest::targetArtifactSelectorsCall {});
677
678        // Insert them into the executor `targeted_abi`.
679        for IInvariantTest::FuzzArtifactSelector { artifact, selectors } in
680            targeted_artifact_selectors
681        {
682            let identifier = self.validate_selected_contract(artifact, &selectors)?;
683            self.artifact_filters.targeted.entry(identifier).or_default().extend(selectors);
684        }
685
686        let targeted_artifacts = self
687            .executor
688            .call_sol_default(invariant_address, &IInvariantTest::targetArtifactsCall {});
689        let excluded_artifacts = self
690            .executor
691            .call_sol_default(invariant_address, &IInvariantTest::excludeArtifactsCall {});
692
693        // Insert `excludeArtifacts` into the executor `excluded_abi`.
694        for contract in excluded_artifacts {
695            let identifier = self.validate_selected_contract(contract, &[])?;
696
697            if !self.artifact_filters.excluded.contains(&identifier) {
698                self.artifact_filters.excluded.push(identifier);
699            }
700        }
701
702        // Exclude any artifact without mutable functions.
703        for (artifact, contract) in self.project_contracts.iter() {
704            if contract
705                .abi
706                .functions()
707                .filter(|func| {
708                    !matches!(
709                        func.state_mutability,
710                        alloy_json_abi::StateMutability::Pure
711                            | alloy_json_abi::StateMutability::View
712                    )
713                })
714                .count()
715                == 0
716                && !self.artifact_filters.excluded.contains(&artifact.identifier())
717            {
718                self.artifact_filters.excluded.push(artifact.identifier());
719            }
720        }
721
722        // Insert `targetArtifacts` into the executor `targeted_abi`, if they have not been seen
723        // before.
724        for contract in targeted_artifacts {
725            let identifier = self.validate_selected_contract(contract, &[])?;
726
727            if !self.artifact_filters.targeted.contains_key(&identifier)
728                && !self.artifact_filters.excluded.contains(&identifier)
729            {
730                self.artifact_filters.targeted.insert(identifier, vec![]);
731            }
732        }
733        Ok(())
734    }
735
736    /// Makes sure that the contract exists in the project. If so, it returns its artifact
737    /// identifier.
738    fn validate_selected_contract(
739        &mut self,
740        contract: String,
741        selectors: &[FixedBytes<4>],
742    ) -> Result<String> {
743        if let Some((artifact, contract_data)) =
744            self.project_contracts.find_by_name_or_identifier(&contract)?
745        {
746            // Check that the selectors really exist for this contract.
747            for selector in selectors {
748                contract_data
749                    .abi
750                    .functions()
751                    .find(|func| func.selector().as_slice() == selector.as_slice())
752                    .wrap_err(format!("{contract} does not have the selector {selector:?}"))?;
753            }
754
755            return Ok(artifact.identifier());
756        }
757        eyre::bail!(
758            "{contract} not found in the project. Allowed format: `contract_name` or `contract_path:contract_name`."
759        );
760    }
761
762    /// Selects senders and contracts based on the contract methods `targetSenders() -> address[]`,
763    /// `targetContracts() -> address[]` and `excludeContracts() -> address[]`.
764    pub fn select_contracts_and_senders(
765        &self,
766        to: Address,
767    ) -> Result<(SenderFilters, FuzzRunIdentifiedContracts)> {
768        let targeted_senders =
769            self.executor.call_sol_default(to, &IInvariantTest::targetSendersCall {});
770        let mut excluded_senders =
771            self.executor.call_sol_default(to, &IInvariantTest::excludeSendersCall {});
772        // Extend with default excluded addresses - https://github.com/foundry-rs/foundry/issues/4163
773        excluded_senders.extend([
774            CHEATCODE_ADDRESS,
775            HARDHAT_CONSOLE_ADDRESS,
776            DEFAULT_CREATE2_DEPLOYER,
777        ]);
778        // Extend with precompiles - https://github.com/foundry-rs/foundry/issues/4287
779        excluded_senders.extend(PRECOMPILES);
780        let sender_filters = SenderFilters::new(targeted_senders, excluded_senders);
781
782        let selected = self.executor.call_sol_default(to, &IInvariantTest::targetContractsCall {});
783        let excluded = self.executor.call_sol_default(to, &IInvariantTest::excludeContractsCall {});
784
785        let contracts = self
786            .setup_contracts
787            .iter()
788            .filter(|&(addr, (identifier, _))| {
789                // Include to address if explicitly set as target.
790                if *addr == to && selected.contains(&to) {
791                    return true;
792                }
793
794                *addr != to
795                    && *addr != CHEATCODE_ADDRESS
796                    && *addr != HARDHAT_CONSOLE_ADDRESS
797                    && (selected.is_empty() || selected.contains(addr))
798                    && (excluded.is_empty() || !excluded.contains(addr))
799                    && self.artifact_filters.matches(identifier)
800            })
801            .map(|(addr, (identifier, abi))| {
802                (*addr, TargetedContract::new(identifier.clone(), abi.clone()))
803            })
804            .collect();
805        let mut contracts = TargetedContracts { inner: contracts };
806
807        self.target_interfaces(to, &mut contracts)?;
808
809        self.select_selectors(to, &mut contracts)?;
810
811        // There should be at least one contract identified as target for fuzz runs.
812        if contracts.is_empty() {
813            eyre::bail!("No contracts to fuzz.");
814        }
815
816        Ok((sender_filters, FuzzRunIdentifiedContracts::new(contracts, selected.is_empty())))
817    }
818
819    /// Extends the contracts and selectors to fuzz with the addresses and ABIs specified in
820    /// `targetInterfaces() -> (address, string[])[]`. Enables targeting of addresses that are
821    /// not deployed during `setUp` such as when fuzzing in a forked environment. Also enables
822    /// targeting of delegate proxies and contracts deployed with `create` or `create2`.
823    pub fn target_interfaces(
824        &self,
825        invariant_address: Address,
826        targeted_contracts: &mut TargetedContracts,
827    ) -> Result<()> {
828        let interfaces = self
829            .executor
830            .call_sol_default(invariant_address, &IInvariantTest::targetInterfacesCall {});
831
832        // Since `targetInterfaces` returns a tuple array there is no guarantee
833        // that the addresses are unique this map is used to merge functions of
834        // the specified interfaces for the same address. For example:
835        // `[(addr1, ["IERC20", "IOwnable"])]` and `[(addr1, ["IERC20"]), (addr1, ("IOwnable"))]`
836        // should be equivalent.
837        let mut combined = TargetedContracts::new();
838
839        // Loop through each address and its associated artifact identifiers.
840        // We're borrowing here to avoid taking full ownership.
841        for IInvariantTest::FuzzInterface { addr, artifacts } in &interfaces {
842            // Identifiers are specified as an array, so we loop through them.
843            for identifier in artifacts {
844                // Try to find the contract by name or identifier in the project's contracts.
845                if let Some(abi) = self.project_contracts.find_abi_by_name_or_identifier(identifier)
846                {
847                    combined
848                        // Check if there's an entry for the given key in the 'combined' map.
849                        .entry(*addr)
850                        // If the entry exists, extends its ABI with the function list.
851                        .and_modify(|entry| {
852                            // Extend the ABI's function list with the new functions.
853                            entry.abi.functions.extend(abi.functions.clone());
854                        })
855                        // Otherwise insert it into the map.
856                        .or_insert_with(|| TargetedContract::new(identifier.to_string(), abi));
857                }
858            }
859        }
860
861        targeted_contracts.extend(combined.inner);
862
863        Ok(())
864    }
865
866    /// Selects the functions to fuzz based on the contract method `targetSelectors()` and
867    /// `targetArtifactSelectors()`.
868    pub fn select_selectors(
869        &self,
870        address: Address,
871        targeted_contracts: &mut TargetedContracts,
872    ) -> Result<()> {
873        for (address, (identifier, _)) in self.setup_contracts {
874            if let Some(selectors) = self.artifact_filters.targeted.get(identifier) {
875                self.add_address_with_functions(*address, selectors, false, targeted_contracts)?;
876            }
877        }
878
879        let mut target_test_selectors = vec![];
880        let mut excluded_test_selectors = vec![];
881
882        // Collect contract functions marked as target for fuzzing campaign.
883        let selectors =
884            self.executor.call_sol_default(address, &IInvariantTest::targetSelectorsCall {});
885        for IInvariantTest::FuzzSelector { addr, selectors } in selectors {
886            if addr == address {
887                target_test_selectors = selectors.clone();
888            }
889            self.add_address_with_functions(addr, &selectors, false, targeted_contracts)?;
890        }
891
892        // Collect contract functions excluded from fuzzing campaign.
893        let excluded_selectors =
894            self.executor.call_sol_default(address, &IInvariantTest::excludeSelectorsCall {});
895        for IInvariantTest::FuzzSelector { addr, selectors } in excluded_selectors {
896            if addr == address {
897                // If fuzz selector address is the test contract, then record selectors to be
898                // later excluded if needed.
899                excluded_test_selectors = selectors.clone();
900            }
901            self.add_address_with_functions(addr, &selectors, true, targeted_contracts)?;
902        }
903
904        if target_test_selectors.is_empty()
905            && let Some(target) = targeted_contracts.get(&address)
906        {
907            // If test contract is marked as a target and no target selector explicitly set, then
908            // include only state-changing functions that are not reserved and selectors that are
909            // not explicitly excluded.
910            let selectors: Vec<_> = target
911                .abi
912                .functions()
913                .filter_map(|func| {
914                    if matches!(
915                        func.state_mutability,
916                        alloy_json_abi::StateMutability::Pure
917                            | alloy_json_abi::StateMutability::View
918                    ) || func.is_reserved()
919                        || excluded_test_selectors.contains(&func.selector())
920                    {
921                        None
922                    } else {
923                        Some(func.selector())
924                    }
925                })
926                .collect();
927            self.add_address_with_functions(address, &selectors, false, targeted_contracts)?;
928        }
929
930        Ok(())
931    }
932
933    /// Adds the address and fuzzed or excluded functions to `TargetedContracts`.
934    fn add_address_with_functions(
935        &self,
936        address: Address,
937        selectors: &[Selector],
938        should_exclude: bool,
939        targeted_contracts: &mut TargetedContracts,
940    ) -> eyre::Result<()> {
941        // Do not add address in target contracts if no function selected.
942        if selectors.is_empty() {
943            return Ok(());
944        }
945
946        let contract = match targeted_contracts.entry(address) {
947            Entry::Occupied(entry) => entry.into_mut(),
948            Entry::Vacant(entry) => {
949                let (identifier, abi) = self.setup_contracts.get(&address).ok_or_else(|| {
950                    eyre::eyre!(
951                        "[{}] address does not have an associated contract: {}",
952                        if should_exclude { "excludeSelectors" } else { "targetSelectors" },
953                        address
954                    )
955                })?;
956                entry.insert(TargetedContract::new(identifier.clone(), abi.clone()))
957            }
958        };
959        contract.add_selectors(selectors.iter().copied(), should_exclude)?;
960        Ok(())
961    }
962}
963
964/// Collects data from call for fuzzing. However, it first verifies that the sender is not an EOA
965/// before inserting it into the dictionary. Otherwise, we flood the dictionary with
966/// randomly generated addresses.
967fn collect_data(
968    invariant_test: &InvariantTest,
969    state_changeset: &mut HashMap<Address, Account>,
970    tx: &BasicTxDetails,
971    call_result: &RawCallResult,
972    run_depth: u32,
973) {
974    // Verify it has no code.
975    let mut has_code = false;
976    if let Some(Some(code)) =
977        state_changeset.get(&tx.sender).map(|account| account.info.code.as_ref())
978    {
979        has_code = !code.is_empty();
980    }
981
982    // We keep the nonce changes to apply later.
983    let mut sender_changeset = None;
984    if !has_code {
985        sender_changeset = state_changeset.remove(&tx.sender);
986    }
987
988    // Collect values from fuzzed call result and add them to fuzz dictionary.
989    invariant_test.fuzz_state.collect_values_from_call(
990        &invariant_test.targeted_contracts,
991        tx,
992        &call_result.result,
993        &call_result.logs,
994        &*state_changeset,
995        run_depth,
996    );
997
998    // Re-add changes
999    if let Some(changed) = sender_changeset {
1000        state_changeset.insert(tx.sender, changed);
1001    }
1002}
1003
1004/// Calls the `afterInvariant()` function on a contract.
1005/// Returns call result and if call succeeded.
1006/// The state after the call is not persisted.
1007pub(crate) fn call_after_invariant_function(
1008    executor: &Executor,
1009    to: Address,
1010) -> Result<(RawCallResult, bool), EvmError> {
1011    let calldata = Bytes::from_static(&IInvariantTest::afterInvariantCall::SELECTOR);
1012    let mut call_result = executor.call_raw(CALLER, to, calldata, U256::ZERO)?;
1013    let success = executor.is_raw_call_mut_success(to, &mut call_result, false);
1014    Ok((call_result, success))
1015}
1016
1017/// Calls the invariant function and returns call result and if succeeded.
1018pub(crate) fn call_invariant_function(
1019    executor: &Executor,
1020    address: Address,
1021    calldata: Bytes,
1022) -> Result<(RawCallResult, bool)> {
1023    let mut call_result = executor.call_raw(CALLER, address, calldata, U256::ZERO)?;
1024    let success = executor.is_raw_call_mut_success(address, &mut call_result, false);
1025    Ok((call_result, success))
1026}