fix: reprocess notes in pxe when a new contract is added

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

@@ -6,7 +6,16 @@ import { Fr } from '@aztec/foundation/fields';
 import { L2Block, MerkleTreeId, NullifierMembershipWitness, PublicDataWitness } from '@aztec/types';
 
 import { NoteData } from '../acvm/index.js';
-import { CommitmentsDB } from '../public/index.js';
+import { CommitmentsDB } from '../public/db.js';
+
+/**
+ * Error thrown when a contract is not found in the database.
+ */
+export class ContractNotFoundError extends Error {
+  constructor(contractAddress: string) {
+    super(`DB has no contract with address ${contractAddress}`);
+  }
+}
 
 /**
  * A function artifact with optional debug metadata

yarn-project/archiver/src/archiver/archiver.ts

@@ -59,7 +59,6 @@ export class Archiver implements L2BlockSource, L2LogsSource, ContractDataSource
    * @param inboxAddress - Ethereum address of the inbox contract.
    * @param registryAddress - Ethereum address of the registry contract.
    * @param contractDeploymentEmitterAddress - Ethereum address of the contractDeploymentEmitter contract.
-   * @param searchStartBlock - The L1 block from which to start searching for new blocks.
    * @param pollingIntervalMs - The interval for polling for L1 logs (in milliseconds).
    * @param store - An archiver data store for storage & retrieval of blocks, encrypted logs & contract data.
    * @param log - A logger.

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

@@ -211,8 +211,11 @@ describe('e2e_2_pxes', () => {
 
     await awaitServerSynchronized(pxeA);
 
-    const storedValue = await getChildStoredValue(childCompleteAddress, pxeB);
-    expect(storedValue).toEqual(newValueToSet);
+    const storedValueOnB = await getChildStoredValue(childCompleteAddress, pxeB);
+    expect(storedValueOnB).toEqual(newValueToSet);
+
+    const storedValueOnA = await getChildStoredValue(childCompleteAddress, pxeA);
+    expect(storedValueOnA).toEqual(newValueToSet);
   });
 
   it('private state is "zero" when Private eXecution Environment (PXE) does not have the account private key', async () => {
@@ -270,4 +273,111 @@ describe('e2e_2_pxes', () => {
     // registering should wait for the account to be synchronized
     await expect(walletOnB.isAccountStateSynchronized(completeAddress.address)).resolves.toBe(true);
   });
+
+  it('permits sending funds to a user before they have registered the contract', async () => {
+    const initialBalance = 987n;
+    const transferAmount1 = 654n;
+
+    const completeTokenAddress = await deployTokenContract(initialBalance, userA.address, pxeA);
+    const tokenAddress = completeTokenAddress.address;
+
+    // Add account B to wallet A
+    await pxeA.registerRecipient(userB);
+    // Add account A to wallet B
+    await pxeB.registerRecipient(userA);
+
+    // Check initial balances and logs are as expected
+    await expectTokenBalance(walletA, tokenAddress, userA.address, initialBalance);
+    // don't check userB yet
+
+    await expectsNumOfEncryptedLogsInTheLastBlockToBe(aztecNode, 1);
+
+    // Transfer funds from A to B via PXE A
+    const contractWithWalletA = await TokenContract.at(tokenAddress, walletA);
+    const receiptAToB = await contractWithWalletA.methods
+      .transfer(userA.address, userB.address, transferAmount1, 0)
+      .send()
+      .wait();
+    expect(receiptAToB.status).toBe(TxStatus.MINED);
+
+    // now add the contract and check balances
+    await pxeB.addContracts([
+      {
+        artifact: TokenContract.artifact,
+        completeAddress: completeTokenAddress,
+        portalContract: EthAddress.ZERO,
+      },
+    ]);
+    await expectTokenBalance(walletA, tokenAddress, userA.address, initialBalance - transferAmount1);
+    await expectTokenBalance(walletB, tokenAddress, userB.address, transferAmount1);
+  });
+
+  it('permits sending funds to a user, and spending them, before they have registered the contract', async () => {
+    const initialBalance = 987n;
+    const transferAmount1 = 654n;
+    const transferAmount2 = 323n;
+
+    // setup an account that is shared across PXEs
+    const sharedPrivateKey = GrumpkinScalar.random();
+    const sharedAccountOnA = getUnsafeSchnorrAccount(pxeA, sharedPrivateKey, Fr.random());
+    const sharedAccountAddress = sharedAccountOnA.getCompleteAddress();
+    const sharedWalletOnA = await sharedAccountOnA.waitDeploy();
+    await expect(sharedWalletOnA.isAccountStateSynchronized(sharedAccountAddress.address)).resolves.toBe(true);
+
+    const sharedAccountOnB = getUnsafeSchnorrAccount(pxeB, sharedPrivateKey, sharedAccountAddress);
+    await sharedAccountOnB.register();
+    const sharedWalletOnB = await sharedAccountOnB.getWallet();
+
+    await pxeA.registerRecipient(userB);
+
+    // deploy the contract on PXE A
+    const completeTokenAddress = await deployTokenContract(initialBalance, userA.address, pxeA);
+    const tokenAddress = completeTokenAddress.address;
+
+    // Transfer funds from A to Shared Wallet via PXE A
+    const contractWithWalletA = await TokenContract.at(tokenAddress, walletA);
+    const receiptAToShared = await contractWithWalletA.methods
+      .transfer(userA.address, sharedAccountAddress.address, transferAmount1, 0)
+      .send()
+      .wait();
+    expect(receiptAToShared.status).toBe(TxStatus.MINED);
+
+    // Now send funds from Shared Wallet to B via PXE A
+    const contractWithSharedWalletA = await TokenContract.at(tokenAddress, sharedWalletOnA);
+    const receiptSharedToB = await contractWithSharedWalletA.methods
+      .transfer(sharedAccountAddress.address, userB.address, transferAmount2, 0)
+      .send()
+      .wait();
+    expect(receiptSharedToB.status).toBe(TxStatus.MINED);
+
+    // check balances from PXE-A's perspective
+    await expectTokenBalance(walletA, tokenAddress, userA.address, initialBalance - transferAmount1);
+    await expectTokenBalance(
+      sharedWalletOnA,
+      tokenAddress,
+      sharedAccountAddress.address,
+      transferAmount1 - transferAmount2,
+    );
+
+    // now add the contract and check balances from PXE-B's perspective.
+    // The process should be:
+    // PXE-B had previously deferred the notes from A -> Shared, and Shared -> B
+    // PXE-B adds the contract
+    // PXE-B reprocesses the deferred notes, and sees the nullifier for A -> Shared
+    await pxeB.addContracts([
+      {
+        artifact: TokenContract.artifact,
+        completeAddress: completeTokenAddress,
+        portalContract: EthAddress.ZERO,
+      },
+    ]);
+    await expectTokenBalance(walletB, tokenAddress, userB.address, transferAmount2);
+    await expect(sharedWalletOnB.isAccountStateSynchronized(sharedAccountAddress.address)).resolves.toBe(true);
+    await expectTokenBalance(
+      sharedWalletOnB,
+      tokenAddress,
+      sharedAccountAddress.address,
+      transferAmount1 - transferAmount2,
+    );
+  });
 });

yarn-project/pxe/src/contract_data_oracle/index.ts

@@ -1,5 +1,6 @@
-import { AztecAddress, MembershipWitness, VK_TREE_HEIGHT } from '@aztec/circuits.js';
-import { FunctionDebugMetadata, FunctionSelector, getFunctionDebugMetadata } from '@aztec/foundation/abi';
+import { ContractNotFoundError } from '@aztec/acir-simulator';
+import { AztecAddress, ContractFunctionDao, MembershipWitness, VK_TREE_HEIGHT } from '@aztec/circuits.js';
+import { FunctionDebugMetadata, FunctionSelector } from '@aztec/foundation/abi';
 import { ContractDatabase, StateInfoProvider } from '@aztec/types';
 
 import { ContractTree } from '../contract_tree/index.js';
@@ -47,20 +48,25 @@ export class ContractDataOracle {
   /**
    * Retrieves the artifact of a specified function within a given contract.
    * The function is identified by its name, which is unique within a contract.
+   * Throws if the contract has not been added to the database.
    *
    * @param contractAddress - The AztecAddress representing the contract containing the function.
    * @param functionName - The name of the function.
-   * @returns The corresponding function's artifact as an object.
+   * @returns The corresponding function's artifact as an object
    */
-  public async getFunctionArtifactByName(contractAddress: AztecAddress, functionName: string) {
-    const contract = await this.db.getContract(contractAddress);
-    return contract?.functions.find(f => f.name === functionName);
+  public async getFunctionArtifactByName(
+    contractAddress: AztecAddress,
+    functionName: string,
+  ): Promise<ContractFunctionDao | undefined> {
+    const tree = await this.getTree(contractAddress);
+    return tree.contract.getFunctionArtifactByName(functionName);
   }
 
   /**
    * Retrieves the debug metadata of a specified function within a given contract.
    * The function is identified by its selector, which is a unique code generated from the function's signature.
    * Returns undefined if the debug metadata for the given function is not found.
+   * Throws if the contract has not been added to the database.
    *
    * @param contractAddress - The AztecAddress representing the contract containing the function.
    * @param selector - The function selector.
@@ -70,14 +76,14 @@ export class ContractDataOracle {
     contractAddress: AztecAddress,
     selector: FunctionSelector,
   ): Promise<FunctionDebugMetadata | undefined> {
-    const contract = await this.db.getContract(contractAddress);
-    const functionArtifact = contract?.functions.find(f => f.selector.equals(selector));
+    const tree = await this.getTree(contractAddress);
+    const functionArtifact = tree.contract.getFunctionArtifact(selector);
 
-    if (!contract || !functionArtifact) {
+    if (!functionArtifact) {
       return undefined;
     }
 
-    return getFunctionDebugMetadata(contract, functionArtifact.name);
+    return tree.contract.getFunctionDebugMetadataByName(functionArtifact.name);
   }
 
   /**
@@ -88,6 +94,7 @@ export class ContractDataOracle {
    * @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.getTree(contractAddress);
@@ -147,12 +154,12 @@ export class ContractDataOracle {
    * @returns A ContractTree instance associated with the specified contract address.
    * @throws An Error if the contract is not found in the ContractDatabase.
    */
-  private async getTree(contractAddress: AztecAddress) {
+  private async getTree(contractAddress: AztecAddress): Promise<ContractTree> {
     let tree = this.trees.find(t => t.contract.completeAddress.address.equals(contractAddress));
     if (!tree) {
       const contract = await this.db.getContract(contractAddress);
       if (!contract) {
-        throw new Error(`Unknown contract: ${contractAddress}`);
+        throw new ContractNotFoundError(contractAddress.toString());
       }
 
       tree = new ContractTree(contract, this.stateProvider);

yarn-project/pxe/src/database/deferred_note_dao.test.ts

@@ -0,0 +1,34 @@
+import { AztecAddress, Fr, Point } from '@aztec/circuits.js';
+import { Note, randomTxHash } from '@aztec/types';
+
+import { DeferredNoteDao } from './deferred_note_dao.js';
+
+export const randomDeferredNoteDao = ({
+  publicKey = Point.random(),
+  note = Note.random(),
+  contractAddress = AztecAddress.random(),
+  txHash = randomTxHash(),
+  storageSlot = Fr.random(),
+  txNullifier = Fr.random(),
+  newCommitments = [Fr.random(), Fr.random()],
+  dataStartIndexForTx = Math.floor(Math.random() * 100),
+}: Partial<DeferredNoteDao> = {}) => {
+  return new DeferredNoteDao(
+    publicKey,
+    note,
+    contractAddress,
+    storageSlot,
+    txHash,
+    txNullifier,
+    newCommitments,
+    dataStartIndexForTx,
+  );
+};
+
+describe('Deferred Note DAO', () => {
+  it('convert to and from buffer', () => {
+    const deferredNote = randomDeferredNoteDao();
+    const buf = deferredNote.toBuffer();
+    expect(DeferredNoteDao.fromBuffer(buf)).toEqual(deferredNote);
+  });
+});

yarn-project/pxe/src/database/deferred_note_dao.ts

@@ -0,0 +1,55 @@
+import { AztecAddress, Fr, Point, PublicKey, Vector } from '@aztec/circuits.js';
+import { serializeToBuffer } from '@aztec/circuits.js/utils';
+import { BufferReader, Note, TxHash } from '@aztec/types';
+
+/**
+ * A note that is intended for us, but we cannot decode it yet because the contract is not yet in our database.
+ *
+ * So keep the state that we need to decode it later.
+ */
+export class DeferredNoteDao {
+  constructor(
+    /** The public key associated with this note */
+    public publicKey: PublicKey,
+    /** The note as emitted from the Noir contract. */
+    public note: Note,
+    /** The contract address this note is created in. */
+    public contractAddress: AztecAddress,
+    /** The specific storage location of the note on the contract. */
+    public storageSlot: Fr,
+    /** The hash of the tx the note was created in. */
+    public txHash: TxHash,
+    /** The first nullifier emitted by the transaction */
+    public txNullifier: Fr,
+    /** New commitments in this transaction, one of which belongs to this note */
+    public newCommitments: Fr[],
+    /** The next available leaf index for the note hash tree for this transaction */
+    public dataStartIndexForTx: number,
+  ) {}
+
+  toBuffer(): Buffer {
+    return serializeToBuffer(
+      this.publicKey.toBuffer(),
+      this.note.toBuffer(),
+      this.contractAddress.toBuffer(),
+      this.storageSlot.toBuffer(),
+      this.txHash.toBuffer(),
+      this.txNullifier.toBuffer(),
+      new Vector(this.newCommitments),
+      this.dataStartIndexForTx,
+    );
+  }
+  static fromBuffer(buffer: Buffer | BufferReader) {
+    const reader = BufferReader.asReader(buffer);
+    return new DeferredNoteDao(
+      reader.readObject(Point),
+      reader.readObject(Note),
+      reader.readObject(AztecAddress),
+      reader.readObject(Fr),
+      reader.readObject(TxHash),
+      reader.readObject(Fr),
+      reader.readVector(Fr),
+      reader.readNumber(),
+    );
+  }
+}