1#![cfg_attr(not(test), warn(unused_crate_dependencies))]
6#![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))]
7
8use alloy_primitives::{Address, Bytes, B256};
9use foundry_compilers::{
10 artifacts::{CompactContractBytecodeCow, Libraries},
11 contracts::ArtifactContracts,
12 Artifact, ArtifactId,
13};
14use semver::Version;
15use std::{
16 collections::{BTreeMap, BTreeSet},
17 path::{Path, PathBuf},
18 str::FromStr,
19};
20
21#[derive(Debug, thiserror::Error)]
23pub enum LinkerError {
24 #[error("wasn't able to find artifact for library {name} at {file}")]
25 MissingLibraryArtifact { file: String, name: String },
26 #[error("target artifact is not present in provided artifacts set")]
27 MissingTargetArtifact,
28 #[error(transparent)]
29 InvalidAddress(<Address as std::str::FromStr>::Err),
30 #[error("cyclic dependency found, can't link libraries via CREATE2")]
31 CyclicDependency,
32}
33
34pub struct Linker<'a> {
35 pub root: PathBuf,
37 pub contracts: ArtifactContracts<CompactContractBytecodeCow<'a>>,
39}
40
41pub struct LinkOutput {
43 pub libraries: Libraries,
46 pub libs_to_deploy: Vec<Bytes>,
49}
50
51impl<'a> Linker<'a> {
52 pub fn new(
53 root: impl Into<PathBuf>,
54 contracts: ArtifactContracts<CompactContractBytecodeCow<'a>>,
55 ) -> Self {
56 Linker { root: root.into(), contracts }
57 }
58
59 fn convert_artifact_id_to_lib_path(&self, id: &ArtifactId) -> (PathBuf, String) {
64 let path = id.source.strip_prefix(self.root.as_path()).unwrap_or(&id.source);
65 let name = id.name.split('.').next().unwrap();
67
68 (path.to_path_buf(), name.to_owned())
69 }
70
71 fn find_artifact_id_by_library_path(
76 &'a self,
77 file: &str,
78 name: &str,
79 version: Option<&Version>,
80 ) -> Option<&'a ArtifactId> {
81 for id in self.contracts.keys() {
82 if let Some(version) = version {
83 if id.version != *version {
84 continue;
85 }
86 }
87 let (artifact_path, artifact_name) = self.convert_artifact_id_to_lib_path(id);
88
89 if artifact_name == *name && artifact_path == Path::new(file) {
90 return Some(id);
91 }
92 }
93
94 None
95 }
96
97 fn collect_dependencies(
99 &'a self,
100 target: &'a ArtifactId,
101 deps: &mut BTreeSet<&'a ArtifactId>,
102 ) -> Result<(), LinkerError> {
103 let contract = self.contracts.get(target).ok_or(LinkerError::MissingTargetArtifact)?;
104
105 let mut references = BTreeMap::new();
106 if let Some(bytecode) = &contract.bytecode {
107 references.extend(bytecode.link_references.clone());
108 }
109 if let Some(deployed_bytecode) = &contract.deployed_bytecode {
110 if let Some(bytecode) = &deployed_bytecode.bytecode {
111 references.extend(bytecode.link_references.clone());
112 }
113 }
114
115 for (file, libs) in &references {
116 for contract in libs.keys() {
117 let id = self
118 .find_artifact_id_by_library_path(file, contract, Some(&target.version))
119 .ok_or_else(|| LinkerError::MissingLibraryArtifact {
120 file: file.to_string(),
121 name: contract.to_string(),
122 })?;
123 if deps.insert(id) {
124 self.collect_dependencies(id, deps)?;
125 }
126 }
127 }
128
129 Ok(())
130 }
131
132 pub fn link_with_nonce_or_address(
142 &'a self,
143 libraries: Libraries,
144 sender: Address,
145 mut nonce: u64,
146 targets: impl IntoIterator<Item = &'a ArtifactId>,
147 ) -> Result<LinkOutput, LinkerError> {
148 let mut libraries = libraries.with_stripped_file_prefixes(self.root.as_path());
151
152 let mut needed_libraries = BTreeSet::new();
153 for target in targets {
154 self.collect_dependencies(target, &mut needed_libraries)?;
155 }
156
157 let mut libs_to_deploy = Vec::new();
158
159 for id in needed_libraries {
162 let (lib_path, lib_name) = self.convert_artifact_id_to_lib_path(id);
163
164 libraries.libs.entry(lib_path).or_default().entry(lib_name).or_insert_with(|| {
165 let address = sender.create(nonce);
166 libs_to_deploy.push((id, address));
167 nonce += 1;
168
169 address.to_checksum(None)
170 });
171 }
172
173 let libs_to_deploy = libs_to_deploy
175 .into_iter()
176 .map(|(id, _)| {
177 Ok(self.link(id, &libraries)?.get_bytecode_bytes().unwrap().into_owned())
178 })
179 .collect::<Result<Vec<_>, LinkerError>>()?;
180
181 Ok(LinkOutput { libraries, libs_to_deploy })
182 }
183
184 pub fn link_with_create2(
185 &'a self,
186 libraries: Libraries,
187 sender: Address,
188 salt: B256,
189 target: &'a ArtifactId,
190 ) -> Result<LinkOutput, LinkerError> {
191 let mut libraries = libraries.with_stripped_file_prefixes(self.root.as_path());
194
195 let mut needed_libraries = BTreeSet::new();
196 self.collect_dependencies(target, &mut needed_libraries)?;
197
198 let mut needed_libraries = needed_libraries
199 .into_iter()
200 .filter(|id| {
201 let (file, name) = self.convert_artifact_id_to_lib_path(id);
203 !libraries.libs.contains_key(&file) || !libraries.libs[&file].contains_key(&name)
204 })
205 .map(|id| {
206 let bytecode = self.link(id, &libraries).unwrap().bytecode.unwrap();
208 (id, bytecode)
209 })
210 .collect::<Vec<_>>();
211
212 let mut libs_to_deploy = Vec::new();
213
214 while !needed_libraries.is_empty() {
217 let deployable = needed_libraries
219 .iter()
220 .enumerate()
221 .find(|(_, (_, bytecode))| !bytecode.object.is_unlinked());
222
223 let Some((index, &(id, _))) = deployable else {
225 return Err(LinkerError::CyclicDependency);
226 };
227 let (_, bytecode) = needed_libraries.swap_remove(index);
228 let code = bytecode.bytes().unwrap();
229 let address = sender.create2_from_code(salt, code);
230 libs_to_deploy.push(code.clone());
231
232 let (file, name) = self.convert_artifact_id_to_lib_path(id);
233
234 for (_, bytecode) in &mut needed_libraries {
235 bytecode.to_mut().link(&file.to_string_lossy(), &name, address);
236 }
237
238 libraries.libs.entry(file).or_default().insert(name, address.to_checksum(None));
239 }
240
241 Ok(LinkOutput { libraries, libs_to_deploy })
242 }
243
244 pub fn link(
246 &self,
247 target: &ArtifactId,
248 libraries: &Libraries,
249 ) -> Result<CompactContractBytecodeCow<'a>, LinkerError> {
250 let mut contract =
251 self.contracts.get(target).ok_or(LinkerError::MissingTargetArtifact)?.clone();
252 for (file, libs) in &libraries.libs {
253 for (name, address) in libs {
254 let address = Address::from_str(address).map_err(LinkerError::InvalidAddress)?;
255 if let Some(bytecode) = contract.bytecode.as_mut() {
256 bytecode.to_mut().link(&file.to_string_lossy(), name, address);
257 }
258 if let Some(deployed_bytecode) =
259 contract.deployed_bytecode.as_mut().and_then(|b| b.to_mut().bytecode.as_mut())
260 {
261 deployed_bytecode.link(&file.to_string_lossy(), name, address);
262 }
263 }
264 }
265 Ok(contract)
266 }
267
268 pub fn get_linked_artifacts(
269 &self,
270 libraries: &Libraries,
271 ) -> Result<ArtifactContracts, LinkerError> {
272 self.contracts.keys().map(|id| Ok((id.clone(), self.link(id, libraries)?))).collect()
273 }
274
275 pub fn get_linked_artifacts_cow(
276 &self,
277 libraries: &Libraries,
278 ) -> Result<ArtifactContracts<CompactContractBytecodeCow<'a>>, LinkerError> {
279 self.contracts.keys().map(|id| Ok((id.clone(), self.link(id, libraries)?))).collect()
280 }
281}
282
283#[cfg(test)]
284mod tests {
285 use super::*;
286 use alloy_primitives::{address, fixed_bytes, map::HashMap};
287 use foundry_compilers::{
288 multi::MultiCompiler,
289 solc::{Solc, SolcCompiler},
290 Project, ProjectCompileOutput, ProjectPathsConfig,
291 };
292
293 struct LinkerTest {
294 project: Project,
295 output: ProjectCompileOutput,
296 dependency_assertions: HashMap<String, Vec<(String, Address)>>,
297 }
298
299 impl LinkerTest {
300 fn new(path: impl Into<PathBuf>, strip_prefixes: bool) -> Self {
301 let path = path.into();
302 let paths = ProjectPathsConfig::builder()
303 .root("../../testdata")
304 .lib("../../testdata/lib")
305 .sources(path.clone())
306 .tests(path)
307 .build()
308 .unwrap();
309
310 let solc = Solc::find_or_install(&Version::new(0, 8, 18)).unwrap();
311 let project = Project::builder()
312 .paths(paths)
313 .ephemeral()
314 .no_artifacts()
315 .build(MultiCompiler { solc: Some(SolcCompiler::Specific(solc)), vyper: None })
316 .unwrap();
317
318 let mut output = project.compile().unwrap();
319
320 if strip_prefixes {
321 output = output.with_stripped_file_prefixes(project.root());
322 }
323
324 Self { project, output, dependency_assertions: HashMap::default() }
325 }
326
327 fn assert_dependencies(
328 mut self,
329 artifact_id: String,
330 deps: Vec<(String, Address)>,
331 ) -> Self {
332 self.dependency_assertions.insert(artifact_id, deps);
333 self
334 }
335
336 fn test_with_sender_and_nonce(self, sender: Address, initial_nonce: u64) {
337 let linker = Linker::new(self.project.root(), self.output.artifact_ids().collect());
338 for (id, identifier) in self.iter_linking_targets(&linker) {
339 let output = linker
340 .link_with_nonce_or_address(Default::default(), sender, initial_nonce, [id])
341 .expect("Linking failed");
342 self.validate_assertions(identifier, output);
343 }
344 }
345
346 fn test_with_create2(self, sender: Address, salt: B256) {
347 let linker = Linker::new(self.project.root(), self.output.artifact_ids().collect());
348 for (id, identifier) in self.iter_linking_targets(&linker) {
349 let output = linker
350 .link_with_create2(Default::default(), sender, salt, id)
351 .expect("Linking failed");
352 self.validate_assertions(identifier, output);
353 }
354 }
355
356 fn iter_linking_targets<'a>(
357 &'a self,
358 linker: &'a Linker<'_>,
359 ) -> impl IntoIterator<Item = (&'a ArtifactId, String)> + 'a {
360 linker.contracts.keys().filter_map(move |id| {
361 let source = id
365 .source
366 .strip_prefix(self.project.root())
367 .unwrap_or(&id.source)
368 .to_string_lossy();
369 let identifier = format!("{source}:{}", id.name);
370
371 if identifier.contains("DSTest") {
374 return None;
375 }
376
377 Some((id, identifier))
378 })
379 }
380
381 fn validate_assertions(&self, identifier: String, output: LinkOutput) {
382 let LinkOutput { libs_to_deploy, libraries } = output;
383
384 let assertions = self
385 .dependency_assertions
386 .get(&identifier)
387 .unwrap_or_else(|| panic!("Unexpected artifact: {identifier}"));
388
389 assert_eq!(
390 libs_to_deploy.len(),
391 assertions.len(),
392 "artifact {identifier} has more/less dependencies than expected ({} vs {}): {:#?}",
393 libs_to_deploy.len(),
394 assertions.len(),
395 libs_to_deploy
396 );
397
398 for (dep_identifier, address) in assertions {
399 let (file, name) = dep_identifier.split_once(':').unwrap();
400 if let Some(lib_address) =
401 libraries.libs.get(Path::new(file)).and_then(|libs| libs.get(name))
402 {
403 assert_eq!(
404 *lib_address,
405 address.to_string(),
406 "incorrect library address for dependency {dep_identifier} of {identifier}"
407 );
408 } else {
409 panic!("Library {dep_identifier} not found");
410 }
411 }
412 }
413 }
414
415 fn link_test(path: impl Into<PathBuf>, test_fn: impl Fn(LinkerTest)) {
416 let path = path.into();
417 test_fn(LinkerTest::new(path.clone(), true));
418 test_fn(LinkerTest::new(path, false));
419 }
420
421 #[test]
422 fn link_simple() {
423 link_test("../../testdata/default/linking/simple", |linker| {
424 linker
425 .assert_dependencies("default/linking/simple/Simple.t.sol:Lib".to_string(), vec![])
426 .assert_dependencies(
427 "default/linking/simple/Simple.t.sol:LibraryConsumer".to_string(),
428 vec![(
429 "default/linking/simple/Simple.t.sol:Lib".to_string(),
430 address!("0x5a443704dd4b594b382c22a083e2bd3090a6fef3"),
431 )],
432 )
433 .assert_dependencies(
434 "default/linking/simple/Simple.t.sol:SimpleLibraryLinkingTest".to_string(),
435 vec![(
436 "default/linking/simple/Simple.t.sol:Lib".to_string(),
437 address!("0x5a443704dd4b594b382c22a083e2bd3090a6fef3"),
438 )],
439 )
440 .test_with_sender_and_nonce(Address::default(), 1);
441 });
442 }
443
444 #[test]
445 fn link_nested() {
446 link_test("../../testdata/default/linking/nested", |linker| {
447 linker
448 .assert_dependencies("default/linking/nested/Nested.t.sol:Lib".to_string(), vec![])
449 .assert_dependencies(
450 "default/linking/nested/Nested.t.sol:NestedLib".to_string(),
451 vec![(
452 "default/linking/nested/Nested.t.sol:Lib".to_string(),
453 address!("0x5a443704dd4b594b382c22a083e2bd3090a6fef3"),
454 )],
455 )
456 .assert_dependencies(
457 "default/linking/nested/Nested.t.sol:LibraryConsumer".to_string(),
458 vec![
459 (
462 "default/linking/nested/Nested.t.sol:Lib".to_string(),
463 Address::from_str("0x5a443704dd4b594b382c22a083e2bd3090a6fef3")
464 .unwrap(),
465 ),
466 (
467 "default/linking/nested/Nested.t.sol:NestedLib".to_string(),
468 Address::from_str("0x47e9Fbef8C83A1714F1951F142132E6e90F5fa5D")
469 .unwrap(),
470 ),
471 ],
472 )
473 .assert_dependencies(
474 "default/linking/nested/Nested.t.sol:NestedLibraryLinkingTest".to_string(),
475 vec![
476 (
477 "default/linking/nested/Nested.t.sol:Lib".to_string(),
478 Address::from_str("0x5a443704dd4b594b382c22a083e2bd3090a6fef3")
479 .unwrap(),
480 ),
481 (
482 "default/linking/nested/Nested.t.sol:NestedLib".to_string(),
483 Address::from_str("0x47e9fbef8c83a1714f1951f142132e6e90f5fa5d")
484 .unwrap(),
485 ),
486 ],
487 )
488 .test_with_sender_and_nonce(Address::default(), 1);
489 });
490 }
491
492 #[test]
493 fn link_duplicate() {
494 link_test("../../testdata/default/linking/duplicate", |linker| {
495 linker
496 .assert_dependencies(
497 "default/linking/duplicate/Duplicate.t.sol:A".to_string(),
498 vec![],
499 )
500 .assert_dependencies(
501 "default/linking/duplicate/Duplicate.t.sol:B".to_string(),
502 vec![],
503 )
504 .assert_dependencies(
505 "default/linking/duplicate/Duplicate.t.sol:C".to_string(),
506 vec![(
507 "default/linking/duplicate/Duplicate.t.sol:A".to_string(),
508 address!("0x5a443704dd4b594b382c22a083e2bd3090a6fef3"),
509 )],
510 )
511 .assert_dependencies(
512 "default/linking/duplicate/Duplicate.t.sol:D".to_string(),
513 vec![(
514 "default/linking/duplicate/Duplicate.t.sol:B".to_string(),
515 address!("0x5a443704dd4b594b382c22a083e2bd3090a6fef3"),
516 )],
517 )
518 .assert_dependencies(
519 "default/linking/duplicate/Duplicate.t.sol:E".to_string(),
520 vec![
521 (
522 "default/linking/duplicate/Duplicate.t.sol:A".to_string(),
523 Address::from_str("0x5a443704dd4b594b382c22a083e2bd3090a6fef3")
524 .unwrap(),
525 ),
526 (
527 "default/linking/duplicate/Duplicate.t.sol:C".to_string(),
528 Address::from_str("0x47e9fbef8c83a1714f1951f142132e6e90f5fa5d")
529 .unwrap(),
530 ),
531 ],
532 )
533 .assert_dependencies(
534 "default/linking/duplicate/Duplicate.t.sol:LibraryConsumer".to_string(),
535 vec![
536 (
537 "default/linking/duplicate/Duplicate.t.sol:A".to_string(),
538 Address::from_str("0x5a443704dd4b594b382c22a083e2bd3090a6fef3")
539 .unwrap(),
540 ),
541 (
542 "default/linking/duplicate/Duplicate.t.sol:B".to_string(),
543 Address::from_str("0x47e9fbef8c83a1714f1951f142132e6e90f5fa5d")
544 .unwrap(),
545 ),
546 (
547 "default/linking/duplicate/Duplicate.t.sol:C".to_string(),
548 Address::from_str("0x8be503bcded90ed42eff31f56199399b2b0154ca")
549 .unwrap(),
550 ),
551 (
552 "default/linking/duplicate/Duplicate.t.sol:D".to_string(),
553 Address::from_str("0x47c5e40890bce4a473a49d7501808b9633f29782")
554 .unwrap(),
555 ),
556 (
557 "default/linking/duplicate/Duplicate.t.sol:E".to_string(),
558 Address::from_str("0x29b2440db4a256b0c1e6d3b4cdcaa68e2440a08f")
559 .unwrap(),
560 ),
561 ],
562 )
563 .assert_dependencies(
564 "default/linking/duplicate/Duplicate.t.sol:DuplicateLibraryLinkingTest"
565 .to_string(),
566 vec![
567 (
568 "default/linking/duplicate/Duplicate.t.sol:A".to_string(),
569 Address::from_str("0x5a443704dd4b594b382c22a083e2bd3090a6fef3")
570 .unwrap(),
571 ),
572 (
573 "default/linking/duplicate/Duplicate.t.sol:B".to_string(),
574 Address::from_str("0x47e9fbef8c83a1714f1951f142132e6e90f5fa5d")
575 .unwrap(),
576 ),
577 (
578 "default/linking/duplicate/Duplicate.t.sol:C".to_string(),
579 Address::from_str("0x8be503bcded90ed42eff31f56199399b2b0154ca")
580 .unwrap(),
581 ),
582 (
583 "default/linking/duplicate/Duplicate.t.sol:D".to_string(),
584 Address::from_str("0x47c5e40890bce4a473a49d7501808b9633f29782")
585 .unwrap(),
586 ),
587 (
588 "default/linking/duplicate/Duplicate.t.sol:E".to_string(),
589 Address::from_str("0x29b2440db4a256b0c1e6d3b4cdcaa68e2440a08f")
590 .unwrap(),
591 ),
592 ],
593 )
594 .test_with_sender_and_nonce(Address::default(), 1);
595 });
596 }
597
598 #[test]
599 fn link_cycle() {
600 link_test("../../testdata/default/linking/cycle", |linker| {
601 linker
602 .assert_dependencies(
603 "default/linking/cycle/Cycle.t.sol:Foo".to_string(),
604 vec![
605 (
606 "default/linking/cycle/Cycle.t.sol:Foo".to_string(),
607 Address::from_str("0x47e9Fbef8C83A1714F1951F142132E6e90F5fa5D")
608 .unwrap(),
609 ),
610 (
611 "default/linking/cycle/Cycle.t.sol:Bar".to_string(),
612 Address::from_str("0x5a443704dd4B594B382c22a083e2BD3090A6feF3")
613 .unwrap(),
614 ),
615 ],
616 )
617 .assert_dependencies(
618 "default/linking/cycle/Cycle.t.sol:Bar".to_string(),
619 vec![
620 (
621 "default/linking/cycle/Cycle.t.sol:Foo".to_string(),
622 Address::from_str("0x47e9Fbef8C83A1714F1951F142132E6e90F5fa5D")
623 .unwrap(),
624 ),
625 (
626 "default/linking/cycle/Cycle.t.sol:Bar".to_string(),
627 Address::from_str("0x5a443704dd4B594B382c22a083e2BD3090A6feF3")
628 .unwrap(),
629 ),
630 ],
631 )
632 .test_with_sender_and_nonce(Address::default(), 1);
633 });
634 }
635
636 #[test]
637 fn link_create2_nested() {
638 link_test("../../testdata/default/linking/nested", |linker| {
639 linker
640 .assert_dependencies("default/linking/nested/Nested.t.sol:Lib".to_string(), vec![])
641 .assert_dependencies(
642 "default/linking/nested/Nested.t.sol:NestedLib".to_string(),
643 vec![(
644 "default/linking/nested/Nested.t.sol:Lib".to_string(),
645 address!("0xddb1Cd2497000DAeA687CEa3dc34Af44084BEa74"),
646 )],
647 )
648 .assert_dependencies(
649 "default/linking/nested/Nested.t.sol:LibraryConsumer".to_string(),
650 vec![
651 (
654 "default/linking/nested/Nested.t.sol:Lib".to_string(),
655 Address::from_str("0xddb1Cd2497000DAeA687CEa3dc34Af44084BEa74")
656 .unwrap(),
657 ),
658 (
659 "default/linking/nested/Nested.t.sol:NestedLib".to_string(),
660 Address::from_str("0xfebE2F30641170642f317Ff6F644Cee60E7Ac369")
661 .unwrap(),
662 ),
663 ],
664 )
665 .assert_dependencies(
666 "default/linking/nested/Nested.t.sol:NestedLibraryLinkingTest".to_string(),
667 vec![
668 (
669 "default/linking/nested/Nested.t.sol:Lib".to_string(),
670 Address::from_str("0xddb1Cd2497000DAeA687CEa3dc34Af44084BEa74")
671 .unwrap(),
672 ),
673 (
674 "default/linking/nested/Nested.t.sol:NestedLib".to_string(),
675 Address::from_str("0xfebE2F30641170642f317Ff6F644Cee60E7Ac369")
676 .unwrap(),
677 ),
678 ],
679 )
680 .test_with_create2(
681 Address::default(),
682 fixed_bytes!(
683 "19bf59b7b67ae8edcbc6e53616080f61fa99285c061450ad601b0bc40c9adfc9"
684 ),
685 );
686 });
687 }
688}