Skip to main content

foundry_common/preprocessor/
deps.rs

1use super::{
2    data::{ContractData, PreprocessorData, deploy_helper_path, reserved_identifiers},
3    span_to_range,
4};
5use crate::fs::{canonicalize_path, normalize_path};
6use alloy_primitives::keccak256;
7use foundry_compilers::{
8    ProjectPathsConfig, Updates,
9    artifacts::{SolcLanguage, remappings::Remapping},
10    project::{NativeDependencies, NativeDependencyState, PreprocessorState},
11};
12use itertools::Itertools;
13use path_slash::PathExt;
14use solar::{
15    ast::{ItemKind, UserDefinableOperator},
16    sema::{
17        Gcx, Hir,
18        builtins::Builtin,
19        hir::{
20            CallArgs, CallArgsKind, CallOptions, Contract, ContractId, ContractKind, Expr,
21            ExprKind, Function, FunctionId, FunctionKind, Modifier, Res, SourceId, StateMutability,
22            Stmt, StmtKind, TypeKind, UsingDirective, UsingEntryKind, Variable, VariableId,
23            Visibility, Visit,
24        },
25        interface::{SourceMap, Symbol, data_structures::Never, source_map::FileName},
26    },
27};
28use std::{
29    collections::{BTreeMap, BTreeSet, HashSet},
30    ops::{ControlFlow, Range},
31    path::{Path, PathBuf},
32};
33
34/// Compiler and source context whose validation would be lost by rewriting construction.
35#[derive(Clone, Copy)]
36pub(super) struct ConstructorContext {
37    pub abi_coder_v2: bool,
38    pub supports_create2: bool,
39}
40
41impl ConstructorContext {
42    pub(super) fn for_source(mut self, gcx: Gcx<'_>, source: SourceId) -> Self {
43        let ast = gcx
44            .sources
45            .get_file(&gcx.hir.source(source).file)
46            .and_then(|(_, source)| source.ast.as_ref());
47        let Some(ast) = ast else {
48            self.abi_coder_v2 = false;
49            return self;
50        };
51        for item in ast.items.iter() {
52            if let ItemKind::Pragma(pragma) = &item.kind
53                && let Some((name, Some(value))) = pragma.tokens.as_name_and_value()
54            {
55                match (name.as_str(), value.as_str()) {
56                    ("abicoder", "v1") => {
57                        self.abi_coder_v2 = false;
58                        return self;
59                    }
60                    ("abicoder", "v2") | ("experimental", "ABIEncoderV2") => {
61                        self.abi_coder_v2 = true;
62                    }
63                    _ => {}
64                }
65            }
66        }
67        self
68    }
69}
70
71/// Holds data about referenced source contracts and bytecode dependencies.
72pub(crate) struct PreprocessorDependencies {
73    // Mapping contract id to preprocess -> contract bytecode dependencies.
74    pub preprocessed_contracts: BTreeMap<ContractId, Vec<BytecodeDependency>>,
75    // Referenced contract ids.
76    pub referenced_contracts: HashSet<ContractId>,
77}
78
79impl PreprocessorDependencies {
80    #[allow(clippy::too_many_arguments)]
81    pub fn new(
82        gcx: Gcx<'_>,
83        constructor_context: ConstructorContext,
84        paths: &[PathBuf],
85        script_paths: &HashSet<PathBuf>,
86        project_paths: &ProjectPathsConfig<SolcLanguage>,
87        source_units: &[PathBuf],
88        mocks: &mut HashSet<PathBuf>,
89        preprocessor_state: &mut PreprocessorState,
90    ) -> Self {
91        let relative_paths = project_paths.paths_relative();
92        let src_dir = &relative_paths.sources;
93        let root_dir = &project_paths.root;
94        let remappings = &project_paths.remappings;
95        let mut preprocessed_contracts = BTreeMap::new();
96        let mut referenced_contracts = HashSet::new();
97        let mut current_mocks = HashSet::new();
98        let mut current_native_dependencies = NativeDependencies::new();
99        let candidate_files =
100            paths.iter().map(|path| normalize_path(&root_dir.join(path))).collect::<HashSet<_>>();
101        let mut conservative_files = HashSet::new();
102        let global_using_dependencies = using_dependency_sources(
103            gcx,
104            gcx.hir
105                .source_ids()
106                .flat_map(|id| gcx.hir.source(id).usings)
107                .filter(|directive| directive.global),
108        );
109
110        // Helper closure for iterating candidate contracts to preprocess (tests and scripts).
111        let candidate_contracts = || {
112            gcx.hir.contract_ids().filter_map(|id| {
113                let contract = gcx.hir.contract(id);
114                let source = gcx.hir.source(contract.source);
115                let FileName::Real(path) = &source.file.name else {
116                    return None;
117                };
118
119                if !paths.contains(path) {
120                    trace!("{} is not test or script", path.display());
121                    return None;
122                }
123
124                Some((id, contract, source, path))
125            })
126        };
127
128        // An internal call can observe return data left by another function. Analyze the whole
129        // owning contract and its reachable helpers before deciding which scopes can be rewritten.
130        let mut native_return_data_contracts = HashSet::new();
131        for (id, _, _, _) in candidate_contracts() {
132            let mut observer = ReturnDataObserver::new(gcx);
133            let _ = observer.visit_nested_contract(id);
134            if observer.observes_return_data {
135                native_return_data_contracts.insert(id);
136                native_return_data_contracts.extend(observer.contracts);
137            }
138        }
139
140        // Collect current mocks.
141        for (_, contract, _, path) in candidate_contracts() {
142            let full_path = normalize_path(&root_dir.join(path));
143            let mut inherited_dependencies = BTreeSet::new();
144            let mut inherits_source_contract = false;
145            for base_id in contract.linearized_bases {
146                let base = gcx.hir.contract(*base_id);
147                if let FileName::Real(base_path) = &gcx.hir.source(base.source).file.name {
148                    let base_path = normalize_path(&root_dir.join(base_path));
149                    let is_source_contract = is_path_in_dir(&base_path, src_dir, root_dir);
150                    inherits_source_contract |= is_source_contract;
151                    if base_path != full_path {
152                        inherited_dependencies.insert(base_path);
153                    }
154                }
155            }
156            if inherits_source_contract {
157                trace!("found mock contract {}", full_path.display());
158                current_mocks.insert(full_path.clone());
159            }
160            if !inherited_dependencies.is_empty() {
161                add_native_dependencies(
162                    &mut current_native_dependencies,
163                    full_path,
164                    inherited_dependencies,
165                );
166            }
167        }
168
169        // Collect dependencies for non-mock test/script contracts.
170        for (contract_id, contract, source, path) in candidate_contracts() {
171            let full_path = normalize_path(&root_dir.join(path));
172
173            if current_mocks.contains(&full_path) {
174                trace!("{} is a mock, skipping", path.display());
175                continue;
176            }
177
178            // Treat the contract as a script when its file lives under the configured script
179            // directory, or when it inherits from a `Script` base (forge-std). The inheritance
180            // check covers atypical layouts where script contracts are placed under `src/`.
181            let is_script = script_paths.contains(path)
182                || contract
183                    .linearized_bases
184                    .iter()
185                    .skip(1)
186                    .any(|base_id| gcx.hir.contract(*base_id).name.as_str() == "Script");
187            let mut deps_collector = BytecodeDependencyCollector::new(
188                gcx,
189                contract_id,
190                path,
191                src_dir,
192                root_dir,
193                is_script,
194                constructor_context.for_source(gcx, contract.source),
195            );
196            deps_collector.preserve_native_bytecode =
197                native_return_data_contracts.contains(&contract_id);
198            let mut using_dependencies = global_using_dependencies.clone();
199            using_dependencies.extend(using_dependency_sources(
200                gcx,
201                source.usings.iter().chain(
202                    contract
203                        .linearized_bases
204                        .iter()
205                        .flat_map(|base_id| gcx.hir.contract(*base_id).usings),
206                ),
207            ));
208            for source_id in using_dependencies {
209                deps_collector.collect_source_dependencies(source_id);
210            }
211            // Analyze current contract.
212            let _ = deps_collector.walk_contract(contract);
213            if deps_collector.has_unresolved_native_dependency {
214                conservative_files.insert(full_path.clone());
215            }
216            if !deps_collector.native_dependencies.is_empty() {
217                add_native_dependencies(
218                    &mut current_native_dependencies,
219                    full_path.clone(),
220                    deps_collector.native_dependencies,
221                );
222            }
223            deps_collector.dependencies.retain(|dependency| {
224                let dependency_id = dependency.referenced_contract;
225                let dependency = gcx.hir.contract(dependency_id);
226                let dependency_source = gcx.hir.source(dependency.source);
227                let FileName::Real(dependency_path) = &dependency_source.file.name else {
228                    conservative_files.insert(full_path.clone());
229                    return false;
230                };
231                let has_constructor_args = dependency
232                    .ctor
233                    .is_some_and(|ctor_id| !gcx.hir.function(ctor_id).parameters.is_empty());
234                if can_rewrite(
235                    dependency_path,
236                    path,
237                    root_dir,
238                    source_units,
239                    remappings,
240                    has_constructor_args,
241                    dependency.name.as_str(),
242                ) {
243                    true
244                } else {
245                    add_native_dependencies(
246                        &mut current_native_dependencies,
247                        full_path.clone(),
248                        [normalize_path(&root_dir.join(dependency_path))],
249                    );
250                    false
251                }
252            });
253            // Ignore empty test contracts declared in source files with other contracts.
254            if !deps_collector.dependencies.is_empty() {
255                preprocessed_contracts.insert(contract_id, deps_collector.dependencies);
256            }
257        }
258
259        for file in conservative_files {
260            current_native_dependencies.insert(file, NativeDependencyState::Conservative);
261        }
262
263        // Replace classifications only for files examined in this compiler job. This clears stale
264        // mocks after a file is refactored while preserving fallback state across narrower jobs.
265        for file in candidate_files {
266            let state = current_native_dependencies.remove(&file);
267            if preprocessor_state.update(file.clone(), state) {
268                mocks.remove(&file);
269            }
270        }
271        mocks.extend(current_mocks);
272
273        for dependencies in preprocessed_contracts.values() {
274            referenced_contracts.extend(dependencies.iter().map(|dep| dep.referenced_contract));
275        }
276
277        Self { preprocessed_contracts, referenced_contracts }
278    }
279}
280
281/// Adds exact dependency paths unless the source is already classified conservatively.
282fn add_native_dependencies(
283    dependencies: &mut NativeDependencies,
284    file: PathBuf,
285    incoming: impl IntoIterator<Item = PathBuf>,
286) {
287    match dependencies.entry(file) {
288        std::collections::btree_map::Entry::Vacant(entry) => {
289            entry.insert(NativeDependencyState::Known(incoming.into_iter().collect()));
290        }
291        std::collections::btree_map::Entry::Occupied(entry) => {
292            if let NativeDependencyState::Known(dependencies) = entry.into_mut() {
293                dependencies.extend(incoming);
294            }
295        }
296    }
297}
298
299/// Returns sources whose code can be embedded through the given `using for` directives.
300fn using_dependency_sources<'gcx>(
301    gcx: Gcx<'gcx>,
302    directives: impl IntoIterator<Item = &'gcx UsingDirective<'gcx>>,
303) -> HashSet<SourceId> {
304    directives
305        .into_iter()
306        .flat_map(|directive| directive.entries)
307        .flat_map(|entry| match entry.kind {
308            UsingEntryKind::Library(contract_id) => vec![gcx.hir.contract(contract_id).source],
309            UsingEntryKind::Functions(function_ids) => function_ids
310                .iter()
311                .map(|function_id| gcx.hir.function(*function_id).source)
312                .collect(),
313            UsingEntryKind::Err(_) => Vec::new(),
314        })
315        .collect()
316}
317
318/// Represents a bytecode dependency kind.
319#[derive(Debug)]
320enum BytecodeDependencyKind {
321    /// `type(Contract).creationCode`
322    CreationCode,
323    /// `type(Contract).runtimeCode`.
324    RuntimeCode,
325    /// `new Contract`.
326    New {
327        /// Contract name.
328        name: String,
329        /// Constructor args length.
330        args_length: usize,
331        /// Constructor call args offset.
332        call_args_offset: usize,
333        /// `msg.value` (if any) used when creating contract.
334        value: Option<String>,
335        /// `salt` (if any) used when creating contract.
336        salt: Option<String>,
337    },
338}
339
340/// Represents a single bytecode dependency.
341#[derive(Debug)]
342pub(crate) struct BytecodeDependency {
343    /// Dependency kind.
344    kind: BytecodeDependencyKind,
345    /// Source map location of this dependency.
346    loc: Range<usize>,
347    /// HIR id of referenced contract.
348    referenced_contract: ContractId,
349    /// The original expression must reach Solc to preserve its validation.
350    preserve_native: bool,
351}
352
353/// Walks over contract HIR and collects [`BytecodeDependency`]s and referenced contracts.
354struct BytecodeDependencyCollector<'gcx, 'src> {
355    /// Source map, used for determining contract item locations.
356    gcx: Gcx<'gcx>,
357    /// Contract whose lexically owned bytecode references may be rewritten.
358    owner_contract: ContractId,
359    /// Constructor validation context of the source being rewritten.
360    constructor_context: ConstructorContext,
361    /// Source path of the current contract.
362    source_path: PathBuf,
363    /// Project source dir, used to determine if referenced contract is a source contract.
364    src_dir: &'src Path,
365    /// Project root, used to compare relative and absolute source paths.
366    root_dir: &'src Path,
367    /// Whether the contract being analyzed lives in a script file.
368    /// Script bytecode references must not be rewritten: native script CREATE/CREATE2 frames
369    /// are handled by the script execution inspector, and `type(Contract).creationCode` must keep
370    /// its native mutability semantics.
371    is_script: bool,
372    /// Whether `type(Contract).creationCode` should keep native Solidity semantics.
373    preserve_native_creation_code: bool,
374    /// Whether bytecode references are being visited outside the owning contract's rewrite scope.
375    preserve_native_bytecode: bool,
376    /// Dependencies collected for current contract.
377    dependencies: Vec<BytecodeDependency>,
378    /// Dependencies that cannot be rewritten and remain embedded in the importer bytecode.
379    native_dependencies: BTreeSet<PathBuf>,
380    /// Whether a native dependency could not be assigned a stable filesystem identity.
381    has_unresolved_native_dependency: bool,
382    /// Functions followed while finding transitively embedded implementation code.
383    visited_functions: HashSet<FunctionId>,
384    /// Imported sources already classified as native dependencies.
385    visited_sources: HashSet<SourceId>,
386    /// Constants followed while finding embedded initializers, including aliases and cycles.
387    visited_variables: HashSet<VariableId>,
388}
389
390impl<'gcx, 'src> BytecodeDependencyCollector<'gcx, 'src> {
391    fn new(
392        gcx: Gcx<'gcx>,
393        owner_contract: ContractId,
394        source_path: &Path,
395        src_dir: &'src Path,
396        root_dir: &'src Path,
397        is_script: bool,
398        constructor_context: ConstructorContext,
399    ) -> Self {
400        Self {
401            gcx,
402            owner_contract,
403            constructor_context,
404            source_path: normalize_path(&root_dir.join(source_path)),
405            src_dir,
406            root_dir,
407            is_script,
408            preserve_native_creation_code: false,
409            preserve_native_bytecode: false,
410            dependencies: vec![],
411            native_dependencies: BTreeSet::new(),
412            has_unresolved_native_dependency: false,
413            visited_functions: HashSet::new(),
414            visited_sources: HashSet::new(),
415            visited_variables: HashSet::new(),
416        }
417    }
418
419    /// Visits embedded implementation code without allowing edits outside the owning contract.
420    fn collect_function_dependency(&mut self, function_id: FunctionId) {
421        let function = self.gcx.hir.function(function_id);
422        if function.contract == Some(self.owner_contract)
423            || !self.visited_functions.insert(function_id)
424        {
425            return;
426        }
427
428        let source = self.gcx.hir.source(function.source);
429        if let FileName::Real(path) = &source.file.name {
430            let path = normalize_path(&self.root_dir.join(path));
431            if path != self.source_path {
432                self.native_dependencies.insert(path);
433            }
434        } else {
435            self.has_unresolved_native_dependency = true;
436        }
437
438        let previous = self.preserve_native_bytecode;
439        self.preserve_native_bytecode = true;
440        let _ = self.visit_function(function);
441        self.preserve_native_bytecode = previous;
442    }
443
444    /// Visits an expression for dependencies without rewriting within its source range.
445    fn collect_native_expr(&mut self, expr: &'gcx Expr<'gcx>) {
446        let previous = self.preserve_native_bytecode;
447        self.preserve_native_bytecode = true;
448        let _ = self.visit_expr(expr);
449        self.preserve_native_bytecode = previous;
450    }
451
452    /// Records a source containing embedded code and all of its transitive imports.
453    fn collect_source_dependencies(&mut self, source_id: SourceId) {
454        if !self.visited_sources.insert(source_id) {
455            return;
456        }
457        let source = self.gcx.hir.source(source_id);
458        if let FileName::Real(path) = &source.file.name {
459            let path = normalize_path(&self.root_dir.join(path));
460            if path != self.source_path {
461                self.native_dependencies.insert(path);
462            }
463        } else {
464            self.has_unresolved_native_dependency = true;
465        }
466        for &(_, imported_source) in source.imports {
467            self.collect_source_dependencies(imported_source);
468        }
469    }
470
471    /// Classifies a bytecode dependency as rewritable or native.
472    fn collect_dependency(&mut self, dependency: BytecodeDependency) {
473        let contract = self.gcx.hir.contract(dependency.referenced_contract);
474        let source = self.gcx.hir.source(contract.source);
475        let FileName::Real(path) = &source.file.name else {
476            self.has_unresolved_native_dependency = true;
477            return;
478        };
479        let native_path = normalize_path(&self.root_dir.join(path));
480
481        if self.preserve_native_bytecode || dependency.preserve_native {
482            self.native_dependencies.insert(native_path);
483            return;
484        }
485
486        if matches!(&dependency.kind, BytecodeDependencyKind::RuntimeCode) {
487            self.native_dependencies.insert(native_path);
488            return;
489        }
490
491        // Script bytecode references must not be rewritten. See field doc on `is_script`.
492        if self.is_script {
493            match &dependency.kind {
494                BytecodeDependencyKind::CreationCode | BytecodeDependencyKind::RuntimeCode => {
495                    trace!("skip creationCode in script");
496                    self.native_dependencies.insert(native_path);
497                    return;
498                }
499                BytecodeDependencyKind::New { .. } => {
500                    trace!("skip new-expression in script");
501                    self.native_dependencies.insert(native_path);
502                    return;
503                }
504            }
505        }
506
507        // `type(Contract).creationCode` has native `pure` semantics. Rewriting it to a `view`
508        // cheatcode call would make valid pure functions fail to compile.
509        if self.preserve_native_creation_code
510            && matches!(&dependency.kind, BytecodeDependencyKind::CreationCode)
511        {
512            trace!("skip creationCode in native creationCode context");
513            self.native_dependencies.insert(native_path);
514            return;
515        }
516
517        let has_constructor_args = contract
518            .ctor
519            .is_some_and(|ctor_id| !self.gcx.hir.function(ctor_id).parameters.is_empty());
520        // Solidity only permits a custom layout on the most-derived contract, so the generated
521        // constructor helper cannot inherit a target that declares one; keep this dependency
522        // native.
523        if contract.layout.is_some() && has_constructor_args {
524            trace!("skip dependency on custom-layout contract");
525            self.native_dependencies.insert(native_path);
526            return;
527        }
528
529        // Constructor parameter types are copied into a derived helper contract. Private
530        // constants used as array dimensions are not accessible in that scope.
531        if constructor_uses_private_constants(self.gcx, contract) {
532            self.native_dependencies.insert(native_path);
533            return;
534        }
535
536        // Remapped imports can have absolute or symlinked paths, while compiler input paths are
537        // relative and configured source directories can be canonicalized.
538        if !is_path_in_dir(path, self.src_dir, self.root_dir) {
539            let path = path.display();
540            trace!("keep external dependency {path} native");
541            self.native_dependencies.insert(native_path);
542            return;
543        }
544
545        self.dependencies.push(dependency);
546    }
547
548    /// Follows constants whose initializer can embed code from another source.
549    fn collect_variable_dependency(&mut self, id: VariableId) {
550        let variable = self.gcx.hir.variable(id);
551        if !variable.is_constant() || !self.visited_variables.insert(id) {
552            return;
553        }
554        if let FileName::Real(path) = &self.gcx.hir.source(variable.source).file.name {
555            let path = normalize_path(&self.root_dir.join(path));
556            if path != self.source_path {
557                self.native_dependencies.insert(path);
558            }
559        } else {
560            self.has_unresolved_native_dependency = true;
561        }
562        if let Some(initializer) = variable.initializer {
563            self.collect_native_expr(initializer);
564        }
565    }
566}
567
568/// Returns whether constructor parameter types reference private constants.
569fn constructor_uses_private_constants(gcx: Gcx<'_>, contract: &Contract<'_>) -> bool {
570    contract.ctor.is_some_and(|ctor| {
571        gcx.hir.function(ctor).parameters.iter().any(|&param| {
572            gcx.hir
573                .variable(param)
574                .ty
575                .visit(&gcx.hir, &mut |ty| {
576                    if let TypeKind::Array(array) = &ty.kind
577                        && let Some(size) = array.size
578                    {
579                        size.visit(&mut |expr| {
580                            if gcx.resolved_variable(expr).is_some_and(|var| {
581                                gcx.hir.variable(var).visibility == Some(Visibility::Private)
582                            }) {
583                                return ControlFlow::Break(());
584                            }
585                            ControlFlow::Continue(())
586                        })?;
587                    }
588                    ControlFlow::Continue(())
589                })
590                .is_break()
591        })
592    })
593}
594
595/// Returns whether generated helper and artifact references preserve the source-unit identity.
596fn can_rewrite(
597    path: &Path,
598    source_path: &Path,
599    root_dir: &Path,
600    source_units: &[PathBuf],
601    remappings: &[Remapping],
602    has_constructor_args: bool,
603    contract_name: &str,
604) -> bool {
605    let generated_path = path.strip_prefix(root_dir).unwrap_or(path);
606    if !source_units.iter().any(|source_unit| source_unit == generated_path)
607        || source_units.iter().filter(|source_unit| source_unit.ends_with(generated_path)).count()
608            != 1
609    {
610        return false;
611    }
612
613    // Runtime artifact lookup uses the running test's context, which can differ from the source
614    // containing an inherited helper. Any remapping matching the generated path is therefore
615    // unsafe unless every possible runtime context is known.
616    if remappings.iter().any(|remapping| remapping_matches_path(remapping, generated_path)) {
617        return false;
618    }
619
620    if !has_constructor_args {
621        return true;
622    }
623
624    let identifier = keccak256(format!("{}:{contract_name}", generated_path.to_slash_lossy()));
625    let helper_path = deploy_helper_path(&format!("{identifier:x}"), source_units);
626    !remappings.iter().any(|remapping| {
627        // The test imports the generated helper, which in turn imports the dependency.
628        remapping_applies(remapping, &helper_path, source_path, root_dir)
629            || remapping_applies(remapping, generated_path, &helper_path, root_dir)
630    })
631}
632
633/// Returns whether `path` resolves within `dir`, accepting relative, absolute, and symlinked paths.
634fn is_path_in_dir(path: &Path, dir: &Path, root_dir: &Path) -> bool {
635    let path = normalize_path(&root_dir.join(path));
636    let dir = normalize_path(&root_dir.join(dir));
637    path.starts_with(&dir)
638        || canonicalize_path(path)
639            .is_ok_and(|path| canonicalize_path(dir).is_ok_and(|dir| path.starts_with(dir)))
640}
641
642/// Returns whether a generated import would be redirected by `remapping`.
643fn remapping_applies(
644    remapping: &Remapping,
645    import_path: &Path,
646    source_unit: &Path,
647    root_dir: &Path,
648) -> bool {
649    let source_unit = source_unit.strip_prefix(root_dir).unwrap_or(source_unit).to_slash_lossy();
650    remapping
651        .context
652        .as_ref()
653        .is_none_or(|context| source_unit.starts_with(Path::new(context).to_slash_lossy().as_ref()))
654        && remapping_matches_path(remapping, import_path)
655}
656
657/// Returns whether `path` has the string prefix selected by `remapping`.
658fn remapping_matches_path(remapping: &Remapping, path: &Path) -> bool {
659    path.to_slash_lossy().starts_with(&remapping.name)
660}
661
662impl<'gcx> Visit<'gcx> for BytecodeDependencyCollector<'gcx, '_> {
663    type BreakValue = Never;
664
665    fn hir(&self) -> &'gcx Hir<'gcx> {
666        &self.gcx.hir
667    }
668
669    fn visit_function(&mut self, func: &'gcx Function<'gcx>) -> ControlFlow<Self::BreakValue> {
670        let previous = self.preserve_native_creation_code;
671        self.preserve_native_creation_code = previous
672            || func.state_mutability == StateMutability::Pure
673            || matches!(func.kind, FunctionKind::Modifier);
674        self.walk_function(func)?;
675        self.preserve_native_creation_code = previous;
676        ControlFlow::Continue(())
677    }
678
679    fn visit_var(&mut self, var: &'gcx Variable<'gcx>) -> ControlFlow<Self::BreakValue> {
680        let previous = self.preserve_native_creation_code;
681        self.preserve_native_creation_code |= var.is_constant();
682        self.walk_var(var)?;
683        self.preserve_native_creation_code = previous;
684        ControlFlow::Continue(())
685    }
686
687    fn visit_expr(&mut self, expr: &'gcx Expr<'gcx>) -> ControlFlow<Self::BreakValue> {
688        #[allow(clippy::collapsible_match)]
689        match &expr.kind {
690            ExprKind::Ident(resolutions) => {
691                for &resolution in *resolutions {
692                    match resolution {
693                        Res::Namespace(source_id) => {
694                            self.collect_source_dependencies(source_id);
695                        }
696                        Res::Item(item) => {
697                            if let Some(function_id) = item.as_function() {
698                                self.collect_function_dependency(function_id);
699                            } else if let Some(variable_id) = item.as_variable() {
700                                self.collect_variable_dependency(variable_id);
701                            }
702                        }
703                        _ => {}
704                    }
705                }
706            }
707            ExprKind::Call(callee, call_args) => {
708                let (call_expr, named_args) = callee.split_call_options();
709                if let Some(dependency) = handle_call_expr(
710                    self.gcx,
711                    self.constructor_context,
712                    expr,
713                    call_expr,
714                    call_args,
715                    &named_args,
716                ) {
717                    self.collect_dependency(dependency);
718                    // Call options are copied into the replacement expression. Keep their
719                    // bytecode references native so edits cannot overlap the outer replacement.
720                    self.visit_expr(call_expr)?;
721                    if let Some(call_options) = named_args {
722                        for arg in call_options.args {
723                            self.collect_native_expr(&arg.value);
724                        }
725                    }
726                    self.visit_call_args(call_args)?;
727                    return ControlFlow::Continue(());
728                }
729                if let Some(function_id) = self.gcx.resolved_function(call_expr) {
730                    self.collect_function_dependency(function_id);
731                }
732            }
733            ExprKind::Member(member_expr, ident) => {
734                // Solar does not resolve `Library.function` member expressions as functions. The
735                // internal library implementation is embedded in the caller, so classify the
736                // library source and its imports as native dependencies directly.
737                if let ExprKind::Ident(resolutions) = member_expr.kind {
738                    for resolution in resolutions {
739                        if let Res::Item(item) = resolution
740                            && let Some(contract_id) = item.as_contract()
741                        {
742                            let contract = self.gcx.hir.contract(contract_id);
743                            if contract.kind == ContractKind::Library {
744                                self.collect_source_dependencies(contract.source);
745                            }
746                        }
747                    }
748                }
749                if let ExprKind::TypeCall(ty) = &member_expr.kind
750                    && let TypeKind::Custom(contract_id) = &ty.kind
751                    && let Some(contract_id) = contract_id.as_contract()
752                    && let kind = match ident.name.as_str() {
753                        "creationCode" => BytecodeDependencyKind::CreationCode,
754                        "runtimeCode" => BytecodeDependencyKind::RuntimeCode,
755                        _ => return self.walk_expr(expr),
756                    }
757                {
758                    self.collect_dependency(BytecodeDependency {
759                        kind,
760                        loc: span_to_range(self.gcx.sess.source_map(), expr.span),
761                        referenced_contract: contract_id,
762                        preserve_native: false,
763                    });
764                }
765            }
766            _ => {}
767        }
768        self.walk_expr(expr)
769    }
770
771    fn visit_stmt(&mut self, stmt: &'gcx Stmt<'gcx>) -> ControlFlow<Self::BreakValue> {
772        if let StmtKind::Try(stmt_try) = stmt.kind
773            && let ExprKind::Call(callee, ..) = &stmt_try.expr.kind
774            && let (call_expr, _) = callee.split_call_options()
775            && matches!(call_expr.kind, ExprKind::New(_))
776        {
777            // Keep try deployments native: a static-context violation halts the current frame,
778            // whereas a deployment cheatcode revert could be caught by an untyped try. Typed
779            // returns also require native creation to preserve the constructor catch boundary.
780            self.collect_native_expr(&stmt_try.expr);
781
782            for clause in stmt_try.clauses {
783                for &var in clause.args {
784                    self.visit_nested_var(var)?;
785                }
786                for stmt in clause.block.stmts {
787                    self.visit_stmt(stmt)?;
788                }
789            }
790            return ControlFlow::Continue(());
791        }
792        self.walk_stmt(stmt)
793    }
794}
795
796/// Helper function to analyze and extract bytecode dependency from a given call expression.
797fn handle_call_expr(
798    gcx: Gcx<'_>,
799    context: ConstructorContext,
800    parent_expr: &Expr<'_>,
801    call_expr: &Expr<'_>,
802    call_args: &CallArgs<'_>,
803    call_options: &Option<&CallOptions<'_>>,
804) -> Option<BytecodeDependency> {
805    if let ExprKind::New(ty_new) = &call_expr.kind
806        && let TypeKind::Custom(item_id) = ty_new.kind
807        && let Some(contract_id) = item_id.as_contract()
808    {
809        let source_map = gcx.sess.source_map();
810        let name = source_map.span_to_snippet(ty_new.span).ok()?;
811
812        // Calculate the offset to remove call options and parentheses between the new type and
813        // constructor arguments. For example, in `new Counter {value: 333} (address(this))`, the
814        // offset is used to replace `{value: 333} (` with `(`. This also removes closing
815        // parentheses around the callee when no call options are present.
816        let call_args_offset = if call_args.is_empty() {
817            0
818        } else {
819            (call_args.span.lo() - ty_new.span.hi()).to_usize()
820        };
821
822        let args_len = parent_expr.span.hi() - ty_new.span.hi();
823        return Some(BytecodeDependency {
824            kind: BytecodeDependencyKind::New {
825                name,
826                args_length: args_len.to_usize(),
827                call_args_offset,
828                value: named_arg(call_options, "value", source_map),
829                salt: named_arg(call_options, "salt", source_map),
830            },
831            // The HIR callee excludes parentheses, so start at the full call expression.
832            loc: span_to_range(source_map, parent_expr.span.with_hi(call_expr.span.hi())),
833            referenced_contract: contract_id,
834            preserve_native: !valid_constructor_call(
835                gcx,
836                context,
837                contract_id,
838                call_args,
839                call_options,
840            ),
841        });
842    }
843    None
844}
845
846/// Helper function to extract value of a given named arg.
847fn named_arg(
848    call_options: &Option<&CallOptions<'_>>,
849    arg: &str,
850    source_map: &SourceMap,
851) -> Option<String> {
852    call_options
853        .map(|options| options.args)
854        .unwrap_or_default()
855        .iter()
856        .find(|named_arg| named_arg.name.as_str() == arg)
857        .and_then(|named_arg| source_map.span_to_snippet(named_arg.value.span).ok())
858}
859
860/// Goes over all test/script files and replaces bytecode dependencies with cheatcode
861/// invocations.
862///
863/// Try deployments remain native to preserve their constructor failure boundaries.
864pub(crate) fn remove_bytecode_dependencies(
865    gcx: Gcx<'_>,
866    deps: &PreprocessorDependencies,
867    data: &PreprocessorData,
868    root_dir: &Path,
869    use_uint160_return: bool,
870) -> Updates {
871    let mut updates = Updates::default();
872    for (contract_id, deps) in &deps.preprocessed_contracts {
873        let contract = gcx.hir.contract(*contract_id);
874        let source = gcx.hir.source(contract.source);
875        let FileName::Real(path) = &source.file.name else {
876            continue;
877        };
878
879        let reserved_identifiers = reserved_identifiers(gcx, contract.source);
880        let identifier = keccak256(format!(
881            "{}:{}",
882            path.strip_prefix(root_dir).unwrap_or(path).to_slash_lossy(),
883            contract.name,
884        ));
885        let updates = updates.entry(path.clone()).or_default();
886        let mut used_helpers = BTreeSet::new();
887
888        let vm_interface_name =
889            unique_identifier(&reserved_identifiers, format!("VmContractHelper{identifier:x}"));
890        // `address(uint160(uint256(keccak256("hevm cheat code"))))`
891        let vm = format!("{vm_interface_name}(0x7109709ECfa91a80626fF3989D68f67F5b1DD12D)");
892        for dep in deps {
893            let Some(ContractData { artifact, constructor_data, .. }) =
894                data.get(&dep.referenced_contract)
895            else {
896                continue;
897            };
898
899            match &dep.kind {
900                BytecodeDependencyKind::CreationCode => {
901                    // for creation code we need to just call getCode
902                    updates.insert((
903                        dep.loc.start,
904                        dep.loc.end,
905                        format!("{vm}.getCode(\"{artifact}\")"),
906                    ));
907                }
908                BytecodeDependencyKind::RuntimeCode => {
909                    unreachable!("runtimeCode is never rewritten")
910                }
911                BytecodeDependencyKind::New {
912                    name,
913                    args_length,
914                    call_args_offset,
915                    value,
916                    salt,
917                } => {
918                    let mut update = format!("{name}(payable(");
919                    if use_uint160_return {
920                        update.push_str("address(");
921                    }
922                    let closing_seq = if use_uint160_return { "}))))" } else { "})))" };
923                    update.push_str(&format!("{vm}.deployCode({{"));
924                    update.push_str(&format!("_artifact: \"{artifact}\""));
925
926                    if let Some(value) = value {
927                        update.push_str(", ");
928                        update.push_str(&format!("_value: {value}"));
929                    }
930
931                    if let Some(salt) = salt {
932                        update.push_str(", ");
933                        update.push_str(&format!("_salt: {salt}"));
934                    }
935
936                    if let Some(constructor_data) = constructor_data {
937                        // Insert our helper.
938                        used_helpers.insert(dep.referenced_contract);
939                        let helper_contract = unique_identifier(
940                            &reserved_identifiers,
941                            constructor_data.helper_contract.clone(),
942                        );
943                        let encode_function = unique_identifier(
944                            &reserved_identifiers,
945                            constructor_data.encode_function.clone(),
946                        );
947
948                        update.push_str(", ");
949                        update.push_str(&format!(
950                            "_args: {}({}.{}",
951                            encode_function, helper_contract, constructor_data.args_struct,
952                        ));
953                        updates.insert((dep.loc.start, dep.loc.end + call_args_offset, update));
954
955                        updates.insert((
956                            dep.loc.end + args_length,
957                            dep.loc.end + args_length,
958                            format!("){closing_seq}"),
959                        ));
960                    } else {
961                        update.push_str(closing_seq);
962                        updates.insert((dep.loc.start, dep.loc.end + args_length, update));
963                    }
964                }
965            };
966        }
967
968        let helper_imports = used_helpers
969            .into_iter()
970            .map(|id| {
971                let constructor = data[&id].constructor_data.as_ref().unwrap();
972                let helper_contract = &constructor.helper_contract;
973                let encode_function = &constructor.encode_function;
974                let local_helper =
975                    unique_identifier(&reserved_identifiers, helper_contract.clone());
976                let local_encoder =
977                    unique_identifier(&reserved_identifiers, encode_function.clone());
978                let helper_import = import_alias(helper_contract, &local_helper);
979                let encoder_import = import_alias(encode_function, &local_encoder);
980                let helper_path = constructor.helper_path.to_slash_lossy();
981                format!("import {{{helper_import}, {encoder_import}}} from \"{helper_path}\";",)
982            })
983            .sorted()
984            .join("\n");
985        let deployment_return_type = if use_uint160_return { "uint160" } else { "address" };
986        updates.insert((
987            source.file.src.len(),
988            source.file.src.len(),
989            format!(
990                r#"
991{helper_imports}
992
993interface {vm_interface_name} {{
994    function deployCode(string memory _artifact) external returns ({deployment_return_type});
995    function deployCode(string memory _artifact, bytes32 _salt) external returns ({deployment_return_type});
996    function deployCode(string memory _artifact, bytes memory _args) external returns ({deployment_return_type});
997    function deployCode(string memory _artifact, bytes memory _args, bytes32 _salt) external returns ({deployment_return_type});
998    function deployCode(string memory _artifact, uint256 _value) external returns ({deployment_return_type});
999    function deployCode(string memory _artifact, uint256 _value, bytes32 _salt) external returns ({deployment_return_type});
1000    function deployCode(string memory _artifact, bytes memory _args, uint256 _value) external returns ({deployment_return_type});
1001    function deployCode(string memory _artifact, bytes memory _args, uint256 _value, bytes32 _salt) external returns ({deployment_return_type});
1002    function getCode(string memory _artifact) external view returns (bytes memory);
1003}}"#
1004            ),
1005        ));
1006    }
1007    updates
1008}
1009
1010/// Returns an identifier that cannot collide with text in the original source.
1011fn unique_identifier(source: &str, mut identifier: String) -> String {
1012    while source.contains(&identifier) {
1013        identifier.push('_');
1014    }
1015    identifier
1016}
1017
1018fn import_alias(identifier: &str, local: &str) -> String {
1019    if identifier == local { identifier.to_string() } else { format!("{identifier} as {local}") }
1020}
1021
1022/// Checks constraints that disappear when `new` is replaced by a cheatcode call.
1023/// The generated argument struct retains type checks, but not source ABI-coder restrictions.
1024/// Keep parameterized ABI-coder-v1 calls native rather than duplicating Solc ABI validation.
1025fn valid_constructor_call(
1026    gcx: Gcx<'_>,
1027    context: ConstructorContext,
1028    id: ContractId,
1029    args: &CallArgs<'_>,
1030    options: &Option<&CallOptions<'_>>,
1031) -> bool {
1032    let contract = gcx.hir.contract(id);
1033    if contract.kind != ContractKind::Contract {
1034        return false;
1035    }
1036    let constructor = contract.ctor.map(|id| gcx.hir.function(id));
1037    let parameters = constructor.map_or(&[][..], |ctor| ctor.parameters);
1038    if args.len() != parameters.len() || (!parameters.is_empty() && !context.abi_coder_v2) {
1039        return false;
1040    }
1041    if let CallArgsKind::Named(args) = args.kind {
1042        let mut names = HashSet::new();
1043        if args.iter().any(|arg| {
1044            !names.insert(arg.name.name)
1045                || !parameters.iter().any(|id| {
1046                    gcx.hir.variable(*id).name.is_some_and(|name| name.name == arg.name.name)
1047                })
1048        }) {
1049            return false;
1050        }
1051    }
1052    if let Some(options) = options {
1053        let mut names = HashSet::new();
1054        if options.args.iter().any(|arg| {
1055            !names.insert(arg.name.name)
1056                || match arg.name.as_str() {
1057                    "salt" => !context.supports_create2,
1058                    "value" => constructor
1059                        .is_none_or(|ctor| ctor.state_mutability != StateMutability::Payable),
1060                    _ => true,
1061                }
1062        }) {
1063            return false;
1064        }
1065    }
1066    true
1067}
1068
1069/// Finds return-buffer observations in a contract and the helpers it can call internally.
1070struct ReturnDataObserver<'gcx> {
1071    gcx: Gcx<'gcx>,
1072    observes_return_data: bool,
1073    functions: HashSet<FunctionId>,
1074    contracts: HashSet<ContractId>,
1075    visited_contracts: HashSet<ContractId>,
1076    member_functions: HashSet<(FunctionId, Symbol)>,
1077    sources: HashSet<SourceId>,
1078    operator_functions: HashSet<(FunctionId, UserDefinableOperator)>,
1079}
1080
1081impl<'gcx> ReturnDataObserver<'gcx> {
1082    fn new(gcx: Gcx<'gcx>) -> Self {
1083        Self {
1084            gcx,
1085            observes_return_data: false,
1086            functions: HashSet::new(),
1087            contracts: HashSet::new(),
1088            visited_contracts: HashSet::new(),
1089            member_functions: HashSet::new(),
1090            sources: HashSet::new(),
1091            operator_functions: HashSet::new(),
1092        }
1093    }
1094
1095    fn collect_member_functions(&mut self, source: SourceId, contract: Option<ContractId>) {
1096        let bases = contract
1097            .into_iter()
1098            .flat_map(|id| self.gcx.hir.contract(id).linearized_bases)
1099            .copied()
1100            .collect::<Vec<_>>();
1101        for &id in &bases {
1102            self.member_functions.extend(
1103                self.gcx
1104                    .hir
1105                    .contract(id)
1106                    .functions()
1107                    .filter_map(|id| self.gcx.hir.function(id).name.map(|name| (id, name.name))),
1108            );
1109        }
1110        let directives = self
1111            .gcx
1112            .hir
1113            .source(source)
1114            .usings
1115            .iter()
1116            .chain(bases.iter().flat_map(|&id| self.gcx.hir.contract(id).usings))
1117            .chain(
1118                self.gcx
1119                    .hir
1120                    .source_ids()
1121                    .flat_map(|id| self.gcx.hir.source(id).usings)
1122                    .filter(|directive| directive.global),
1123            );
1124        for directive in directives {
1125            for entry in directive.entries {
1126                match entry.kind {
1127                    UsingEntryKind::Library(id) => self.member_functions.extend(
1128                        self.gcx.hir.contract(id).functions().filter_map(|id| {
1129                            self.gcx.hir.function(id).name.map(|name| (id, name.name))
1130                        }),
1131                    ),
1132                    UsingEntryKind::Functions(ids) => {
1133                        if let Some(operator) = entry.operator {
1134                            self.operator_functions.extend(ids.iter().map(|&id| (id, operator)));
1135                        }
1136                        self.member_functions.extend(ids.iter().copied().filter_map(|id| {
1137                            entry
1138                                .name
1139                                .or_else(|| self.gcx.hir.function(id).name.map(|name| name.name))
1140                                .map(|name| (id, name))
1141                        }))
1142                    }
1143                    UsingEntryKind::Err(_) => {}
1144                }
1145            }
1146        }
1147    }
1148}
1149
1150impl<'gcx> Visit<'gcx> for ReturnDataObserver<'gcx> {
1151    type BreakValue = Never;
1152
1153    fn hir(&self) -> &'gcx Hir<'gcx> {
1154        &self.gcx.hir
1155    }
1156
1157    fn visit_nested_source(&mut self, id: SourceId) -> ControlFlow<Self::BreakValue> {
1158        if self.sources.insert(id) {
1159            let source = self.gcx.hir.source(id);
1160            self.walk_nested_source(id)?;
1161            for &(_, id) in source.imports {
1162                self.visit_nested_source(id)?;
1163            }
1164        }
1165        ControlFlow::Continue(())
1166    }
1167
1168    fn visit_nested_contract(&mut self, id: ContractId) -> ControlFlow<Self::BreakValue> {
1169        if self.visited_contracts.insert(id) {
1170            let contract = self.gcx.hir.contract(id);
1171            self.contracts.insert(id);
1172            // State initializers run before functions are visited, but can already call methods
1173            // supplied by contract-scoped using directives.
1174            self.collect_member_functions(contract.source, Some(id));
1175            // An inherited observer can call a derived override that produces return data. Treat
1176            // the complete inheritance hierarchy as one execution scope before rewriting it.
1177            for &base in contract.linearized_bases {
1178                if base != id {
1179                    self.visit_nested_contract(base)?;
1180                }
1181            }
1182            self.walk_contract(contract)?;
1183        }
1184        ControlFlow::Continue(())
1185    }
1186
1187    fn visit_nested_function(&mut self, id: FunctionId) -> ControlFlow<Self::BreakValue> {
1188        if self.functions.insert(id) {
1189            let function = self.gcx.hir.function(id);
1190            self.collect_member_functions(function.source, function.contract);
1191            if let Some(id) = function.contract {
1192                self.contracts.insert(id);
1193            }
1194            self.walk_function(function)?;
1195        }
1196        ControlFlow::Continue(())
1197    }
1198
1199    fn visit_modifier(&mut self, modifier: &'gcx Modifier<'gcx>) -> ControlFlow<Self::BreakValue> {
1200        if let Some(id) = modifier.id.as_function() {
1201            self.visit_nested_function(id)?;
1202        }
1203        self.walk_modifier(modifier)
1204    }
1205
1206    fn visit_expr(&mut self, expr: &'gcx Expr<'gcx>) -> ControlFlow<Self::BreakValue> {
1207        let operator = match &expr.kind {
1208            ExprKind::Unary(op, _) => UserDefinableOperator::from_unop(op.kind),
1209            ExprKind::Binary(_, op, _) => UserDefinableOperator::from_binop(op.kind),
1210            _ => None,
1211        };
1212        if let Some(operator) = operator {
1213            // Type checking has not selected an overload yet, so visit every visible binding.
1214            let functions = self
1215                .operator_functions
1216                .iter()
1217                .filter_map(|&(id, bound)| (bound == operator).then_some(id))
1218                .collect::<Vec<_>>();
1219            for id in functions {
1220                self.visit_nested_function(id)?;
1221            }
1222        }
1223        match &expr.kind {
1224            ExprKind::Ident(resolutions) => {
1225                for resolution in *resolutions {
1226                    match resolution {
1227                        Res::Builtin(Builtin::YulReturndatasize | Builtin::YulReturndatacopy) => {
1228                            self.observes_return_data = true;
1229                        }
1230                        Res::Item(item) => {
1231                            if let Some(id) = item.as_function() {
1232                                self.visit_nested_function(id)?;
1233                            }
1234                        }
1235                        _ => {}
1236                    }
1237                }
1238            }
1239            ExprKind::Call(callee, _) => {
1240                let (callee, _) = callee.split_call_options();
1241                if let Some(id) = self.gcx.resolved_function(callee) {
1242                    self.visit_nested_function(id)?;
1243                }
1244            }
1245            ExprKind::Member(member, name) => {
1246                // Resolve the complete namespace path, including renamed re-exports and library
1247                // members, without visiting unrelated declarations in the imported sources.
1248                let mut names = vec![*name];
1249                let mut root = member.peel_parens();
1250                while let ExprKind::Member(parent, name) = &root.kind {
1251                    names.push(*name);
1252                    root = parent.peel_parens();
1253                }
1254                if let ExprKind::Ident(resolutions) = &root.kind {
1255                    names.reverse();
1256                    for resolution in *resolutions {
1257                        if let Res::Namespace(source) = resolution {
1258                            if let Some(resolutions) =
1259                                self.gcx.source_path_resolutions(&names, *source, None)
1260                                && let Some(targets) = resolutions.last()
1261                            {
1262                                for target in targets {
1263                                    if let Res::Item(item) = target {
1264                                        if let Some(id) = item.as_function() {
1265                                            self.visit_nested_function(id)?;
1266                                        } else if let Some(id) = item.as_variable() {
1267                                            self.visit_nested_var(id)?;
1268                                        }
1269                                    }
1270                                }
1271                            } else {
1272                                // Retain the conservative fallback when name resolution is
1273                                // incomplete.
1274                                self.visit_nested_source(*source)?;
1275                            }
1276                        }
1277                    }
1278                }
1279                // Type checking has not run, so include every visible overload of an inherited
1280                // or using-for method with this name.
1281                let functions = self
1282                    .member_functions
1283                    .iter()
1284                    .copied()
1285                    .filter_map(|(id, attached_name)| (attached_name == name.name).then_some(id))
1286                    .collect::<Vec<_>>();
1287                for id in functions {
1288                    self.visit_nested_function(id)?;
1289                }
1290                if let ExprKind::Ident(resolutions) = member.peel_parens().kind {
1291                    for resolution in resolutions {
1292                        if let Res::Item(item) = resolution
1293                            && let Some(id) = item.as_contract()
1294                            && self.gcx.hir.contract(id).kind == ContractKind::Library
1295                        {
1296                            for id in self.gcx.hir.contract(id).functions() {
1297                                if self
1298                                    .gcx
1299                                    .hir
1300                                    .function(id)
1301                                    .name
1302                                    .is_some_and(|ident| ident.name == name.name)
1303                                {
1304                                    self.visit_nested_function(id)?;
1305                                }
1306                            }
1307                        }
1308                    }
1309                }
1310            }
1311            _ => {}
1312        }
1313        self.walk_expr(expr)
1314    }
1315}