@aztec/pxe 0.0.1-commit.c7c42ec → 0.0.1-commit.f295ac2

package/src/pxe.ts

@@ -1,4 +1,5 @@
 import type { PrivateEventFilter } from '@aztec/aztec.js/wallet';
+import { BlockNumber } from '@aztec/foundation/branded-types';
 import { Fr } from '@aztec/foundation/curves/bn254';
 import { type Logger, createLogger } from '@aztec/foundation/log';
 import { SerialQueue } from '@aztec/foundation/queue';
@@ -17,16 +18,15 @@ import {
 } from '@aztec/stdlib/abi';
 import type { AuthWitness } from '@aztec/stdlib/auth-witness';
 import type { AztecAddress } from '@aztec/stdlib/aztec-address';
+import { L2BlockHash } from '@aztec/stdlib/block';
 import {
   CompleteAddress,
-  type ContractClassWithId,
   type ContractInstanceWithAddress,
   type PartialAddress,
   computeContractAddressFromInstance,
   getContractClassFromArtifact,
 } from '@aztec/stdlib/contract';
 import { SimulationError } from '@aztec/stdlib/errors';
-import { computeProtocolNullifier, siloNullifier } from '@aztec/stdlib/hash';
 import type { AztecNode, PrivateKernelProver } from '@aztec/stdlib/interfaces/client';
 import type {
   PrivateExecutionStep,
@@ -54,6 +54,7 @@ import { inspect } from 'util';
 
 import { BlockSynchronizer } from './block_synchronizer/index.js';
 import type { PXEConfig } from './config/index.js';
+import { BenchmarkedNodeFactory } from './contract_function_simulator/benchmarked_node.js';
 import {
   ContractFunctionSimulator,
   generateSimulatedProvingResult,
@@ -63,11 +64,12 @@ import { ProxiedContractStoreFactory } from './contract_function_simulator/proxi
 import { PXEDebugUtils } from './debug/pxe_debug_utils.js';
 import { enrichPublicSimulationError, enrichSimulationError } from './error_enriching.js';
 import { PrivateEventFilterValidator } from './events/private_event_filter_validator.js';
+import { JobCoordinator } from './job_coordinator/job_coordinator.js';
 import {
   PrivateKernelExecutionProver,
   type PrivateKernelExecutionProverConfig,
 } from './private_kernel/private_kernel_execution_prover.js';
-import { PrivateKernelOracleImpl } from './private_kernel/private_kernel_oracle_impl.js';
+import { PrivateKernelOracle } from './private_kernel/private_kernel_oracle.js';
 import { AddressStore } from './storage/address_store/address_store.js';
 import { AnchorBlockStore } from './storage/anchor_block_store/anchor_block_store.js';
 import { CapsuleStore } from './storage/capsule_store/capsule_store.js';
@@ -107,6 +109,7 @@ export class PXE {
     private protocolContractsProvider: ProtocolContractsProvider,
     private log: Logger,
     private jobQueue: SerialQueue,
+    private jobCoordinator: JobCoordinator,
     public debug: PXEDebugUtils,
   ) {}
 
@@ -143,7 +146,19 @@ export class PXE {
     const capsuleStore = new CapsuleStore(store);
     const keyStore = new KeyStore(store);
     const tipsStore = new L2TipsKVStore(store, 'pxe');
-    const synchronizer = new BlockSynchronizer(node, anchorBlockStore, noteStore, tipsStore, config, loggerOrSuffix);
+    const synchronizer = new BlockSynchronizer(
+      node,
+      store,
+      anchorBlockStore,
+      noteStore,
+      privateEventStore,
+      tipsStore,
+      config,
+      loggerOrSuffix,
+    );
+
+    const jobCoordinator = new JobCoordinator(store);
+    jobCoordinator.registerStores([capsuleStore, senderTaggingStore, recipientTaggingStore, privateEventStore]);
 
     const debugUtils = new PXEDebugUtils(contractStore, noteStore);
 
@@ -168,6 +183,7 @@ export class PXE {
       protocolContractsProvider,
       log,
       jobQueue,
+      jobCoordinator,
       debugUtils,
     );
 
@@ -191,7 +207,7 @@ export class PXE {
       this.noteStore,
       this.keyStore,
       this.addressStore,
-      this.node,
+      BenchmarkedNodeFactory.create(this.node),
       this.anchorBlockStore,
       this.senderTaggingStore,
       this.recipientTaggingStore,
@@ -222,7 +238,7 @@ export class PXE {
    *
    * Useful for tasks that cannot run concurrently, such as contract function simulation.
    */
-  #putInJobQueue<T>(fn: () => Promise<T>): Promise<T> {
+  #putInJobQueue<T>(fn: (jobId: string) => Promise<T>): Promise<T> {
     // TODO(#12636): relax the conditions under which we forbid concurrency.
     if (this.jobQueue.length() != 0) {
       this.log.warn(
@@ -230,7 +246,22 @@ export class PXE {
       );
     }
 
-    return this.jobQueue.put(fn);
+    return this.jobQueue.put(async () => {
+      const jobId = this.jobCoordinator.beginJob();
+      this.log.verbose(`Beginning job ${jobId}`);
+
+      try {
+        const result = await fn(jobId);
+        this.log.verbose(`Committing job ${jobId}`);
+
+        await this.jobCoordinator.commitJob(jobId);
+        return result;
+      } catch (err) {
+        this.log.verbose(`Aborting job ${jobId}`);
+        await this.jobCoordinator.abortJob(jobId);
+        throw err;
+      }
+    });
   }
 
   async #registerProtocolContracts() {
@@ -245,25 +276,13 @@ export class PXE {
     this.log.verbose(`Registered protocol contracts in pxe`, registered);
   }
 
-  async #isContractClassPubliclyRegistered(id: Fr): Promise<boolean> {
-    return !!(await this.node.getContractClass(id));
-  }
-
-  async #isContractPublished(address: AztecAddress): Promise<boolean> {
-    return !!(await this.node.getContract(address));
-  }
-
-  async #isContractInitialized(address: AztecAddress): Promise<boolean> {
-    const initNullifier = await siloNullifier(address, address.toField());
-    return !!(await this.node.getNullifierMembershipWitness('latest', initNullifier));
-  }
-
   // Executes the entrypoint private function, as well as all nested private
   // functions that might arise.
   async #executePrivate(
     contractFunctionSimulator: ContractFunctionSimulator,
     txRequest: TxExecutionRequest,
-    scopes?: AztecAddress[],
+    scopes: AztecAddress[] | undefined,
+    jobId: string,
   ): Promise<PrivateExecutionResult> {
     const { origin: contractAddress, functionSelector } = txRequest;
 
@@ -280,6 +299,7 @@ export class PXE {
         // contract entrypoint
         undefined, // senderForTags
         scopes,
+        jobId,
       );
       this.log.debug(`Private simulation completed for ${contractAddress.toString()}:${functionSelector}`);
       return result;
@@ -298,17 +318,19 @@ export class PXE {
    * @param authWitnesses - Authentication witnesses required for the function call.
    * @param scopes - Optional array of account addresses whose notes can be accessed in this call. Defaults to all
    * accounts if not specified.
+   * @param jobId - The job ID for staged writes.
    * @returns The simulation result containing the outputs of the utility function.
    */
   async #simulateUtility(
     contractFunctionSimulator: ContractFunctionSimulator,
     call: FunctionCall,
-    authWitnesses?: AuthWitness[],
-    scopes?: AztecAddress[],
+    authWitnesses: AuthWitness[] | undefined,
+    scopes: AztecAddress[] | undefined,
+    jobId: string,
   ) {
     try {
       const anchorBlockHeader = await this.anchorBlockStore.getBlockHeader();
-      return contractFunctionSimulator.runUtility(call, authWitnesses ?? [], anchorBlockHeader, scopes);
+      return contractFunctionSimulator.runUtility(call, authWitnesses ?? [], anchorBlockHeader, scopes, jobId);
     } catch (err) {
       if (err instanceof SimulationError) {
         await enrichSimulationError(err, this.contractStore, this.log);
@@ -361,13 +383,9 @@ export class PXE {
     privateExecutionResult: PrivateExecutionResult,
     config: PrivateKernelExecutionProverConfig,
   ): Promise<PrivateKernelExecutionProofOutput<PrivateKernelTailCircuitPublicInputs>> {
-    const simulationAnchorBlock = privateExecutionResult.getSimulationAnchorBlockNumber();
-    const kernelOracle = new PrivateKernelOracleImpl(
-      this.contractStore,
-      this.keyStore,
-      this.node,
-      simulationAnchorBlock,
-    );
+    const anchorBlockHeader = await this.anchorBlockStore.getBlockHeader();
+    const anchorBlockHash = L2BlockHash.fromField(await anchorBlockHeader.hash());
+    const kernelOracle = new PrivateKernelOracle(this.contractStore, this.keyStore, this.node, anchorBlockHash);
     const kernelTraceProver = new PrivateKernelExecutionProver(kernelOracle, proofCreator, !this.proverEnabled);
     this.log.debug(`Executing kernel trace prover (${JSON.stringify(config)})...`);
     return await kernelTraceProver.proveWithKernels(txExecutionRequest.toTxRequest(), privateExecutionResult, config);
@@ -380,66 +398,12 @@ export class PXE {
   }
 
   /**
-   * Returns the contract class metadata given a contract class id.
-   * The metadata consists of its contract class, whether it has been publicly registered, and its artifact.
-   * @remark - it queries the node to check whether the contract class with the given id has been publicly registered.
-   * @param id - Identifier of the class.
-   * @param includeArtifact - Identifier of the class.
-   * @returns - It returns the contract class metadata, with the artifact field being optional, and will only be returned if true is passed in
-   * for `includeArtifact`
-   * TODO(@spalladino): The PXE actually holds artifacts and not classes, what should we return? Also,
-   * should the pxe query the node for contract public info, and merge it with its own definitions?
-   * TODO(@spalladino): This method is strictly needed to decide whether to publicly register a class or not
-   * during a public deployment. We probably want a nicer and more general API for this, but it'll have to
-   * do for the time being.
+   * Returns the contract artifact for a given contract class id, if it's registered in the PXE.
+   * @param id - Identifier of the contract class.
+   * @returns The contract artifact if found, undefined otherwise.
    */
-  public async getContractClassMetadata(
-    id: Fr,
-    includeArtifact: boolean = false,
-  ): Promise<{
-    contractClass: ContractClassWithId | undefined;
-    isContractClassPubliclyRegistered: boolean;
-    artifact: ContractArtifact | undefined;
-  }> {
-    const artifact = await this.contractStore.getContractArtifact(id);
-    if (!artifact) {
-      this.log.warn(`No artifact found for contract class ${id.toString()} when looking for its metadata`);
-    }
-
-    return {
-      contractClass: artifact && (await getContractClassFromArtifact(artifact)),
-      isContractClassPubliclyRegistered: await this.#isContractClassPubliclyRegistered(id),
-      artifact: includeArtifact ? artifact : undefined,
-    };
-  }
-
-  /**
-   * Returns the contract metadata given an address.
-   * The metadata consists of its contract instance, which includes the contract class identifier,
-   * initialization hash, deployment salt, and public keys hash; whether the contract instance has been initialized;
-   * and whether the contract instance with the given address has been publicly deployed.
-   * @remark - it queries the node to check whether the contract instance has been initialized / publicly deployed through a node.
-   * This query is not dependent on the PXE.
-   * @param address - The address that the contract instance resides at.
-   * @returns - It returns the contract metadata
-   * TODO(@spalladino): Should we return the public keys in plain as well here?
-   */
-  public async getContractMetadata(address: AztecAddress): Promise<{
-    contractInstance: ContractInstanceWithAddress | undefined;
-    isContractInitialized: boolean;
-    isContractPublished: boolean;
-  }> {
-    let instance;
-    try {
-      instance = await this.contractStore.getContractInstance(address);
-    } catch {
-      this.log.warn(`No instance found for contract ${address.toString()} when looking for its metadata`);
-    }
-    return {
-      contractInstance: instance,
-      isContractInitialized: await this.#isContractInitialized(address),
-      isContractPublished: await this.#isContractPublished(address),
-    };
+  public async getContractArtifact(id: Fr): Promise<ContractArtifact | undefined> {
+    return await this.contractStore.getContractArtifact(id);
   }
 
   /**
@@ -612,13 +576,7 @@ export class PXE {
 
       const header = await this.anchorBlockStore.getBlockHeader();
 
-      const currentClassId = await readCurrentClassId(
-        contractAddress,
-        currentInstance,
-        this.node,
-        header.globalVariables.blockNumber,
-        header.globalVariables.timestamp,
-      );
+      const currentClassId = await readCurrentClassId(contractAddress, currentInstance, this.node, header);
       if (!contractClass.id.equals(currentClassId)) {
         throw new Error('Could not update contract to a class different from the current one.');
       }
@@ -657,14 +615,14 @@ export class PXE {
     let privateExecutionResult: PrivateExecutionResult;
     // We disable proving concurrently mostly out of caution, since it accesses some of our stores. Proving is so
     // computationally demanding that it'd be rare for someone to try to do it concurrently regardless.
-    return this.#putInJobQueue(async () => {
+    return this.#putInJobQueue(async jobId => {
       const totalTimer = new Timer();
       try {
         const syncTimer = new Timer();
         await this.blockStateSynchronizer.sync();
         const syncTime = syncTimer.ms();
         const contractFunctionSimulator = this.#getSimulatorForTx();
-        privateExecutionResult = await this.#executePrivate(contractFunctionSimulator, txRequest);
+        privateExecutionResult = await this.#executePrivate(contractFunctionSimulator, txRequest, undefined, jobId);
 
         const {
           publicInputs,
@@ -711,7 +669,7 @@ export class PXE {
           // TODO(benesjan): The following is an expensive operation. Figure out a way to avoid it.
           const txHash = (await txProvingResult.toTx()).txHash;
 
-          await this.senderTaggingStore.storePendingIndexes(preTagsUsedInTheTx, txHash);
+          await this.senderTaggingStore.storePendingIndexes(preTagsUsedInTheTx, txHash, jobId);
           this.log.debug(`Stored used pre-tags as sender for the tx`, {
             preTagsUsedInTheTx,
           });
@@ -741,7 +699,7 @@ export class PXE {
     skipProofGeneration: boolean = true,
   ): Promise<TxProfileResult> {
     // We disable concurrent profiles for consistency with simulateTx.
-    return this.#putInJobQueue(async () => {
+    return this.#putInJobQueue(async jobId => {
       const totalTimer = new Timer();
       try {
         const txInfo = {
@@ -761,7 +719,12 @@ export class PXE {
         const syncTime = syncTimer.ms();
 
         const contractFunctionSimulator = this.#getSimulatorForTx();
-        const privateExecutionResult = await this.#executePrivate(contractFunctionSimulator, txRequest);
+        const privateExecutionResult = await this.#executePrivate(
+          contractFunctionSimulator,
+          txRequest,
+          undefined,
+          jobId,
+        );
 
         const { executionSteps, timings: { proving } = {} } = await this.#prove(
           txRequest,
@@ -841,7 +804,7 @@ export class PXE {
     // We disable concurrent simulations since those might execute oracles which read and write to the PXE stores (e.g.
     // to the capsules), and we need to prevent concurrent runs from interfering with one another (e.g. attempting to
     // delete the same read value, or reading values that another simulation is currently modifying).
-    return this.#putInJobQueue(async () => {
+    return this.#putInJobQueue(async jobId => {
       try {
         const totalTimer = new Timer();
         const txInfo = {
@@ -867,20 +830,14 @@ export class PXE {
         const skipKernels = overrides?.contracts !== undefined && Object.keys(overrides.contracts ?? {}).length > 0;
 
         // Execution of private functions only; no proving, and no kernel logic.
-        const privateExecutionResult = await this.#executePrivate(contractFunctionSimulator, txRequest, scopes);
+        const privateExecutionResult = await this.#executePrivate(contractFunctionSimulator, txRequest, scopes, jobId);
 
         let publicInputs: PrivateKernelTailCircuitPublicInputs | undefined;
         let executionSteps: PrivateExecutionStep[] = [];
 
         if (skipKernels) {
-          // According to the protocol rules, the nonce generator for the note hashes
-          // can either be the first nullifier in the tx or the protocol nullifier if there are none.
-          const nonceGenerator = privateExecutionResult.firstNullifier.equals(Fr.ZERO)
-            ? await computeProtocolNullifier(await txRequest.toTxRequest().hash())
-            : privateExecutionResult.firstNullifier;
           ({ publicInputs, executionSteps } = await generateSimulatedProvingResult(
             privateExecutionResult,
-            nonceGenerator,
             this.contractStore,
           ));
         } else {
@@ -982,7 +939,7 @@ export class PXE {
     // We disable concurrent simulations since those might execute oracles which read and write to the PXE stores (e.g.
     // to the capsules), and we need to prevent concurrent runs from interfering with one another (e.g. attempting to
     // delete the same read value, or reading values that another simulation is currently modifying).
-    return this.#putInJobQueue(async () => {
+    return this.#putInJobQueue(async jobId => {
       try {
         const totalTimer = new Timer();
         const syncTimer = new Timer();
@@ -990,7 +947,18 @@ export class PXE {
         const syncTime = syncTimer.ms();
         const functionTimer = new Timer();
         const contractFunctionSimulator = this.#getSimulatorForTx();
-        const executionResult = await this.#simulateUtility(contractFunctionSimulator, call, authwits ?? [], scopes);
+
+        await this.contractStore.syncPrivateState(call.to, call.selector, privateSyncCall =>
+          this.#simulateUtility(contractFunctionSimulator, privateSyncCall, [], undefined, jobId),
+        );
+
+        const executionResult = await this.#simulateUtility(
+          contractFunctionSimulator,
+          call,
+          authwits ?? [],
+          scopes,
+          jobId,
+        );
         const functionTime = functionTimer.ms();
 
         const totalTime = totalTimer.ms();
@@ -1035,11 +1003,25 @@ export class PXE {
     eventSelector: EventSelector,
     filter: PrivateEventFilter,
   ): Promise<PackedPrivateEvent[]> {
-    // We need to manually trigger private state sync to have a guarantee that all the events are available.
-    const call = await this.contractStore.getFunctionCall('sync_private_state', [], filter.contractAddress);
-    await this.simulateUtility(call);
+    let anchorBlockNumber: BlockNumber;
+
+    await this.#putInJobQueue(async jobId => {
+      await this.blockStateSynchronizer.sync();
+
+      anchorBlockNumber = (await this.anchorBlockStore.getBlockHeader()).getBlockNumber();
+
+      const contractFunctionSimulator = this.#getSimulatorForTx();
+
+      await this.contractStore.syncPrivateState(
+        filter.contractAddress,
+        null,
+        async privateSyncCall =>
+          await this.#simulateUtility(contractFunctionSimulator, privateSyncCall, [], undefined, jobId),
+      );
+    });
 
-    const sanitizedFilter = await new PrivateEventFilterValidator(this.anchorBlockStore).validate(filter);
+    // anchorBlockNumber is set during the job and fixed to whatever it is after a block sync
+    const sanitizedFilter = new PrivateEventFilterValidator(anchorBlockNumber!).validate(filter);
 
     this.log.debug(
       `Getting private events for ${sanitizedFilter.contractAddress.toString()} from ${sanitizedFilter.fromBlock} to ${sanitizedFilter.toBlock}`,