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, CompactDeployedBytecode,
12 ConfigurableContractArtifact, ContractBytecodeSome, Offsets,
13 },
14 utils::canonicalized,
15};
16use std::{
17 collections::BTreeMap,
18 ops::Deref,
19 path::{Path, PathBuf},
20 str::FromStr,
21 sync::Arc,
22};
23
24const CALL_PROTECTION_BYTECODE_PREFIX: [u8; 21] =
29 hex!("730000000000000000000000000000000000000000");
30
31#[expect(missing_docs)]
33#[derive(Debug, Clone)]
34pub struct BytecodeData {
35 pub object: Option<BytecodeObject>,
36 pub link_references: BTreeMap<String, BTreeMap<String, Vec<Offsets>>>,
37 pub immutable_references: BTreeMap<String, Vec<Offsets>>,
38}
39
40impl BytecodeData {
41 fn bytes(&self) -> Option<&Bytes> {
42 self.object.as_ref().and_then(|b| b.as_bytes())
43 }
44}
45
46impl From<CompactBytecode> for BytecodeData {
47 fn from(bytecode: CompactBytecode) -> Self {
48 Self {
49 object: Some(bytecode.object),
50 link_references: bytecode.link_references,
51 immutable_references: BTreeMap::new(),
52 }
53 }
54}
55
56impl From<CompactDeployedBytecode> for BytecodeData {
57 fn from(bytecode: CompactDeployedBytecode) -> Self {
58 let (object, link_references) = if let Some(compact) = bytecode.bytecode {
59 (Some(compact.object), compact.link_references)
60 } else {
61 (None, BTreeMap::new())
62 };
63 Self { object, link_references, immutable_references: bytecode.immutable_references }
64 }
65}
66
67#[derive(Debug)]
69pub struct ContractData {
70 pub name: String,
72 pub abi: JsonAbi,
74 pub bytecode: Option<BytecodeData>,
76 pub deployed_bytecode: Option<BytecodeData>,
78}
79
80impl ContractData {
81 pub fn bytecode(&self) -> Option<&Bytes> {
83 self.bytecode.as_ref()?.bytes().filter(|b| !b.is_empty())
84 }
85
86 pub fn deployed_bytecode(&self) -> Option<&Bytes> {
88 self.deployed_bytecode.as_ref()?.bytes().filter(|b| !b.is_empty())
89 }
90
91 pub fn bytecode_without_placeholders(&self) -> Option<Bytes> {
93 strip_bytecode_placeholders(self.bytecode.as_ref()?.object.as_ref()?)
94 }
95
96 pub fn deployed_bytecode_without_placeholders(&self) -> Option<Bytes> {
98 strip_bytecode_placeholders(self.deployed_bytecode.as_ref()?.object.as_ref()?)
99 }
100}
101
102type ArtifactWithContractRef<'a> = (&'a ArtifactId, &'a ContractData);
103
104#[derive(Clone, Default, Debug)]
106pub struct ContractsByArtifact(Arc<BTreeMap<ArtifactId, ContractData>>);
107
108impl ContractsByArtifact {
109 pub fn new(artifacts: impl IntoIterator<Item = (ArtifactId, CompactContractBytecode)>) -> Self {
111 let map = artifacts
112 .into_iter()
113 .filter_map(|(id, artifact)| {
114 let name = id.name.clone();
115 let CompactContractBytecode { abi, bytecode, deployed_bytecode } = artifact;
116 Some((
117 id,
118 ContractData {
119 name,
120 abi: abi?,
121 bytecode: bytecode.map(Into::into),
122 deployed_bytecode: deployed_bytecode.map(Into::into),
123 },
124 ))
125 })
126 .collect();
127 Self(Arc::new(map))
128 }
129
130 pub fn clear(&mut self) {
132 *self = Self::default();
133 }
134
135 pub fn find_by_creation_code(&self, code: &[u8]) -> Option<ArtifactWithContractRef<'_>> {
137 self.find_by_code(code, 0.1, true, ContractData::bytecode)
138 }
139
140 pub fn find_by_deployed_code(&self, code: &[u8]) -> Option<ArtifactWithContractRef<'_>> {
142 self.find_by_code(code, 0.15, false, ContractData::deployed_bytecode)
143 }
144
145 fn find_by_code(
148 &self,
149 code: &[u8],
150 accepted_score: f64,
151 strip_ctor_args: bool,
152 get: impl Fn(&ContractData) -> Option<&Bytes>,
153 ) -> Option<ArtifactWithContractRef<'_>> {
154 self.iter()
155 .filter_map(|(id, contract)| {
156 if let Some(deployed_bytecode) = get(contract) {
157 let mut code = code;
158 if strip_ctor_args && code.len() > deployed_bytecode.len() {
159 if let Some(constructor) = contract.abi.constructor() {
161 let constructor_args = &code[deployed_bytecode.len()..];
162 if constructor.abi_decode_input(constructor_args).is_ok() {
163 code = &code[..deployed_bytecode.len()]
166 }
167 }
168 };
169
170 let score = bytecode_diff_score(deployed_bytecode.as_ref(), code);
171 (score <= accepted_score).then_some((score, (id, contract)))
172 } else {
173 None
174 }
175 })
176 .min_by(|(score1, _), (score2, _)| score1.total_cmp(score2))
177 .map(|(_, data)| data)
178 }
179
180 pub fn find_by_deployed_code_exact(&self, code: &[u8]) -> Option<ArtifactWithContractRef<'_>> {
183 if code.is_empty() {
185 return None;
186 }
187
188 self.iter().find(|(_, contract)| {
189 let Some(deployed_bytecode) = &contract.deployed_bytecode else {
190 return false;
191 };
192 let Some(deployed_code) = &deployed_bytecode.object else {
193 return false;
194 };
195
196 let len = match deployed_code {
197 BytecodeObject::Bytecode(bytes) => bytes.len(),
198 BytecodeObject::Unlinked(bytes) => bytes.len() / 2,
199 };
200
201 if len != code.len() {
202 return false;
203 }
204
205 let mut ignored = deployed_bytecode
207 .immutable_references
208 .values()
209 .chain(deployed_bytecode.link_references.values().flat_map(|v| v.values()))
210 .flatten()
211 .cloned()
212 .collect::<Vec<_>>();
213
214 let has_call_protection = match deployed_code {
219 BytecodeObject::Bytecode(bytes) => {
220 bytes.starts_with(&CALL_PROTECTION_BYTECODE_PREFIX)
221 }
222 BytecodeObject::Unlinked(bytes) => {
223 if let Ok(bytes) =
224 Bytes::from_str(&bytes[..CALL_PROTECTION_BYTECODE_PREFIX.len() * 2])
225 {
226 bytes.starts_with(&CALL_PROTECTION_BYTECODE_PREFIX)
227 } else {
228 false
229 }
230 }
231 };
232
233 if has_call_protection {
234 ignored.push(Offsets { start: 1, length: 20 });
235 }
236
237 ignored.sort_by_key(|o| o.start);
238
239 let mut left = 0;
240 for offset in ignored {
241 let right = offset.start as usize;
242
243 let matched = match deployed_code {
244 BytecodeObject::Bytecode(bytes) => bytes[left..right] == code[left..right],
245 BytecodeObject::Unlinked(bytes) => {
246 if let Ok(bytes) = Bytes::from_str(&bytes[left * 2..right * 2]) {
247 bytes == code[left..right]
248 } else {
249 false
250 }
251 }
252 };
253
254 if !matched {
255 return false;
256 }
257
258 left = right + offset.length as usize;
259 }
260
261 if left < code.len() {
262 match deployed_code {
263 BytecodeObject::Bytecode(bytes) => bytes[left..] == code[left..],
264 BytecodeObject::Unlinked(bytes) => {
265 if let Ok(bytes) = Bytes::from_str(&bytes[left * 2..]) {
266 bytes == code[left..]
267 } else {
268 false
269 }
270 }
271 }
272 } else {
273 true
274 }
275 })
276 }
277
278 pub fn find_by_name_or_identifier(
281 &self,
282 id: &str,
283 ) -> Result<Option<ArtifactWithContractRef<'_>>> {
284 let contracts = self
285 .iter()
286 .filter(|(artifact, _)| artifact.name == id || artifact.identifier() == id)
287 .collect::<Vec<_>>();
288
289 if contracts.len() > 1 {
290 eyre::bail!("{id} has more than one implementation.");
291 }
292
293 Ok(contracts.first().copied())
294 }
295
296 pub fn find_abi_by_name_or_identifier(&self, id: &str) -> Option<JsonAbi> {
298 self.iter()
299 .find(|(artifact, _)| {
300 artifact.name.split(".").next().unwrap() == id || artifact.identifier() == id
301 })
302 .map(|(_, contract)| contract.abi.clone())
303 }
304
305 pub fn find_abi_by_name_or_src_path(&self, name_or_path: &str) -> Option<(JsonAbi, String)> {
309 self.iter()
310 .find(|(artifact, _)| {
311 artifact.name == name_or_path || artifact.source == PathBuf::from(name_or_path)
312 })
313 .map(|(_, contract)| (contract.abi.clone(), contract.name.clone()))
314 }
315
316 pub fn flatten(&self) -> (BTreeMap<Selector, Function>, BTreeMap<B256, Event>, JsonAbi) {
318 let mut funcs = BTreeMap::new();
319 let mut events = BTreeMap::new();
320 let mut errors_abi = JsonAbi::new();
321 for (_name, contract) in self.iter() {
322 for func in contract.abi.functions() {
323 funcs.insert(func.selector(), func.clone());
324 }
325 for event in contract.abi.events() {
326 events.insert(event.selector(), event.clone());
327 }
328 for error in contract.abi.errors() {
329 errors_abi.errors.entry(error.name.clone()).or_default().push(error.clone());
330 }
331 }
332 (funcs, events, errors_abi)
333 }
334}
335
336impl From<ProjectCompileOutput> for ContractsByArtifact {
337 fn from(value: ProjectCompileOutput) -> Self {
338 Self::new(value.into_artifacts().map(|(id, ar)| {
339 (
340 id,
341 CompactContractBytecode {
342 abi: ar.abi,
343 bytecode: ar.bytecode,
344 deployed_bytecode: ar.deployed_bytecode,
345 },
346 )
347 }))
348 }
349}
350
351impl Deref for ContractsByArtifact {
352 type Target = BTreeMap<ArtifactId, ContractData>;
353
354 fn deref(&self) -> &Self::Target {
355 &self.0
356 }
357}
358
359pub type ContractsByAddress = BTreeMap<Address, (String, JsonAbi)>;
361
362pub fn bytecode_diff_score<'a>(mut a: &'a [u8], mut b: &'a [u8]) -> f64 {
366 if a.len() < b.len() {
368 std::mem::swap(&mut a, &mut b);
369 }
370
371 let mut n_different_bytes = a.len() - b.len();
373
374 if n_different_bytes > 32 && n_different_bytes * 10 > a.len() {
379 return 1.0;
380 }
381
382 n_different_bytes += unsafe { count_different_bytes(a, b) };
385
386 n_different_bytes as f64 / a.len() as f64
387}
388
389unsafe fn count_different_bytes(a: &[u8], b: &[u8]) -> usize {
395 let a_ptr = a.as_ptr();
400 let b_ptr = b.as_ptr();
401 let len = b.len();
402
403 let mut sum = 0;
404 let mut i = 0;
405 while i < len {
406 sum += unsafe { *a_ptr.add(i) != *b_ptr.add(i) } as usize;
408 i += 1;
409 }
410 sum
411}
412
413pub fn get_contract_name(id: &str) -> &str {
432 id.rsplit(':').next().unwrap_or(id)
433}
434
435pub fn get_file_name(id: &str) -> &str {
447 id.split(':').next().unwrap_or(id)
448}
449
450pub fn compact_to_contract(contract: CompactContractBytecode) -> Result<ContractBytecodeSome> {
452 Ok(ContractBytecodeSome {
453 abi: contract.abi.ok_or_else(|| eyre::eyre!("No contract abi"))?,
454 bytecode: contract.bytecode.ok_or_else(|| eyre::eyre!("No contract bytecode"))?.into(),
455 deployed_bytecode: contract
456 .deployed_bytecode
457 .ok_or_else(|| eyre::eyre!("No contract deployed bytecode"))?
458 .into(),
459 })
460}
461
462pub fn find_target_path(project: &Project, identifier: &PathOrContractInfo) -> Result<PathBuf> {
464 match identifier {
465 PathOrContractInfo::Path(path) => Ok(canonicalized(project.root().join(path))),
466 PathOrContractInfo::ContractInfo(info) => {
467 if let Some(path) = info.path.as_ref() {
468 let path = canonicalized(project.root().join(path));
469 let sources = project.sources()?;
470 let contract_path = sources
471 .iter()
472 .find_map(|(src_path, _)| {
473 if **src_path == path {
474 return Some(src_path.clone());
475 }
476 None
477 })
478 .ok_or_else(|| {
479 eyre::eyre!(
480 "Could not find source file for contract `{}` at {}",
481 info.name,
482 path.strip_prefix(project.root()).unwrap().display()
483 )
484 })?;
485 return Ok(contract_path);
486 }
487 let path = project.find_contract_path(&info.name)?;
491 Ok(path)
492 }
493 }
494}
495
496pub fn find_matching_contract_artifact(
498 output: &mut ProjectCompileOutput,
499 target_path: &Path,
500 target_name: Option<&str>,
501) -> eyre::Result<ConfigurableContractArtifact> {
502 if let Some(name) = target_name {
503 output
504 .remove(target_path, name)
505 .ok_or_eyre(format!("Could not find artifact `{name}` in the compiled artifacts"))
506 } else {
507 let possible_targets = output
508 .artifact_ids()
509 .filter(|(id, _artifact)| id.source == target_path)
510 .collect::<Vec<_>>();
511
512 if possible_targets.is_empty() {
513 eyre::bail!(
514 "Could not find artifact linked to source `{target_path:?}` in the compiled artifacts"
515 );
516 }
517
518 let (target_id, target_artifact) = possible_targets[0].clone();
519 if possible_targets.len() == 1 {
520 return Ok(target_artifact.clone());
521 }
522
523 if !target_id.name.contains(".")
527 && possible_targets.iter().any(|(id, _)| id.name != target_id.name)
528 {
529 eyre::bail!(
530 "Multiple contracts found in the same file, please specify the target <path>:<contract> or <contract>"
531 );
532 }
533
534 let artifact = possible_targets
537 .iter()
538 .find_map(|(id, artifact)| if id.profile == "default" { Some(*artifact) } else { None })
539 .unwrap_or(target_artifact);
540
541 Ok(artifact.clone())
542 }
543}
544#[cfg(test)]
545mod tests {
546 use super::*;
547
548 #[test]
549 fn bytecode_diffing() {
550 assert_eq!(bytecode_diff_score(b"a", b"a"), 0.0);
551 assert_eq!(bytecode_diff_score(b"a", b"b"), 1.0);
552
553 let a_100 = &b"a".repeat(100)[..];
554 assert_eq!(bytecode_diff_score(a_100, &b"b".repeat(100)), 1.0);
555 assert_eq!(bytecode_diff_score(a_100, &b"b".repeat(99)), 1.0);
556 assert_eq!(bytecode_diff_score(a_100, &b"b".repeat(101)), 1.0);
557 assert_eq!(bytecode_diff_score(a_100, &b"b".repeat(120)), 1.0);
558 assert_eq!(bytecode_diff_score(a_100, &b"b".repeat(1000)), 1.0);
559
560 let a_99 = &b"a".repeat(99)[..];
561 assert!(bytecode_diff_score(a_100, a_99) <= 0.01);
562 }
563
564 #[test]
565 fn find_by_deployed_code_exact_with_empty_deployed() {
566 let contracts = ContractsByArtifact::new(vec![]);
567
568 assert!(contracts.find_by_deployed_code_exact(&[]).is_none());
569 }
570}