foundry_common/preprocessor/
data.rs

1use super::span_to_range;
2use foundry_compilers::artifacts::{Source, Sources};
3use path_slash::PathExt;
4use solar::sema::{
5    Gcx,
6    hir::{Contract, ContractId},
7    interface::source_map::FileName,
8};
9use std::{
10    collections::{BTreeMap, HashSet},
11    path::{Path, PathBuf},
12};
13
14/// Keeps data about project contracts definitions referenced from tests and scripts.
15/// Contract id -> Contract data definition mapping.
16pub type PreprocessorData = BTreeMap<ContractId, ContractData>;
17
18/// Keeps data about a contract constructor.
19#[derive(Debug)]
20pub struct ContractConstructorData {
21    /// ABI encoded args.
22    pub abi_encode_args: String,
23    /// Constructor struct fields.
24    pub struct_fields: String,
25    /// Generated helper contract identifier.
26    pub helper_contract: String,
27    /// Generated constructor argument struct identifier.
28    pub args_struct: String,
29    /// Generated ABI encoding function identifier.
30    pub encode_function: String,
31    /// Generated helper source-unit path.
32    pub helper_path: PathBuf,
33}
34
35/// Keeps data about a single contract definition.
36#[derive(Debug)]
37pub(crate) struct ContractData {
38    /// Path of the source file.
39    path: PathBuf,
40    /// Name of the contract
41    name: String,
42    /// Constructor parameters, if any.
43    pub constructor_data: Option<ContractConstructorData>,
44    /// Artifact string to pass into cheatcodes.
45    pub artifact: String,
46}
47
48impl ContractData {
49    fn new(
50        gcx: Gcx<'_>,
51        contract_id: ContractId,
52        contract: &Contract<'_>,
53        path: &Path,
54        source: &solar::sema::hir::Source<'_>,
55        reserved_identifiers: &str,
56        source_units: &[PathBuf],
57    ) -> Self {
58        let artifact =
59            solidity_string_content(&format!("{}:{}", path.to_slash_lossy(), contract.name));
60
61        // Process data for contracts with constructor and parameters.
62        let constructor_data = contract
63            .ctor
64            .map(|ctor_id| gcx.hir.function(ctor_id))
65            .filter(|ctor| !ctor.parameters.is_empty())
66            .map(|ctor| {
67                let contract_id = contract_id.index();
68                let mut abi_encode_args = vec![];
69                let mut struct_fields = vec![];
70                let mut arg_index = 0;
71                for param_id in ctor.parameters {
72                    let param = gcx.hir.variable(*param_id);
73                    let loc = span_to_range(gcx.sess.source_map(), param.ty.span);
74                    let ty = &source.file.src[loc];
75                    let name = if let Some(ident) = param.name {
76                        ident.name.to_string()
77                    } else {
78                        // Generate a unique name if the constructor arg does not have one.
79                        arg_index += 1;
80                        unique_identifier(
81                            reserved_identifiers,
82                            format!("foundry_pp_ctor_arg{arg_index}"),
83                        )
84                    };
85                    abi_encode_args.push(format!("args.{name}"));
86                    struct_fields.push(format!("{ty} {name}"));
87                }
88
89                ContractConstructorData {
90                    abi_encode_args: abi_encode_args.join(", "),
91                    struct_fields: struct_fields.join("; "),
92                    helper_contract: unique_identifier(
93                        reserved_identifiers,
94                        format!("DeployHelper{contract_id}"),
95                    ),
96                    args_struct: unique_identifier(
97                        reserved_identifiers,
98                        "FoundryPpConstructorArgs".to_string(),
99                    ),
100                    encode_function: unique_identifier(
101                        reserved_identifiers,
102                        format!("encodeArgs{contract_id}"),
103                    ),
104                    helper_path: deploy_helper_path(contract_id, source_units),
105                }
106            });
107
108        Self {
109            path: path.to_path_buf(),
110            name: contract.name.to_string(),
111            constructor_data,
112            artifact,
113        }
114    }
115
116    /// If contract has a non-empty constructor, generates a helper source file for it containing a
117    /// helper to encode constructor arguments.
118    ///
119    /// This is needed because current preprocessing wraps the arguments, leaving them unchanged.
120    /// This allows us to handle nested new expressions correctly. However, this requires us to have
121    /// a way to wrap both named and unnamed arguments. i.e you can't do abi.encode({arg: val}).
122    ///
123    /// This function produces a helper struct + a helper function to encode the arguments. The
124    /// struct is defined in scope of an abstract contract inheriting the contract containing the
125    /// constructor. This is done as a hack to allow us to inherit the same scope of definitions.
126    ///
127    /// The resulted helper looks like this:
128    /// ```solidity
129    /// import "lib/openzeppelin-contracts/contracts/token/ERC20.sol";
130    ///
131    /// abstract contract DeployHelper335 is ERC20 {
132    ///     struct FoundryPpConstructorArgs {
133    ///         string name;
134    ///         string symbol;
135    ///     }
136    /// }
137    ///
138    /// function encodeArgs335(DeployHelper335.FoundryPpConstructorArgs memory args) pure returns (bytes memory) {
139    ///     return abi.encode(args.name, args.symbol);
140    /// }
141    /// ```
142    ///
143    /// Example usage:
144    /// ```solidity
145    /// new ERC20(name, symbol)
146    /// ```
147    /// becomes
148    /// ```solidity
149    /// vm.deployCode("artifact path", encodeArgs335(DeployHelper335.FoundryPpConstructorArgs(name, symbol)))
150    /// ```
151    /// With named arguments:
152    /// ```solidity
153    /// new ERC20({name: name, symbol: symbol})
154    /// ```
155    /// becomes
156    /// ```solidity
157    /// vm.deployCode("artifact path", encodeArgs335(DeployHelper335.FoundryPpConstructorArgs({name: name, symbol: symbol})))
158    /// ```
159    pub fn build_helper(&self) -> Option<String> {
160        let Self { path, name, constructor_data, artifact: _ } = self;
161
162        let Some(constructor_details) = constructor_data else { return None };
163        let struct_fields = &constructor_details.struct_fields;
164        let abi_encode_args = &constructor_details.abi_encode_args;
165        let helper_contract = &constructor_details.helper_contract;
166        let args_struct = &constructor_details.args_struct;
167        let encode_function = &constructor_details.encode_function;
168
169        let path = solidity_string_content(path.to_slash_lossy().as_ref());
170        let helper = format!(
171            r#"
172// SPDX-License-Identifier: MIT
173pragma solidity >=0.4.0;
174
175import "{path}";
176
177abstract contract {helper_contract} is {name} {{
178    struct {args_struct} {{
179        {struct_fields};
180    }}
181}}
182
183function {encode_function}({helper_contract}.{args_struct} memory args) pure returns (bytes memory) {{
184    return abi.encode({abi_encode_args});
185}}
186        "#,
187        );
188
189        Some(helper)
190    }
191}
192
193fn unique_identifier(source: &str, mut identifier: String) -> String {
194    while source.contains(&identifier) {
195        identifier.push('_');
196    }
197    identifier
198}
199
200fn solidity_string_content(value: &str) -> String {
201    value.chars().fold(String::new(), |mut escaped, char| {
202        match char {
203            '\\' => escaped.push_str("\\\\"),
204            '"' => escaped.push_str("\\\""),
205            '\n' => escaped.push_str("\\n"),
206            '\r' => escaped.push_str("\\r"),
207            '\t' => escaped.push_str("\\t"),
208            char => escaped.push(char),
209        }
210        escaped
211    })
212}
213
214/// Collects preprocessor data from referenced contracts.
215pub(crate) fn collect_preprocessor_data(
216    gcx: Gcx<'_>,
217    referenced_contracts: &HashSet<ContractId>,
218    root_dir: &Path,
219    source_units: &[PathBuf],
220) -> PreprocessorData {
221    let mut data = PreprocessorData::default();
222    let reserved_identifiers = gcx
223        .hir
224        .source_ids()
225        .map(|source_id| gcx.hir.source(source_id).file.src.as_str())
226        .collect::<Vec<_>>()
227        .join("\n");
228    for contract_id in referenced_contracts {
229        let contract = gcx.hir.contract(*contract_id);
230        let source = gcx.hir.source(contract.source);
231
232        let FileName::Real(path) = &source.file.name else {
233            continue;
234        };
235
236        // Match the compiler input paths in generated imports and artifact references.
237        let path = path.strip_prefix(root_dir).unwrap_or(path);
238        let contract_data = ContractData::new(
239            gcx,
240            *contract_id,
241            contract,
242            path,
243            source,
244            &reserved_identifiers,
245            source_units,
246        );
247        data.insert(*contract_id, contract_data);
248    }
249    data
250}
251
252/// Creates helper libraries for contracts with a non-empty constructor.
253///
254/// See [`ContractData::build_helper`] for more details.
255pub(crate) fn create_deploy_helpers(data: &BTreeMap<ContractId, ContractData>) -> Sources {
256    let mut deploy_helpers = Sources::new();
257    for contract in data.values() {
258        if let Some(code) = contract.build_helper() {
259            let path = contract.constructor_data.as_ref().unwrap().helper_path.clone();
260            deploy_helpers.insert(path, Source::new(code));
261        }
262    }
263    deploy_helpers
264}
265
266/// Directory of the deploy helper sources generated by the preprocessor.
267const DEPLOY_HELPERS_DIR: &str = "foundry-pp";
268
269/// Returns a generated helper path that cannot overwrite an existing source unit.
270pub(crate) fn deploy_helper_path(contract_id: usize, source_units: &[PathBuf]) -> PathBuf {
271    let mut stem = format!("DeployHelper{contract_id}");
272    loop {
273        let path = PathBuf::from(format!("{DEPLOY_HELPERS_DIR}/{stem}.sol"));
274        if !source_units.iter().any(|source_unit| source_unit == &path) {
275            return path;
276        }
277        stem.push('_');
278    }
279}
280
281/// Returns whether the path points to a deploy helper source generated by the preprocessor.
282///
283/// These sources only exist in the compiler input and are never written to disk.
284pub fn is_deploy_helper_path(path: &Path) -> bool {
285    path.parent().and_then(Path::file_name).is_some_and(|dir| dir == DEPLOY_HELPERS_DIR)
286        && path
287            .file_stem()
288            .and_then(|stem| stem.to_str())
289            .is_some_and(|stem| stem.starts_with("DeployHelper"))
290}
291
292#[cfg(test)]
293mod tests {
294    use super::*;
295
296    #[test]
297    fn deploy_helper_path_does_not_replace_source_units() {
298        let source_units = [
299            PathBuf::from("foundry-pp/DeployHelper7.sol"),
300            PathBuf::from("foundry-pp/DeployHelper7_.sol"),
301        ];
302
303        assert_eq!(
304            deploy_helper_path(7, &source_units),
305            PathBuf::from("foundry-pp/DeployHelper7__.sol")
306        );
307    }
308
309    #[test]
310    fn is_deploy_helper_path_matches_generated_helpers() {
311        assert!(is_deploy_helper_path(Path::new("foundry-pp/DeployHelper7.sol")));
312        assert!(is_deploy_helper_path(Path::new("/project/foundry-pp/DeployHelper7_.sol")));
313        assert!(!is_deploy_helper_path(Path::new("/project/src/DeployHelper7.sol")));
314        assert!(!is_deploy_helper_path(Path::new("/project/foundry-pp/Counter.sol")));
315    }
316}
foundry_common/preprocessor/data.rs

foundry_common/preprocessor/
data.rs
