1use 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
25const CALL_PROTECTION_BYTECODE_PREFIX: [u8; 21] =
30 hex!("730000000000000000000000000000000000000000");
31
32#[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#[derive(Debug)]
70pub struct ContractData {
71 pub name: String,
73 pub abi: JsonAbi,
75 pub bytecode: Option<BytecodeData>,
77 pub deployed_bytecode: Option<BytecodeData>,
79 pub storage_layout: Option<Arc<StorageLayout>>,
81}
82
83impl ContractData {
84 pub fn bytecode(&self) -> Option<&Bytes> {
86 self.bytecode.as_ref()?.bytes().filter(|b| !b.is_empty())
87 }
88
89 pub fn deployed_bytecode(&self) -> Option<&Bytes> {
91 self.deployed_bytecode.as_ref()?.bytes().filter(|b| !b.is_empty())
92 }
93
94 pub fn bytecode_without_placeholders(&self) -> Option<Bytes> {
96 strip_bytecode_placeholders(self.bytecode.as_ref()?.object.as_ref()?)
97 }
98
99 pub fn deployed_bytecode_without_placeholders(&self) -> Option<Bytes> {
101 strip_bytecode_placeholders(self.deployed_bytecode.as_ref()?.object.as_ref()?)
102 }
103}
104
105pub struct ContractsByArtifactBuilder<'a> {
108 artifacts: BTreeMap<ArtifactId, CompactContractBytecodeCow<'a>>,
110 storage_layouts: BTreeMap<ArtifactId, StorageLayout>,
112}
113
114impl<'a> ContractsByArtifactBuilder<'a> {
115 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 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 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#[derive(Clone, Default, Debug)]
161pub struct ContractsByArtifact(Arc<BTreeMap<ArtifactId, ContractData>>);
162
163impl ContractsByArtifact {
164 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 pub fn clear(&mut self) {
188 *self = Self::default();
189 }
190
191 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 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 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 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 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 pub fn find_by_deployed_code_exact(&self, code: &[u8]) -> Option<ArtifactWithContractRef<'_>> {
239 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 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 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 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 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 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 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
415pub type ContractsByAddress = BTreeMap<Address, (String, JsonAbi)>;
417
418pub fn bytecode_diff_score<'a>(mut a: &'a [u8], mut b: &'a [u8]) -> f64 {
422 if a.len() < b.len() {
424 std::mem::swap(&mut a, &mut b);
425 }
426
427 let mut n_different_bytes = a.len() - b.len();
429
430 if n_different_bytes > 32 && n_different_bytes * 10 > a.len() {
435 return 1.0;
436 }
437
438 n_different_bytes += unsafe { count_different_bytes(a, b) };
441
442 n_different_bytes as f64 / a.len() as f64
443}
444
445unsafe fn count_different_bytes(a: &[u8], b: &[u8]) -> usize {
451 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 sum += unsafe { *a_ptr.add(i) != *b_ptr.add(i) } as usize;
464 i += 1;
465 }
466 sum
467}
468
469pub fn get_contract_name(id: &str) -> &str {
488 id.rsplit(':').next().unwrap_or(id)
489}
490
491pub fn get_file_name(id: &str) -> &str {
503 id.split(':').next().unwrap_or(id)
504}
505
506pub 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
518pub 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 let path = project.find_contract_path(&info.name)?;
547 Ok(path)
548 }
549 }
550}
551
552pub 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 !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 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}