feat(world): rework types for callFrom viem action (#3414)

.changeset/thick-peas-burn.md

@@ -0,0 +1,5 @@
+---
+"@latticexyz/world": patch
+---
+
+Reworked `callFrom` action to use `getAction` internally, rather than a decorated Viem client, and updated types to better match Viem.

packages/world/ts/actions/callFrom.ts

@@ -6,13 +6,12 @@ import {
   type Account,
   type Hex,
   type WalletActions,
-  type WriteContractReturnType,
   type EncodeFunctionDataParameters,
-  type PublicClient,
   Client,
+  PublicActions,
 } from "viem";
 import { getAction, encodeFunctionData } from "viem/utils";
-import { writeContract } from "viem/actions";
+import { readContract, writeContract as viem_writeContract } from "viem/actions";
 import { readHex } from "@latticexyz/common";
 import {
   getKeySchema,
@@ -24,19 +23,11 @@ import {
 import worldConfig from "../../mud.config";
 import IStoreReadAbi from "../../out/IStoreRead.sol/IStoreRead.abi.json";
 
-// Accepts either `worldFunctionToSystemFunction` or `publicClient`, but not both.
-type CallFromParameters = CallFromFunctionParameters | CallFromClientParameters;
-type CallFromBaseParameters = {
+type CallFromParameters = {
   worldAddress: Hex;
   delegatorAddress: Hex;
-};
-type CallFromFunctionParameters = CallFromBaseParameters & {
-  worldFunctionToSystemFunction: (worldFunctionSelector: Hex) => Promise<SystemFunction>;
-  publicClient?: never;
-};
-type CallFromClientParameters = CallFromBaseParameters & {
-  worldFunctionToSystemFunction?: never;
-  publicClient: PublicClient;
+  worldFunctionToSystemFunction?: (worldFunctionSelector: Hex) => Promise<SystemFunction>;
+  publicClient?: Client;
 };
 
 type SystemFunction = { systemId: Hex; systemFunctionSelector: Hex };
@@ -50,20 +41,24 @@ type SystemFunction = { systemId: Hex; systemFunctionSelector: Hex };
 // If `publicClient` is provided instead, this function retrieves the corresponding system function from the World contract.
 //
 // The function mapping is cached to avoid redundant retrievals for the same World function.
-export function callFrom<TChain extends Chain, TAccount extends Account>(
+export function callFrom(
   params: CallFromParameters,
-): (client: Client<Transport, TChain, TAccount>) => Pick<WalletActions<TChain, TAccount>, "writeContract"> {
+): <chain extends Chain, account extends Account | undefined>(
+  client: Client<Transport, chain, account>,
+) => Pick<WalletActions<chain, account>, "writeContract"> {
   return (client) => ({
     // Applies to: `client.writeContract`, `getContract(client, ...).write`
-    writeContract: async (writeArgs): Promise<WriteContractReturnType> => {
+    async writeContract(writeArgs) {
+      const _writeContract = getAction(client, viem_writeContract, "writeContract");
+
       // Skip if the contract isn't the World or the function called should not be redirected through `callFrom`.
       if (
         writeArgs.address !== params.worldAddress ||
         writeArgs.functionName === "call" ||
         writeArgs.functionName === "callFrom" ||
         writeArgs.functionName === "callWithSignature"
       ) {
-        return getAction(client, writeContract, "writeContract")(writeArgs);
+        return _writeContract(writeArgs);
       }
 
       // Encode the World's calldata (which includes the World's function selector).
@@ -77,7 +72,11 @@ export function callFrom<TChain extends Chain, TAccount extends Account>(
       const worldFunctionSelector = slice(worldCalldata, 0, 4);
 
       // Get the systemId and System's function selector.
-      const { systemId, systemFunctionSelector } = await worldFunctionToSystemFunction(params, worldFunctionSelector);
+      const { systemId, systemFunctionSelector } = await worldFunctionToSystemFunction({
+        ...params,
+        publicClient: params.publicClient ?? client,
+        worldFunctionSelector,
+      });
 
       // Construct the System's calldata by replacing the World's function selector with the System's.
       // Use `readHex` instead of `slice` to prevent out-of-bounds errors with calldata that has no args.
@@ -91,35 +90,38 @@ export function callFrom<TChain extends Chain, TAccount extends Account>(
       };
 
       // Call `writeContract` with the new args.
-      return getAction(client, writeContract, "writeContract")(callFromArgs);
+      return _writeContract(callFromArgs);
     },
   });
 }
 
 const systemFunctionCache = new Map<Hex, SystemFunction>();
 
-async function worldFunctionToSystemFunction(
-  params: CallFromParameters,
-  worldFunctionSelector: Hex,
-): Promise<SystemFunction> {
-  const cacheKey = concat([params.worldAddress, worldFunctionSelector]);
+async function worldFunctionToSystemFunction(params: {
+  worldAddress: Hex;
+  delegatorAddress: Hex;
+  worldFunctionSelector: Hex;
+  worldFunctionToSystemFunction?: (worldFunctionSelector: Hex) => Promise<SystemFunction>;
+  publicClient: Client;
+}): Promise<SystemFunction> {
+  const cacheKey = concat([params.worldAddress, params.worldFunctionSelector]);
 
   // Use cache if the function has been called previously.
   const cached = systemFunctionCache.get(cacheKey);
   if (cached) return cached;
 
   // If a mapping function is provided, use it. Otherwise, call the World contract.
   const systemFunction = params.worldFunctionToSystemFunction
-    ? await params.worldFunctionToSystemFunction(worldFunctionSelector)
-    : await retrieveSystemFunctionFromContract(params.publicClient, params.worldAddress, worldFunctionSelector);
+    ? await params.worldFunctionToSystemFunction(params.worldFunctionSelector)
+    : await retrieveSystemFunctionFromContract(params.publicClient, params.worldAddress, params.worldFunctionSelector);
 
   systemFunctionCache.set(cacheKey, systemFunction);
 
   return systemFunction;
 }
 
 async function retrieveSystemFunctionFromContract(
-  publicClient: PublicClient,
+  publicClient: Client,
   worldAddress: Hex,
   worldFunctionSelector: Hex,
 ): Promise<SystemFunction> {
@@ -128,7 +130,9 @@ async function retrieveSystemFunctionFromContract(
   const keySchema = getSchemaTypes(getKeySchema(table));
   const valueSchema = getSchemaTypes(getValueSchema(table));
 
-  const [staticData, encodedLengths, dynamicData] = await publicClient.readContract({
+  const _readContract = getAction(publicClient, readContract, "readContract") as PublicActions["readContract"];
+
+  const [staticData, encodedLengths, dynamicData] = await _readContract({
     address: worldAddress,
     abi: IStoreReadAbi,
     functionName: "getRecord",