WIP: suggestioin refactor bytecode resolution

src/artifact/internal/getBuildArtifact.ts

@@ -1,68 +1,8 @@
 import path from "path";
 import { readdirSync, readFileSync, statSync } from "fs";
 
-import { BuildArtifact } from "../../types";
-import { getCreate2Address, keccak256 } from "ethers";
-
-/**
- * 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 replaceLibraryReferences(
-  bytecode: string,
-  linkReferences: Record<string, any>,
-  libraryAddresses: Record<string, string>
-): string {
-  // Iterate through each source file in the linkReferences object
-  for (const sourceFile in linkReferences) {
-    // Iterate through each library name in the source file
-    for (const libName in linkReferences[sourceFile]) {
-      const libAddress = libraryAddresses[libName];
-
-      // Ensure that the library address is provided
-      if (!libAddress) {
-        throw new Error(`Library address for "${libName}" not found.`);
-      }
-
-      // Ensure the address is properly formatted (without '0x' and exactly 40 characters long)
-      const addressHex = libAddress.toLowerCase().replace("0x", "");
-      if (addressHex.length !== 40) {
-        throw new Error(
-          `Invalid library address for "${libName}": ${libAddress}`
-        );
-      }
-
-      // Iterate through each reference of the library within the bytecode
-      linkReferences[sourceFile][libName].forEach(
-        (ref: { start: number; length: number }) => {
-          // Extract the placeholder in the bytecode corresponding to the library
-          const placeholder = bytecode.slice(
-            ref.start * 2,
-            (ref.start + ref.length) * 2
-          );
-
-          // Replace the placeholder with the correctly padded address (if necessary)
-          const paddedAddress = addressHex.padEnd(ref.length * 2, "0");
-          bytecode = bytecode.replace(placeholder, paddedAddress);
-        }
-      );
-    }
-  }
-
-  // Remove any trailing underscores or erroneous placeholders
-
-  return bytecode.replace("__", "");
-}
+import { BuildArtifact, MastercopyArtifact } from "../../types";
+import assert from "assert";
 
 /**
  * Retrieves the build artifact for a specified contract.
@@ -74,41 +14,15 @@ function replaceLibraryReferences(
  */
 export default function getBuildArtifact(
   _contractName: string,
-  buildDirPath: string,
-  factory: string,
-  salt: string
+  buildDirPath: string
 ): BuildArtifact {
   const { artifactPath, buildInfoPath } = resolvePaths(
     _contractName,
     buildDirPath
   );
 
-  const {
-    sourceName,
-    contractName,
-    bytecode: artifactBytecode,
-    abi,
-    linkReferences,
-  } = JSON.parse(readFileSync(artifactPath, "utf8"));
-  let bytecode = artifactBytecode;
-  if (linkReferences && Object.keys(linkReferences).length > 0) {
-    let libraryAddresses: Record<string, string> = {};
-    Object.keys(linkReferences).forEach((sourceFile) => {
-      Object.keys(linkReferences[sourceFile]).forEach((libName) => {
-        console.log(`Processing library: ${libName}`);
-        const { artifactPath } = resolvePaths(libName, buildDirPath);
-        const { bytecode } = JSON.parse(readFileSync(artifactPath, "utf8"));
-        const address = getCreate2Address(factory, salt, keccak256(bytecode));
-        libraryAddresses = { ...libraryAddresses, [libName]: address };
-      });
-    });
-
-    bytecode = replaceLibraryReferences(
-      artifactBytecode,
-      linkReferences,
-      libraryAddresses
-    );
-  }
+  const { sourceName, contractName, bytecode, abi, linkReferences } =
+    JSON.parse(readFileSync(artifactPath, "utf8"));
 
   const { solcLongVersion, input } = JSON.parse(
     readFileSync(buildInfoPath, "utf8")
@@ -119,11 +33,73 @@ export default function getBuildArtifact(
     sourceName,
     compilerVersion: `v${solcLongVersion}`,
     compilerInput: input,
-    abi,
     bytecode,
+    abi,
+    linkReferences,
   };
 }
 
+/**
+ * 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];
+
+      libraryAddress = libraryAddress.toLowerCase().replace(/^0x/, "");
+
+      assert(libraryAddress.length == 40);
+
+      for (const { length, start } of artifact.linkReferences[libraryPath][
+        libraryName
+      ]) {
+        assert(length == 20);
+        const left = bytecode.slice(0, start * 2);
+        const right = bytecode.slice((start + length) * 2);
+        bytecode = `${left}${libraryAddress}${right}`;
+
+        console.log(
+          `Replaced library reference at ${start} with address ${libraryAddress}`
+        );
+      }
+    }
+  }
+
+  return bytecode.replace("__", "");
+}
+
 /**
  * Resolves the paths to the artifact and build info files for a specified contract.
  *

src/artifact/writeMastercopyFromBuild.ts

@@ -8,7 +8,9 @@ import {
   defaultBuildDir,
   defaultMastercopyArtifactsFile,
 } from "./internal/paths";
-import getBuildArtifact from "./internal/getBuildArtifact";
+import getBuildArtifact, {
+  resolveLinksInBytecode,
+} from "./internal/getBuildArtifact";
 
 import { MastercopyArtifact } from "../types";
 
@@ -48,12 +50,7 @@ export default function writeMastercopyFromBuild({
   buildDirPath?: string;
   mastercopyArtifactsFile?: string;
 }) {
-  const buildArtifact = getBuildArtifact(
-    contractName,
-    buildDirPath,
-    factory,
-    salt
-  );
+  const buildArtifact = getBuildArtifact(contractName, buildDirPath);
 
   const mastercopies = existsSync(mastercopyArtifactsFile)
     ? JSON.parse(readFileSync(mastercopyArtifactsFile, "utf8"))
@@ -62,7 +59,13 @@ export default function writeMastercopyFromBuild({
   if (mastercopies[contractVersion]) {
     console.warn(`Warning: overriding artifact for ${contractVersion}`);
   }
-  console.log("buildArtifact.bytecode", buildArtifact.bytecode);
+
+  const bytecode = resolveLinksInBytecode(
+    contractVersion,
+    buildArtifact,
+    mastercopies
+  );
+
   const mastercopyArtifact: MastercopyArtifact = {
     contractName,
     sourceName: buildArtifact.sourceName,
@@ -71,11 +74,11 @@ export default function writeMastercopyFromBuild({
     factory,
     address: predictSingletonAddress({
       factory,
-      bytecode: buildArtifact.bytecode,
+      bytecode,
       constructorArgs,
       salt,
     }),
-    bytecode: buildArtifact.bytecode,
+    bytecode,
     constructorArgs,
     salt,
     abi: buildArtifact.abi,

src/types.ts

@@ -20,9 +20,13 @@ export type BuildArtifact = {
   compilerInput: any;
   bytecode: string;
   abi: any;
+  linkReferences: Record<
+    string,
+    Record<string, { length: number; start: number }[]>
+  >;
 };
 
-export type MastercopyArtifact = BuildArtifact & {
+export type MastercopyArtifact = Omit<BuildArtifact, "linkReferences"> & {
   contractVersion: string;
   factory: string;
   constructorArgs: {