refactor(stdlib): consolidate find-function-by-selector helpers

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);

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<FunctionAbi | undefined> {
     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<FunctionArtifact | undefined> {
-    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<FunctionAbi | undefined> {
-    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<FunctionCall> {
     const contract = await this.getContract(to);
     if (!contract) {

yarn-project/stdlib/src/abi/abi.ts

@@ -418,20 +418,10 @@ export async function getFunctionArtifact(
   artifact: ContractArtifact,
   functionNameOrSelector: string | FunctionSelector,
 ): Promise<FunctionArtifactWithContractName> {
-  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<FunctionArtifact | undefined> {
+  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<FunctionAbi | undefined> {
+  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 || []);