Skip to main content

foundry_evm_fuzz/invariant/
mod.rs

1use alloy_json_abi::{Event, Function, JsonAbi};
2use alloy_primitives::{Address, B256, Bytes, Selector, map::HashMap};
3use foundry_compilers::artifacts::StorageLayout;
4use itertools::Either;
5use serde::{Deserialize, Serialize};
6use std::{
7    cell::{Cell, Ref, RefCell},
8    collections::BTreeMap,
9    fmt,
10    rc::Rc,
11    sync::Arc,
12};
13
14mod call_override;
15pub use call_override::RandomCallGenerator;
16
17mod filters;
18use crate::BasicTxDetails;
19pub use filters::{ArtifactFilters, SenderFilters};
20use foundry_common::{ContractsByAddress, ContractsByArtifact};
21use foundry_evm_core::utils::StateChangeset;
22
23type DynamicTargetArtifactMatchCache =
24    Rc<RefCell<HashMap<(Address, B256), Option<CachedTargetContract>>>>;
25type FuzzedFunction = (Address, Function);
26type FunctionLookup = HashMap<Selector, Function>;
27
28/// Contracts identified as targets during a fuzz run.
29///
30/// During execution, any newly created contract is added as target and used through the rest of
31/// the fuzz run if the collection is updatable (no `targetContract` specified in `setUp`).
32#[derive(Clone, Debug)]
33pub struct FuzzRunIdentifiedContracts {
34    /// Contracts identified as targets during a fuzz run.
35    targets: Rc<RefCell<TargetedContracts>>,
36    /// Flat cache of all currently fuzzable target functions.
37    fuzzed_functions: Rc<RefCell<Vec<FuzzedFunction>>>,
38    /// Generation counter for cached fuzzed functions.
39    fuzzed_functions_generation: Rc<Cell<u64>>,
40    /// Whether target contracts are updatable or not.
41    pub is_updatable: bool,
42    artifact_matches: DynamicTargetArtifactMatchCache,
43}
44
45impl FuzzRunIdentifiedContracts {
46    /// Creates a new `FuzzRunIdentifiedContracts` instance.
47    pub fn new(targets: TargetedContracts, is_updatable: bool) -> Self {
48        let fuzzed_functions = Self::flatten_fuzzed_functions(&targets);
49        Self {
50            targets: Rc::new(RefCell::new(targets)),
51            fuzzed_functions: Rc::new(RefCell::new(fuzzed_functions)),
52            fuzzed_functions_generation: Rc::new(Cell::new(0)),
53            is_updatable,
54            artifact_matches: Rc::new(RefCell::new(HashMap::default())),
55        }
56    }
57
58    /// Borrows the current targeted contracts.
59    pub fn targets(&self) -> Ref<'_, TargetedContracts> {
60        self.targets.borrow()
61    }
62
63    /// Borrows the current flat list of fuzzed target functions.
64    pub fn fuzzed_functions(&self) -> Ref<'_, [FuzzedFunction]> {
65        Ref::map(self.fuzzed_functions.borrow(), Vec::as_slice)
66    }
67
68    /// Returns the current fuzzed-functions generation.
69    pub fn fuzzed_functions_generation(&self) -> u64 {
70        self.fuzzed_functions_generation.get()
71    }
72
73    fn refresh_fuzzed_functions(&self) {
74        let fuzzed_functions = {
75            let targets = self.targets.borrow();
76            Self::flatten_fuzzed_functions(&targets)
77        };
78        *self.fuzzed_functions.borrow_mut() = fuzzed_functions;
79        self.fuzzed_functions_generation.set(self.fuzzed_functions_generation.get() + 1);
80    }
81
82    fn flatten_fuzzed_functions(targets: &TargetedContracts) -> Vec<FuzzedFunction> {
83        targets.fuzzed_functions().map(|(address, function)| (*address, function.clone())).collect()
84    }
85
86    /// If targets are updatable, collect all contracts created during an invariant run (which
87    /// haven't been discovered yet).
88    pub fn collect_created_contracts(
89        &self,
90        state_changeset: &StateChangeset,
91        project_contracts: &ContractsByArtifact,
92        setup_contracts: &ContractsByAddress,
93        artifact_filters: &ArtifactFilters,
94        created_contracts: &mut Vec<Address>,
95    ) -> eyre::Result<()> {
96        if !self.is_updatable {
97            return Ok(());
98        }
99
100        let mut targets_changed = false;
101        for (address, account) in state_changeset {
102            if setup_contracts.contains_key(address) {
103                continue;
104            }
105            if !account.is_touched() {
106                continue;
107            }
108            let Some(code) = &account.info.code else {
109                continue;
110            };
111            if code.is_empty() {
112                continue;
113            }
114            let code_hash = code.hash_slow();
115            let code = code.original_byte_slice();
116            let Some(contract) = self.target_contract_for_code(
117                *address,
118                code_hash,
119                code,
120                project_contracts,
121                artifact_filters,
122            )?
123            else {
124                continue;
125            };
126            created_contracts.push(*address);
127            self.targets.borrow_mut().insert(*address, contract.into_targeted_contract());
128            targets_changed = true;
129        }
130        if targets_changed {
131            self.refresh_fuzzed_functions();
132        }
133        Ok(())
134    }
135
136    fn target_contract_for_code(
137        &self,
138        address: Address,
139        code_hash: B256,
140        code: &[u8],
141        project_contracts: &ContractsByArtifact,
142        artifact_filters: &ArtifactFilters,
143    ) -> eyre::Result<Option<CachedTargetContract>> {
144        let cache_key = (address, code_hash);
145        if let Some(cached_match) = self.artifact_matches.borrow().get(&cache_key) {
146            return Ok(cached_match.clone());
147        }
148
149        let cached_match = if let Some((artifact, contract_data)) =
150            project_contracts.find_by_deployed_code(code)
151        {
152            artifact_filters.get_targeted_functions(artifact, &contract_data.abi)?.map(
153                |targeted_functions| CachedTargetContract {
154                    identifier: artifact.name.clone(),
155                    abi: contract_data.abi.clone(),
156                    targeted_functions,
157                    storage_layout: contract_data.storage_layout.as_ref().map(Arc::clone),
158                    event_lookup: Arc::new(TargetedContractEvents::new(&contract_data.abi)),
159                },
160            )
161        } else {
162            None
163        };
164        self.artifact_matches.borrow_mut().insert(cache_key, cached_match.clone());
165        Ok(cached_match)
166    }
167
168    /// Clears targeted contracts created during an invariant run.
169    pub fn clear_created_contracts(&self, created_contracts: Vec<Address>) {
170        let mut targets_changed = false;
171        if !created_contracts.is_empty() {
172            let mut targets = self.targets.borrow_mut();
173            for addr in &created_contracts {
174                targets_changed |= targets.remove(addr).is_some();
175            }
176        }
177        if targets_changed {
178            self.refresh_fuzzed_functions();
179        }
180    }
181}
182
183#[derive(Clone, Debug)]
184struct CachedTargetContract {
185    identifier: String,
186    abi: JsonAbi,
187    targeted_functions: Vec<Function>,
188    storage_layout: Option<Arc<StorageLayout>>,
189    event_lookup: Arc<TargetedContractEvents>,
190}
191
192impl CachedTargetContract {
193    fn into_targeted_contract(self) -> TargetedContract {
194        TargetedContract::from_parts(
195            self.identifier,
196            self.abi,
197            self.targeted_functions,
198            Vec::new(),
199            self.storage_layout,
200            self.event_lookup,
201        )
202    }
203}
204
205/// A collection of contracts identified as targets for invariant testing.
206#[derive(Clone, Debug, Default)]
207pub struct TargetedContracts {
208    /// The inner map of targeted contracts.
209    pub inner: BTreeMap<Address, TargetedContract>,
210}
211
212impl TargetedContracts {
213    /// Returns a new `TargetedContracts` instance.
214    pub fn new() -> Self {
215        Self::default()
216    }
217
218    /// Returns fuzzed contract and fuzzed function from address and provided calldata.
219    ///
220    /// Used to decode return values and logs in order to add values into fuzz dictionary.
221    pub fn fuzzed_artifacts(
222        &self,
223        tx: &BasicTxDetails,
224    ) -> (Option<&TargetedContract>, Option<&Function>) {
225        match self.inner.get(&tx.call_details.target) {
226            Some(c) => {
227                let function = tx
228                    .call_details
229                    .calldata
230                    .get(..4)
231                    .and_then(|selector| <[u8; 4]>::try_from(selector).ok())
232                    .map(Selector::from)
233                    .and_then(|selector| c.function_by_selector(selector));
234                (Some(c), function)
235            }
236            None => (None, None),
237        }
238    }
239
240    /// Returns flatten target contract address and functions to be fuzzed.
241    /// Includes contract targeted functions if specified, else all mutable contract functions.
242    pub fn fuzzed_functions(&self) -> impl Iterator<Item = (&Address, &Function)> {
243        self.inner
244            .iter()
245            .filter(|(_, c)| !c.abi.functions.is_empty())
246            .flat_map(|(contract, c)| c.abi_fuzzed_functions().map(move |f| (contract, f)))
247    }
248
249    /// Returns whether the given transaction can be replayed or not with known contracts.
250    pub fn can_replay(&self, tx: &BasicTxDetails) -> bool {
251        match self.inner.get(&tx.call_details.target) {
252            Some(c) => tx
253                .call_details
254                .calldata
255                .get(..4)
256                .and_then(|selector| <[u8; 4]>::try_from(selector).ok())
257                .map(Selector::from)
258                .is_some_and(|selector| c.fuzzed_function_by_selector(selector).is_some()),
259            None => false,
260        }
261    }
262
263    /// Identifies fuzzed contract and function based on given tx details and returns unique metric
264    /// key composed from contract identifier and function name.
265    pub fn fuzzed_metric_key(&self, tx: &BasicTxDetails) -> Option<String> {
266        tx.call_details
267            .calldata
268            .get(..4)
269            .and_then(|selector| <[u8; 4]>::try_from(selector).ok())
270            .map(Selector::from)
271            .and_then(|selector| {
272                self.fuzzed_metric_key_for_selector(tx.call_details.target, selector)
273            })
274    }
275
276    /// Identifies fuzzed contract and function from target and selector and returns unique metric
277    /// key composed from contract identifier and function name.
278    pub fn fuzzed_metric_key_for_selector(
279        &self,
280        target: Address,
281        selector: Selector,
282    ) -> Option<String> {
283        self.inner.get(&target).and_then(|contract| {
284            contract
285                .function_by_selector(selector)
286                .map(|function| format!("{}.{}", contract.identifier.as_str(), function.name))
287        })
288    }
289}
290
291impl std::ops::Deref for TargetedContracts {
292    type Target = BTreeMap<Address, TargetedContract>;
293
294    fn deref(&self) -> &Self::Target {
295        &self.inner
296    }
297}
298
299impl std::ops::DerefMut for TargetedContracts {
300    fn deref_mut(&mut self) -> &mut Self::Target {
301        &mut self.inner
302    }
303}
304
305/// A contract identified as target for invariant testing.
306#[derive(Clone, Debug)]
307pub struct TargetedContract {
308    /// The contract identifier. This is only used in error messages.
309    pub identifier: String,
310    /// The contract's ABI.
311    pub abi: JsonAbi,
312    /// The targeted functions of the contract.
313    pub targeted_functions: Vec<Function>,
314    /// The excluded functions of the contract.
315    pub excluded_functions: Vec<Function>,
316    /// The contract's storage layout, if available.
317    pub storage_layout: Option<Arc<StorageLayout>>,
318    /// Contract events indexed by topic0 and indexed-topic count for log dictionary decoding.
319    pub event_lookup: Arc<TargetedContractEvents>,
320    functions_by_selector: FunctionLookup,
321    fuzzed_functions_by_selector: FunctionLookup,
322}
323
324impl TargetedContract {
325    /// Returns a new `TargetedContract` instance.
326    pub fn new(identifier: String, abi: JsonAbi) -> Self {
327        let event_lookup = Arc::new(TargetedContractEvents::new(&abi));
328        Self::from_parts(identifier, abi, Vec::new(), Vec::new(), None, event_lookup)
329    }
330
331    fn from_parts(
332        identifier: String,
333        abi: JsonAbi,
334        targeted_functions: Vec<Function>,
335        excluded_functions: Vec<Function>,
336        storage_layout: Option<Arc<StorageLayout>>,
337        event_lookup: Arc<TargetedContractEvents>,
338    ) -> Self {
339        let mut contract = Self {
340            identifier,
341            abi,
342            targeted_functions,
343            excluded_functions,
344            storage_layout,
345            event_lookup,
346            functions_by_selector: FunctionLookup::default(),
347            fuzzed_functions_by_selector: FunctionLookup::default(),
348        };
349        contract.rebuild_function_lookups();
350        contract
351    }
352
353    /// Determines contract storage layout from project contracts. Needs `storageLayout` to be
354    /// enabled as extra output in project configuration.
355    pub fn with_project_contracts(mut self, project_contracts: &ContractsByArtifact) -> Self {
356        if let Some((src, name)) = self.identifier.split_once(':')
357            && let Some((_, contract_data)) = project_contracts.iter().find(|(artifact, _)| {
358                artifact.name == name && artifact.source.as_path().ends_with(src)
359            })
360        {
361            self.storage_layout = contract_data.storage_layout.as_ref().map(Arc::clone);
362        }
363        self
364    }
365
366    /// Helper to retrieve functions to fuzz for specified abi.
367    /// Returns specified targeted functions if any, else mutable abi functions, always skipping
368    /// functions marked as excluded.
369    pub fn abi_fuzzed_functions(&self) -> impl Iterator<Item = &Function> {
370        if self.targeted_functions.is_empty() {
371            Either::Right(self.abi.functions().filter(|&func| {
372                !matches!(
373                    func.state_mutability,
374                    alloy_json_abi::StateMutability::Pure | alloy_json_abi::StateMutability::View
375                ) && !self.excluded_functions.contains(func)
376            }))
377        } else {
378            Either::Left(
379                self.targeted_functions
380                    .iter()
381                    .filter(|func| !self.excluded_functions.contains(func)),
382            )
383        }
384    }
385
386    pub fn rebuild_function_lookups(&mut self) {
387        let functions_by_selector =
388            self.abi.functions().fold(FunctionLookup::default(), |mut functions, function| {
389                functions.entry(function.selector()).or_insert_with(|| function.clone());
390                functions
391            });
392        let fuzzed_functions_by_selector = self.abi_fuzzed_functions().fold(
393            FunctionLookup::default(),
394            |mut functions, function| {
395                functions.entry(function.selector()).or_insert_with(|| function.clone());
396                functions
397            },
398        );
399        self.functions_by_selector = functions_by_selector;
400        self.fuzzed_functions_by_selector = fuzzed_functions_by_selector;
401    }
402
403    /// Returns any ABI function for the given selector.
404    pub fn function_by_selector(&self, selector: Selector) -> Option<&Function> {
405        self.functions_by_selector.get(&selector)
406    }
407
408    /// Returns a fuzzable function for the given selector.
409    pub fn fuzzed_function_by_selector(&self, selector: Selector) -> Option<&Function> {
410        self.fuzzed_functions_by_selector.get(&selector)
411    }
412
413    /// Returns the function for the given selector.
414    pub fn get_function(&self, selector: Selector) -> eyre::Result<&Function> {
415        self.function_by_selector(selector)
416            .ok_or_else(|| eyre::eyre!("{} does not have the selector {selector}", self.identifier))
417    }
418
419    /// Adds the specified selectors to the targeted functions.
420    pub fn add_selectors(
421        &mut self,
422        selectors: impl IntoIterator<Item = Selector>,
423        should_exclude: bool,
424    ) -> eyre::Result<()> {
425        for selector in selectors {
426            if should_exclude {
427                self.excluded_functions.push(self.get_function(selector)?.clone());
428            } else {
429                self.targeted_functions.push(self.get_function(selector)?.clone());
430            }
431        }
432        self.rebuild_function_lookups();
433        Ok(())
434    }
435}
436
437/// Events for a targeted contract, pre-indexed for log dictionary decoding.
438#[derive(Clone, Debug, Default)]
439pub struct TargetedContractEvents {
440    by_topic: HashMap<(B256, usize), Vec<TargetedContractEvent>>,
441    anonymous: Vec<TargetedContractEvent>,
442}
443
444impl TargetedContractEvents {
445    fn new(abi: &JsonAbi) -> Self {
446        let mut events = Self::default();
447        for (order, event) in abi.events().enumerate() {
448            let event = TargetedContractEvent { order, event: event.clone() };
449            if event.event.anonymous {
450                events.anonymous.push(event);
451            } else {
452                let indexed_count = event.event.inputs.iter().filter(|input| input.indexed).count();
453                events
454                    .by_topic
455                    .entry((event.event.selector(), indexed_count))
456                    .or_default()
457                    .push(event);
458            }
459        }
460        events
461    }
462
463    pub fn by_topic(
464        &self,
465        selector: &B256,
466        indexed_count: usize,
467    ) -> Option<&[TargetedContractEvent]> {
468        self.by_topic.get(&(*selector, indexed_count)).map(Vec::as_slice)
469    }
470
471    pub fn anonymous(&self) -> &[TargetedContractEvent] {
472        &self.anonymous
473    }
474}
475
476/// Event with its flattened ABI order for preserving log decode priority.
477#[derive(Clone, Debug)]
478pub struct TargetedContractEvent {
479    event: Event,
480    order: usize,
481}
482
483impl TargetedContractEvent {
484    pub const fn order(&self) -> usize {
485        self.order
486    }
487
488    pub const fn event(&self) -> &Event {
489        &self.event
490    }
491}
492
493/// Test contract which is testing its invariants.
494#[derive(Clone, Debug)]
495pub struct InvariantContract<'a> {
496    /// Address of the test contract.
497    pub address: Address,
498    /// Name of the test contract.
499    pub name: &'a str,
500    /// Invariant functions to assert against, paired with their `fail_on_revert` config.
501    /// Stored in **source declaration order** so failure-event attribution and report
502    /// rendering match user expectations.
503    pub invariant_fns: Vec<(&'a Function, bool)>,
504    /// Index into [`Self::invariant_fns`] of the stable campaign anchor. Boolean invariant
505    /// suites use a deterministic contract-local anchor so test filters do not affect
506    /// corpus/failure namespaces.
507    pub anchor_idx: usize,
508    /// If true, `afterInvariant` function is called after each invariant run.
509    pub call_after_invariant: bool,
510    /// ABI of the test contract.
511    pub abi: &'a JsonAbi,
512    /// ABI calldata for each invariant function.
513    pub invariant_calldata: Vec<Bytes>,
514}
515
516impl<'a> InvariantContract<'a> {
517    /// Creates a new invariant contract.
518    ///
519    /// Caller must ensure `invariant_fns` is non-empty and `anchor_idx < invariant_fns.len()`.
520    pub fn new(
521        address: Address,
522        name: &'a str,
523        invariant_fns: Vec<(&'a Function, bool)>,
524        anchor_idx: usize,
525        call_after_invariant: bool,
526        abi: &'a JsonAbi,
527    ) -> Self {
528        let invariant_calldata =
529            invariant_fns.iter().map(|(func, _)| func.selector().to_vec().into()).collect();
530        Self {
531            address,
532            name,
533            invariant_fns,
534            anchor_idx,
535            call_after_invariant,
536            abi,
537            invariant_calldata,
538        }
539    }
540
541    /// Returns the stable campaign anchor.
542    pub fn anchor(&self) -> &'a Function {
543        self.invariant_fns[self.anchor_idx].0
544    }
545
546    /// Returns cached calldata for the invariant at `idx`.
547    pub fn invariant_calldata(&self, idx: usize) -> Bytes {
548        self.invariant_calldata[idx].clone()
549    }
550
551    /// Returns cached calldata for the stable campaign anchor.
552    pub fn anchor_calldata(&self) -> Bytes {
553        self.invariant_calldata(self.anchor_idx)
554    }
555
556    /// Returns true if this is an optimization mode invariant (returns int256).
557    pub fn is_optimization(&self) -> bool {
558        is_optimization_invariant(self.anchor())
559    }
560}
561
562/// Settings that determine the validity of a persisted invariant counterexample.
563///
564/// When a counterexample is replayed, it's only valid if the same contracts, selectors,
565/// senders, and fail_on_revert settings are used. Changes to unrelated code (e.g., adding
566/// a log statement) should not invalidate the counterexample.
567#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
568pub struct InvariantSettings {
569    /// Target contracts with their addresses and identifiers.
570    pub target_contracts: BTreeMap<Address, String>,
571    /// Target selectors per contract address.
572    pub target_selectors: BTreeMap<Address, Vec<Selector>>,
573    /// Target senders for the invariant test.
574    pub target_senders: Vec<Address>,
575    /// Excluded senders for the invariant test.
576    pub excluded_senders: Vec<Address>,
577    /// Whether the test should fail on any revert.
578    pub fail_on_revert: bool,
579}
580
581impl InvariantSettings {
582    /// Creates new invariant settings from the given components.
583    pub fn new(
584        targeted_contracts: &TargetedContracts,
585        sender_filters: &SenderFilters,
586        fail_on_revert: bool,
587    ) -> Self {
588        let mut target_contracts = BTreeMap::new();
589        let mut target_selectors = BTreeMap::new();
590        for (addr, contract) in &targeted_contracts.inner {
591            target_contracts.insert(*addr, contract.identifier.clone());
592            target_selectors
593                .insert(*addr, contract.abi_fuzzed_functions().map(|f| f.selector()).collect());
594        }
595
596        let mut target_senders = sender_filters.targeted.clone();
597        target_senders.sort_unstable();
598
599        let mut excluded_senders = sender_filters.excluded.clone();
600        excluded_senders.sort_unstable();
601
602        Self {
603            target_contracts,
604            target_selectors,
605            target_senders,
606            excluded_senders,
607            fail_on_revert,
608        }
609    }
610
611    /// Compares these settings with another and returns a description of what changed.
612    /// Returns `None` if the settings are equivalent.
613    pub fn diff(&self, other: &Self) -> Option<String> {
614        let mut changes = Vec::new();
615
616        if self.target_contracts != other.target_contracts {
617            let added: Vec<_> = other
618                .target_contracts
619                .iter()
620                .filter(|(addr, _)| !self.target_contracts.contains_key(*addr))
621                .map(|(_, name)| name.as_str())
622                .collect();
623            let removed: Vec<_> = self
624                .target_contracts
625                .iter()
626                .filter(|(addr, _)| !other.target_contracts.contains_key(*addr))
627                .map(|(_, name)| name.as_str())
628                .collect();
629
630            if !added.is_empty() {
631                changes.push(format!("added target contracts: {}", added.join(", ")));
632            }
633            if !removed.is_empty() {
634                changes.push(format!("removed target contracts: {}", removed.join(", ")));
635            }
636        }
637
638        if self.target_selectors != other.target_selectors {
639            changes.push("target selectors changed".to_string());
640        }
641
642        if self.target_senders != other.target_senders {
643            changes.push("target senders changed".to_string());
644        }
645
646        if self.excluded_senders != other.excluded_senders {
647            changes.push("excluded senders changed".to_string());
648        }
649
650        if self.fail_on_revert != other.fail_on_revert {
651            changes.push(format!(
652                "fail_on_revert changed from {} to {}",
653                self.fail_on_revert, other.fail_on_revert
654            ));
655        }
656
657        if changes.is_empty() { None } else { Some(changes.join(", ")) }
658    }
659}
660
661impl fmt::Display for InvariantSettings {
662    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
663        write!(
664            f,
665            "targets: {}, selectors: {}, senders: {}, excluded: {}, fail_on_revert: {}",
666            self.target_contracts.len(),
667            self.target_selectors.values().map(|v| v.len()).sum::<usize>(),
668            self.target_senders.len(),
669            self.excluded_senders.len(),
670            self.fail_on_revert,
671        )
672    }
673}
674
675/// Returns true if the function returns `int256`, indicating optimization mode.
676/// In optimization mode, the fuzzer maximizes the return value instead of checking invariants.
677pub fn is_optimization_invariant(func: &Function) -> bool {
678    func.outputs.len() == 1 && func.outputs[0].ty == "int256"
679}
680
681#[cfg(test)]
682mod tests {
683    use super::*;
684    use crate::CallDetails;
685    use alloy_primitives::U256;
686    use foundry_compilers::{
687        ArtifactId,
688        artifacts::{
689            BytecodeObject, CompactBytecode, CompactContractBytecode, CompactDeployedBytecode,
690        },
691    };
692    use revm::{bytecode::Bytecode, state::Account};
693
694    fn abi_with_functions(functions: &[&str]) -> JsonAbi {
695        let mut abi = JsonAbi::new();
696        for function in functions {
697            let function = Function::parse(function).unwrap();
698            abi.functions.entry(function.name.clone()).or_default().push(function);
699        }
700        abi
701    }
702
703    fn targeted_contracts_with_functions(target: Address, functions: &[&str]) -> TargetedContracts {
704        let mut targets = TargetedContracts::new();
705        targets.inner.insert(
706            target,
707            TargetedContract::new("Target".to_string(), abi_with_functions(functions)),
708        );
709        targets
710    }
711
712    fn targeted_contracts_with_function(target: Address, function: Function) -> TargetedContracts {
713        let mut abi = JsonAbi::new();
714        abi.functions.entry(function.name.clone()).or_default().push(function);
715        let mut targets = TargetedContracts::new();
716        targets.inner.insert(target, TargetedContract::new("Target".to_string(), abi));
717        targets
718    }
719
720    fn tx(target: Address, calldata: impl Into<Bytes>) -> BasicTxDetails {
721        BasicTxDetails {
722            warp: None,
723            roll: None,
724            sender: Address::ZERO,
725            call_details: CallDetails { target, calldata: calldata.into(), value: None },
726        }
727    }
728
729    fn artifact_id(name: &str) -> ArtifactId {
730        ArtifactId {
731            path: format!("{name}.json").into(),
732            name: name.to_string(),
733            source: format!("{name}.sol").into(),
734            version: "0.8.30".parse().unwrap(),
735            build_id: "test".to_string(),
736            profile: "test".to_string(),
737        }
738    }
739
740    fn project_contracts_with_runtime_code_and_abi(
741        name: &str,
742        code: Bytes,
743        abi: JsonAbi,
744    ) -> ContractsByArtifact {
745        let deployed_bytecode = CompactDeployedBytecode {
746            bytecode: Some(CompactBytecode {
747                object: BytecodeObject::Bytecode(code),
748                source_map: None,
749                link_references: Default::default(),
750            }),
751            immutable_references: Default::default(),
752        };
753        let artifact = CompactContractBytecode {
754            abi: Some(abi),
755            bytecode: None,
756            deployed_bytecode: Some(deployed_bytecode),
757        };
758        ContractsByArtifact::new([(artifact_id(name), artifact)])
759    }
760
761    fn touched_account_with_code(code: Bytes) -> Account {
762        let mut account = Account::default();
763        account.info.balance = U256::ZERO;
764        account.info.code = Some(Bytecode::new_raw(code));
765        account.mark_touch();
766        account
767    }
768
769    #[test]
770    fn targeted_contracts_short_calldata_is_not_replayable_or_decodable() {
771        let target = Address::from([0x42; 20]);
772        let targets = targeted_contracts_with_function(target, Function::parse("foo()").unwrap());
773        let tx = tx(target, vec![0xde, 0xad, 0xbe]);
774
775        assert!(!targets.can_replay(&tx));
776        assert!(targets.fuzzed_artifacts(&tx).1.is_none());
777        assert!(targets.fuzzed_metric_key(&tx).is_none());
778    }
779
780    #[test]
781    fn abi_fuzzed_functions_filters_excluded_targeted_functions() {
782        let allowed = Function::parse("allowed()").unwrap();
783        let excluded = Function::parse("excluded()").unwrap();
784        let mut contract = TargetedContract::new("Target".to_string(), JsonAbi::new());
785        contract.targeted_functions = vec![allowed.clone(), excluded.clone()];
786        contract.excluded_functions = vec![excluded];
787
788        let selectors = contract.abi_fuzzed_functions().map(Function::selector).collect::<Vec<_>>();
789
790        assert_eq!(selectors, vec![allowed.selector()]);
791    }
792
793    #[test]
794    fn targeted_contracts_refresh_selector_lookup_after_filters() {
795        let target = Address::from([0x42; 20]);
796        let foo = Function::parse("foo()").unwrap();
797        let bar = Function::parse("bar()").unwrap();
798
799        let mut excluded = targeted_contracts_with_functions(target, &["foo()", "bar()"]);
800        excluded.inner.get_mut(&target).unwrap().add_selectors([foo.selector()], true).unwrap();
801        assert!(!excluded.can_replay(&tx(target, foo.selector().to_vec())));
802        assert!(excluded.can_replay(&tx(target, bar.selector().to_vec())));
803        assert_eq!(
804            excluded.fuzzed_artifacts(&tx(target, foo.selector().to_vec())).1.unwrap().name,
805            "foo"
806        );
807
808        let mut targeted = targeted_contracts_with_functions(target, &["foo()", "bar()"]);
809        targeted.inner.get_mut(&target).unwrap().add_selectors([foo.selector()], false).unwrap();
810        assert!(targeted.can_replay(&tx(target, foo.selector().to_vec())));
811        assert!(!targeted.can_replay(&tx(target, bar.selector().to_vec())));
812        assert_eq!(
813            targeted.fuzzed_metric_key_for_selector(target, bar.selector()).unwrap(),
814            "Target.bar"
815        );
816    }
817
818    #[test]
819    fn fuzz_run_identified_contracts_cache_fuzzed_functions_in_target_order() {
820        let first = Address::from([0x01; 20]);
821        let second = Address::from([0x02; 20]);
822        let mut targets = targeted_contracts_with_functions(second, &["bar()", "baz(uint256)"]);
823        targets.inner.insert(
824            first,
825            TargetedContract::new(
826                "First".to_string(),
827                abi_with_functions(&["foo()", "qux(address)"]),
828            ),
829        );
830        let expected = targets
831            .fuzzed_functions()
832            .map(|(address, function)| (*address, function.selector()))
833            .collect::<Vec<_>>();
834
835        let identified = FuzzRunIdentifiedContracts::new(targets, true);
836        let actual = identified
837            .fuzzed_functions()
838            .iter()
839            .map(|(address, function)| (*address, function.selector()))
840            .collect::<Vec<_>>();
841
842        assert_eq!(actual, expected);
843    }
844
845    #[test]
846    fn collect_created_contracts_caches_deployed_code_matches() {
847        let existing = Address::from([0x42; 20]);
848        let created = Address::from([0x43; 20]);
849        let setup = Address::from([0x44; 20]);
850        let untouched = Address::from([0x45; 20]);
851        let runtime_code = Bytes::from_static(&[0x60, 0x00, 0x56]);
852        let project_contracts = project_contracts_with_runtime_code_and_abi(
853            "DynamicTarget",
854            runtime_code.clone(),
855            JsonAbi::new(),
856        );
857        let mut targets = TargetedContracts::new();
858        for address in [existing, setup, untouched] {
859            targets.inner.insert(
860                address,
861                TargetedContract::new("AlreadyTargeted".to_string(), JsonAbi::new()),
862            );
863        }
864        let identified = FuzzRunIdentifiedContracts::new(targets, true);
865
866        let mut state_changeset = StateChangeset::default();
867        state_changeset.insert(existing, touched_account_with_code(runtime_code.clone()));
868        state_changeset.insert(setup, touched_account_with_code(runtime_code.clone()));
869        state_changeset.insert(untouched, Account::default());
870        state_changeset.insert(created, touched_account_with_code(runtime_code));
871        let mut created_contracts = Vec::new();
872        let setup_contracts =
873            ContractsByAddress::from([(setup, ("Setup".to_string(), JsonAbi::new()))]);
874
875        identified
876            .collect_created_contracts(
877                &state_changeset,
878                &project_contracts,
879                &setup_contracts,
880                &ArtifactFilters::default(),
881                &mut created_contracts,
882            )
883            .unwrap();
884
885        created_contracts.sort_unstable();
886        assert_eq!(created_contracts, vec![existing, created]);
887        let targets = identified.targets();
888        assert_eq!(targets[&existing].identifier, "DynamicTarget");
889        assert_eq!(targets[&created].identifier, "DynamicTarget");
890        assert_eq!(targets[&setup].identifier, "AlreadyTargeted");
891        assert_eq!(targets[&untouched].identifier, "AlreadyTargeted");
892        drop(targets);
893
894        identified
895            .collect_created_contracts(
896                &state_changeset,
897                &Default::default(),
898                &setup_contracts,
899                &ArtifactFilters::default(),
900                &mut created_contracts,
901            )
902            .unwrap();
903
904        created_contracts.sort_unstable();
905        assert_eq!(created_contracts, vec![existing, existing, created, created]);
906    }
907
908    #[test]
909    fn collect_and_clear_created_contracts_refresh_fuzzed_function_cache() {
910        let existing = Address::from([0x42; 20]);
911        let created = Address::from([0x43; 20]);
912        let runtime_code = Bytes::from_static(&[0x60, 0x00, 0x56]);
913        let project_contracts = project_contracts_with_runtime_code_and_abi(
914            "DynamicTarget",
915            runtime_code.clone(),
916            abi_with_functions(&["dynamic(uint256)"]),
917        );
918        let identified = FuzzRunIdentifiedContracts::new(
919            targeted_contracts_with_functions(existing, &["existing()"]),
920            true,
921        );
922
923        let initial = identified
924            .fuzzed_functions()
925            .iter()
926            .map(|(address, function)| (*address, function.selector()))
927            .collect::<Vec<_>>();
928        assert_eq!(initial, vec![(existing, Function::parse("existing()").unwrap().selector())]);
929        assert_eq!(identified.fuzzed_functions_generation(), 0);
930
931        let mut state_changeset = StateChangeset::default();
932        state_changeset.insert(created, touched_account_with_code(runtime_code));
933        let mut created_contracts = Vec::new();
934
935        identified
936            .collect_created_contracts(
937                &state_changeset,
938                &project_contracts,
939                &ContractsByAddress::default(),
940                &ArtifactFilters::default(),
941                &mut created_contracts,
942            )
943            .unwrap();
944
945        let with_created = identified
946            .fuzzed_functions()
947            .iter()
948            .map(|(address, function)| (*address, function.selector()))
949            .collect::<Vec<_>>();
950        assert_eq!(identified.fuzzed_functions_generation(), 1);
951        assert_eq!(
952            with_created,
953            vec![
954                (existing, Function::parse("existing()").unwrap().selector()),
955                (created, Function::parse("dynamic(uint256)").unwrap().selector()),
956            ]
957        );
958
959        identified.clear_created_contracts(created_contracts);
960        let cleared = identified
961            .fuzzed_functions()
962            .iter()
963            .map(|(address, function)| (*address, function.selector()))
964            .collect::<Vec<_>>();
965        assert_eq!(cleared, initial);
966        assert_eq!(identified.fuzzed_functions_generation(), 2);
967    }
968}