From f223560f276c69231c6f8100c2e6c801b21c1b3c Mon Sep 17 00:00:00 2001 From: dbanks12 Date: Wed, 6 May 2026 19:58:18 +0000 Subject: [PATCH] refactor(stdlib): consolidate find-function-by-selector helpers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The same artifact lookup loop was duplicated three times — as ContractStore's private #findFunctionArtifactBySelector and #findFunctionAbiBySelector, and as a local helper inside ProxiedContractStoreFactory. Hoist the two variants into stdlib/abi as findFunctionArtifactBySelector and findFunctionAbiBySelector so both ContractStore and the proxied store call the same primitive. Also have the existing throwing getFunctionArtifact helper reuse findFunctionArtifactBySelector for its selector branch. --- .../proxied_contract_data_source.ts | 26 +++------- .../storage/contract_store/contract_store.ts | 35 +++---------- yarn-project/stdlib/src/abi/abi.ts | 52 ++++++++++++++----- 3 files changed, 52 insertions(+), 61 deletions(-) diff --git a/yarn-project/pxe/src/contract_function_simulator/proxied_contract_data_source.ts b/yarn-project/pxe/src/contract_function_simulator/proxied_contract_data_source.ts index a4c1c6e2744b..f7f9ce20da5f 100644 --- a/yarn-project/pxe/src/contract_function_simulator/proxied_contract_data_source.ts +++ b/yarn-project/pxe/src/contract_function_simulator/proxied_contract_data_source.ts @@ -1,4 +1,4 @@ -import { FunctionSelector } from '@aztec/stdlib/abi'; +import { type FunctionSelector, findFunctionArtifactBySelector } from '@aztec/stdlib/abi'; import type { AztecAddress } from '@aztec/stdlib/aztec-address'; import type { ContractOverrides } from '@aztec/stdlib/tx'; @@ -22,22 +22,6 @@ export class ProxiedContractStoreFactory { return contractStore; } - const findFunctionInArtifact = async ( - artifact: { name: string; functions: { name: string; parameters: any[] }[] }, - selector: FunctionSelector, - contractAddress: AztecAddress, - ) => { - for (const fn of artifact.functions) { - const fnSelector = await FunctionSelector.fromNameAndParameters(fn.name, fn.parameters); - if (fnSelector.equals(selector)) { - return { ...fn, contractName: artifact.name } as any; - } - } - throw new Error( - `Function with selector ${selector} not found in stub artifact for overridden contract at ${contractAddress}.`, - ); - }; - return new Proxy(contractStore, { get(target, prop: keyof ContractStore) { if (prop === 'getContractInstance') { @@ -59,7 +43,13 @@ export class ProxiedContractStoreFactory { `at ${contractAddress}. Register it via pxe.registerContractClass(...) before simulating.`, ); } - return findFunctionInArtifact(artifact, selector, contractAddress); + const fn = await findFunctionArtifactBySelector(artifact, selector); + if (!fn) { + throw new Error( + `Function with selector ${selector} not found in stub artifact for overridden contract at ${contractAddress}.`, + ); + } + return { ...fn, contractName: artifact.name }; }; } const value = Reflect.get(target, prop); diff --git a/yarn-project/pxe/src/storage/contract_store/contract_store.ts b/yarn-project/pxe/src/storage/contract_store/contract_store.ts index 3c9cccccb864..3e7fbf41c0c6 100644 --- a/yarn-project/pxe/src/storage/contract_store/contract_store.ts +++ b/yarn-project/pxe/src/storage/contract_store/contract_store.ts @@ -7,7 +7,6 @@ import type { AztecAsyncKVStore, AztecAsyncMap } from '@aztec/kv-store'; import { type ContractArtifact, type FunctionAbi, - type FunctionArtifact, type FunctionArtifactWithContractName, FunctionCall, type FunctionDebugMetadata, @@ -16,6 +15,8 @@ import { contractArtifactFromBuffer, contractArtifactToBuffer, encodeArguments, + findFunctionAbiBySelector, + findFunctionArtifactBySelector, getFunctionDebugMetadata, } from '@aztec/stdlib/abi'; import { AztecAddress } from '@aztec/stdlib/aztec-address'; @@ -290,7 +291,7 @@ export class ContractStore { if (!artifact) { return undefined; } - const fn = await this.#findFunctionArtifactBySelector(artifact, selector); + const fn = await findFunctionArtifactBySelector(artifact, selector); return fn && { ...fn, contractName: artifact.name }; } @@ -322,7 +323,7 @@ export class ContractStore { selector: FunctionSelector, ): Promise { const artifact = await this.#getArtifactByAddress(contractAddress); - return artifact && (await this.#findFunctionAbiBySelector(artifact, selector)); + return artifact && (await findFunctionAbiBySelector(artifact, selector)); } /** @@ -340,7 +341,7 @@ export class ContractStore { if (!artifact) { return undefined; } - const fn = await this.#findFunctionArtifactBySelector(artifact, selector); + const fn = await findFunctionArtifactBySelector(artifact, selector); return fn && getFunctionDebugMetadata(artifact, fn); } @@ -374,34 +375,10 @@ export class ContractStore { public async getDebugFunctionName(contractAddress: AztecAddress, selector: FunctionSelector) { const artifact = await this.#getArtifactByAddress(contractAddress); - const fn = artifact && (await this.#findFunctionAbiBySelector(artifact, selector)); + const fn = artifact && (await findFunctionAbiBySelector(artifact, selector)); return `${artifact?.name ?? contractAddress}:${fn?.name ?? selector}`; } - async #findFunctionArtifactBySelector( - artifact: ContractArtifact, - selector: FunctionSelector, - ): Promise { - for (const fn of artifact.functions) { - const fnSelector = await FunctionSelector.fromNameAndParameters(fn.name, fn.parameters); - if (fnSelector.equals(selector)) { - return fn; - } - } - } - - async #findFunctionAbiBySelector( - artifact: ContractArtifact, - selector: FunctionSelector, - ): Promise { - for (const fn of [...artifact.functions, ...(artifact.nonDispatchPublicFunctions ?? [])]) { - const fnSelector = await FunctionSelector.fromNameAndParameters(fn.name, fn.parameters); - if (fnSelector.equals(selector)) { - return fn; - } - } - } - public async getFunctionCall(functionName: string, args: any[], to: AztecAddress): Promise { const contract = await this.getContract(to); if (!contract) { diff --git a/yarn-project/stdlib/src/abi/abi.ts b/yarn-project/stdlib/src/abi/abi.ts index 631b84c0dcef..c2328d262ebf 100644 --- a/yarn-project/stdlib/src/abi/abi.ts +++ b/yarn-project/stdlib/src/abi/abi.ts @@ -418,20 +418,10 @@ export async function getFunctionArtifact( artifact: ContractArtifact, functionNameOrSelector: string | FunctionSelector, ): Promise { - let functionArtifact; - if (typeof functionNameOrSelector === 'string') { - functionArtifact = artifact.functions.find(f => f.name === functionNameOrSelector); - } else { - const functionsAndSelectors = await Promise.all( - artifact.functions.map(async fn => ({ - fn, - selector: await FunctionSelector.fromNameAndParameters(fn.name, fn.parameters), - })), - ); - functionArtifact = functionsAndSelectors.find(fnAndSelector => - functionNameOrSelector.equals(fnAndSelector.selector), - )?.fn; - } + const functionArtifact = + typeof functionNameOrSelector === 'string' + ? artifact.functions.find(f => f.name === functionNameOrSelector) + : await findFunctionArtifactBySelector(artifact, functionNameOrSelector); if (!functionArtifact) { throw new Error(`Unknown function ${functionNameOrSelector}`); } @@ -441,6 +431,40 @@ export async function getFunctionArtifact( return { ...functionArtifact, debug: debugMetadata, contractName: artifact.name }; } +/** + * Finds the function artifact within `artifact.functions` whose selector matches `selector`. + * Returns `undefined` if no match is found. + */ +export async function findFunctionArtifactBySelector( + artifact: ContractArtifact, + selector: FunctionSelector, +): Promise { + const fnsAndSelectors = await Promise.all( + artifact.functions.map(async fn => ({ + fn, + selector: await FunctionSelector.fromNameAndParameters(fn.name, fn.parameters), + })), + ); + return fnsAndSelectors.find(({ selector: s }) => s.equals(selector))?.fn; +} + +/** + * Finds the function abi (across both `functions` and `nonDispatchPublicFunctions`) whose selector + * matches `selector`. Returns `undefined` if no match is found. + */ +export async function findFunctionAbiBySelector( + artifact: ContractArtifact, + selector: FunctionSelector, +): Promise { + const fnsAndSelectors = await Promise.all( + getAllFunctionAbis(artifact).map(async fn => ({ + fn, + selector: await FunctionSelector.fromNameAndParameters(fn.name, fn.parameters), + })), + ); + return fnsAndSelectors.find(({ selector: s }) => s.equals(selector))?.fn; +} + /** Gets all function abis */ export function getAllFunctionAbis(artifact: ContractArtifact): FunctionAbi[] { return artifact.functions.map(f => f as FunctionAbi).concat(artifact.nonDispatchPublicFunctions || []);