refactor: move linking code into its own function

gnosisguild · Sep 13, 2024 · a9524e9 · a9524e9
1 parent 627dde8
commit a9524e9
Show file tree

Hide file tree

Showing 3 changed files with 137 additions and 89 deletions.
diff --git a/src/artifact/internal/getBuildArtifact.ts b/src/artifact/internal/getBuildArtifact.ts
@@ -40,68 +40,6 @@ export default function getBuildArtifact(
   };
 }
 
-/**
- * Replaces library references in the bytecode with actual deployed addresses.
- *
- * This function scans the bytecode and replaces placeholder references
- * to libraries with their actual on-chain addresses. It ensures that
- * the library addresses are valid and properly formatted.
- *
- * @param {string} bytecode - The bytecode that may contain library references.
- * @param {Record<string, any>} linkReferences - References to libraries, as returned by the compiler.
- * @param {Record<string, string>} libraryAddresses - A map of library names to their deployed addresses.
- * @returns {string} - The updated bytecode with library references replaced by actual addresses.
- *
- * @throws {Error} - Throws if a library address is missing or incorrectly formatted.
- */
-export function resolveLinksInBytecode(
-  contractVersion: string,
-  artifact: BuildArtifact,
-  mastercopies: Record<string, Record<string, MastercopyArtifact>>
-): string {
-  let bytecode = artifact.bytecode;
-
-  for (const libraryPath of Object.keys(artifact.linkReferences)) {
-    for (const libraryName of Object.keys(
-      artifact.linkReferences[libraryPath]
-    )) {
-      console.log(`libraryPath ${libraryPath} libraryName ${libraryName}`);
-
-      if (
-        !mastercopies[libraryName] ||
-        !mastercopies[libraryName][contractVersion]
-      ) {
-        throw new Error(
-          `Could not link ${libraryName} for ${artifact.contractName}`
-        );
-      }
-
-      let { address: libraryAddress } =
-        mastercopies[libraryName][contractVersion];
-
-      assert(isAddress(libraryAddress));
-
-      for (const { length, start: offset } of artifact.linkReferences[
-        libraryPath
-      ][libraryName]) {
-        assert(length == 20);
-
-        // the offset is in bytes, and does not account for the trailing 0x
-        const left = 2 + offset * 2;
-        const right = left + length * 2;
-
-        bytecode = `${bytecode.slice(0, left)}${libraryAddress.slice(2).toLowerCase()}${bytecode.slice(right)}`;
-
-        console.log(
-          `Replaced library reference at ${offset} with address ${libraryAddress}`
-        );
-      }
-    }
-  }
-
-  return bytecode;
-}
-
 /**
  * Resolves the paths to the artifact and build info files for a specified contract.
  *

diff --git a/src/artifact/internal/linkBuildArtifact.ts b/src/artifact/internal/linkBuildArtifact.ts
@@ -0,0 +1,127 @@
+import assert from "assert";
+import { isAddress } from "ethers";
+
+import { BuildArtifact, MastercopyArtifact } from "../../types";
+
+/**
+ * Resolves library links in a build artifact
+ *
+ */
+export default function linkBuildArtifact({
+  artifact,
+  contractVersion,
+  minimalCompilerInput,
+  mastercopies,
+}: {
+  artifact: BuildArtifact;
+  contractVersion: string;
+  minimalCompilerInput?: string;
+  mastercopies: Record<string, Record<string, MastercopyArtifact>>;
+}): BuildArtifact {
+  const bytecode = linkBytecode(artifact, contractVersion, mastercopies);
+  const compilerInput = linkCompilerInput(
+    artifact,
+    contractVersion,
+    minimalCompilerInput || artifact.compilerInput,
+    mastercopies
+  );
+
+  return {
+    ...artifact,
+    bytecode,
+    compilerInput,
+  };
+}
+
+/**
+ * Replaces library references in the bytecode with actual deployed addresses.
+ *
+ * This function scans the bytecode and replaces placeholder references
+ * to libraries with their actual on-chain addresses. It ensures that
+ * the library addresses are valid and properly formatted.
+ *
+ * @param {string} bytecode - The bytecode that may contain library references.
+ * @param {Record<string, any>} linkReferences - References to libraries, as returned by the compiler.
+ * @param {Record<string, string>} libraryAddresses - A map of library names to their deployed addresses.
+ * @returns {string} - The updated bytecode with library references replaced by actual addresses.
+ *
+ * @throws {Error} - Throws if a library address is missing or incorrectly formatted.
+ */
+function linkBytecode(
+  artifact: BuildArtifact,
+  contractVersion: string,
+  mastercopies: Record<string, Record<string, MastercopyArtifact>>
+): string {
+  let bytecode = artifact.bytecode;
+
+  for (const libraryPath of Object.keys(artifact.linkReferences)) {
+    for (const libraryName of Object.keys(
+      artifact.linkReferences[libraryPath]
+    )) {
+      console.log(`libraryPath ${libraryPath} libraryName ${libraryName}`);
+
+      if (
+        !mastercopies[libraryName] ||
+        !mastercopies[libraryName][contractVersion]
+      ) {
+        throw new Error(
+          `Could not link ${libraryName} for ${artifact.contractName}`
+        );
+      }
+
+      let { address: libraryAddress } =
+        mastercopies[libraryName][contractVersion];
+
+      assert(isAddress(libraryAddress));
+
+      for (const { length, start: offset } of artifact.linkReferences[
+        libraryPath
+      ][libraryName]) {
+        assert(length == 20);
+
+        // the offset is in bytes, and does not account for the trailing 0x
+        const left = 2 + offset * 2;
+        const right = left + length * 2;
+
+        bytecode = `${bytecode.slice(0, left)}${libraryAddress.slice(2).toLowerCase()}${bytecode.slice(right)}`;
+
+        console.log(
+          `Replaced library reference at ${offset} with address ${libraryAddress}`
+        );
+      }
+    }
+  }
+
+  return bytecode;
+}
+
+function linkCompilerInput(
+  artifact: BuildArtifact,
+  contractVersion: string,
+  compilerInput: any,
+  mastercopies: Record<string, Record<string, MastercopyArtifact>>
+): any {
+  for (const libraryPath of Object.keys(artifact.linkReferences)) {
+    compilerInput.settings.libraries[libraryPath] =
+      compilerInput.settings.libraries[libraryPath] || {};
+    for (const libraryName of Object.keys(
+      artifact.linkReferences[libraryPath]
+    )) {
+      const libraryAddress =
+        mastercopies[libraryName]?.[contractVersion]?.address;
+      if (!libraryAddress) {
+        continue;
+      }
+
+      assert(isAddress(libraryAddress));
+
+      compilerInput.settings = {
+        ...compilerInput.settings,
+        libraries: {
+          ...compilerInput.settings.libraries,
+          [libraryPath]: { [libraryName]: libraryAddress },
+        },
+      };
+    }
+  }
+}
diff --git a/src/artifact/writeMastercopyFromBuild.ts b/src/artifact/writeMastercopyFromBuild.ts
@@ -8,9 +8,8 @@ import {
   defaultBuildDir,
   defaultMastercopyArtifactsFile,
 } from "./internal/paths";
-import getBuildArtifact, {
-  resolveLinksInBytecode,
-} from "./internal/getBuildArtifact";
+import getBuildArtifact from "./internal/getBuildArtifact";
+import linkBuildArtifact from "./internal/linkBuildArtifact";
 
 import { MastercopyArtifact } from "../types";
 
@@ -60,29 +59,13 @@ export default function writeMastercopyFromBuild({
     console.warn(`Warning: overriding artifact for ${contractVersion}`);
   }
 
-  const bytecode = resolveLinksInBytecode(
+  const artifact = linkBuildArtifact({
+    artifact: buildArtifact,
     contractVersion,
-    buildArtifact,
-    mastercopies
-  );
+    minimalCompilerInput,
+    mastercopies,
+  });
 
-  const compilerInput = minimalCompilerInput || buildArtifact.compilerInput;
-  compilerInput.settings = compilerInput.settings || {};
-  compilerInput.settings.libraries = compilerInput.settings.libraries || {};
-  for (const libraryPath of Object.keys(buildArtifact.linkReferences)) {
-    compilerInput.settings.libraries[libraryPath] =
-      compilerInput.settings.libraries[libraryPath] || {};
-    for (const libraryName of Object.keys(
-      buildArtifact.linkReferences[libraryPath]
-    )) {
-      const libraryAddress =
-        mastercopies[libraryName]?.[contractVersion]?.address;
-      if (libraryAddress) {
-        compilerInput.settings.libraries[libraryPath][libraryName] =
-          libraryAddress;
-      }
-    }
-  }
   const mastercopyArtifact: MastercopyArtifact = {
     contractName,
     sourceName: buildArtifact.sourceName,
@@ -91,15 +74,15 @@ export default function writeMastercopyFromBuild({
     factory,
     address: predictSingletonAddress({
       factory,
-      bytecode,
+      bytecode: artifact.bytecode,
       constructorArgs,
       salt,
     }),
-    bytecode,
+    bytecode: artifact.bytecode,
     constructorArgs,
     salt,
     abi: buildArtifact.abi,
-    compilerInput,
+    compilerInput: artifact.compilerInput,
   };
 
   const nextMastercopies = {