Recursive public fns in ACIR public simulator

yarn-project/acir-simulator/src/acvm/acvm.ts

@@ -25,6 +25,7 @@ export interface ACIRCallback {
   notifyCreatedNote(params: ACVMField[]): Promise<[ACVMField]>;
   notifyNullifiedNote(params: ACVMField[]): Promise<[ACVMField]>;
   callPrivateFunction(params: ACVMField[]): Promise<ACVMField[]>;
+  callPublicFunction(params: ACVMField[]): Promise<ACVMField[]>;
   storageRead(params: ACVMField[]): Promise<[ACVMField]>;
   storageWrite(params: ACVMField[]): Promise<[ACVMField]>;
   viewNotesPage(params: ACVMField[]): Promise<ACVMField[]>;

yarn-project/acir-simulator/src/client/private_execution.ts

@@ -157,6 +157,7 @@ export class PrivateFunctionExecution {
       viewNotesPage: notAvailable,
       storageRead: notAvailable,
       storageWrite: notAvailable,
+      callPublicFunction: notAvailable,
     });
 
     const publicInputs = extractPublicInputs(partialWitness, acir);

yarn-project/acir-simulator/src/client/unconstrained_execution.ts

@@ -58,6 +58,7 @@ export class UnconstrainedFunctionExecution {
       notifyCreatedNote: notAvailable,
       notifyNullifiedNote: notAvailable,
       callPrivateFunction: notAvailable,
+      callPublicFunction: notAvailable,
       storageRead: notAvailable,
       storageWrite: notAvailable,
     });

yarn-project/acir-simulator/src/public/db.ts

@@ -1,10 +1,11 @@
+import { EthAddress } from '@aztec/circuits.js';
 import { AztecAddress } from '@aztec/foundation/aztec-address';
 import { Fr } from '@aztec/foundation/fields';
 
 /**
- * The database interface for public functions.
+ * Database interface for providing access to public state.
  */
-export interface PublicDB {
+export interface PublicStateDB {
   /**
    * Reads a value from public storage, returning zero if none.
    * @param contract - Owner of the storage.
@@ -13,3 +14,23 @@ export interface PublicDB {
    */
   storageRead(contract: AztecAddress, slot: Fr): Promise<Fr>;
 }
+
+/**
+ * Database interface for providing access to public contract data.
+ */
+export interface PublicContractsDB {
+  /**
+   * Returns the brillig (public bytecode) of a function.
+   * @param address - The contract address that owns this function.
+   * @param functionSelector - The selector for the function.
+   * @returns The bytecode or undefined if not found.
+   */
+  getBytecode(address: AztecAddress, functionSelector: Buffer): Promise<Buffer | undefined>;
+
+  /**
+   * Returns the portal contract address for an L2 address.
+   * @param address - The L2 contract address.
+   * @returns The portal contract address or undefined if not found.
+   */
+  getPortalContractAddress(address: AztecAddress): Promise<EthAddress | undefined>;
+}

yarn-project/acir-simulator/src/public/execution.ts

@@ -1,12 +1,4 @@
-import { CallContext, FunctionData, StateRead, StateTransition, TxRequest } from '@aztec/circuits.js';
-import { createDebugLogger } from '@aztec/foundation/log';
-import { select_return_flattened as selectPublicWitnessFlattened } from '@noir-lang/noir_util_wasm';
-import { acvm, fromACVMField, toACVMField, toACVMWitness } from '../acvm/index.js';
-import { PublicDB } from './db.js';
-import { StateActionsCollector } from './state_actions.js';
-import { Fr } from '@aztec/foundation/fields';
-import { AztecAddress } from '@aztec/foundation/aztec-address';
-import { EthAddress } from '@aztec/foundation/eth-address';
+import { AztecAddress, CallContext, Fr, FunctionData, StateRead, StateTransition, TxRequest } from '@aztec/circuits.js';
 
 /**
  * The public function execution result.
@@ -18,113 +10,30 @@ export interface PublicExecutionResult {
   stateReads: StateRead[];
   /** The state transitions performed by the function. */
   stateTransitions: StateTransition[];
+  /** The results of nested calls. */
+  nestedExecutions: this[];
 }
 
 /**
- * Generates the initial witness for a public function.
- * @param args - The arguments to the function.
- * @param callContext - The call context of the function.
- * @param witnessStartIndex - The index where to start inserting the parameters.
- * @returns The initial witness.
+ * The execution of a public function.
  */
-function getInitialWitness(args: Fr[], callContext: CallContext, witnessStartIndex = 1) {
-  return toACVMWitness(witnessStartIndex, [
-    callContext.isContractDeployment,
-    callContext.isDelegateCall,
-    callContext.isStaticCall,
-    callContext.msgSender,
-    callContext.portalContractAddress,
-    callContext.storageContractAddress,
-    ...args,
-  ]);
+export interface PublicExecution {
+  /** Address of the contract being executed. */
+  contractAddress: AztecAddress;
+  /** Function of the contract being called. */
+  functionData: FunctionData;
+  /** Arguments for the call. */
+  args: Fr[];
+  /** Context of the call. */
+  callContext: CallContext;
 }
 
 /**
- * The public function execution class.
+ * Returns whether the input is a public execution.
+ * @param input - Input to check.
+ * @returns Whether it's a public execution.
  */
-export class PublicExecution {
-  constructor(
-    /** The public database. */
-    public readonly db: PublicDB,
-    /** The ACIR bytecode of the public function. */
-    public readonly publicFunctionBytecode: Buffer,
-    /** The address of the contract to execute. */
-    public readonly contractAddress: AztecAddress,
-    /** The function data of the function to execute. */
-    public readonly functionData: FunctionData,
-    /** The arguments of the function to execute. */
-    public readonly args: Fr[],
-    /** The call context of the execution. */
-    public readonly callContext: CallContext,
-
-    private log = createDebugLogger('aztec:simulator:public-execution'),
-  ) {}
-
-  /**
-   * Creates a public function execution from a transaction request.
-   * @param db - The public database.
-   * @param request - The transaction request.
-   * @param bytecode - The bytecode of the public function.
-   * @param portalContractAddress - The address of the portal contract.
-   * @returns The public function execution.
-   */
-  static fromTransactionRequest(db: PublicDB, request: TxRequest, bytecode: Buffer, portalContractAddress: EthAddress) {
-    const contractAddress = request.to;
-    const callContext: CallContext = new CallContext(
-      request.from,
-      request.to,
-      portalContractAddress,
-      false,
-      false,
-      false,
-    );
-    return new this(db, bytecode, contractAddress, request.functionData, request.args, callContext);
-  }
-
-  /**
-   * Executes the public function.
-   * @returns The execution result.
-   */
-  public async run(): Promise<PublicExecutionResult> {
-    const selectorHex = this.functionData.functionSelector.toString('hex');
-    this.log(`Executing public external function ${this.contractAddress.toShortString()}:${selectorHex}`);
-
-    const acir = this.publicFunctionBytecode;
-    const initialWitness = getInitialWitness(this.args, this.callContext);
-    const stateActions = new StateActionsCollector(this.db, this.contractAddress);
-
-    const notAvailable = () => Promise.reject(`Built-in not available for public execution simulation`);
-
-    const { partialWitness } = await acvm(acir, initialWitness, {
-      getSecretKey: notAvailable,
-      getNotes2: notAvailable,
-      getRandomField: notAvailable,
-      notifyCreatedNote: notAvailable,
-      notifyNullifiedNote: notAvailable,
-      callPrivateFunction: notAvailable,
-      viewNotesPage: notAvailable,
-      storageRead: async ([slot]) => {
-        const storageSlot = fromACVMField(slot);
-        const value = await stateActions.read(storageSlot);
-        this.log(`Oracle storage read: slot=${storageSlot.toShortString()} value=${value.toString()}`);
-        return [toACVMField(value)];
-      },
-      storageWrite: async ([slot, value]) => {
-        const storageSlot = fromACVMField(slot);
-        const newValue = fromACVMField(value);
-        await stateActions.write(storageSlot, newValue);
-        this.log(`Oracle storage write: slot=${storageSlot.toShortString()} value=${value.toString()}`);
-        return [toACVMField(newValue)];
-      },
-    });
-
-    const returnValues = selectPublicWitnessFlattened(acir, partialWitness).map(fromACVMField);
-    const [stateReads, stateTransitions] = stateActions.collect();
-
-    return {
-      stateReads,
-      stateTransitions,
-      returnValues,
-    };
-  }
+export function isPublicExecution(input: PublicExecution | TxRequest): input is PublicExecution {
+  const execution = input as PublicExecution;
+  return !!execution.callContext && !!execution.args && !!execution.contractAddress && !!execution.functionData;
 }

yarn-project/acir-simulator/src/public/executor.ts

@@ -0,0 +1,149 @@
+import { AztecAddress, CallContext, EthAddress, Fr, FunctionData, TxRequest } from '@aztec/circuits.js';
+import { padArrayEnd } from '@aztec/foundation/collection';
+import { createDebugLogger } from '@aztec/foundation/log';
+import { select_return_flattened as selectPublicWitnessFlattened } from '@noir-lang/noir_util_wasm';
+import { acvm, frToAztecAddress, frToSelector, fromACVMField, toACVMField, toACVMWitness } from '../acvm/index.js';
+import { PublicContractsDB, PublicStateDB } from './db.js';
+import { PublicExecution, PublicExecutionResult } from './execution.js';
+import { StateActionsCollector } from './state_actions.js';
+
+// Copied from crate::abi at noir-contracts/src/contracts/noir-aztec3/src/abi.nr
+const NOIR_MAX_RETURN_VALUES = 4;
+
+/**
+ * Handles execution of public functions.
+ */
+export class PublicExecutor {
+  constructor(
+    private readonly stateDb: PublicStateDB,
+    private readonly contractsDb: PublicContractsDB,
+
+    private log = createDebugLogger('aztec:simulator:public-executor'),
+  ) {}
+
+  /**
+   * Executes a public execution request.
+   * @param execution - The execution to run.
+   * @returns The result of the run plus all nested runs.
+   */
+  public async execute(execution: PublicExecution): Promise<PublicExecutionResult> {
+    const selectorHex = execution.functionData.functionSelector.toString('hex');
+    this.log(`Executing public external function ${execution.contractAddress.toShortString()}:${selectorHex}`);
+
+    const selector = execution.functionData.functionSelector;
+    const acir = await this.contractsDb.getBytecode(execution.contractAddress, selector);
+    if (!acir) throw new Error(`Bytecode not found for ${execution.contractAddress.toShortString()}:${selectorHex}`);
+
+    const initialWitness = getInitialWitness(execution.args, execution.callContext);
+    const stateActions = new StateActionsCollector(this.stateDb, execution.contractAddress);
+    const nestedExecutions: PublicExecutionResult[] = [];
+
+    const notAvailable = () => Promise.reject(`Built-in not available for public execution simulation`);
+
+    const { partialWitness } = await acvm(acir, initialWitness, {
+      getSecretKey: notAvailable,
+      getNotes2: notAvailable,
+      getRandomField: notAvailable,
+      notifyCreatedNote: notAvailable,
+      notifyNullifiedNote: notAvailable,
+      callPrivateFunction: notAvailable,
+      viewNotesPage: notAvailable,
+      storageRead: async ([slot]) => {
+        const storageSlot = fromACVMField(slot);
+        const value = await stateActions.read(storageSlot);
+        this.log(`Oracle storage read: slot=${storageSlot.toShortString()} value=${value.toString()}`);
+        return [toACVMField(value)];
+      },
+      storageWrite: async ([slot, value]) => {
+        const storageSlot = fromACVMField(slot);
+        const newValue = fromACVMField(value);
+        await stateActions.write(storageSlot, newValue);
+        this.log(`Oracle storage write: slot=${storageSlot.toShortString()} value=${value.toString()}`);
+        return [toACVMField(newValue)];
+      },
+      callPublicFunction: async ([address, functionSelector, ...args]) => {
+        this.log(`Public function call: addr=${address} selector=${functionSelector} args=${args.join(',')}`);
+        const childExecutionResult = await this.callPublicFunction(
+          frToAztecAddress(fromACVMField(address)),
+          frToSelector(fromACVMField(functionSelector)),
+          args.map(f => fromACVMField(f)),
+          execution.callContext,
+        );
+
+        nestedExecutions.push(childExecutionResult);
+        this.log(`Returning from nested call: ret=${childExecutionResult.returnValues.join(', ')}`);
+        return padArrayEnd(childExecutionResult.returnValues, Fr.ZERO, NOIR_MAX_RETURN_VALUES).map(fr => fr.toString());
+      },
+    });
+
+    const returnValues = selectPublicWitnessFlattened(acir, partialWitness).map(fromACVMField);
+    const [stateReads, stateTransitions] = stateActions.collect();
+
+    return {
+      stateReads,
+      stateTransitions,
+      returnValues,
+      nestedExecutions,
+    };
+  }
+
+  /**
+   * Creates a PublicExecution out of a TxRequest to a public function.
+   * @param input - The TxRequest calling a public function.
+   * @returns A PublicExecution object that can be run via execute.
+   */
+  public async getPublicExecution(input: TxRequest): Promise<PublicExecution> {
+    const contractAddress = input.to;
+    const portalContractAddress = (await this.contractsDb.getPortalContractAddress(contractAddress)) ?? EthAddress.ZERO;
+    const callContext: CallContext = new CallContext(input.from, input.to, portalContractAddress, false, false, false);
+
+    return { callContext, contractAddress, functionData: input.functionData, args: input.args };
+  }
+
+  private async callPublicFunction(
+    targetContractAddress: AztecAddress,
+    targetFunctionSelector: Buffer,
+    targetArgs: Fr[],
+    callerContext: CallContext,
+  ) {
+    const portalAddress = (await this.contractsDb.getPortalContractAddress(targetContractAddress)) ?? EthAddress.ZERO;
+    const functionData = new FunctionData(targetFunctionSelector, false, false);
+
+    const callContext = CallContext.from({
+      msgSender: callerContext.storageContractAddress,
+      portalContractAddress: portalAddress,
+      storageContractAddress: targetContractAddress,
+      isContractDeployment: false,
+      isDelegateCall: false,
+      isStaticCall: false,
+    });
+
+    const nestedExecution: PublicExecution = {
+      args: targetArgs,
+      contractAddress: targetContractAddress,
+      functionData,
+      callContext,
+    };
+
+    return this.execute(nestedExecution);
+  }
+}
+
+/**
+ * Generates the initial witness for a public function.
+ * @param args - The arguments to the function.
+ * @param callContext - The call context of the function.
+ * @param witnessStartIndex - The index where to start inserting the parameters.
+ * @returns The initial witness.
+ */
+function getInitialWitness(args: Fr[], callContext: CallContext, witnessStartIndex = 1) {
+  return toACVMWitness(witnessStartIndex, [
+    callContext.isContractDeployment,
+    callContext.isDelegateCall,
+    callContext.isStaticCall,
+    callContext.msgSender,
+    callContext.portalContractAddress,
+    callContext.storageContractAddress,
+    ...args,
+  ]);
+}