fix(test:zk): proper stripping during link

feat(link:script): register target & deps

Author: Karrq

crates/forge/tests/it/zk/linking.rs

@@ -8,6 +8,8 @@ use crate::{
     test_helpers::{deploy_zk_contract, run_zk_script_test, TEST_DATA_DEFAULT},
 };
 
+const ZKSOLC_MIN_LINKING_VERSION: Version = Version::new(1, 5, 9);
+
 #[tokio::test(flavor = "multi_thread")]
 async fn test_zk_deploy_time_linking() {
     let runner = TEST_DATA_DEFAULT.runner_zksync();
@@ -20,7 +22,7 @@ forgetest_async!(
     #[should_panic = "no bytecode for contract; is it abstract or unlinked?"]
     script_zk_fails_indirect_reference_to_unlinked,
     |prj, cmd| {
-        setup_libs_prj(&mut prj, &mut cmd, Some(Version::new(1, 5, 9)));
+        setup_libs_prj(&mut prj, &mut cmd, Some(ZKSOLC_MIN_LINKING_VERSION));
         run_zk_script_test(
             prj.root(),
             &mut cmd,
@@ -35,7 +37,7 @@ forgetest_async!(
 );
 
 forgetest_async!(script_zk_deploy_time_linking, |prj, cmd| {
-    setup_libs_prj(&mut prj, &mut cmd, Some(Version::new(1, 5, 9)));
+    setup_libs_prj(&mut prj, &mut cmd, Some(ZKSOLC_MIN_LINKING_VERSION));
     run_zk_script_test(
         prj.root(),
         &mut cmd,
@@ -54,7 +56,10 @@ forgetest_async!(
     #[should_panic = "deploy-time linking not supported"]
     script_zk_deploy_time_linking_fails_older_version,
     |prj, cmd| {
-        setup_libs_prj(&mut prj, &mut cmd, Some(Version::new(1, 5, 8)));
+        let mut version = ZKSOLC_MIN_LINKING_VERSION;
+        version.patch -= 1;
+
+        setup_libs_prj(&mut prj, &mut cmd, Some(version));
         run_zk_script_test(
             prj.root(),
             &mut cmd,

crates/linking/src/zksync.rs

@@ -115,10 +115,11 @@ impl<'a> ZkLinker<'a> {
 
     /// Performs DFS on the graph of link references, and populates `deps` with all found libraries,
     /// including ones of factory deps.
-    fn zk_collect_dependencies(
+    pub fn zk_collect_dependencies(
         &'a self,
         target: &'a ArtifactId,
-        deps: &mut BTreeSet<&'a ArtifactId>,
+        libraries: &mut BTreeSet<&'a ArtifactId>,
+        factory_deps: Option<&mut BTreeSet<&'a ArtifactId>>,
     ) -> Result<(), LinkerError> {
         let (_, artifact) = self
             .zk_artifacts()
@@ -131,10 +132,11 @@ impl<'a> ZkLinker<'a> {
         }
 
         // find all nested factory deps's link references
-        let mut factory_deps = BTreeSet::new();
-        self.zk_collect_factory_deps(target, &mut factory_deps)?;
+        let mut fdeps_default = BTreeSet::new();
+        let factory_deps = factory_deps.unwrap_or(&mut fdeps_default);
+        self.zk_collect_factory_deps(target, factory_deps)?;
 
-        for (_, fdep) in factory_deps.into_iter().filter_map(|target| {
+        for (_, fdep) in factory_deps.iter().filter_map(|target| {
             self.zk_artifacts().find(|(id, _)| id.source == target.source && id.name == target.name)
         }) {
             if let Some(bytecode) = &fdep.bytecode {
@@ -151,8 +153,8 @@ impl<'a> ZkLinker<'a> {
                         file: file.to_string(),
                         name: contract.to_string(),
                     })?;
-                if deps.insert(id) {
-                    self.zk_collect_dependencies(id, deps)?;
+                if libraries.insert(id) {
+                    self.zk_collect_dependencies(id, libraries, Some(factory_deps))?;
                 }
             }
         }
@@ -182,7 +184,7 @@ impl<'a> ZkLinker<'a> {
 
         let mut needed_libraries = BTreeSet::new();
         for target in targets {
-            self.zk_collect_dependencies(target, &mut needed_libraries)?;
+            self.zk_collect_dependencies(target, &mut needed_libraries, None)?;
         }
 
         let mut libs_to_deploy = Vec::new();
@@ -225,7 +227,7 @@ impl<'a> ZkLinker<'a> {
         let mut contracts = self.linker.contracts.clone();
 
         let mut needed_libraries = BTreeSet::new();
-        self.zk_collect_dependencies(target, &mut needed_libraries)?;
+        self.zk_collect_dependencies(target, &mut needed_libraries, None)?;
 
         let attempt_link = |contracts: &mut ArtifactContracts<CompactContractBytecodeCow<'a>>,
                             id,

crates/script/src/build.rs

@@ -134,16 +134,14 @@ impl BuildData {
             output.libraries
         };
 
+        let mut factory_deps = Default::default();
+        let mut libs = Default::default();
+        linker.zk_collect_dependencies(target, &mut libs, Some(&mut factory_deps)).expect("able to enumerate all factory deps");
+
         let linked_contracts = linker
             .zk_get_linked_artifacts(
-                input
-                    .artifact_ids()
-                    .filter(|(_, artifact)| artifact.is_unlinked())
-                    // we can't use the `id` directly here
-                    // becuase we expect an iterator of references
-                    .map(|(id, _)| {
-                        linker.linker.contracts.get_key_value(&id).expect("id to be present").0
-                    }),
+                // only retrieve target and its deps
+                factory_deps.into_iter().chain(libs.into_iter()).chain([target]),
                 &libraries,
             )
             .context("retrieving all fully linked contracts")?;

crates/strategy/zksync/src/executor/runner.rs

@@ -144,6 +144,9 @@ impl ExecutorStrategyRunner for ZksyncExecutorStrategyRunner {
             return Err(LinkerError::MissingTargetArtifact);
         };
 
+        // we don't strip here unlinke upstream due to
+        // `input` being used later during linking
+        // and that is unstripped
         let contracts: ArtifactContracts<CompactContractBytecodeCow<'_>> =
             input.artifact_ids().collect();
 
@@ -194,33 +197,25 @@ impl ExecutorStrategyRunner for ZksyncExecutorStrategyRunner {
             .map_err(zk_linker_error_to_linker)?;
 
         let linked_contracts = linker
-            .zk_get_linked_artifacts(
-                input
-                    .artifact_ids()
-                    .filter(|(_, artifact)| artifact.is_unlinked())
-                    // we can't use the `id` directly here
-                    // becuase we expect an iterator of references
-                    .map(|(id, _)| {
-                        linker.linker.contracts.get_key_value(&id).expect("id to be present").0
-                    }),
-                &libraries,
-            )
+            .zk_get_linked_artifacts(linker.linker.contracts.keys(), &libraries)
             .map_err(zk_linker_error_to_linker)?;
 
         let newly_linked_dual_compiled_contracts = linked_contracts
             .iter()
             .flat_map(|(needle, zk)| {
+                // match EVM linking's prefix stripping
+                let stripped = needle.clone().with_stripped_file_prefixes(&root);
                 evm_link
                     .linked_contracts
                     .iter()
-                    .find(|(id, _)| id.source == needle.source && id.name == needle.name)
-                    .map(|(_, evm)| (needle, zk, evm))
+                    .find(|(id, _)| id.source == stripped.source && id.name == stripped.name)
+                    .map(|(_, evm)| (needle, stripped, zk, evm))
             })
-            .filter(|(_, zk, evm)| zk.bytecode.is_some() && evm.bytecode.is_some())
-            .map(|(id, linked_zk, evm)| {
+            .filter(|(_, _, zk, evm)| zk.bytecode.is_some() && evm.bytecode.is_some())
+            .map(|(unstripped_id, id, linked_zk, evm)| {
                 let (_, unlinked_zk_artifact) = input
                     .artifact_ids()
-                    .find(|(contract_id, _)| contract_id == id)
+                    .find(|(contract_id, _)| contract_id == unstripped_id)
                     .expect("unable to find original (pre-linking) artifact");
                 let zk_bytecode =
                     linked_zk.get_bytecode_bytes().expect("no EraVM bytecode (or unlinked)");
@@ -363,7 +358,7 @@ impl ExecutorStrategyRunner for ZksyncExecutorStrategyRunner {
         // also mark this library as persistent, this will ensure that the state of the library is
         // persistent across fork swaps in forking mode
         executor.backend_mut().add_persistent_account(address);
-        debug!(%address, "deployed contract with create2");
+        debug!(%address, "deployed contract");
 
         let mut request = TransactionMaybeSigned::new(Default::default());
         let unsigned = request.as_unsigned_mut().unwrap();