feat(cli): add the batch registration feature to systems, accesses, functions and tables

packages/cli/src/deploy/ensureFunctions.ts

@@ -1,10 +1,11 @@
-import { Client, Transport, Chain, Account, Hex } from "viem";
+import { Client, Transport, Chain, Account, Hex, encodeFunctionData } from "viem";
 import { hexToResource, writeContract } from "@latticexyz/common";
 import { WorldDeploy, WorldFunction, worldAbi } from "./common";
 import { debug } from "./debug";
 import { getFunctions } from "./getFunctions";
 import pRetry from "p-retry";
 import { wait } from "@latticexyz/common/utils";
+import { getBatchCallData } from "../utils/getBatchCallData";
 
 export async function ensureFunctions({
   client,
@@ -36,47 +37,57 @@ export async function ensureFunctions({
 
   if (!toAdd.length) return [];
 
-  debug("registering functions:", toAdd.map((func) => func.signature).join(", "));
+  debug("Batch registering functions:", toAdd.map((func) => func.signature).join(", "));
 
-  return Promise.all(
-    toAdd.map((func) => {
-      const { namespace } = hexToResource(func.systemId);
+  const toAddBySystem = toAdd.reduce((acc: { [key: Hex]: WorldFunction[] }, func) => {
+    if (acc[func.systemId]) {
+      acc[func.systemId].push(func);
+    } else {
+      acc[func.systemId] = [func];
+    }
+    return acc;
+  }, {});
+
+  const encodedFunctionDataBySystem = Object.entries(toAddBySystem).reduce((acc, [systemId, functions]) => {
+    const { namespace } = hexToResource(systemId as Hex);
+    const encodedFunctionDataList = functions.map((func) => {
+      let encodedFunctionData: Hex;
       if (namespace === "") {
-        return pRetry(
-          () =>
-            writeContract(client, {
-              chain: client.chain ?? null,
-              address: worldDeploy.address,
-              abi: worldAbi,
-              // TODO: replace with batchCall (https://github.com/latticexyz/mud/issues/1645)
-              functionName: "registerRootFunctionSelector",
-              args: [func.systemId, func.systemFunctionSignature, func.systemFunctionSelector],
-            }),
-          {
-            retries: 3,
-            onFailedAttempt: async (error) => {
-              const delay = error.attemptNumber * 500;
-              debug(`failed to register function ${func.signature}, retrying in ${delay}ms...`);
-              await wait(delay);
-            },
-          }
-        );
+        encodedFunctionData = encodeFunctionData({
+          abi: worldAbi,
+          functionName: "registerRootFunctionSelector",
+          args: [systemId as Hex, func.systemFunctionSignature, func.systemFunctionSelector],
+        });
+      } else {
+        encodedFunctionData = encodeFunctionData({
+          abi: worldAbi,
+          functionName: "registerFunctionSelector",
+          args: [systemId as Hex, func.systemFunctionSignature],
+        });
       }
+      return encodedFunctionData;
+    });
+    acc[systemId] = encodedFunctionDataList;
+    return acc;
+  }, {} as Record<string, Hex[]>);
+
+  return Promise.all(
+    Object.entries(encodedFunctionDataBySystem).map(([systemId, encodedFunctionData]) => {
       return pRetry(
         () =>
           writeContract(client, {
             chain: client.chain ?? null,
             address: worldDeploy.address,
             abi: worldAbi,
-            // TODO: replace with batchCall (https://github.com/latticexyz/mud/issues/1645)
-            functionName: "registerFunctionSelector",
-            args: [func.systemId, func.systemFunctionSignature],
+            functionName: "batchCall",
+            args: [getBatchCallData(encodedFunctionData)],
           }),
         {
           retries: 3,
           onFailedAttempt: async (error) => {
             const delay = error.attemptNumber * 500;
-            debug(`failed to register function ${func.signature}, retrying in ${delay}ms...`);
+            //TODO: Replace the systemId with system label
+            debug(`failed to register function ${systemId}, retrying in ${delay}ms...`);
             await wait(delay);
           },
         }

packages/cli/src/deploy/ensureSystems.ts

@@ -1,4 +1,4 @@
-import { Client, Transport, Chain, Account, Hex, getAddress } from "viem";
+import { Client, Transport, Chain, Account, Hex, getAddress, encodeFunctionData } from "viem";
 import { writeContract } from "@latticexyz/common";
 import { System, WorldDeploy, worldAbi } from "./common";
 import { debug } from "./debug";
@@ -7,6 +7,7 @@ import { getSystems } from "./getSystems";
 import { getResourceAccess } from "./getResourceAccess";
 import { uniqueBy, wait } from "@latticexyz/common/utils";
 import pRetry from "p-retry";
+import { getBatchCallData } from "../utils/getBatchCallData";
 import { ensureContractsDeployed } from "./ensureContractsDeployed";
 
 export async function ensureSystems({
@@ -22,29 +23,42 @@ export async function ensureSystems({
     getSystems({ client, worldDeploy }),
     getResourceAccess({ client, worldDeploy }),
   ]);
+
+  // Prepare access changes
+
   const systemIds = systems.map((system) => system.systemId);
+
   const currentAccess = worldAccess.filter(({ resourceId }) => systemIds.includes(resourceId));
+
   const desiredAccess = systems.flatMap((system) =>
     system.allowedAddresses.map((address) => ({ resourceId: system.systemId, address }))
   );
 
-  const accessToAdd = desiredAccess.filter(
-    (access) =>
-      !currentAccess.some(
-        ({ resourceId, address }) =>
-          resourceId === access.resourceId && getAddress(address) === getAddress(access.address)
-      )
-  );
+  type accessType = {
+    resourceId: Hex;
+    address: Hex;
+    isRemove: boolean;
+  };
 
-  const accessToRemove = currentAccess.filter(
-    (access) =>
-      !desiredAccess.some(
-        ({ resourceId, address }) =>
-          resourceId === access.resourceId && getAddress(address) === getAddress(access.address)
-      )
-  );
+  const accessToAdd: accessType[] = desiredAccess
+    .filter(
+      (access) =>
+        !currentAccess.some(
+          ({ resourceId, address }) =>
+            resourceId === access.resourceId && getAddress(address) === getAddress(access.address)
+        )
+    )
+    .map((access) => ({ ...access, isRemove: false }));
 
-  // TODO: move each system access+registration to batch call to be atomic
+  const accessToRemove: accessType[] = currentAccess
+    .filter(
+      (access) =>
+        !desiredAccess.some(
+          ({ resourceId, address }) =>
+            resourceId === access.resourceId && getAddress(address) === getAddress(access.address)
+        )
+    )
+    .map((access) => ({ ...access, isRemove: true }));
 
   if (accessToRemove.length) {
     debug("revoking", accessToRemove.length, "access grants");
@@ -53,48 +67,20 @@ export async function ensureSystems({
     debug("adding", accessToAdd.length, "access grants");
   }
 
-  const accessTxs = [
-    ...accessToRemove.map((access) =>
-      pRetry(
-        () =>
-          writeContract(client, {
-            chain: client.chain ?? null,
-            address: worldDeploy.address,
-            abi: worldAbi,
-            functionName: "revokeAccess",
-            args: [access.resourceId, access.address],
-          }),
-        {
-          retries: 3,
-          onFailedAttempt: async (error) => {
-            const delay = error.attemptNumber * 500;
-            debug(`failed to revoke access, retrying in ${delay}ms...`);
-            await wait(delay);
-          },
-        }
-      )
-    ),
-    ...accessToAdd.map((access) =>
-      pRetry(
-        () =>
-          writeContract(client, {
-            chain: client.chain ?? null,
-            address: worldDeploy.address,
-            abi: worldAbi,
-            functionName: "grantAccess",
-            args: [access.resourceId, access.address],
-          }),
-        {
-          retries: 3,
-          onFailedAttempt: async (error) => {
-            const delay = error.attemptNumber * 500;
-            debug(`failed to grant access, retrying in ${delay}ms...`);
-            await wait(delay);
-          },
-        }
-      )
-    ),
-  ];
+  const accessChanges = accessToAdd.concat(accessToRemove);
+
+  const accessChangeMap = accessChanges.reduce((map, access) => {
+    const systemId = access.resourceId;
+    const accesses = map[systemId];
+    if (accesses) {
+      accesses.push(access);
+    } else {
+      map[systemId] = [access];
+    }
+    return map;
+  }, {} as Record<string, accessType[]>);
+
+  // Prepare Systems to register
 
   const existingSystems = systems.filter((system) =>
     worldSystems.some(
@@ -135,27 +121,68 @@ export async function ensureSystems({
     })),
   });
 
-  const registerTxs = missingSystems.map((system) =>
-    pRetry(
-      () =>
-        writeContract(client, {
-          chain: client.chain ?? null,
-          address: worldDeploy.address,
-          abi: worldAbi,
-          // TODO: replace with batchCall (https://github.com/latticexyz/mud/issues/1645)
-          functionName: "registerSystem",
-          args: [system.systemId, system.address, system.allowAll],
-        }),
-      {
-        retries: 3,
-        onFailedAttempt: async (error) => {
-          const delay = error.attemptNumber * 500;
-          debug(`failed to register system ${resourceLabel(system)}, retrying in ${delay}ms...`);
-          await wait(delay);
-        },
+  const missingSystemsMap = missingSystems.reduce((map, system) => {
+    map[system.systemId] = system;
+    return map;
+  }, {} as Record<string, System>);
+
+  if (Object.keys(missingSystemsMap).length !== missingSystems.length) {
+    throw new Error("Duplicate systemId found in systems");
+  }
+
+  // Combine systems and access changes
+  const registrationAndAccess: [System, accessType[]][] = Object.entries(missingSystemsMap).map(
+    ([systemId, system]) => {
+      if (systemId in accessChangeMap) {
+        return [system, accessChangeMap[systemId]];
+      } else {
+        return [system, [] as accessType[]];
       }
-    )
+    }
   );
 
-  return await Promise.all([...accessTxs, ...registerTxs]);
+  const encodedFunctionDataBySystem = registrationAndAccess.reduce((acc, [system, access]) => {
+    const registerSystemFunctionData: Hex = encodeFunctionData({
+      abi: worldAbi,
+      functionName: "registerSystem",
+      args: [system.systemId, system.address, system.allowAll],
+    });
+
+    const accessFunctionDataList: Hex[] = access.map((access: any) => {
+      const functionName = access.isRemove ? "revokeAccess" : "grantAccess";
+      const encodedData = encodeFunctionData({
+        abi: worldAbi,
+        functionName: functionName,
+        args: [access.resourceId, access.address],
+      });
+      return encodedData;
+    });
+    const systemLabel = resourceLabel(system);
+    acc[systemLabel] = [registerSystemFunctionData, ...accessFunctionDataList];
+    return acc;
+  }, {} as Record<string, Hex[]>);
+
+  const systemTxs = [
+    ...Object.entries(encodedFunctionDataBySystem).map(([systemLabel, encodedFunctionData]) =>
+      pRetry(
+        () =>
+          writeContract(client, {
+            chain: client.chain ?? null,
+            address: worldDeploy.address,
+            abi: worldAbi,
+            functionName: "batchCall",
+            args: [getBatchCallData(encodedFunctionData)],
+          }),
+        {
+          retries: 3,
+          onFailedAttempt: async (error) => {
+            const delay = error.attemptNumber * 500;
+            debug(`failed to register or change access permission at ${systemLabel}, retrying in ${delay}ms...`);
+            await wait(delay);
+          },
+        }
+      )
+    ),
+  ];
+  return await Promise.all([...systemTxs]);
 }

packages/cli/src/deploy/ensureTables.ts

@@ -1,4 +1,4 @@
-import { Client, Transport, Chain, Account, Hex } from "viem";
+import { Client, Transport, Chain, Account, Hex, encodeFunctionData } from "viem";
 import { Table } from "./configToTables";
 import { writeContract } from "@latticexyz/common";
 import { WorldDeploy, worldAbi } from "./common";
@@ -8,6 +8,7 @@ import { resourceLabel } from "./resourceLabel";
 import { getTables } from "./getTables";
 import pRetry from "p-retry";
 import { wait } from "@latticexyz/common/utils";
+import { getBatchCallData } from "../utils/getBatchCallData";
 
 export async function ensureTables({
   client,
@@ -27,32 +28,48 @@ export async function ensureTables({
   }
 
   const missingTables = tables.filter((table) => !worldTableIds.includes(table.tableId));
+  const encodedFunctionDataList = missingTables.map((table) => {
+    const encodedFunctionData = encodeFunctionData({
+      abi: worldAbi,
+      functionName: "registerTable",
+      args: [
+        table.tableId,
+        valueSchemaToFieldLayoutHex(table.valueSchema),
+        keySchemaToHex(table.keySchema),
+        valueSchemaToHex(table.valueSchema),
+        Object.keys(table.keySchema),
+        Object.keys(table.valueSchema),
+      ],
+    });
+    return encodedFunctionData;
+  });
+
+  // Slice register calls to avoid hitting the block gas limit
+  // TODO: the batchSize can be configurable
+  const batchSize = 10;
+  const iterations = Math.ceil(encodedFunctionDataList.length / batchSize);
+  const slicedFunctionDataList = Array.from({ length: iterations }, (_, i) => {
+    return encodedFunctionDataList.slice(i * batchSize, (i + 1) * batchSize);
+  });
+
   if (missingTables.length) {
-    debug("registering tables", missingTables.map(resourceLabel).join(", "));
+    debug("Batch registering tables", missingTables.map(resourceLabel).join(", "));
     return await Promise.all(
-      missingTables.map((table) =>
+      slicedFunctionDataList.map((encodedFunctionDataList) =>
         pRetry(
           () =>
             writeContract(client, {
               chain: client.chain ?? null,
               address: worldDeploy.address,
               abi: worldAbi,
-              // TODO: replace with batchCall (https://github.com/latticexyz/mud/issues/1645)
-              functionName: "registerTable",
-              args: [
-                table.tableId,
-                valueSchemaToFieldLayoutHex(table.valueSchema),
-                keySchemaToHex(table.keySchema),
-                valueSchemaToHex(table.valueSchema),
-                Object.keys(table.keySchema),
-                Object.keys(table.valueSchema),
-              ],
+              functionName: "batchCall",
+              args: [getBatchCallData(encodedFunctionDataList)],
             }),
           {
             retries: 3,
             onFailedAttempt: async (error) => {
               const delay = error.attemptNumber * 500;
-              debug(`failed to register table ${resourceLabel(table)}, retrying in ${delay}ms...`);
+              debug(`failed to register tables, retrying in ${delay}ms...`);
               await wait(delay);
             },
           }