1use crate::{compile::PathOrContractInfo, find_metadata_start, strip_bytecode_placeholders};
4use alloy_dyn_abi::JsonAbiExt;
5use alloy_json_abi::{Event, Function, JsonAbi};
6use alloy_primitives::{Address, B256, Bytes, Selector, address, 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
25const CALL_PROTECTION_BYTECODE_PREFIX: [u8; 21] =
30 hex!("730000000000000000000000000000000000000000");
31
32pub const LIBRARY_DEPLOYER: Address = address!("0x1F95D37F27EA0dEA9C252FC09D5A6eaA97647353");
36
37#[expect(missing_docs)]
39#[derive(Debug, Clone)]
40pub struct BytecodeData {
41 pub object: Option<BytecodeObject>,
42 pub link_references: BTreeMap<String, BTreeMap<String, Vec<Offsets>>>,
43 pub immutable_references: BTreeMap<String, Vec<Offsets>>,
44}
45
46impl BytecodeData {
47 fn bytes(&self) -> Option<&Bytes> {
48 self.object.as_ref().and_then(|b| b.as_bytes())
49 }
50}
51
52impl From<CompactBytecode> for BytecodeData {
53 fn from(bytecode: CompactBytecode) -> Self {
54 Self {
55 object: Some(bytecode.object),
56 link_references: bytecode.link_references,
57 immutable_references: BTreeMap::new(),
58 }
59 }
60}
61
62impl From<CompactDeployedBytecode> for BytecodeData {
63 fn from(bytecode: CompactDeployedBytecode) -> Self {
64 let (object, link_references) = if let Some(compact) = bytecode.bytecode {
65 (Some(compact.object), compact.link_references)
66 } else {
67 (None, BTreeMap::new())
68 };
69 Self { object, link_references, immutable_references: bytecode.immutable_references }
70 }
71}
72
73#[derive(Debug)]
75pub struct ContractData {
76 pub name: String,
78 pub abi: JsonAbi,
80 pub bytecode: Option<BytecodeData>,
82 pub deployed_bytecode: Option<BytecodeData>,
84 pub storage_layout: Option<Arc<StorageLayout>>,
86}
87
88impl ContractData {
89 pub fn bytecode(&self) -> Option<&Bytes> {
91 self.bytecode.as_ref()?.bytes().filter(|b| !b.is_empty())
92 }
93
94 pub fn deployed_bytecode(&self) -> Option<&Bytes> {
96 self.deployed_bytecode.as_ref()?.bytes().filter(|b| !b.is_empty())
97 }
98
99 pub fn bytecode_without_placeholders(&self) -> Option<Bytes> {
101 strip_bytecode_placeholders(self.bytecode.as_ref()?.object.as_ref()?)
102 }
103
104 pub fn deployed_bytecode_without_placeholders(&self) -> Option<Bytes> {
106 strip_bytecode_placeholders(self.deployed_bytecode.as_ref()?.object.as_ref()?)
107 }
108}
109
110pub struct ContractsByArtifactBuilder<'a> {
113 artifacts: BTreeMap<ArtifactId, CompactContractBytecodeCow<'a>>,
115 storage_layouts: BTreeMap<ArtifactId, StorageLayout>,
117}
118
119impl<'a> ContractsByArtifactBuilder<'a> {
120 pub fn new(
122 artifacts: impl IntoIterator<Item = (ArtifactId, CompactContractBytecodeCow<'a>)>,
123 ) -> Self {
124 Self { artifacts: artifacts.into_iter().collect(), storage_layouts: BTreeMap::new() }
125 }
126
127 pub fn with_output(self, output: &ProjectCompileOutput, base: &Path) -> Self {
129 self.with_storage_layouts(output.artifact_ids().filter_map(|(id, artifact)| {
130 artifact
131 .storage_layout
132 .as_ref()
133 .map(|layout| (id.with_stripped_file_prefixes(base), layout.clone()))
134 }))
135 }
136
137 pub fn with_storage_layouts(
139 mut self,
140 layouts: impl IntoIterator<Item = (ArtifactId, StorageLayout)>,
141 ) -> Self {
142 self.storage_layouts.extend(layouts);
143 self
144 }
145
146 pub fn build(self) -> ContractsByArtifact {
148 let map = self
149 .artifacts
150 .into_iter()
151 .filter_map(|(id, artifact)| {
152 let name = id.name.clone();
153 let CompactContractBytecodeCow { abi, bytecode, deployed_bytecode } = artifact;
154
155 Some((
156 id.clone(),
157 ContractData {
158 name,
159 abi: abi?.into_owned(),
160 bytecode: bytecode.map(|b| b.into_owned().into()),
161 deployed_bytecode: deployed_bytecode.map(|b| b.into_owned().into()),
162 storage_layout: self.storage_layouts.get(&id).map(|l| Arc::new(l.clone())),
163 },
164 ))
165 })
166 .collect();
167
168 ContractsByArtifact(Arc::new(map))
169 }
170}
171
172type ArtifactWithContractRef<'a> = (&'a ArtifactId, &'a ContractData);
173
174#[derive(Clone, Default, Debug)]
176pub struct ContractsByArtifact(Arc<BTreeMap<ArtifactId, ContractData>>);
177
178impl ContractsByArtifact {
179 pub fn new(artifacts: impl IntoIterator<Item = (ArtifactId, CompactContractBytecode)>) -> Self {
181 let map = artifacts
182 .into_iter()
183 .filter_map(|(id, artifact)| {
184 let name = id.name.clone();
185 let CompactContractBytecode { abi, bytecode, deployed_bytecode } = artifact;
186 Some((
187 id,
188 ContractData {
189 name,
190 abi: abi?,
191 bytecode: bytecode.map(Into::into),
192 deployed_bytecode: deployed_bytecode.map(Into::into),
193 storage_layout: None,
194 },
195 ))
196 })
197 .collect();
198 Self(Arc::new(map))
199 }
200
201 pub fn clear(&mut self) {
203 *self = Self::default();
204 }
205
206 pub fn find_by_creation_code(&self, code: &[u8]) -> Option<ArtifactWithContractRef<'_>> {
208 self.find_by_code(code, 0.1, true, ContractData::bytecode)
209 }
210
211 pub fn find_by_deployed_code(&self, code: &[u8]) -> Option<ArtifactWithContractRef<'_>> {
213 self.find_by_code(code, 0.15, false, ContractData::deployed_bytecode)
214 }
215
216 fn find_by_code(
219 &self,
220 code: &[u8],
221 accepted_score: f64,
222 strip_ctor_args: bool,
223 get: impl Fn(&ContractData) -> Option<&Bytes>,
224 ) -> Option<ArtifactWithContractRef<'_>> {
225 self.iter()
226 .filter_map(|(id, contract)| {
227 if let Some(deployed_bytecode) = get(contract) {
228 let mut code = code;
229 if strip_ctor_args && code.len() > deployed_bytecode.len() {
230 if let Some(constructor) = contract.abi.constructor() {
232 let constructor_args = &code[deployed_bytecode.len()..];
233 if constructor.abi_decode_input(constructor_args).is_ok() {
234 code = &code[..deployed_bytecode.len()]
237 }
238 }
239 };
240
241 let score = bytecode_diff_score(deployed_bytecode.as_ref(), code);
242 (score <= accepted_score).then_some((score, (id, contract)))
243 } else {
244 None
245 }
246 })
247 .min_by(|(score1, _), (score2, _)| score1.total_cmp(score2))
248 .map(|(_, data)| data)
249 }
250
251 pub fn find_by_deployed_code_exact(&self, code: &[u8]) -> Option<ArtifactWithContractRef<'_>> {
254 self.find_by_deployed_code_exact_inner(code, false)
255 }
256
257 pub fn find_by_deployed_code_exact_unique(
259 &self,
260 code: &[u8],
261 ) -> Option<ArtifactWithContractRef<'_>> {
262 self.find_by_deployed_code_exact_inner(code, true)
263 }
264
265 fn find_by_deployed_code_exact_inner(
266 &self,
267 code: &[u8],
268 unique: bool,
269 ) -> Option<ArtifactWithContractRef<'_>> {
270 if code.is_empty() {
272 return None;
273 }
274
275 let mut partial_match = None;
276 let mut unique_match = None;
277 let matched = self.iter().find(|(id, contract)| {
278 let Some(deployed_bytecode) = &contract.deployed_bytecode else {
279 return false;
280 };
281 let Some(deployed_code) = &deployed_bytecode.object else {
282 return false;
283 };
284
285 let len = match deployed_code {
286 BytecodeObject::Bytecode(bytes) => bytes.len(),
287 BytecodeObject::Unlinked(bytes) => bytes.len() / 2,
288 };
289
290 if len != code.len() {
291 return false;
292 }
293
294 let mut ignored = deployed_bytecode
296 .immutable_references
297 .values()
298 .chain(deployed_bytecode.link_references.values().flat_map(|v| v.values()))
299 .flatten()
300 .cloned()
301 .collect::<Vec<_>>();
302
303 let has_call_protection = match deployed_code {
308 BytecodeObject::Bytecode(bytes) => {
309 bytes.starts_with(&CALL_PROTECTION_BYTECODE_PREFIX)
310 }
311 BytecodeObject::Unlinked(bytes) => {
312 if let Ok(bytes) =
313 Bytes::from_str(&bytes[..CALL_PROTECTION_BYTECODE_PREFIX.len() * 2])
314 {
315 bytes.starts_with(&CALL_PROTECTION_BYTECODE_PREFIX)
316 } else {
317 false
318 }
319 }
320 };
321
322 if has_call_protection {
323 ignored.push(Offsets { start: 1, length: 20 });
324 }
325
326 let metadata_start = find_metadata_start(code);
327
328 if let Some(metadata) = metadata_start {
329 ignored.push(Offsets {
330 start: metadata as u32,
331 length: (code.len() - metadata) as u32,
332 });
333 }
334
335 ignored.sort_by_key(|o| o.start);
336
337 let mut left = 0;
338 for offset in ignored {
339 let right = offset.start as usize;
340
341 let matched = match deployed_code {
342 BytecodeObject::Bytecode(bytes) => bytes[left..right] == code[left..right],
343 BytecodeObject::Unlinked(bytes) => {
344 if let Ok(bytes) = Bytes::from_str(&bytes[left * 2..right * 2]) {
345 bytes == code[left..right]
346 } else {
347 false
348 }
349 }
350 };
351
352 if !matched {
353 return false;
354 }
355
356 left = right + offset.length as usize;
357 }
358
359 let is_partial = if left < code.len() {
360 match deployed_code {
361 BytecodeObject::Bytecode(bytes) => bytes[left..] == code[left..],
362 BytecodeObject::Unlinked(bytes) => {
363 if let Ok(bytes) = Bytes::from_str(&bytes[left * 2..]) {
364 bytes == code[left..]
365 } else {
366 false
367 }
368 }
369 }
370 } else {
371 true
372 };
373
374 if !is_partial {
375 return false;
376 }
377
378 let Some(metadata) = metadata_start else {
379 if unique && unique_match.is_none() {
380 unique_match = Some((*id, *contract));
381 return false;
382 }
383 return true;
384 };
385
386 let exact_match = match deployed_code {
387 BytecodeObject::Bytecode(bytes) => bytes[metadata..] == code[metadata..],
388 BytecodeObject::Unlinked(bytes) => {
389 if let Ok(bytes) = Bytes::from_str(&bytes[metadata * 2..]) {
390 bytes == code[metadata..]
391 } else {
392 false
393 }
394 }
395 };
396
397 if exact_match {
398 if unique && unique_match.is_none() {
399 unique_match = Some((*id, *contract));
400 false
401 } else {
402 true
403 }
404 } else {
405 partial_match = Some((*id, *contract));
406 false
407 }
408 });
409
410 if unique {
411 matched.is_none().then_some(unique_match).flatten()
412 } else {
413 matched.or(partial_match)
414 }
415 }
416
417 pub fn find_by_name_or_identifier(
420 &self,
421 id: &str,
422 ) -> Result<Option<ArtifactWithContractRef<'_>>> {
423 let mut iter =
424 self.iter().filter(|(artifact, _)| artifact.name == id || artifact.identifier() == id);
425 let first = iter.next();
426 if first.is_some() && iter.next().is_some() {
427 eyre::bail!("{id} has more than one implementation.");
428 }
429
430 Ok(first)
431 }
432
433 pub fn find_abi_by_name_or_src_path(&self, name_or_path: &str) -> Option<(JsonAbi, String)> {
437 self.iter()
438 .find(|(artifact, _)| {
439 artifact.name == name_or_path || artifact.source == Path::new(name_or_path)
440 })
441 .map(|(_, contract)| (contract.abi.clone(), contract.name.clone()))
442 }
443
444 pub fn flatten(&self) -> (BTreeMap<Selector, Function>, BTreeMap<B256, Event>, JsonAbi) {
446 let mut funcs = BTreeMap::new();
447 let mut events = BTreeMap::new();
448 let mut errors_abi = JsonAbi::new();
449 for contract in self.values() {
450 for func in contract.abi.functions() {
451 funcs.insert(func.selector(), func.clone());
452 }
453 for event in contract.abi.events() {
454 events.insert(event.selector(), event.clone());
455 }
456 for error in contract.abi.errors() {
457 errors_abi.errors.entry(error.name.clone()).or_default().push(error.clone());
458 }
459 }
460 (funcs, events, errors_abi)
461 }
462}
463
464impl From<ProjectCompileOutput> for ContractsByArtifact {
465 fn from(value: ProjectCompileOutput) -> Self {
466 Self::new(value.into_artifacts().map(|(id, ar)| {
467 (
468 id,
469 CompactContractBytecode {
470 abi: ar.abi,
471 bytecode: ar.bytecode,
472 deployed_bytecode: ar.deployed_bytecode,
473 },
474 )
475 }))
476 }
477}
478
479impl Deref for ContractsByArtifact {
480 type Target = BTreeMap<ArtifactId, ContractData>;
481
482 fn deref(&self) -> &Self::Target {
483 &self.0
484 }
485}
486
487pub type ContractsByAddress = BTreeMap<Address, (String, JsonAbi)>;
489
490pub fn matches_contract_creation(contract: &ContractData, creation_code: &[u8]) -> bool {
496 let Some(bytecode) = contract.bytecode() else { return false };
497 let Some(arguments) = creation_code.strip_prefix(bytecode.as_ref()) else { return false };
498 match contract.abi.constructor() {
499 Some(constructor) => constructor
500 .abi_decode_input(arguments)
501 .ok()
502 .and_then(|values| constructor.abi_encode_input(&values).ok())
503 .is_some_and(|encoded| encoded == arguments),
504 None => arguments.is_empty(),
505 }
506}
507
508pub fn bytecode_diff_score<'a>(mut a: &'a [u8], mut b: &'a [u8]) -> f64 {
509 if a.len() < b.len() {
511 std::mem::swap(&mut a, &mut b);
512 }
513
514 let mut n_different_bytes = a.len() - b.len();
516
517 if n_different_bytes > 32 && n_different_bytes * 10 > a.len() {
522 return 1.0;
523 }
524
525 n_different_bytes += unsafe { count_different_bytes(a, b) };
528
529 n_different_bytes as f64 / a.len() as f64
530}
531
532const unsafe fn count_different_bytes(a: &[u8], b: &[u8]) -> usize {
538 let a_ptr = a.as_ptr();
543 let b_ptr = b.as_ptr();
544 let len = b.len();
545
546 let mut sum = 0;
547 let mut i = 0;
548 while i < len {
549 sum += unsafe { *a_ptr.add(i) != *b_ptr.add(i) } as usize;
551 i += 1;
552 }
553 sum
554}
555
556pub fn get_contract_name(id: &str) -> &str {
575 id.rsplit(':').next().unwrap_or(id)
576}
577
578pub fn get_file_name(id: &str) -> &str {
590 id.split(':').next().unwrap_or(id)
591}
592
593pub fn compact_to_contract(contract: CompactContractBytecode) -> Result<ContractBytecodeSome> {
595 Ok(ContractBytecodeSome {
596 abi: contract.abi.ok_or_else(|| eyre::eyre!("No contract abi"))?,
597 bytecode: contract.bytecode.ok_or_else(|| eyre::eyre!("No contract bytecode"))?.into(),
598 deployed_bytecode: contract
599 .deployed_bytecode
600 .ok_or_else(|| eyre::eyre!("No contract deployed bytecode"))?
601 .into(),
602 })
603}
604
605pub fn find_target_path(project: &Project, identifier: &PathOrContractInfo) -> Result<PathBuf> {
607 match identifier {
608 PathOrContractInfo::Path(path) => Ok(canonicalized(project.root().join(path))),
609 PathOrContractInfo::ContractInfo(info) => {
610 if let Some(path) = info.path.as_ref() {
611 let path = canonicalized(project.root().join(path));
612 if !path.is_file() {
613 eyre::bail!(
614 "Could not find source file for contract `{}` at {}",
615 info.name,
616 path.strip_prefix(project.root()).unwrap_or(&path).display()
617 );
618 }
619 return Ok(path);
620 }
621 let path = project.find_contract_path(&info.name)?;
625 Ok(path)
626 }
627 }
628}
629
630pub fn find_matching_contract_artifact(
632 output: &mut ProjectCompileOutput,
633 target_path: &Path,
634 target_name: Option<&str>,
635) -> eyre::Result<ConfigurableContractArtifact> {
636 if let Some(name) = target_name {
637 if let Some(artifact) = output.remove(target_path, name) {
638 return Ok(artifact);
639 }
640
641 let target_path = canonicalized(target_path);
642 let matching_source = output.artifact_ids().find_map(|(id, _artifact)| {
643 (id.name == name && canonicalized(&id.source) == target_path).then(|| id.source.clone())
644 });
645
646 matching_source
647 .and_then(|source| output.remove(&source, name))
648 .ok_or_eyre(format!("Could not find artifact `{name}` in the compiled artifacts"))
649 } else {
650 let possible_targets = output
651 .artifact_ids()
652 .filter(|(id, _artifact)| id.source == target_path)
653 .collect::<Vec<_>>();
654
655 if possible_targets.is_empty() {
656 eyre::bail!(
657 "Could not find artifact linked to source `{target_path:?}` in the compiled artifacts"
658 );
659 }
660
661 let (target_id, target_artifact) = possible_targets[0].clone();
662 if possible_targets.len() == 1 {
663 return Ok(target_artifact.clone());
664 }
665
666 if !target_id.name.contains('.')
670 && possible_targets.iter().any(|(id, _)| id.name != target_id.name)
671 {
672 eyre::bail!(
673 "Multiple contracts found in the same file, please specify the target <path>:<contract> or <contract>"
674 );
675 }
676
677 let artifact = possible_targets
680 .iter()
681 .find_map(|(id, artifact)| (id.profile == "default").then_some(*artifact))
682 .unwrap_or(target_artifact);
683
684 Ok(artifact.clone())
685 }
686}
687
688#[cfg(test)]
689mod tests {
690 use super::*;
691 use alloy_dyn_abi::DynSolValue;
692 use alloy_primitives::U256;
693 use semver::Version;
694
695 fn deployed_artifact(name: &str, code: Bytes) -> (ArtifactId, CompactContractBytecode) {
696 (
697 ArtifactId {
698 path: format!("out/{name}.json").into(),
699 name: name.to_owned(),
700 source: format!("src/{name}.sol").into(),
701 version: Version::new(0, 8, 30),
702 build_id: String::new(),
703 profile: "default".to_owned(),
704 },
705 CompactContractBytecode {
706 abi: Some(Default::default()),
707 bytecode: None,
708 deployed_bytecode: Some(CompactDeployedBytecode {
709 bytecode: Some(CompactBytecode {
710 object: BytecodeObject::Bytecode(code),
711 source_map: None,
712 link_references: Default::default(),
713 }),
714 immutable_references: Default::default(),
715 }),
716 },
717 )
718 }
719
720 #[test]
721 fn exact_creation_match_requires_canonical_constructor_suffix() {
722 let abi = JsonAbi::parse(["constructor(uint256 value)"]).unwrap();
723 let bytecode = Bytes::from_static(&[0x60, 0x00]);
724 let contract = ContractData {
725 name: "C".to_owned(),
726 abi,
727 bytecode: Some(BytecodeData {
728 object: Some(BytecodeObject::Bytecode(bytecode.clone())),
729 link_references: BTreeMap::new(),
730 immutable_references: BTreeMap::new(),
731 }),
732 deployed_bytecode: None,
733 storage_layout: None,
734 };
735 let arguments = contract
736 .abi
737 .constructor()
738 .unwrap()
739 .abi_encode_input(&[DynSolValue::Uint(U256::from(1), 256)])
740 .unwrap();
741 let creation = [bytecode.as_ref(), &arguments].concat();
742
743 assert!(matches_contract_creation(&contract, &creation));
744 assert!(!matches_contract_creation(&contract, &creation[..creation.len() - 1]));
745 assert!(!matches_contract_creation(&contract, &[creation, vec![0]].concat()));
746 }
747
748 #[test]
749 fn bytecode_diffing() {
750 assert_eq!(bytecode_diff_score(b"a", b"a"), 0.0);
751 assert_eq!(bytecode_diff_score(b"a", b"b"), 1.0);
752
753 let a_100 = &b"a".repeat(100)[..];
754 assert_eq!(bytecode_diff_score(a_100, &b"b".repeat(100)), 1.0);
755 assert_eq!(bytecode_diff_score(a_100, &b"b".repeat(99)), 1.0);
756 assert_eq!(bytecode_diff_score(a_100, &b"b".repeat(101)), 1.0);
757 assert_eq!(bytecode_diff_score(a_100, &b"b".repeat(120)), 1.0);
758 assert_eq!(bytecode_diff_score(a_100, &b"b".repeat(1000)), 1.0);
759
760 let a_99 = &b"a".repeat(99)[..];
761 assert!(bytecode_diff_score(a_100, a_99) <= 0.01);
762 }
763
764 #[test]
765 fn find_by_deployed_code_exact_with_empty_deployed() {
766 let contracts = ContractsByArtifact::new(vec![]);
767
768 assert!(contracts.find_by_deployed_code_exact(&[]).is_none());
769 }
770
771 #[test]
772 fn find_by_deployed_code_exact_unique_rejects_ambiguity() {
773 let code = Bytes::from_static(&[0x60, 0x00]);
774 let contracts = ContractsByArtifact::new([
775 deployed_artifact("A", code.clone()),
776 deployed_artifact("B", code.clone()),
777 ]);
778
779 assert!(contracts.find_by_deployed_code_exact_unique(&code).is_none());
780
781 let contracts = ContractsByArtifact::new([deployed_artifact("A", code.clone())]);
782 assert_eq!(contracts.find_by_deployed_code_exact_unique(&code).unwrap().0.name, "A");
783 }
784
785 #[test]
786 fn find_by_deployed_code_exact_unique_rejects_partial_metadata_match() {
787 let artifact_code = Bytes::from_static(&[0x60, 0x00, 0xa0, 0x00, 0x01]);
788 let deployed_code = Bytes::from_static(&[0x60, 0x00, 0xf6, 0x00, 0x01]);
789 let contracts = ContractsByArtifact::new([deployed_artifact("A", artifact_code)]);
790
791 assert!(contracts.find_by_deployed_code_exact(&deployed_code).is_some());
792 assert!(contracts.find_by_deployed_code_exact_unique(&deployed_code).is_none());
793 }
794}