feat: sync cache invalidation oracle (backport #21459)

noir-projects/aztec-nr/aztec/src/messages/processing/offchain.nr

@@ -5,6 +5,7 @@ use crate::{
         encoding::MESSAGE_CIPHERTEXT_LEN,
         processing::{MessageContext, MessageTxContext, OffchainMessageWithContext, resolve_message_contexts},
     },
+    oracle::contract_sync::invalidate_contract_sync_cache,
     protocol::{
         address::AztecAddress,
         constants::MAX_TX_LIFETIME,
@@ -104,6 +105,9 @@ pub unconstrained fn receive(
     messages: BoundedVec<OffchainMessage, MAX_OFFCHAIN_MESSAGES_PER_RECEIVE_CALL>,
 ) {
     let inbox: CapsuleArray<PendingOffchainMsg> = CapsuleArray::at(contract_address, OFFCHAIN_INBOX_SLOT);
+    // May contain duplicates if multiple messages target the same recipient. This is harmless since
+    // cache invalidation on the TS side is idempotent (deleting an already-deleted key is a no-op).
+    let mut scopes: BoundedVec<AztecAddress, MAX_OFFCHAIN_MESSAGES_PER_RECEIVE_CALL> = BoundedVec::new();
     let mut i = 0;
     let messages_len = messages.len();
     while i < messages_len {
@@ -121,10 +125,11 @@ pub unconstrained fn receive(
                 anchor_block_timestamp: msg.anchor_block_timestamp,
             },
         );
+        scopes.push(msg.recipient);
         i += 1;
     }
 
-    // TODO: Invoke an oracle to invalidate contract sync state cache
+    invalidate_contract_sync_cache(contract_address, scopes);
 }
 
 /// Returns offchain-delivered messages to process during sync.

noir-projects/aztec-nr/aztec/src/oracle/contract_sync.nr

@@ -0,0 +1,18 @@
+use crate::protocol::address::AztecAddress;
+
+#[oracle(aztec_utl_invalidateContractSyncCache)]
+unconstrained fn invalidate_contract_sync_cache_oracle<let N: u32>(
+    contract_address: AztecAddress,
+    scopes: BoundedVec<AztecAddress, N>,
+) {}
+
+/// Forces the PXE to re-sync the given contract for a set of scopes on the next query.
+///
+/// Call this after writing data (e.g. offchain messages) that the contract's `sync_state` function needs to discover.
+/// Without invalidation, the sync cache would skip re-running `sync_state` until the next block.
+pub unconstrained fn invalidate_contract_sync_cache<let N: u32>(
+    contract_address: AztecAddress,
+    scopes: BoundedVec<AztecAddress, N>,
+) {
+    invalidate_contract_sync_cache_oracle(contract_address, scopes);
+}

noir-projects/aztec-nr/aztec/src/oracle/mod.nr

@@ -10,6 +10,7 @@ pub mod auth_witness;
 pub mod block_header;
 pub mod call_private_function;
 pub mod capsules;
+pub mod contract_sync;
 pub mod public_call;
 pub mod tx_phase;
 pub mod execution;

noir-projects/aztec-nr/aztec/src/oracle/version.nr

@@ -4,7 +4,7 @@
 ///
 /// @dev Whenever a contract function or Noir test is run, the `aztec_utl_assertCompatibleOracleVersion` oracle is
 /// called and if the oracle version is incompatible an error is thrown.
-pub global ORACLE_VERSION: Field = 17;
+pub global ORACLE_VERSION: Field = 18;
 
 /// Asserts that the version of the oracle is compatible with the version expected by the contract.
 pub fn assert_compatible_oracle_version() {

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

@@ -27,6 +27,22 @@ describe('e2e_2_pxes', () => {
   let teardownA: () => Promise<void>;
   let teardownB: () => Promise<void>;
 
+  async function setupSecondaryPXE(
+    node: AztecNode,
+    fundedAccounts: InitialAccountData[],
+    accountIndex: number,
+    pxeName: string,
+  ) {
+    const { wallet, teardown } = await setupPXEAndGetWallet(node, {}, undefined, pxeName);
+    const accountManager = await wallet.createSchnorrAccount(
+      fundedAccounts[accountIndex].secret,
+      fundedAccounts[accountIndex].salt,
+    );
+    const deployMethod = await accountManager.getDeployMethod();
+    await deployMethod.send({ from: AztecAddress.ZERO });
+    return { wallet, address: accountManager.address, teardown };
+  }
+
   beforeEach(async () => {
     ({
       aztecNode,
@@ -37,17 +53,11 @@ describe('e2e_2_pxes', () => {
       teardown: teardownA,
     } = await setup(1, { numberOfInitialFundedAccounts: 3 }));
 
-    // Account A is already deployed in setup
-
-    // Deploy accountB via walletB.
-    ({ wallet: walletB, teardown: teardownB } = await setupPXEAndGetWallet(aztecNode, {}, undefined, 'pxe-1'));
-    const accountBManager = await walletB.createSchnorrAccount(
-      initialFundedAccounts[1].secret,
-      initialFundedAccounts[1].salt,
-    );
-    accountBAddress = accountBManager.address;
-    const accountBDeployMethod = await accountBManager.getDeployMethod();
-    await accountBDeployMethod.send({ from: AztecAddress.ZERO });
+    ({
+      wallet: walletB,
+      address: accountBAddress,
+      teardown: teardownB,
+    } = await setupSecondaryPXE(aztecNode, initialFundedAccounts, 1, 'pxe-b'));
 
     await walletA.registerSender(accountBAddress, 'accountB');
     await walletB.registerSender(accountAAddress, 'accountA');
@@ -211,4 +221,32 @@ describe('e2e_2_pxes', () => {
     await expectTokenBalance(walletB, token, accountBAddress, transferAmount2, logger);
     await expectTokenBalance(walletB, token, sharedAccountAddress, transferAmount1 - transferAmount2, logger);
   });
+
+  it('balance updates automatically after sender is registered', async () => {
+    const initialBalance = 500n;
+    const transferAmount = 200n;
+
+    const { contract: token, instance } = await deployToken(walletA, accountAAddress, initialBalance, logger);
+
+    // Set up a third PXE (C) that does NOT have sender A registered
+    const {
+      wallet: walletC,
+      address: accountCAddress,
+      teardown: teardownC,
+    } = await setupSecondaryPXE(aztecNode, initialFundedAccounts, 2, 'pxe-c');
+    await walletC.registerContract(instance, TokenContract.artifact);
+
+    // Transfer from A to C
+    const contractWithWalletA = TokenContract.at(token.address, walletA);
+    await contractWithWalletA.methods.transfer(accountCAddress, transferAmount).send({ from: accountAAddress });
+
+    // Balance is 0 because PXE C doesn't know about sender A yet
+    await expectTokenBalance(walletC, token, accountCAddress, 0n, logger);
+
+    // Register sender A on PXE C -- cache invalidation makes balance visible immediately
+    await walletC.registerSender(accountAAddress, 'accountA');
+    await expectTokenBalance(walletC, token, accountCAddress, transferAmount, logger);
+
+    await teardownC();
+  });
 });

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

@@ -109,9 +109,6 @@ describe('e2e_offchain_payment', () => {
       ])
       .simulate({ from: bob });
 
-    // Force an empty block so the PXE re-syncs and discovers the offchain-delivered notes.
-    await forceEmptyBlock();
-
     // TODO(F-413): we need to implement scopes on capsules so we can check Alice's balance too here. This is not
     // possible right now because the offchain inbox is shared for all accounts using this contract in the same PXE,
     // which is bad.
@@ -155,14 +152,6 @@ describe('e2e_offchain_payment', () => {
       ])
       .simulate({ from: bob });
 
-    // TODO: revisit this. The call to offchain_receive is a utility and as such it causes the contract to sync, which,
-    // in combination with our caching policies, means subsequent utility calls won't trigger a re-sync.
-    // Given we're hooking the offchain sync process to the general sync process, this means we won't process any new
-    // offchain messages until at least one block passes.
-    // A potential escape hatch for this is to remove the check that forbids external invocation of `sync_state`.
-    // That would let users trigger syncs manually to circumvent caching issues like this.
-    await forceEmptyBlock();
-
     // Check that Bob got the payment before a re-org
     const { result: bobBalance } = await contract.methods.get_balance(bob).simulate({ from: bob });
     expect(bobBalance).toBe(paymentAmount);

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

@@ -343,6 +343,7 @@ export class ContractFunctionSimulator {
       capsuleStore: this.capsuleStore,
       privateEventStore: this.privateEventStore,
       messageContextService: this.messageContextService,
+      contractSyncService: this.contractSyncService,
       jobId,
       scopes,
     });

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

@@ -143,6 +143,7 @@ export interface IUtilityExecutionOracle {
   copyCapsule(contractAddress: AztecAddress, srcKey: Fr, dstKey: Fr, numEntries: number): Promise<void>;
   aes128Decrypt(ciphertext: Buffer, iv: Buffer, symKey: Buffer): Promise<Buffer>;
   getSharedSecret(address: AztecAddress, ephPk: Point): Promise<Point>;
+  invalidateContractSyncCache(contractAddress: AztecAddress, scopes: AztecAddress[]): void;
   emitOffchainEffect(data: Fr[]): Promise<void>;
 }
 

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

@@ -634,6 +634,20 @@ export class Oracle {
     return secret.toFields().map(toACVMField);
   }
 
+  // eslint-disable-next-line camelcase
+  aztec_utl_invalidateContractSyncCache(
+    [contractAddress]: ACVMField[],
+    scopes: ACVMField[],
+    [scopeCount]: ACVMField[],
+  ): Promise<ACVMField[]> {
+    const scopeAddresses = scopes.slice(0, +scopeCount).map(s => AztecAddress.fromField(Fr.fromString(s)));
+    this.handlerAsUtility().invalidateContractSyncCache(
+      AztecAddress.fromField(Fr.fromString(contractAddress)),
+      scopeAddresses,
+    );
+    return Promise.resolve([]);
+  }
+
   // eslint-disable-next-line camelcase
   async aztec_utl_emitOffchainEffect(data: ACVMField[]) {
     await this.handlerAsUtility().emitOffchainEffect(data.map(Fr.fromString));

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

@@ -26,7 +26,6 @@ import {
 } from '@aztec/stdlib/tx';
 
 import type { AccessScopes } from '../../access_scopes.js';
-import type { ContractSyncService } from '../../contract_sync/contract_sync_service.js';
 import { NoteService } from '../../notes/note_service.js';
 import type { SenderTaggingStore } from '../../storage/tagging_store/sender_tagging_store.js';
 import { syncSenderTaggingIndexes } from '../../tagging/index.js';
@@ -49,7 +48,6 @@ export type PrivateExecutionOracleArgs = Omit<UtilityExecutionOracleArgs, 'contr
   noteCache: ExecutionNoteCache;
   taggingIndexCache: ExecutionTaggingIndexCache;
   senderTaggingStore: SenderTaggingStore;
-  contractSyncService: ContractSyncService;
   totalPublicCalldataCount?: number;
   sideEffectCounter?: number;
   senderForTags?: AztecAddress;
@@ -83,7 +81,6 @@ export class PrivateExecutionOracle extends UtilityExecutionOracle implements IP
   private readonly noteCache: ExecutionNoteCache;
   private readonly taggingIndexCache: ExecutionTaggingIndexCache;
   private readonly senderTaggingStore: SenderTaggingStore;
-  private readonly contractSyncService: ContractSyncService;
   private totalPublicCalldataCount: number;
   protected sideEffectCounter: number;
   private senderForTags?: AztecAddress;
@@ -103,7 +100,6 @@ export class PrivateExecutionOracle extends UtilityExecutionOracle implements IP
     this.noteCache = args.noteCache;
     this.taggingIndexCache = args.taggingIndexCache;
     this.senderTaggingStore = args.senderTaggingStore;
-    this.contractSyncService = args.contractSyncService;
     this.totalPublicCalldataCount = args.totalPublicCalldataCount ?? 0;
     this.sideEffectCounter = args.sideEffectCounter ?? 0;
     this.senderForTags = args.senderForTags;

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

@@ -232,6 +232,7 @@ describe('Utility Execution test suite', () => {
         capsuleStore,
         privateEventStore,
         messageContextService,
+        contractSyncService,
         jobId: 'test-job-id',
         scopes: 'ALL_SCOPES',
       });
@@ -245,6 +246,24 @@ describe('Utility Execution test suite', () => {
       });
     });
 
+    describe('invalidateContractSyncCache', () => {
+      it('throws when contract address does not match', async () => {
+        const otherAddress = await AztecAddress.random();
+        const scope = await AztecAddress.random();
+        expect(() => utilityExecutionOracle.invalidateContractSyncCache(otherAddress, [scope])).toThrow(
+          `Contract ${contractAddress} cannot invalidate sync cache of ${otherAddress}`,
+        );
+        expect(contractSyncService.invalidateContractForScopes).not.toHaveBeenCalled();
+      });
+
+      it('invalidates cache for the given scopes', async () => {
+        const scopeA = await AztecAddress.random();
+        const scopeB = await AztecAddress.random();
+        utilityExecutionOracle.invalidateContractSyncCache(contractAddress, [scopeA, scopeB]);
+        expect(contractSyncService.invalidateContractForScopes).toHaveBeenCalledWith(contractAddress, [scopeA, scopeB]);
+      });
+    });
+
     describe('resolveMessageContexts', () => {
       const requestSlot = Fr.random();
       const responseSlot = Fr.random();

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

@@ -23,6 +23,7 @@ import type { BlockHeader, Capsule, OffchainEffect } from '@aztec/stdlib/tx';
 
 import type { AccessScopes } from '../../access_scopes.js';
 import { createContractLogger, logContractMessage } from '../../contract_logging.js';
+import type { ContractSyncService } from '../../contract_sync/contract_sync_service.js';
 import { EventService } from '../../events/event_service.js';
 import { LogService } from '../../logs/log_service.js';
 import { MessageContextService } from '../../messages/message_context_service.js';
@@ -62,6 +63,7 @@ export type UtilityExecutionOracleArgs = {
   capsuleStore: CapsuleStore;
   privateEventStore: PrivateEventStore;
   messageContextService: MessageContextService;
+  contractSyncService: ContractSyncService;
   jobId: string;
   log?: ReturnType<typeof createLogger>;
   scopes: AccessScopes;
@@ -91,6 +93,7 @@ export class UtilityExecutionOracle implements IMiscOracle, IUtilityExecutionOra
   protected readonly capsuleStore: CapsuleStore;
   protected readonly privateEventStore: PrivateEventStore;
   protected readonly messageContextService: MessageContextService;
+  protected readonly contractSyncService: ContractSyncService;
   protected readonly jobId: string;
   protected logger: ReturnType<typeof createLogger>;
   protected readonly scopes: AccessScopes;
@@ -110,6 +113,7 @@ export class UtilityExecutionOracle implements IMiscOracle, IUtilityExecutionOra
     this.capsuleStore = args.capsuleStore;
     this.privateEventStore = args.privateEventStore;
     this.messageContextService = args.messageContextService;
+    this.contractSyncService = args.contractSyncService;
     this.jobId = args.jobId;
     this.logger = args.log ?? createLogger('simulator:client_view_context');
     this.scopes = args.scopes;
@@ -651,6 +655,17 @@ export class UtilityExecutionOracle implements IMiscOracle, IUtilityExecutionOra
     return this.capsuleStore.copyCapsule(this.contractAddress, srcSlot, dstSlot, numEntries, this.jobId);
   }
 
+  /**
+   * Clears cached sync state for a contract for a set of scopes, forcing re-sync on the next query so that newly
+   * stored notes or events are discovered.
+   */
+  public invalidateContractSyncCache(contractAddress: AztecAddress, scopes: AztecAddress[]): void {
+    if (!contractAddress.equals(this.contractAddress)) {
+      throw new Error(`Contract ${this.contractAddress} cannot invalidate sync cache of ${contractAddress}`);
+    }
+    this.contractSyncService.invalidateContractForScopes(contractAddress, scopes);
+  }
+
   // TODO(#11849): consider replacing this oracle with a pure Noir implementation of aes decryption.
   public aes128Decrypt(ciphertext: Buffer, iv: Buffer, symKey: Buffer): Promise<Buffer> {
     const aes128 = new Aes128();