From 701b78a6ad61f2a002fb510ea23eff246b66d4a8 Mon Sep 17 00:00:00 2001
From: alvrs <alvarius@lattice.xyz>
Date: Thu, 19 Sep 2024 13:11:43 +0100
Subject: [PATCH] change library map interface

---
 .../cli/src/deploy/createPrepareDeploy.ts     | 14 ++------
 packages/cli/src/deploy/getLibraryMap.ts      | 35 +++++++++++++++----
 2 files changed, 32 insertions(+), 17 deletions(-)

diff --git a/packages/cli/src/deploy/createPrepareDeploy.ts b/packages/cli/src/deploy/createPrepareDeploy.ts
index 8ec3e75579..7d536b309e 100644
--- a/packages/cli/src/deploy/createPrepareDeploy.ts
+++ b/packages/cli/src/deploy/createPrepareDeploy.ts
@@ -1,7 +1,7 @@
 import { DeterministicContract, LibraryPlaceholder, salt } from "./common";
 import { spliceHex } from "@latticexyz/common";
 import { Hex, getCreate2Address, Address } from "viem";
-import { LibraryMap, getLibraryKey } from "./getLibraryMap";
+import { LibraryMap } from "./getLibraryMap";
 
 export function createPrepareDeploy(
   bytecodeWithPlaceholders: Hex,
@@ -10,16 +10,8 @@ export function createPrepareDeploy(
   return function prepareDeploy(deployer: Address, libraryMap: LibraryMap) {
     let bytecode = bytecodeWithPlaceholders;
     for (const placeholder of placeholders) {
-      const libraryKey = getLibraryKey(placeholder);
-      const library = libraryMap[libraryKey];
-      if (!library) {
-        throw new Error(`Could not find library for bytecode placeholder ${placeholder.path}:${placeholder.name}`);
-      }
-
-      // Store the prepared address in the library map to avoid preparing the same library twice
-      library.address ??= library.prepareDeploy(deployer, libraryMap).address;
-
-      bytecode = spliceHex(bytecode, placeholder.start, placeholder.length, library.address);
+      const address = libraryMap.getAddress({ name: placeholder.name, path: placeholder.path, deployer });
+      bytecode = spliceHex(bytecode, placeholder.start, placeholder.length, address);
     }
     return {
       bytecode,
diff --git a/packages/cli/src/deploy/getLibraryMap.ts b/packages/cli/src/deploy/getLibraryMap.ts
index 766f2a0f72..49153b7671 100644
--- a/packages/cli/src/deploy/getLibraryMap.ts
+++ b/packages/cli/src/deploy/getLibraryMap.ts
@@ -1,12 +1,35 @@
-import { Hex } from "viem";
+import { Address } from "viem";
 import { Library } from "./common";
 
-export type LibraryMap = { [key: string]: Library & { address?: Hex } };
+export type LibraryMap = {
+  getAddress: (opts: { path: string; name: string; deployer: Address }) => Address;
+};
 
-export function getLibraryMap(libraries: readonly Library[]): LibraryMap {
-  return Object.fromEntries(libraries.map((library) => [getLibraryKey(library), library]));
+function getLibraryKey({ path, name }: { path: string; name: string }): string {
+  return `${path}:${name}`;
 }
 
-export function getLibraryKey({ path, name }: { path: string; name: string }): string {
-  return `${path}:${name}`;
+type LibraryCache = {
+  [key: string]: Library & {
+    address?: {
+      [deployer: Address]: Address;
+    };
+  };
+};
+
+export function getLibraryMap(libraries: readonly Library[]): LibraryMap {
+  const cache: LibraryCache = Object.fromEntries(libraries.map((library) => [getLibraryKey(library), library]));
+  const libraryMap = {
+    getAddress: ({ path, name, deployer }) => {
+      const library = cache[getLibraryKey({ path, name })];
+      if (!library) {
+        throw new Error(`Could not find library for bytecode placeholder ${path}:${name}`);
+      }
+      library.address ??= {};
+      // Store the prepared address in the library cache to avoid preparing the same library twice
+      library.address[deployer] ??= library.prepareDeploy(deployer, libraryMap).address;
+      return library.address[deployer];
+    },
+  } satisfies LibraryMap;
+  return libraryMap;
 }