refactor(PXE): sync from PXE instead of from contracts

noir-projects/aztec-nr/aztec/src/macros/internals_functions_generation/external/helpers.nr

@@ -25,21 +25,6 @@ pub(crate) comptime fn get_abi_relevant_attributes(f: FunctionDefinition) -> Quo
     attributes
 }
 
-/// Injects a call to `aztec::messages::discovery::discover_new_messages`, causing for new notes to be added to PXE and made
-/// available for the current execution.
-pub(crate) comptime fn create_message_discovery_call() -> Quoted {
-    quote {
-        /// Safety: message discovery returns nothing and is performed solely for its side-effects. It is therefore
-        /// always safe to call.
-        unsafe {
-            dep::aztec::messages::discovery::discover_new_messages(
-                self.address,
-                _compute_note_hash_and_nullifier,
-            );
-        };
-    }
-}
-
 /// Injects an authwit verification check of the form:
 /// ```
 ///   if (!from.eq(context.msg_sender().unwrap())) {

noir-projects/aztec-nr/aztec/src/macros/internals_functions_generation/external/private.nr

@@ -1,8 +1,7 @@
 use crate::macros::{
     internals_functions_generation::external::helpers::{
-        create_authorize_once_check, create_message_discovery_call, get_abi_relevant_attributes,
+        create_authorize_once_check, get_abi_relevant_attributes,
     },
-    notes::NOTES,
     utils::{
         fn_has_authorize_once, fn_has_noinitcheck, fn_has_nophasecheck, is_fn_initializer,
         is_fn_only_self, is_fn_view, module_has_initializer, module_has_storage,
@@ -120,16 +119,6 @@ pub(crate) comptime fn generate_private_external(f: FunctionDefinition) -> Quote
         }
     };
 
-    // All private functions perform message discovery, since they may need to access notes. This is slightly
-    // inefficient and could be improved by only doing it once we actually attempt to read any. Note that the message
-    // discovery call syncs private events as well. We do not sync those here if there are no notes because we don't
-    // have an API that would access events from private functions.
-    let message_discovery_call = if NOTES.len() > 0 {
-        create_message_discovery_call()
-    } else {
-        quote {}
-    };
-
     // Inject the authwit check if the function is marked with #[authorize_once].
     let authorize_once_check = if fn_has_authorize_once(f) {
         create_authorize_once_check(f, true)
@@ -192,7 +181,6 @@ pub(crate) comptime fn generate_private_external(f: FunctionDefinition) -> Quote
         $init_check
         $internal_check
         $view_check
-        $message_discovery_call
         $authorize_once_check
     };
 

noir-projects/aztec-nr/aztec/src/macros/internals_functions_generation/external/utility.nr

@@ -1,7 +1,4 @@
-use crate::macros::{
-    internals_functions_generation::external::helpers::create_message_discovery_call,
-    utils::module_has_storage,
-};
+use crate::macros::utils::module_has_storage;
 
 pub(crate) comptime fn generate_utility_external(f: FunctionDefinition) -> Quoted {
     // Initialize Storage if module has storage
@@ -28,16 +25,10 @@ pub(crate) comptime fn generate_utility_external(f: FunctionDefinition) -> Quote
         };
     };
 
-    // All utility functions perform message discovery, since they may need to access private notes that would be
-    // found during this process or they may be used to sync private events from TypeScript
-    // (`sync_private_state` function gets invoked by PXE::getPrivateEvents function).
-    let message_discovery_call = create_message_discovery_call();
-
     // A quote to be injected at the beginning of the function body.
     let to_prepend = quote {
         dep::aztec::oracle::version::assert_compatible_oracle_version();
         $contract_self_creation
-        $message_discovery_call
     };
 
     let original_function_name = f.name();

yarn-project/end-to-end/src/e2e_multiple_blobs.test.ts

@@ -54,15 +54,19 @@ describe('e2e_multiple_blobs', () => {
     const utilityFunctions = contractArtifact.functions.filter(fn => fn.functionType == FunctionType.UTILITY);
 
     // Increase the minimum number of txs per block so that all txs will be mined in the same block.
-    const TX_COUNT = 5;
+    // We need enough function broadcasts to produce more than FIELDS_PER_BLOB (4096) fields.
+    const TX_COUNT = 8;
     await aztecNodeAdmin.setConfig({ minTxsPerBlock: TX_COUNT });
 
     const provenTxs = [
       // 1 contract deployment tx.
       await publishContractClass(wallet, AvmTestContract.artifact),
-      // 2 private function broadcast txs.
+      // 5 private function broadcast txs.
       await broadcastFunction(privateFunctions[0]),
       await broadcastFunction(privateFunctions[1]),
+      await broadcastFunction(privateFunctions[2]),
+      await broadcastFunction(privateFunctions[3]),
+      await broadcastFunction(privateFunctions[4]),
       // 1 utility function broadcast tx.
       await broadcastFunction(utilityFunctions[0]),
       // 1 tx to emit note hash, nullifier, l2_to_l1_message, private log and public log.

yarn-project/pxe/src/contract_function_simulator/contract_function_simulator.ts

@@ -23,7 +23,7 @@ import { type Logger, createLogger } from '@aztec/foundation/log';
 import { Timer } from '@aztec/foundation/timer';
 import type { KeyStore } from '@aztec/key-store';
 import { getVKTreeRoot } from '@aztec/noir-protocol-circuits-types/vk-tree';
-import { protocolContractsHash } from '@aztec/protocol-contracts';
+import { isProtocolContract, protocolContractsHash } from '@aztec/protocol-contracts';
 import {
   type CircuitSimulator,
   ExecutionError,
@@ -134,6 +134,13 @@ export class ContractFunctionSimulator {
   ): Promise<PrivateExecutionResult> {
     const simulatorSetupTimer = new Timer();
 
+    // Protocol contracts don't have private state to sync
+    if (!isProtocolContract(contractAddress)) {
+      await this.contractDataProvider.syncPrivateState(contractAddress, selector, privateSyncCall =>
+        this.runUtility(privateSyncCall, [], anchorBlockHeader, scopes),
+      );
+    }
+
     await verifyCurrentClassId(contractAddress, this.aztecNode, this.contractDataProvider, anchorBlockHeader);
 
     const entryPointArtifact = await this.contractDataProvider.getFunctionArtifactWithDebugMetadata(
@@ -170,6 +177,9 @@ export class ContractFunctionSimulator {
       request.txContext,
       callContext,
       anchorBlockHeader,
+      async call => {
+        await this.runUtility(call, [], anchorBlockHeader, scopes);
+      },
       request.authWitnesses,
       request.capsules,
       HashedValuesCache.create(request.argsOfCalls),

yarn-project/pxe/src/contract_function_simulator/oracle/private_execution.test.ts

@@ -416,6 +416,26 @@ describe('Private Execution test suite', () => {
       }
       return { ...artifact, debug: undefined };
     });
+    contractDataProvider.getFunctionCall.mockImplementation(async (functionName, args, to) => {
+      const contract = contracts[to.toString()];
+      if (!contract) {
+        throw new Error(`Contract not found: ${to}`);
+      }
+      const functionArtifact = getFunctionArtifactByName(contract, functionName);
+      return {
+        name: functionArtifact.name,
+        args: encodeArguments(functionArtifact, args),
+        selector: await FunctionSelector.fromNameAndParameters(functionArtifact.name, functionArtifact.parameters),
+        type: functionArtifact.functionType,
+        to,
+        hideMsgSender: false,
+        isStatic: functionArtifact.isStatic,
+        returnTypes: functionArtifact.returnTypes,
+      };
+    });
+    contractDataProvider.syncPrivateState.mockImplementation(async (contractAddress, _functionSelector, _executor) => {
+      await contractDataProvider.getFunctionCall('sync_private_state', [], contractAddress);
+    });
 
     capsuleDataProvider.loadCapsule.mockImplementation((_, __) => Promise.resolve(null));
 
@@ -671,18 +691,40 @@ describe('Private Execution test suite', () => {
         artifact: ParentContractArtifact,
         anchorBlockHeader,
         functionName: 'entry_point',
+        contractAddress: parentAddress,
       });
 
       expect(result.returnValues).toEqual([new Fr(privateIncrement)]);
 
-      // First fetch of the function artifact is the parent contract
-      expect(contractDataProvider.getFunctionArtifact.mock.calls[1]).toEqual([childAddress, childSelector]);
+      expect(contractDataProvider.getFunctionArtifact).toHaveBeenCalledWith(childAddress, childSelector);
       expect(result.nestedExecutionResults).toHaveLength(1);
       expect(result.nestedExecutionResults[0].returnValues).toEqual([new Fr(privateIncrement)]);
       expect(result.publicInputs.privateCallRequests.array[0].callContext).toEqual(
         result.nestedExecutionResults[0].publicInputs.callContext,
       );
     });
+
+    it('syncs private state for child in nested calls', async () => {
+      const childArtifact = getFunctionArtifactByName(ChildContractArtifact, 'value');
+      const parentAddress = await AztecAddress.random();
+      const childAddress = await AztecAddress.random();
+      const childSelector = await FunctionSelector.fromNameAndParameters(childArtifact.name, childArtifact.parameters);
+
+      await mockContractInstance(ChildContractArtifact, childAddress);
+
+      contractDataProvider.getFunctionCall.mockClear();
+
+      const args = [childAddress, childSelector];
+      await runSimulator({
+        args,
+        artifact: ParentContractArtifact,
+        anchorBlockHeader,
+        functionName: 'entry_point',
+        contractAddress: parentAddress,
+      });
+
+      expect(contractDataProvider.getFunctionCall).toHaveBeenCalledWith('sync_private_state', [], childAddress);
+    });
   });
 
   describe('consuming messages', () => {

yarn-project/pxe/src/contract_function_simulator/oracle/private_execution_oracle.ts

@@ -3,10 +3,12 @@ import { Fr } from '@aztec/foundation/curves/bn254';
 import { createLogger } from '@aztec/foundation/log';
 import { Timer } from '@aztec/foundation/timer';
 import type { KeyStore } from '@aztec/key-store';
+import { isProtocolContract } from '@aztec/protocol-contracts';
 import { type CircuitSimulator, toACVMWitness } from '@aztec/simulator/client';
 import {
   type FunctionAbi,
   type FunctionArtifact,
+  type FunctionCall,
   FunctionSelector,
   type NoteSelector,
   countArgumentsSize,
@@ -83,6 +85,8 @@ export class PrivateExecutionOracle extends UtilityExecutionOracle implements IP
     private readonly callContext: CallContext,
     /** Header of a block whose state is used during private execution (not the block the transaction is included in). */
     protected override readonly anchorBlockHeader: BlockHeader,
+    /** Needed to trigger contract synchronization before nested calls */
+    private readonly utilityExecutor: (call: FunctionCall) => Promise<void>,
     /** List of transient auth witnesses to be used during this simulation */
     authWitnesses: AuthWitness[],
     capsules: Capsule[],
@@ -553,6 +557,11 @@ export class PrivateExecutionOracle extends UtilityExecutionOracle implements IP
       this.anchorBlockHeader,
     );
 
+    // Protocol contracts don't have private state to sync
+    if (!isProtocolContract(targetContractAddress)) {
+      await this.contractDataProvider.syncPrivateState(targetContractAddress, functionSelector, this.utilityExecutor);
+    }
+
     const targetArtifact = await this.contractDataProvider.getFunctionArtifactWithDebugMetadata(
       targetContractAddress,
       functionSelector,
@@ -567,6 +576,7 @@ export class PrivateExecutionOracle extends UtilityExecutionOracle implements IP
       derivedTxContext,
       derivedCallContext,
       this.anchorBlockHeader,
+      this.utilityExecutor,
       this.authWitnesses,
       this.capsules,
       this.executionCache,

yarn-project/pxe/src/pxe.ts

@@ -6,7 +6,7 @@ import { Timer } from '@aztec/foundation/timer';
 import { KeyStore } from '@aztec/key-store';
 import type { AztecAsyncKVStore } from '@aztec/kv-store';
 import { L2TipsKVStore } from '@aztec/kv-store/stores';
-import { type ProtocolContractsProvider, protocolContractNames } from '@aztec/protocol-contracts';
+import { type ProtocolContractsProvider, isProtocolContract, protocolContractNames } from '@aztec/protocol-contracts';
 import type { CircuitSimulator } from '@aztec/simulator/client';
 import {
   type ContractArtifact,
@@ -996,6 +996,14 @@ export class PXE {
         const syncTime = syncTimer.ms();
         const functionTimer = new Timer();
         const contractFunctionSimulator = this.#getSimulatorForTx();
+
+        // Protocol contracts don't have private state to sync
+        if (!isProtocolContract(call.to)) {
+          await this.contractDataProvider.syncPrivateState(call.to, call.selector, privateSyncCall =>
+            this.#simulateUtility(contractFunctionSimulator, privateSyncCall),
+          );
+        }
+
         const executionResult = await this.#simulateUtility(contractFunctionSimulator, call, authwits ?? [], scopes);
         const functionTime = functionTimer.ms();
 
@@ -1037,21 +1045,27 @@ export class PXE {
    *    Defaults to the latest known block to PXE + 1.
    * @returns - The packed events with block and tx metadata.
    */
-  public async getPrivateEvents(
-    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.contractDataProvider.getFunctionCall('sync_private_state', [], filter.contractAddress);
-    await this.simulateUtility(call);
-
-    const sanitizedFilter = await new PrivateEventFilterValidator(this.anchorBlockDataProvider).validate(filter);
-
-    this.log.debug(
-      `Getting private events for ${sanitizedFilter.contractAddress.toString()} from ${sanitizedFilter.fromBlock} to ${sanitizedFilter.toBlock}`,
-    );
+  public getPrivateEvents(eventSelector: EventSelector, filter: PrivateEventFilter): Promise<PackedPrivateEvent[]> {
+    return this.#putInJobQueue(async () => {
+      await this.blockStateSynchronizer.sync();
+      const contractFunctionSimulator = this.#getSimulatorForTx();
+      // Protocol contracts don't have private state to sync
+      if (!isProtocolContract(filter.contractAddress)) {
+        await this.contractDataProvider.syncPrivateState(
+          filter.contractAddress,
+          null,
+          async privateSyncCall => await this.#simulateUtility(contractFunctionSimulator, privateSyncCall),
+        );
+      }
+
+      const sanitizedFilter = await new PrivateEventFilterValidator(this.anchorBlockDataProvider).validate(filter);
 
-    return this.privateEventDataProvider.getPrivateEvents(eventSelector, sanitizedFilter);
+      this.log.debug(
+        `Getting private events for ${sanitizedFilter.contractAddress.toString()} from ${sanitizedFilter.fromBlock} to ${sanitizedFilter.toBlock}`,
+      );
+
+      return this.privateEventDataProvider.getPrivateEvents(eventSelector, sanitizedFilter);
+    });
   }
 
   /**

yarn-project/pxe/src/storage/contract_data_provider/contract_data_provider.ts

@@ -316,4 +316,20 @@ export class ContractDataProvider {
       returnTypes: functionDao.returnTypes,
     };
   }
+
+  // Synchronize target contract data
+  public async syncPrivateState(
+    contractAddress: AztecAddress,
+    functionToInvokeAfterSync: FunctionSelector | null,
+    utilityExecutor: (privateSyncCall: FunctionCall) => Promise<any>,
+  ) {
+    const syncPrivateStateFunctionCall = await this.getFunctionCall('sync_private_state', [], contractAddress);
+    if (functionToInvokeAfterSync && functionToInvokeAfterSync.equals(syncPrivateStateFunctionCall.selector)) {
+      throw new Error(
+        'Forbidden `sync_private_state` invocation. `sync_private_state` can only be invoked by PXE, manual execution can lead to inconsistencies.',
+      );
+    }
+
+    return utilityExecutor(syncPrivateStateFunctionCall);
+  }
 }