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/// Collects preprocessor data from referenced contracts.
19pub(crate) fn collect_preprocessor_data(
20    gcx: Gcx<'_>,
21    referenced_contracts: &HashSet<ContractId>,
22) -> PreprocessorData {
23    let mut data = PreprocessorData::default();
24    for contract_id in referenced_contracts {
25        let contract = gcx.hir.contract(*contract_id);
26        let source = gcx.hir.source(contract.source);
27
28        let FileName::Real(path) = &source.file.name else {
29            continue;
30        };
31
32        let contract_data = ContractData::new(gcx, *contract_id, contract, path, source);
33        data.insert(*contract_id, contract_data);
34    }
35    data
36}
37
38/// Creates helper libraries for contracts with a non-empty constructor.
39///
40/// See [`ContractData::build_helper`] for more details.
41pub(crate) fn create_deploy_helpers(data: &BTreeMap<ContractId, ContractData>) -> Sources {
42    let mut deploy_helpers = Sources::new();
43    for (contract_id, contract) in data {
44        if let Some(code) = contract.build_helper() {
45            let path = format!("foundry-pp/DeployHelper{}.sol", contract_id.get());
46            deploy_helpers.insert(path.into(), Source::new(code));
47        }
48    }
49    deploy_helpers
50}
51
52/// Keeps data about a contract constructor.
53#[derive(Debug)]
54pub struct ContractConstructorData {
55    /// ABI encoded args.
56    pub abi_encode_args: String,
57    /// Constructor struct fields.
58    pub struct_fields: String,
59}
60
61/// Keeps data about a single contract definition.
62#[derive(Debug)]
63pub(crate) struct ContractData {
64    /// HIR Id of the contract.
65    contract_id: ContractId,
66    /// Path of the source file.
67    path: PathBuf,
68    /// Name of the contract
69    name: String,
70    /// Constructor parameters, if any.
71    pub constructor_data: Option<ContractConstructorData>,
72    /// Artifact string to pass into cheatcodes.
73    pub artifact: String,
74}
75
76impl ContractData {
77    fn new(
78        gcx: Gcx<'_>,
79        contract_id: ContractId,
80        contract: &Contract<'_>,
81        path: &Path,
82        source: &solar::sema::hir::Source<'_>,
83    ) -> Self {
84        let artifact = format!("{}:{}", path.to_slash_lossy(), contract.name);
85
86        // Process data for contracts with constructor and parameters.
87        let constructor_data = contract
88            .ctor
89            .map(|ctor_id| gcx.hir.function(ctor_id))
90            .filter(|ctor| !ctor.parameters.is_empty())
91            .map(|ctor| {
92                let mut abi_encode_args = vec![];
93                let mut struct_fields = vec![];
94                let mut arg_index = 0;
95                for param_id in ctor.parameters {
96                    let src = source.file.src.as_str();
97                    let loc =
98                        span_to_range(gcx.sess.source_map(), gcx.hir.variable(*param_id).span);
99                    let mut new_src = src[loc].replace(" memory ", " ").replace(" calldata ", " ");
100                    if let Some(ident) = gcx.hir.variable(*param_id).name {
101                        abi_encode_args.push(format!("args.{}", ident.name));
102                    } else {
103                        // Generate an unique name if constructor arg doesn't have one.
104                        arg_index += 1;
105                        abi_encode_args.push(format!("args.foundry_pp_ctor_arg{arg_index}"));
106                        new_src.push_str(&format!(" foundry_pp_ctor_arg{arg_index}"));
107                    }
108                    struct_fields.push(new_src);
109                }
110
111                ContractConstructorData {
112                    abi_encode_args: abi_encode_args.join(", "),
113                    struct_fields: struct_fields.join("; "),
114                }
115            });
116
117        Self {
118            contract_id,
119            path: path.to_path_buf(),
120            name: contract.name.to_string(),
121            constructor_data,
122            artifact,
123        }
124    }
125
126    /// If contract has a non-empty constructor, generates a helper source file for it containing a
127    /// helper to encode constructor arguments.
128    ///
129    /// This is needed because current preprocessing wraps the arguments, leaving them unchanged.
130    /// This allows us to handle nested new expressions correctly. However, this requires us to have
131    /// a way to wrap both named and unnamed arguments. i.e you can't do abi.encode({arg: val}).
132    ///
133    /// This function produces a helper struct + a helper function to encode the arguments. The
134    /// struct is defined in scope of an abstract contract inheriting the contract containing the
135    /// constructor. This is done as a hack to allow us to inherit the same scope of definitions.
136    ///
137    /// The resulted helper looks like this:
138    /// ```solidity
139    /// import "lib/openzeppelin-contracts/contracts/token/ERC20.sol";
140    ///
141    /// abstract contract DeployHelper335 is ERC20 {
142    ///     struct FoundryPpConstructorArgs {
143    ///         string name;
144    ///         string symbol;
145    ///     }
146    /// }
147    ///
148    /// function encodeArgs335(DeployHelper335.FoundryPpConstructorArgs memory args) pure returns (bytes memory) {
149    ///     return abi.encode(args.name, args.symbol);
150    /// }
151    /// ```
152    ///
153    /// Example usage:
154    /// ```solidity
155    /// new ERC20(name, symbol)
156    /// ```
157    /// becomes
158    /// ```solidity
159    /// vm.deployCode("artifact path", encodeArgs335(DeployHelper335.FoundryPpConstructorArgs(name, symbol)))
160    /// ```
161    /// With named arguments:
162    /// ```solidity
163    /// new ERC20({name: name, symbol: symbol})
164    /// ```
165    /// becomes
166    /// ```solidity
167    /// vm.deployCode("artifact path", encodeArgs335(DeployHelper335.FoundryPpConstructorArgs({name: name, symbol: symbol})))
168    /// ```
169    pub fn build_helper(&self) -> Option<String> {
170        let Self { contract_id, path, name, constructor_data, artifact: _ } = self;
171
172        let Some(constructor_details) = constructor_data else { return None };
173        let contract_id = contract_id.get();
174        let struct_fields = &constructor_details.struct_fields;
175        let abi_encode_args = &constructor_details.abi_encode_args;
176
177        let helper = format!(
178            r#"
179// SPDX-License-Identifier: MIT
180pragma solidity >=0.4.0;
181
182import "{path}";
183
184abstract contract DeployHelper{contract_id} is {name} {{
185    struct FoundryPpConstructorArgs {{
186        {struct_fields};
187    }}
188}}
189
190function encodeArgs{contract_id}(DeployHelper{contract_id}.FoundryPpConstructorArgs memory args) pure returns (bytes memory) {{
191    return abi.encode({abi_encode_args});
192}}
193        "#,
194            path = path.to_slash_lossy(),
195        );
196
197        Some(helper)
198    }
199}