From a09bf251a85b04549ddf244396ab23be3b6bb98b Mon Sep 17 00:00:00 2001 From: tash-2s <81064017+tash-2s@users.noreply.github.com> Date: Mon, 25 Mar 2024 14:41:30 -0500 Subject: [PATCH] feat(world): add viem actions for delegation (#2366) --- .changeset/nervous-nails-notice.md | 25 +++++ packages/world/package.json | 1 + packages/world/ts/actions/callFrom.ts | 141 ++++++++++++++++++++++++++ packages/world/ts/exports/internal.ts | 2 + pnpm-lock.yaml | 3 + 5 files changed, 172 insertions(+) create mode 100644 .changeset/nervous-nails-notice.md create mode 100644 packages/world/ts/actions/callFrom.ts diff --git a/.changeset/nervous-nails-notice.md b/.changeset/nervous-nails-notice.md new file mode 100644 index 0000000000..cab1c47d4b --- /dev/null +++ b/.changeset/nervous-nails-notice.md @@ -0,0 +1,25 @@ +--- +"@latticexyz/world": patch +--- + +This is an internal feature and is not ready for stable consumption yet. + +Added viem custom client actions for delegation. By extending viem clients with this function after delegation, the delegation is automatically applied to World contract writes. This means that these writes are made on behalf of the delegator. Internally, it transforms the write arguments to use `callFrom`. + +Usage example: + +```ts +walletClient.extend( + callFrom({ + worldAddress, + delegatorAddress, + publicClient, // Instead of using `publicClient`, you can pass a mapping function as shown below. This allows you to use your client store and avoid read requests. + // worldFunctionToSystemFunction: async (worldFunctionSelector) => { + // const systemFunction = useStore + // .getState() + // .getValue(tables.FunctionSelectors, { worldFunctionSelector })!; + // return { systemId: systemFunction.systemId, systemFunctionSelector: systemFunction.systemFunctionSelector }; + // }, + }), +); +``` diff --git a/packages/world/package.json b/packages/world/package.json index c04423c650..701c7340df 100644 --- a/packages/world/package.json +++ b/packages/world/package.json @@ -57,6 +57,7 @@ "@arktype/util": "0.0.27", "@latticexyz/common": "workspace:*", "@latticexyz/config": "workspace:*", + "@latticexyz/protocol-parser": "workspace:*", "@latticexyz/schema-type": "workspace:*", "@latticexyz/store": "workspace:*", "abitype": "1.0.0", diff --git a/packages/world/ts/actions/callFrom.ts b/packages/world/ts/actions/callFrom.ts new file mode 100644 index 0000000000..8261b729a0 --- /dev/null +++ b/packages/world/ts/actions/callFrom.ts @@ -0,0 +1,141 @@ +import { + slice, + concat, + type WalletClient, + type Transport, + type Chain, + type Account, + type Hex, + type WalletActions, + type WriteContractReturnType, + type EncodeFunctionDataParameters, + type PublicClient, +} from "viem"; +import { getAction, encodeFunctionData } from "viem/utils"; +import { writeContract } from "viem/actions"; +import { readHex } from "@latticexyz/common"; +import { + getKeySchema, + getValueSchema, + getSchemaTypes, + decodeValueArgs, + encodeKey, +} from "@latticexyz/protocol-parser/internal"; +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 = { + worldAddress: Hex; + delegatorAddress: Hex; +}; +type CallFromFunctionParameters = CallFromBaseParameters & { + worldFunctionToSystemFunction: (worldFunctionSelector: Hex) => Promise; + publicClient?: never; +}; +type CallFromClientParameters = CallFromBaseParameters & { + worldFunctionToSystemFunction?: never; + publicClient: PublicClient; +}; + +type SystemFunction = { systemId: Hex; systemFunctionSelector: Hex }; + +// By extending viem clients with this function after delegation, the delegation is automatically applied to World contract writes. +// This means that these writes are made on behalf of the delegator. +// Internally, it transforms the write arguments to use `callFrom`. +// +// Accepts either `worldFunctionToSystemFunction` or `publicClient` as an argument. +// `worldFunctionToSystemFunction` allows manually providing the mapping function, thus users can utilize their client store for the lookup. +// 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( + params: CallFromParameters, +): (client: WalletClient) => Pick, "writeContract"> { + return (client) => ({ + // Applies to: `client.writeContract`, `getContract(client, ...).write` + writeContract: async (writeArgs): Promise => { + // Skip if the contract isn't the World. + if (writeArgs.address !== params.worldAddress) { + return getAction(client, writeContract, "writeContract")(writeArgs); + } + + // Encode the World's calldata (which includes the World's function selector). + const worldCalldata = encodeFunctionData({ + abi: writeArgs.abi, + functionName: writeArgs.functionName, + args: writeArgs.args, + } as unknown as EncodeFunctionDataParameters); + + // The first 4 bytes of calldata represent the function selector. + const worldFunctionSelector = slice(worldCalldata, 0, 4); + + // Get the systemId and System's function selector. + const { systemId, systemFunctionSelector } = await worldFunctionToSystemFunction(params, 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. + const systemCalldata = concat([systemFunctionSelector, readHex(worldCalldata, 4)]); + + // Construct args for `callFrom`. + const callFromArgs: typeof writeArgs = { + ...writeArgs, + functionName: "callFrom", + args: [params.delegatorAddress, systemId, systemCalldata], + }; + + // Call `writeContract` with the new args. + return getAction(client, writeContract, "writeContract")(callFromArgs); + }, + }); +} + +const systemFunctionCache = new Map(); + +async function worldFunctionToSystemFunction( + params: CallFromParameters, + worldFunctionSelector: Hex, +): Promise { + const cacheKey = concat([params.worldAddress, 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); + + systemFunctionCache.set(cacheKey, systemFunction); + + return systemFunction; +} + +async function retrieveSystemFunctionFromContract( + publicClient: PublicClient, + worldAddress: Hex, + worldFunctionSelector: Hex, +): Promise { + const table = worldConfig.tables.world__FunctionSelectors; + + const keySchema = getSchemaTypes(getKeySchema(table)); + const valueSchema = getSchemaTypes(getValueSchema(table)); + + const [staticData, encodedLengths, dynamicData] = await publicClient.readContract({ + address: worldAddress, + abi: IStoreReadAbi, + functionName: "getRecord", + args: [table.tableId, encodeKey(keySchema, { worldFunctionSelector })], + }); + + const decoded = decodeValueArgs(valueSchema, { staticData, encodedLengths, dynamicData }); + + const systemFunction: SystemFunction = { + systemId: decoded.systemId, + systemFunctionSelector: decoded.systemFunctionSelector, + }; + + return systemFunction; +} diff --git a/packages/world/ts/exports/internal.ts b/packages/world/ts/exports/internal.ts index 9866fe55e4..a07919cb95 100644 --- a/packages/world/ts/exports/internal.ts +++ b/packages/world/ts/exports/internal.ts @@ -7,3 +7,5 @@ export * from "../encodeSystemCall"; export * from "../encodeSystemCallFrom"; export * from "../encodeSystemCalls"; export * from "../encodeSystemCallsFrom"; + +export * from "../actions/callFrom"; diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index a107c187a6..820b6fa11d 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -1036,6 +1036,9 @@ importers: '@latticexyz/config': specifier: workspace:* version: link:../config + '@latticexyz/protocol-parser': + specifier: workspace:* + version: link:../protocol-parser '@latticexyz/schema-type': specifier: workspace:* version: link:../schema-type