@aztec/pxe 0.82.0 → 0.82.1

package/src/pxe_service/error_enriching.ts

@@ -1,5 +1,3 @@
-import { PUBLIC_DISPATCH_SELECTOR } from '@aztec/constants';
-import { Fr } from '@aztec/foundation/fields';
 import type { Logger } from '@aztec/foundation/log';
 import { resolveAssertionMessageFromRevertData, resolveOpcodeLocations } from '@aztec/simulator/client';
 import { FunctionSelector } from '@aztec/stdlib/abi';
@@ -66,24 +64,20 @@ export async function enrichPublicSimulationError(
 ) {
   const callStack = err.getCallStack();
   const originalFailingFunction = callStack[callStack.length - 1];
-  // TODO(https://github.com/AztecProtocol/aztec-packages/issues/8985): Properly fix this.
-  // To be able to resolve the assertion message, we need to use the information from the public dispatch function,
-  // no matter what the call stack selector points to (since we've modified it to point to the target function).
-  // We should remove this because the AVM (or public protocol) shouldn't be aware of the public dispatch calling convention.
 
-  const artifact = await contractDataProvider.getFunctionArtifact(
-    originalFailingFunction.contractAddress,
-    FunctionSelector.fromField(new Fr(PUBLIC_DISPATCH_SELECTOR)),
-  );
+  const artifact = await contractDataProvider.getPublicFunctionArtifact(originalFailingFunction.contractAddress);
+  if (!artifact) {
+    throw new Error(
+      `Artifact not found when enriching public simulation error. Contract address: ${originalFailingFunction.contractAddress}.`,
+    );
+  }
+
   const assertionMessage = resolveAssertionMessageFromRevertData(err.revertData, artifact);
   if (assertionMessage) {
     err.setOriginalMessage(err.getOriginalMessage() + `${assertionMessage}`);
   }
 
-  const debugInfo = await contractDataProvider.getFunctionDebugMetadata(
-    originalFailingFunction.contractAddress,
-    FunctionSelector.fromField(new Fr(PUBLIC_DISPATCH_SELECTOR)),
-  );
+  const debugInfo = await contractDataProvider.getPublicFunctionDebugMetadata(originalFailingFunction.contractAddress);
 
   const noirCallStack = err.getNoirCallStack();
   if (debugInfo) {

package/src/pxe_service/pxe_service.ts

@@ -450,10 +450,8 @@ export class PXEService implements PXE {
     isContractClassPubliclyRegistered: boolean;
     artifact: ContractArtifact | undefined;
   }> {
-    let artifact;
-    try {
-      artifact = await this.contractDataProvider.getContractArtifact(id);
-    } catch {
+    const artifact = await this.contractDataProvider.getContractArtifact(id);
+    if (!artifact) {
       this.log.warn(`No artifact found for contract class ${id.toString()} when looking for its metadata`);
     }
 
@@ -572,6 +570,11 @@ export class PXEService implements PXE {
     } else {
       // Otherwise, make sure there is an artifact already registered for that class id
       artifact = await this.contractDataProvider.getContractArtifact(instance.currentContractClassId);
+      if (!artifact) {
+        throw new Error(
+          `Artifact not found when registering an instance. Contract class: ${instance.currentContractClassId}.`,
+        );
+      }
     }
 
     await this.contractDataProvider.addContractInstance(instance);
@@ -585,6 +588,9 @@ export class PXEService implements PXE {
     // class while we're simulating it.
     return this.#putInJobQueue(async () => {
       const currentInstance = await this.contractDataProvider.getContractInstance(contractAddress);
+      if (!currentInstance) {
+        throw new Error(`Instance not found when updating a contract. Contract address: ${contractAddress}.`);
+      }
       const contractClass = await getContractClassFromArtifact(artifact);
       await this.synchronizer.sync();
 

package/src/storage/contract_data_provider/contract_data_provider.ts

@@ -2,9 +2,9 @@ import type { Fr } from '@aztec/foundation/fields';
 import { toArray } from '@aztec/foundation/iterable';
 import type { MembershipWitness } from '@aztec/foundation/trees';
 import type { AztecAsyncKVStore, AztecAsyncMap } from '@aztec/kv-store';
-import { ContractClassNotFoundError, ContractNotFoundError } from '@aztec/simulator/client';
 import {
   type ContractArtifact,
+  type FunctionAbi,
   type FunctionArtifact,
   type FunctionArtifactWithContractName,
   type FunctionDebugMetadata,
@@ -19,6 +19,7 @@ import {
   type ContractClass,
   type ContractInstanceWithAddress,
   SerializableContractInstance,
+  getContractClassFromArtifact,
 } from '@aztec/stdlib/contract';
 
 import type { DataProvider } from '../data_provider.js';
@@ -33,7 +34,11 @@ import { PrivateFunctionsTree } from './private_functions_tree.js';
  */
 export class ContractDataProvider implements DataProvider {
   /** Map from contract class id to private function tree. */
-  private contractClassesCache: Map<string, PrivateFunctionsTree> = new Map();
+  // TODO: Update it to be LRU cache so that it doesn't keep all the data all the time.
+  #privateFunctionTrees: Map<string, PrivateFunctionsTree> = new Map();
+
+  /** Map from contract address to contract class id */
+  #contractClassIdMap: Map<string, Fr> = new Map();
 
   #contractArtifacts: AztecAsyncMap<string, Buffer>;
   #contractInstances: AztecAsyncMap<string, Buffer>;
@@ -66,6 +71,8 @@ export class ContractDataProvider implements DataProvider {
   }
 
   async addContractInstance(contract: ContractInstanceWithAddress): Promise<void> {
+    this.#contractClassIdMap.set(contract.address.toString(), contract.currentContractClassId);
+
     await this.#contractInstances.set(
       contract.address.toString(),
       new SerializableContractInstance(contract).toBuffer(),
@@ -74,15 +81,15 @@ export class ContractDataProvider implements DataProvider {
 
   // Private getters
 
-  async #getContractInstance(address: AztecAddress): Promise<ContractInstanceWithAddress | undefined> {
-    const contract = await this.#contractInstances.getAsync(address.toString());
-    return contract && SerializableContractInstance.fromBuffer(contract).withAddress(address);
-  }
-
-  async #getContractArtifact(id: Fr): Promise<ContractArtifact | undefined> {
-    const contract = await this.#contractArtifacts.getAsync(id.toString());
-    // TODO(@spalladino): AztecAsyncMap lies and returns Uint8Arrays instead of Buffers, hence the extra Buffer.from.
-    return contract && contractArtifactFromBuffer(Buffer.from(contract));
+  async #getContractClassId(contractAddress: AztecAddress): Promise<Fr | undefined> {
+    if (!this.#contractClassIdMap.has(contractAddress.toString())) {
+      const instance = await this.getContractInstance(contractAddress);
+      if (!instance) {
+        return;
+      }
+      this.#contractClassIdMap.set(contractAddress.toString(), instance.currentContractClassId);
+    }
+    return this.#contractClassIdMap.get(contractAddress.toString());
   }
 
   /**
@@ -95,31 +102,21 @@ export class ContractDataProvider implements DataProvider {
    * @returns A ContractTree instance associated with the specified contract address.
    * @throws An Error if the contract is not found in the ContractDatabase.
    */
-  private async getTreeForClassId(classId: Fr): Promise<PrivateFunctionsTree> {
-    if (!this.contractClassesCache.has(classId.toString())) {
-      const artifact = await this.#getContractArtifact(classId);
+  async #getPrivateFunctionTreeForClassId(classId: Fr): Promise<PrivateFunctionsTree | undefined> {
+    if (!this.#privateFunctionTrees.has(classId.toString())) {
+      const artifact = await this.getContractArtifact(classId);
       if (!artifact) {
-        throw new ContractClassNotFoundError(classId.toString());
+        return;
       }
       const tree = await PrivateFunctionsTree.create(artifact);
-      this.contractClassesCache.set(classId.toString(), tree);
+      this.#privateFunctionTrees.set(classId.toString(), tree);
     }
-    return this.contractClassesCache.get(classId.toString())!;
+    return this.#privateFunctionTrees.get(classId.toString())!;
   }
 
-  /**
-   * Retrieve or create a ContractTree instance based on the provided AztecAddress.
-   * If an existing tree with the same contract address is found in the cache, it will be returned.
-   * Otherwise, a new ContractTree instance will be created using the contract data from the database
-   * and added to the cache before returning.
-   *
-   * @param contractAddress - The AztecAddress of the contract for which the ContractTree is required.
-   * @returns A ContractTree instance associated with the specified contract address.
-   * @throws An Error if the contract is not found in the ContractDatabase.
-   */
-  private async getTreeForAddress(contractAddress: AztecAddress): Promise<PrivateFunctionsTree> {
-    const instance = await this.getContractInstance(contractAddress);
-    return this.getTreeForClassId(instance.currentContractClassId);
+  async #getContractArtifactByAddress(contractAddress: AztecAddress): Promise<ContractArtifact | undefined> {
+    const contractClassId = await this.#getContractClassId(contractAddress);
+    return contractClassId && this.getContractArtifact(contractClassId);
   }
 
   // Public getters
@@ -130,32 +127,33 @@ export class ContractDataProvider implements DataProvider {
   }
 
   /** Returns a contract instance for a given address. Throws if not found. */
-  public async getContractInstance(contractAddress: AztecAddress): Promise<ContractInstanceWithAddress> {
-    const instance = await this.#getContractInstance(contractAddress);
-    if (!instance) {
-      throw new ContractNotFoundError(contractAddress.toString());
-    }
-    return instance;
+  public async getContractInstance(contractAddress: AztecAddress): Promise<ContractInstanceWithAddress | undefined> {
+    const contract = await this.#contractInstances.getAsync(contractAddress.toString());
+    return contract && SerializableContractInstance.fromBuffer(contract).withAddress(contractAddress);
   }
 
-  public async getContractArtifact(contractClassId: Fr): Promise<ContractArtifact> {
-    const tree = await this.getTreeForClassId(contractClassId);
-    return tree.getArtifact();
+  public async getContractArtifact(contractClassId: Fr): Promise<ContractArtifact | undefined> {
+    const contract = await this.#contractArtifacts.getAsync(contractClassId.toString());
+    // TODO(@spalladino): AztecAsyncMap lies and returns Uint8Arrays instead of Buffers, hence the extra Buffer.from.
+    return contract && contractArtifactFromBuffer(Buffer.from(contract));
   }
 
   /** Returns a contract class for a given class id. Throws if not found. */
-  public async getContractClass(contractClassId: Fr): Promise<ContractClass> {
-    const tree = await this.getTreeForClassId(contractClassId);
-    return tree.getContractClass();
+  public async getContractClass(contractClassId: Fr): Promise<ContractClass | undefined> {
+    const artifact = await this.getContractArtifact(contractClassId);
+    return artifact && getContractClassFromArtifact(artifact);
   }
 
   public async getContract(
     address: AztecAddress,
   ): Promise<(ContractInstanceWithAddress & ContractArtifact) | undefined> {
     const instance = await this.getContractInstance(address);
-    const artifact = instance && (await this.getContractArtifact(instance?.currentContractClassId));
-    if (!instance || !artifact) {
-      return undefined;
+    if (!instance) {
+      return;
+    }
+    const artifact = await this.getContractArtifact(instance.currentContractClassId);
+    if (!artifact) {
+      return;
     }
     return { ...instance, ...artifact };
   }
@@ -172,14 +170,18 @@ export class ContractDataProvider implements DataProvider {
   public async getFunctionArtifact(
     contractAddress: AztecAddress,
     selector: FunctionSelector,
-  ): Promise<FunctionArtifactWithContractName> {
-    const tree = await this.getTreeForAddress(contractAddress);
-    const contractArtifact = tree.getArtifact();
-    const functionArtifact = await tree.getFunctionArtifact(selector);
-    return {
-      ...functionArtifact,
-      contractName: contractArtifact.name,
-    };
+  ): Promise<FunctionArtifactWithContractName | undefined> {
+    const artifact = await this.#getContractArtifactByAddress(contractAddress);
+    const fnArtifact = artifact && (await this.#findFunctionArtifactBySelector(artifact, selector));
+    return fnArtifact && { ...fnArtifact, contractName: artifact.name };
+  }
+
+  public async getPublicFunctionArtifact(
+    contractAddress: AztecAddress,
+  ): Promise<FunctionArtifactWithContractName | undefined> {
+    const artifact = await this.#getContractArtifactByAddress(contractAddress);
+    const fnArtifact = artifact && artifact.functions.find(fn => fn.functionType === FunctionType.PUBLIC);
+    return fnArtifact && { ...fnArtifact, contractName: artifact.name };
   }
 
   /**
@@ -195,8 +197,16 @@ export class ContractDataProvider implements DataProvider {
     contractAddress: AztecAddress,
     functionName: string,
   ): Promise<FunctionArtifact | undefined> {
-    const tree = await this.getTreeForAddress(contractAddress);
-    return tree.getArtifact().functions.find(f => f.name === functionName);
+    const artifact = await this.#getContractArtifactByAddress(contractAddress);
+    return artifact?.functions.find(fn => fn.name === functionName);
+  }
+
+  public async getFunctionAbi(
+    contractAddress: AztecAddress,
+    selector: FunctionSelector,
+  ): Promise<FunctionAbi | undefined> {
+    const artifact = await this.#getContractArtifactByAddress(contractAddress);
+    return artifact && (await this.#findFunctionAbiBySelector(artifact, selector));
   }
 
   /**
@@ -213,24 +223,17 @@ export class ContractDataProvider implements DataProvider {
     contractAddress: AztecAddress,
     selector: FunctionSelector,
   ): Promise<FunctionDebugMetadata | undefined> {
-    const tree = await this.getTreeForAddress(contractAddress);
-    const artifact = await tree.getFunctionArtifact(selector);
-    return getFunctionDebugMetadata(tree.getArtifact(), artifact);
+    const artifact = await this.#getContractArtifactByAddress(contractAddress);
+    const fnArtifact = artifact && (await this.#findFunctionArtifactBySelector(artifact, selector));
+    return fnArtifact && getFunctionDebugMetadata(artifact, fnArtifact);
   }
 
-  /**
-   * Retrieve the bytecode of a specific function in a contract at the given address.
-   * The returned bytecode is required for executing and verifying the function's behavior
-   * in the Aztec network. Throws an error if the contract or function cannot be found.
-   *
-   * @param contractAddress - The contract's address.
-   * @param selector - The function selector.
-   * @returns A Promise that resolves to a Buffer containing the bytecode of the specified function.
-   * @throws Error if the contract address is unknown or not found.
-   */
-  public async getBytecode(contractAddress: AztecAddress, selector: FunctionSelector) {
-    const tree = await this.getTreeForAddress(contractAddress);
-    return tree.getBytecode(selector);
+  public async getPublicFunctionDebugMetadata(
+    contractAddress: AztecAddress,
+  ): Promise<FunctionDebugMetadata | undefined> {
+    const artifact = await this.#getContractArtifactByAddress(contractAddress);
+    const fnArtifact = artifact && artifact.functions.find(fn => fn.functionType === FunctionType.PUBLIC);
+    return fnArtifact && getFunctionDebugMetadata(artifact, fnArtifact);
   }
 
   /**
@@ -245,21 +248,20 @@ export class ContractDataProvider implements DataProvider {
   public async getFunctionMembershipWitness(
     contractClassId: Fr,
     selector: FunctionSelector,
-  ): Promise<MembershipWitness<5>> {
-    const tree = await this.getTreeForClassId(contractClassId);
-    return tree.getFunctionMembershipWitness(selector);
+  ): Promise<MembershipWitness<5> | undefined> {
+    const tree = await this.#getPrivateFunctionTreeForClassId(contractClassId);
+    return tree?.getFunctionMembershipWitness(selector);
   }
 
   public async getDebugContractName(contractAddress: AztecAddress) {
-    const tree = await this.getTreeForAddress(contractAddress);
-    return tree.getArtifact().name;
+    const artifact = await this.#getContractArtifactByAddress(contractAddress);
+    return artifact?.name;
   }
 
   public async getDebugFunctionName(contractAddress: AztecAddress, selector: FunctionSelector) {
-    const tree = await this.getTreeForAddress(contractAddress);
-    const { name: contractName } = tree.getArtifact();
-    const { name: functionName } = await tree.getFunctionArtifact(selector);
-    return `${contractName}:${functionName}`;
+    const artifact = await this.#getContractArtifactByAddress(contractAddress);
+    const fnArtifact = artifact && (await this.#findFunctionAbiBySelector(artifact, selector));
+    return `${artifact?.name ?? contractAddress}:${fnArtifact?.name ?? selector}`;
   }
 
   public async getSize() {
@@ -267,4 +269,32 @@ export class ContractDataProvider implements DataProvider {
       .concat(await toArray(this.#contractArtifacts.valuesAsync()))
       .reduce((sum, value) => sum + value.length, 0);
   }
+
+  async #findFunctionArtifactBySelector(
+    artifact: ContractArtifact,
+    selector: FunctionSelector,
+  ): Promise<FunctionArtifact | undefined> {
+    const functions = artifact.functions;
+    for (let i = 0; i < functions.length; i++) {
+      const fn = functions[i];
+      const fnSelector = await FunctionSelector.fromNameAndParameters(fn.name, fn.parameters);
+      if (fnSelector.equals(selector)) {
+        return fn;
+      }
+    }
+  }
+
+  async #findFunctionAbiBySelector(
+    artifact: ContractArtifact,
+    selector: FunctionSelector,
+  ): Promise<FunctionAbi | undefined> {
+    const functions = [...artifact.functions, ...(artifact.nonDispatchPublicFunctions ?? [])];
+    for (let i = 0; i < functions.length; i++) {
+      const fn = functions[i];
+      const fnSelector = await FunctionSelector.fromNameAndParameters(fn.name, fn.parameters);
+      if (fnSelector.equals(selector)) {
+        return fn;
+      }
+    }
+  }
 }

package/src/storage/contract_data_provider/index.ts

@@ -1,2 +1 @@
 export { ContractDataProvider } from './contract_data_provider.js';
-export { PrivateFunctionsTree } from './private_functions_tree.js';

package/src/storage/contract_data_provider/private_functions_tree.ts

@@ -2,12 +2,12 @@ import { FUNCTION_TREE_HEIGHT } from '@aztec/constants';
 import { Fr } from '@aztec/foundation/fields';
 import { assertLength } from '@aztec/foundation/serialize';
 import { MembershipWitness, type MerkleTree } from '@aztec/foundation/trees';
-import { type ContractArtifact, type FunctionArtifact, FunctionSelector } from '@aztec/stdlib/abi';
+import { type ContractArtifact, FunctionSelector, FunctionType } from '@aztec/stdlib/abi';
 import {
-  type ContractClassWithId,
+  type PrivateFunction,
   computePrivateFunctionLeaf,
   computePrivateFunctionsTree,
-  getContractClassFromArtifact,
+  getContractClassPrivateFunctionFromArtifact,
 } from '@aztec/stdlib/contract';
 
 /**
@@ -19,78 +19,15 @@ import {
 export class PrivateFunctionsTree {
   private tree?: MerkleTree;
 
-  private constructor(private readonly artifact: ContractArtifact, private contractClass: ContractClassWithId) {}
+  private constructor(private readonly privateFunctions: PrivateFunction[]) {}
 
   static async create(artifact: ContractArtifact) {
-    const contractClass = await getContractClassFromArtifact(artifact);
-    return new PrivateFunctionsTree(artifact, contractClass);
-  }
-
-  /**
-   * Retrieve the artifact of a given function.
-   * The function is identified by its selector, which represents a unique identifier for the function's signature.
-   * Throws an error if the function with the provided selector is not found in the contract.
-   *
-   * @param selector - The function selector.
-   * @returns The artifact object containing relevant information about the targeted function.
-   */
-  public async getFunctionArtifact(selector: FunctionSelector): Promise<FunctionArtifact> {
-    const functionsAndSelectors = await Promise.all(
-      this.artifact.functions.map(async f => ({
-        f,
-        selector: await FunctionSelector.fromNameAndParameters(f.name, f.parameters),
-      })),
+    const privateFunctions = await Promise.all(
+      artifact.functions
+        .filter(fn => fn.functionType === FunctionType.PRIVATE)
+        .map(getContractClassPrivateFunctionFromArtifact),
     );
-    const artifact = functionsAndSelectors.find(f => selector.equals(f.selector))?.f;
-    if (!artifact) {
-      throw new Error(
-        `Unknown function. Selector ${selector.toString()} not found in the artifact ${
-          this.artifact.name
-        } with class ${this.getContractClassId().toString()}.`,
-      );
-    }
-    return artifact;
-  }
-
-  /**
-   * Retrieve the bytecode of a function in the contract by its function selector.
-   * The function selector is a unique identifier for each function in a contract.
-   * Throws an error if the function with the given selector is not found in the contract.
-   *
-   * @param selector - The selector of a function to get bytecode for.
-   * @returns The bytecode of the function as a string.
-   */
-  public async getBytecode(selector: FunctionSelector) {
-    const artifact = await this.getFunctionArtifact(selector);
-    return artifact.bytecode;
-  }
-
-  /**
-   * Calculate and return the root of the function tree for the current contract.
-   * This root is a cryptographic commitment to the set of constrained functions within the contract,
-   * which is used in the Aztec node's proof system. The root will be cached after the first call.
-   *
-   * @returns A promise that resolves to the Fr (finite field element) representation of the function tree root.
-   */
-  public getFunctionTreeRoot() {
-    return this.getTree();
-  }
-
-  /** Returns the contract class object. */
-  public getContractClass() {
-    return this.contractClass;
-  }
-
-  /** Returns the contract artifact. */
-  public getArtifact() {
-    return this.artifact;
-  }
-
-  /**
-   * Returns the contract class identifier for the given artifact.
-   */
-  public getContractClassId() {
-    return this.getContractClass().id;
+    return new PrivateFunctionsTree(privateFunctions);
   }
 
   /**
@@ -105,7 +42,7 @@ export class PrivateFunctionsTree {
   public async getFunctionMembershipWitness(
     selector: FunctionSelector,
   ): Promise<MembershipWitness<typeof FUNCTION_TREE_HEIGHT>> {
-    const fn = this.getContractClass().privateFunctions.find(f => f.selector.equals(selector));
+    const fn = this.privateFunctions.find(f => f.selector.equals(selector));
     if (!fn) {
       throw new Error(`Private function with selector ${selector.toString()} not found in contract class.`);
     }
@@ -123,8 +60,7 @@ export class PrivateFunctionsTree {
 
   private async getTree() {
     if (!this.tree) {
-      const fns = this.getContractClass().privateFunctions;
-      this.tree = await computePrivateFunctionsTree(fns);
+      this.tree = await computePrivateFunctionsTree(this.privateFunctions);
     }
     return this.tree;
   }

package/src/test/pxe_test_suite.ts

@@ -81,7 +81,9 @@ export const pxeTestSuite = (testName: string, pxeSetup: () => Promise<PXE>) =>
 
     it('refuses to register a contract with a class that has not been registered', async () => {
       const instance = await randomContractInstanceWithAddress();
-      await expect(pxe.registerContract({ instance })).rejects.toThrow(/DB has no contract class with id/i);
+      await expect(pxe.registerContract({ instance })).rejects.toThrow(
+        /Artifact not found when registering an instance/,
+      );
     });
 
     it('refuses to register a contract with an artifact with mismatching class id', async () => {