foundry_common/
contracts.rs

1//! Commonly used contract types and functions.
2
3use crate::{compile::PathOrContractInfo, strip_bytecode_placeholders};
4use alloy_dyn_abi::JsonAbiExt;
5use alloy_json_abi::{Event, Function, JsonAbi};
6use alloy_primitives::{Address, B256, Bytes, Selector, hex};
7use eyre::{OptionExt, Result};
8use foundry_compilers::{
9    ArtifactId, Project, ProjectCompileOutput,
10    artifacts::{
11        BytecodeObject, CompactBytecode, CompactContractBytecode, CompactContractBytecodeCow,
12        CompactDeployedBytecode, ConfigurableContractArtifact, ContractBytecodeSome, Offsets,
13        StorageLayout,
14    },
15    utils::canonicalized,
16};
17use std::{
18    collections::BTreeMap,
19    ops::Deref,
20    path::{Path, PathBuf},
21    str::FromStr,
22    sync::Arc,
23};
24
25/// Libraries' runtime code always starts with the following instruction:
26/// `PUSH20 0x0000000000000000000000000000000000000000`
27///
28/// See: <https://docs.soliditylang.org/en/latest/contracts.html#call-protection-for-libraries>
29const CALL_PROTECTION_BYTECODE_PREFIX: [u8; 21] =
30    hex!("730000000000000000000000000000000000000000");
31
32/// Subset of [CompactBytecode] excluding sourcemaps.
33#[expect(missing_docs)]
34#[derive(Debug, Clone)]
35pub struct BytecodeData {
36    pub object: Option<BytecodeObject>,
37    pub link_references: BTreeMap<String, BTreeMap<String, Vec<Offsets>>>,
38    pub immutable_references: BTreeMap<String, Vec<Offsets>>,
39}
40
41impl BytecodeData {
42    fn bytes(&self) -> Option<&Bytes> {
43        self.object.as_ref().and_then(|b| b.as_bytes())
44    }
45}
46
47impl From<CompactBytecode> for BytecodeData {
48    fn from(bytecode: CompactBytecode) -> Self {
49        Self {
50            object: Some(bytecode.object),
51            link_references: bytecode.link_references,
52            immutable_references: BTreeMap::new(),
53        }
54    }
55}
56
57impl From<CompactDeployedBytecode> for BytecodeData {
58    fn from(bytecode: CompactDeployedBytecode) -> Self {
59        let (object, link_references) = if let Some(compact) = bytecode.bytecode {
60            (Some(compact.object), compact.link_references)
61        } else {
62            (None, BTreeMap::new())
63        };
64        Self { object, link_references, immutable_references: bytecode.immutable_references }
65    }
66}
67
68/// Container for commonly used contract data.
69#[derive(Debug)]
70pub struct ContractData {
71    /// Contract name.
72    pub name: String,
73    /// Contract ABI.
74    pub abi: JsonAbi,
75    /// Contract creation code.
76    pub bytecode: Option<BytecodeData>,
77    /// Contract runtime code.
78    pub deployed_bytecode: Option<BytecodeData>,
79    /// Contract storage layout, if available.
80    pub storage_layout: Option<Arc<StorageLayout>>,
81}
82
83impl ContractData {
84    /// Returns reference to bytes of contract creation code, if present.
85    pub fn bytecode(&self) -> Option<&Bytes> {
86        self.bytecode.as_ref()?.bytes().filter(|b| !b.is_empty())
87    }
88
89    /// Returns reference to bytes of contract deployed code, if present.
90    pub fn deployed_bytecode(&self) -> Option<&Bytes> {
91        self.deployed_bytecode.as_ref()?.bytes().filter(|b| !b.is_empty())
92    }
93
94    /// Returns the bytecode without placeholders, if present.
95    pub fn bytecode_without_placeholders(&self) -> Option<Bytes> {
96        strip_bytecode_placeholders(self.bytecode.as_ref()?.object.as_ref()?)
97    }
98
99    /// Returns the deployed bytecode without placeholders, if present.
100    pub fn deployed_bytecode_without_placeholders(&self) -> Option<Bytes> {
101        strip_bytecode_placeholders(self.deployed_bytecode.as_ref()?.object.as_ref()?)
102    }
103}
104
105/// Builder for creating a `ContractsByArtifact` instance, optionally including storage layouts
106/// from project compile output.
107pub struct ContractsByArtifactBuilder<'a> {
108    /// All compiled artifact bytecodes (borrowed).
109    artifacts: BTreeMap<ArtifactId, CompactContractBytecodeCow<'a>>,
110    /// Optionally collected storage layouts for matching artifact IDs.
111    storage_layouts: BTreeMap<ArtifactId, StorageLayout>,
112}
113
114impl<'a> ContractsByArtifactBuilder<'a> {
115    /// Creates a new builder from artifacts with present bytecode iterator.
116    pub fn new(
117        artifacts: impl IntoIterator<Item = (ArtifactId, CompactContractBytecodeCow<'a>)>,
118    ) -> Self {
119        Self { artifacts: artifacts.into_iter().collect(), storage_layouts: BTreeMap::new() }
120    }
121
122    /// Adds storage layouts from `ProjectCompileOutput` to known artifacts.
123    pub fn with_storage_layouts(mut self, output: ProjectCompileOutput) -> Self {
124        self.storage_layouts = output
125            .into_artifacts()
126            .filter_map(|(id, artifact)| artifact.storage_layout.map(|layout| (id, layout)))
127            .collect();
128        self
129    }
130
131    /// Builds `ContractsByArtifact`.
132    pub fn build(self) -> ContractsByArtifact {
133        let map = self
134            .artifacts
135            .into_iter()
136            .filter_map(|(id, artifact)| {
137                let name = id.name.clone();
138                let CompactContractBytecodeCow { abi, bytecode, deployed_bytecode } = artifact;
139
140                Some((
141                    id.clone(),
142                    ContractData {
143                        name,
144                        abi: abi?.into_owned(),
145                        bytecode: bytecode.map(|b| b.into_owned().into()),
146                        deployed_bytecode: deployed_bytecode.map(|b| b.into_owned().into()),
147                        storage_layout: self.storage_layouts.get(&id).map(|l| Arc::new(l.clone())),
148                    },
149                ))
150            })
151            .collect();
152
153        ContractsByArtifact(Arc::new(map))
154    }
155}
156
157type ArtifactWithContractRef<'a> = (&'a ArtifactId, &'a ContractData);
158
159/// Wrapper type that maps an artifact to a contract ABI and bytecode.
160#[derive(Clone, Default, Debug)]
161pub struct ContractsByArtifact(Arc<BTreeMap<ArtifactId, ContractData>>);
162
163impl ContractsByArtifact {
164    /// Creates a new instance by collecting all artifacts with present bytecode from an iterator.
165    pub fn new(artifacts: impl IntoIterator<Item = (ArtifactId, CompactContractBytecode)>) -> Self {
166        let map = artifacts
167            .into_iter()
168            .filter_map(|(id, artifact)| {
169                let name = id.name.clone();
170                let CompactContractBytecode { abi, bytecode, deployed_bytecode } = artifact;
171                Some((
172                    id,
173                    ContractData {
174                        name,
175                        abi: abi?,
176                        bytecode: bytecode.map(Into::into),
177                        deployed_bytecode: deployed_bytecode.map(Into::into),
178                        storage_layout: None,
179                    },
180                ))
181            })
182            .collect();
183        Self(Arc::new(map))
184    }
185
186    /// Clears all contracts.
187    pub fn clear(&mut self) {
188        *self = Self::default();
189    }
190
191    /// Finds a contract which has a similar bytecode as `code`.
192    pub fn find_by_creation_code(&self, code: &[u8]) -> Option<ArtifactWithContractRef<'_>> {
193        self.find_by_code(code, 0.1, true, ContractData::bytecode)
194    }
195
196    /// Finds a contract which has a similar deployed bytecode as `code`.
197    pub fn find_by_deployed_code(&self, code: &[u8]) -> Option<ArtifactWithContractRef<'_>> {
198        self.find_by_code(code, 0.15, false, ContractData::deployed_bytecode)
199    }
200
201    /// Finds a contract based on provided bytecode and accepted match score.
202    /// If strip constructor args flag is true then removes args from bytecode to compare.
203    fn find_by_code(
204        &self,
205        code: &[u8],
206        accepted_score: f64,
207        strip_ctor_args: bool,
208        get: impl Fn(&ContractData) -> Option<&Bytes>,
209    ) -> Option<ArtifactWithContractRef<'_>> {
210        self.iter()
211            .filter_map(|(id, contract)| {
212                if let Some(deployed_bytecode) = get(contract) {
213                    let mut code = code;
214                    if strip_ctor_args && code.len() > deployed_bytecode.len() {
215                        // Try to decode ctor args with contract abi.
216                        if let Some(constructor) = contract.abi.constructor() {
217                            let constructor_args = &code[deployed_bytecode.len()..];
218                            if constructor.abi_decode_input(constructor_args).is_ok() {
219                                // If we can decode args with current abi then remove args from
220                                // code to compare.
221                                code = &code[..deployed_bytecode.len()]
222                            }
223                        }
224                    };
225
226                    let score = bytecode_diff_score(deployed_bytecode.as_ref(), code);
227                    (score <= accepted_score).then_some((score, (id, contract)))
228                } else {
229                    None
230                }
231            })
232            .min_by(|(score1, _), (score2, _)| score1.total_cmp(score2))
233            .map(|(_, data)| data)
234    }
235
236    /// Finds a contract which deployed bytecode exactly matches the given code. Accounts for link
237    /// references and immutables.
238    pub fn find_by_deployed_code_exact(&self, code: &[u8]) -> Option<ArtifactWithContractRef<'_>> {
239        // Immediately return None if the code is empty.
240        if code.is_empty() {
241            return None;
242        }
243
244        self.iter().find(|(_, contract)| {
245            let Some(deployed_bytecode) = &contract.deployed_bytecode else {
246                return false;
247            };
248            let Some(deployed_code) = &deployed_bytecode.object else {
249                return false;
250            };
251
252            let len = match deployed_code {
253                BytecodeObject::Bytecode(bytes) => bytes.len(),
254                BytecodeObject::Unlinked(bytes) => bytes.len() / 2,
255            };
256
257            if len != code.len() {
258                return false;
259            }
260
261            // Collect ignored offsets by chaining link and immutable references.
262            let mut ignored = deployed_bytecode
263                .immutable_references
264                .values()
265                .chain(deployed_bytecode.link_references.values().flat_map(|v| v.values()))
266                .flatten()
267                .cloned()
268                .collect::<Vec<_>>();
269
270            // For libraries solidity adds a call protection prefix to the bytecode. We need to
271            // ignore it as it includes library address determined at runtime.
272            // See https://docs.soliditylang.org/en/latest/contracts.html#call-protection-for-libraries and
273            // https://github.com/NomicFoundation/hardhat/blob/af7807cf38842a4f56e7f4b966b806e39631568a/packages/hardhat-verify/src/internal/solc/bytecode.ts#L172
274            let has_call_protection = match deployed_code {
275                BytecodeObject::Bytecode(bytes) => {
276                    bytes.starts_with(&CALL_PROTECTION_BYTECODE_PREFIX)
277                }
278                BytecodeObject::Unlinked(bytes) => {
279                    if let Ok(bytes) =
280                        Bytes::from_str(&bytes[..CALL_PROTECTION_BYTECODE_PREFIX.len() * 2])
281                    {
282                        bytes.starts_with(&CALL_PROTECTION_BYTECODE_PREFIX)
283                    } else {
284                        false
285                    }
286                }
287            };
288
289            if has_call_protection {
290                ignored.push(Offsets { start: 1, length: 20 });
291            }
292
293            ignored.sort_by_key(|o| o.start);
294
295            let mut left = 0;
296            for offset in ignored {
297                let right = offset.start as usize;
298
299                let matched = match deployed_code {
300                    BytecodeObject::Bytecode(bytes) => bytes[left..right] == code[left..right],
301                    BytecodeObject::Unlinked(bytes) => {
302                        if let Ok(bytes) = Bytes::from_str(&bytes[left * 2..right * 2]) {
303                            bytes == code[left..right]
304                        } else {
305                            false
306                        }
307                    }
308                };
309
310                if !matched {
311                    return false;
312                }
313
314                left = right + offset.length as usize;
315            }
316
317            if left < code.len() {
318                match deployed_code {
319                    BytecodeObject::Bytecode(bytes) => bytes[left..] == code[left..],
320                    BytecodeObject::Unlinked(bytes) => {
321                        if let Ok(bytes) = Bytes::from_str(&bytes[left * 2..]) {
322                            bytes == code[left..]
323                        } else {
324                            false
325                        }
326                    }
327                }
328            } else {
329                true
330            }
331        })
332    }
333
334    /// Finds a contract which has the same contract name or identifier as `id`. If more than one is
335    /// found, return error.
336    pub fn find_by_name_or_identifier(
337        &self,
338        id: &str,
339    ) -> Result<Option<ArtifactWithContractRef<'_>>> {
340        let contracts = self
341            .iter()
342            .filter(|(artifact, _)| artifact.name == id || artifact.identifier() == id)
343            .collect::<Vec<_>>();
344
345        if contracts.len() > 1 {
346            eyre::bail!("{id} has more than one implementation.");
347        }
348
349        Ok(contracts.first().copied())
350    }
351
352    /// Finds abi for contract which has the same contract name or identifier as `id`.
353    pub fn find_abi_by_name_or_identifier(&self, id: &str) -> Option<JsonAbi> {
354        self.iter()
355            .find(|(artifact, _)| {
356                artifact.name.split(".").next().unwrap() == id || artifact.identifier() == id
357            })
358            .map(|(_, contract)| contract.abi.clone())
359    }
360
361    /// Finds abi by name or source path
362    ///
363    /// Returns the abi and the contract name.
364    pub fn find_abi_by_name_or_src_path(&self, name_or_path: &str) -> Option<(JsonAbi, String)> {
365        self.iter()
366            .find(|(artifact, _)| {
367                artifact.name == name_or_path || artifact.source == Path::new(name_or_path)
368            })
369            .map(|(_, contract)| (contract.abi.clone(), contract.name.clone()))
370    }
371
372    /// Flattens the contracts into functions, events and errors.
373    pub fn flatten(&self) -> (BTreeMap<Selector, Function>, BTreeMap<B256, Event>, JsonAbi) {
374        let mut funcs = BTreeMap::new();
375        let mut events = BTreeMap::new();
376        let mut errors_abi = JsonAbi::new();
377        for (_name, contract) in self.iter() {
378            for func in contract.abi.functions() {
379                funcs.insert(func.selector(), func.clone());
380            }
381            for event in contract.abi.events() {
382                events.insert(event.selector(), event.clone());
383            }
384            for error in contract.abi.errors() {
385                errors_abi.errors.entry(error.name.clone()).or_default().push(error.clone());
386            }
387        }
388        (funcs, events, errors_abi)
389    }
390}
391
392impl From<ProjectCompileOutput> for ContractsByArtifact {
393    fn from(value: ProjectCompileOutput) -> Self {
394        Self::new(value.into_artifacts().map(|(id, ar)| {
395            (
396                id,
397                CompactContractBytecode {
398                    abi: ar.abi,
399                    bytecode: ar.bytecode,
400                    deployed_bytecode: ar.deployed_bytecode,
401                },
402            )
403        }))
404    }
405}
406
407impl Deref for ContractsByArtifact {
408    type Target = BTreeMap<ArtifactId, ContractData>;
409
410    fn deref(&self) -> &Self::Target {
411        &self.0
412    }
413}
414
415/// Wrapper type that maps an address to a contract identifier and contract ABI.
416pub type ContractsByAddress = BTreeMap<Address, (String, JsonAbi)>;
417
418/// Very simple fuzzy matching of contract bytecode.
419///
420/// Returns a value between `0.0` (identical) and `1.0` (completely different).
421pub fn bytecode_diff_score<'a>(mut a: &'a [u8], mut b: &'a [u8]) -> f64 {
422    // Make sure `a` is the longer one.
423    if a.len() < b.len() {
424        std::mem::swap(&mut a, &mut b);
425    }
426
427    // Account for different lengths.
428    let mut n_different_bytes = a.len() - b.len();
429
430    // If the difference is more than 32 bytes and more than 10% of the total length,
431    // we assume the bytecodes are completely different.
432    // This is a simple heuristic to avoid checking every byte when the lengths are very different.
433    // 32 is chosen to be a reasonable minimum as it's the size of metadata hashes and one EVM word.
434    if n_different_bytes > 32 && n_different_bytes * 10 > a.len() {
435        return 1.0;
436    }
437
438    // Count different bytes.
439    // SAFETY: `a` is longer than `b`.
440    n_different_bytes += unsafe { count_different_bytes(a, b) };
441
442    n_different_bytes as f64 / a.len() as f64
443}
444
445/// Returns the amount of different bytes between two slices.
446///
447/// # Safety
448///
449/// `a` must be at least as long as `b`.
450unsafe fn count_different_bytes(a: &[u8], b: &[u8]) -> usize {
451    // This could've been written as `std::iter::zip(a, b).filter(|(x, y)| x != y).count()`,
452    // however this function is very hot, and has been written to be as primitive as
453    // possible for lower optimization levels.
454
455    let a_ptr = a.as_ptr();
456    let b_ptr = b.as_ptr();
457    let len = b.len();
458
459    let mut sum = 0;
460    let mut i = 0;
461    while i < len {
462        // SAFETY: `a` is at least as long as `b`, and `i` is in bound of `b`.
463        sum += unsafe { *a_ptr.add(i) != *b_ptr.add(i) } as usize;
464        i += 1;
465    }
466    sum
467}
468
469/// Returns contract name for a given contract identifier.
470///
471/// Artifact/Contract identifier can take the following form:
472/// `<artifact file name>:<contract name>`, the `artifact file name` is the name of the json file of
473/// the contract's artifact and the contract name is the name of the solidity contract, like
474/// `SafeTransferLibTest.json:SafeTransferLibTest`
475///
476/// This returns the `contract name` part
477///
478/// # Example
479///
480/// ```
481/// use foundry_common::*;
482/// assert_eq!(
483///     "SafeTransferLibTest",
484///     get_contract_name("SafeTransferLibTest.json:SafeTransferLibTest")
485/// );
486/// ```
487pub fn get_contract_name(id: &str) -> &str {
488    id.rsplit(':').next().unwrap_or(id)
489}
490
491/// This returns the `file name` part, See [`get_contract_name`]
492///
493/// # Example
494///
495/// ```
496/// use foundry_common::*;
497/// assert_eq!(
498///     "SafeTransferLibTest.json",
499///     get_file_name("SafeTransferLibTest.json:SafeTransferLibTest")
500/// );
501/// ```
502pub fn get_file_name(id: &str) -> &str {
503    id.split(':').next().unwrap_or(id)
504}
505
506/// Helper function to convert CompactContractBytecode ~> ContractBytecodeSome
507pub fn compact_to_contract(contract: CompactContractBytecode) -> Result<ContractBytecodeSome> {
508    Ok(ContractBytecodeSome {
509        abi: contract.abi.ok_or_else(|| eyre::eyre!("No contract abi"))?,
510        bytecode: contract.bytecode.ok_or_else(|| eyre::eyre!("No contract bytecode"))?.into(),
511        deployed_bytecode: contract
512            .deployed_bytecode
513            .ok_or_else(|| eyre::eyre!("No contract deployed bytecode"))?
514            .into(),
515    })
516}
517
518/// Returns the canonicalized target path for the given identifier.
519pub fn find_target_path(project: &Project, identifier: &PathOrContractInfo) -> Result<PathBuf> {
520    match identifier {
521        PathOrContractInfo::Path(path) => Ok(canonicalized(project.root().join(path))),
522        PathOrContractInfo::ContractInfo(info) => {
523            if let Some(path) = info.path.as_ref() {
524                let path = canonicalized(project.root().join(path));
525                let sources = project.sources()?;
526                let contract_path = sources
527                    .iter()
528                    .find_map(|(src_path, _)| {
529                        if **src_path == path {
530                            return Some(src_path.clone());
531                        }
532                        None
533                    })
534                    .ok_or_else(|| {
535                        eyre::eyre!(
536                            "Could not find source file for contract `{}` at {}",
537                            info.name,
538                            path.strip_prefix(project.root()).unwrap().display()
539                        )
540                    })?;
541                return Ok(contract_path);
542            }
543            // If ContractInfo.path hasn't been provided we try to find the contract using the name.
544            // This will fail if projects have multiple contracts with the same name. In that case,
545            // path must be specified.
546            let path = project.find_contract_path(&info.name)?;
547            Ok(path)
548        }
549    }
550}
551
552/// Returns the target artifact given the path and name.
553pub fn find_matching_contract_artifact(
554    output: &mut ProjectCompileOutput,
555    target_path: &Path,
556    target_name: Option<&str>,
557) -> eyre::Result<ConfigurableContractArtifact> {
558    if let Some(name) = target_name {
559        output
560            .remove(target_path, name)
561            .ok_or_eyre(format!("Could not find artifact `{name}` in the compiled artifacts"))
562    } else {
563        let possible_targets = output
564            .artifact_ids()
565            .filter(|(id, _artifact)| id.source == target_path)
566            .collect::<Vec<_>>();
567
568        if possible_targets.is_empty() {
569            eyre::bail!(
570                "Could not find artifact linked to source `{target_path:?}` in the compiled artifacts"
571            );
572        }
573
574        let (target_id, target_artifact) = possible_targets[0].clone();
575        if possible_targets.len() == 1 {
576            return Ok(target_artifact.clone());
577        }
578
579        // If all artifact_ids in `possible_targets` have the same name (without ".", indicates
580        // additional compiler profiles), it means that there are multiple contracts in the
581        // same file.
582        if !target_id.name.contains(".")
583            && possible_targets.iter().any(|(id, _)| id.name != target_id.name)
584        {
585            eyre::bail!(
586                "Multiple contracts found in the same file, please specify the target <path>:<contract> or <contract>"
587            );
588        }
589
590        // Otherwise, we're dealing with additional compiler profiles wherein `id.source` is the
591        // same but `id.path` is different.
592        let artifact = possible_targets
593            .iter()
594            .find_map(|(id, artifact)| if id.profile == "default" { Some(*artifact) } else { None })
595            .unwrap_or(target_artifact);
596
597        Ok(artifact.clone())
598    }
599}
600
601#[cfg(test)]
602mod tests {
603    use super::*;
604
605    #[test]
606    fn bytecode_diffing() {
607        assert_eq!(bytecode_diff_score(b"a", b"a"), 0.0);
608        assert_eq!(bytecode_diff_score(b"a", b"b"), 1.0);
609
610        let a_100 = &b"a".repeat(100)[..];
611        assert_eq!(bytecode_diff_score(a_100, &b"b".repeat(100)), 1.0);
612        assert_eq!(bytecode_diff_score(a_100, &b"b".repeat(99)), 1.0);
613        assert_eq!(bytecode_diff_score(a_100, &b"b".repeat(101)), 1.0);
614        assert_eq!(bytecode_diff_score(a_100, &b"b".repeat(120)), 1.0);
615        assert_eq!(bytecode_diff_score(a_100, &b"b".repeat(1000)), 1.0);
616
617        let a_99 = &b"a".repeat(99)[..];
618        assert!(bytecode_diff_score(a_100, a_99) <= 0.01);
619    }
620
621    #[test]
622    fn find_by_deployed_code_exact_with_empty_deployed() {
623        let contracts = ContractsByArtifact::new(vec![]);
624
625        assert!(contracts.find_by_deployed_code_exact(&[]).is_none());
626    }
627}