reverting stupid changes

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

@@ -1,4 +1,4 @@
-import { CompleteAddress, HistoricBlockData } from '@aztec/circuits.js';
+import { CompleteAddress, HistoricBlockData, PublicKey } from '@aztec/circuits.js';
 import { AztecAddress } from '@aztec/foundation/aztec-address';
 import { Fr } from '@aztec/foundation/fields';
 import { ContractDatabase, MerkleTreeId, NoteFilter } from '@aztec/types';
@@ -50,10 +50,10 @@ export interface Database extends ContractDatabase {
    * Remove nullified notes associated with the given account and nullifiers.
    *
    * @param nullifiers - An array of Fr instances representing nullifiers to be matched.
-   * @param account - An account for which the records are being removed.
+   * @param account - A PublicKey instance representing the account for which the records are being removed.
    * @returns Removed notes.
    */
-  removeNullifiedNotes(nullifiers: Fr[], account: AztecAddress): Promise<NoteDao[]>;
+  removeNullifiedNotes(nullifiers: Fr[], account: PublicKey): Promise<NoteDao[]>;
 
   /**
    * Retrieve the stored Merkle tree roots from the database.

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

@@ -32,8 +32,8 @@ describe('Memory DB', () => {
 
       for (let i = 0; i < notes.length; ++i) {
         const result = await db.getNotes({
-          contractAddress: notes[i].extendedNote.contractAddress,
-          storageSlot: notes[i].extendedNote.storageSlot,
+          contractAddress: notes[i].contractAddress,
+          storageSlot: notes[i].storageSlot,
         });
         expect(result).toEqual([notes[i]]);
       }
@@ -45,8 +45,8 @@ describe('Memory DB', () => {
 
       for (let i = 0; i < notes.length; ++i) {
         const result = await db.getNotes({
-          contractAddress: notes[i].extendedNote.contractAddress,
-          storageSlot: notes[i].extendedNote.storageSlot,
+          contractAddress: notes[i].contractAddress,
+          storageSlot: notes[i].storageSlot,
         });
         expect(result).toEqual([notes[i]]);
       }

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

@@ -1,4 +1,4 @@
-import { CompleteAddress, HistoricBlockData } from '@aztec/circuits.js';
+import { CompleteAddress, HistoricBlockData, PublicKey } from '@aztec/circuits.js';
 import { AztecAddress } from '@aztec/foundation/aztec-address';
 import { Fr } from '@aztec/foundation/fields';
 import { createDebugLogger } from '@aztec/foundation/log';
@@ -54,23 +54,31 @@ export class MemoryDB extends MemoryContractDatabase implements Database {
     return Promise.resolve();
   }
 
-  public getNotes(filter: NoteFilter): Promise<NoteDao[]> {
-    const notes = this.notesTable.filter(
-      noteDao =>
-        (filter.contractAddress == undefined || noteDao.extendedNote.contractAddress.equals(filter.contractAddress)) &&
-        (filter.txHash == undefined || noteDao.extendedNote.txHash.equals(filter.txHash)) &&
-        (filter.storageSlot == undefined || noteDao.extendedNote.storageSlot.equals(filter.storageSlot!)) &&
-        (filter.owner == undefined || noteDao.extendedNote.owner.equals(filter.owner!)),
+  public async getNotes(filter: NoteFilter): Promise<NoteDao[]> {
+    let ownerPublicKey: PublicKey | undefined;
+    if (filter.owner !== undefined) {
+      const ownerCompleteAddress = await this.getCompleteAddress(filter.owner);
+      if (ownerCompleteAddress === undefined) {
+        throw new Error(`Owner ${filter.owner.toString()} not found in memory database`);
+      }
+      ownerPublicKey = ownerCompleteAddress.publicKey;
+    }
+
+    return this.notesTable.filter(
+      note =>
+        (filter.contractAddress == undefined || note.contractAddress.equals(filter.contractAddress)) &&
+        (filter.txHash == undefined || note.txHash.equals(filter.txHash)) &&
+        (filter.storageSlot == undefined || note.storageSlot.equals(filter.storageSlot!)) &&
+        (ownerPublicKey == undefined || note.publicKey.equals(ownerPublicKey!)),
     );
-    return Promise.resolve(notes);
   }
 
-  public removeNullifiedNotes(nullifiers: Fr[], account: AztecAddress) {
+  public removeNullifiedNotes(nullifiers: Fr[], account: PublicKey) {
     const nullifierSet = new Set(nullifiers.map(nullifier => nullifier.toString()));
     const [remaining, removed] = this.notesTable.reduce(
       (acc: [NoteDao[], NoteDao[]], note) => {
         const nullifier = note.siloedNullifier.toString();
-        if (note.extendedNote.owner.equals(account) && nullifierSet.has(nullifier)) {
+        if (note.publicKey.equals(account) && nullifierSet.has(nullifier)) {
           acc[1].push(note);
         } else {
           acc[0].push(note);

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

@@ -1,16 +1,30 @@
-import { Fr } from '@aztec/circuits.js';
-import { ExtendedNote, randomExtendedNote } from '@aztec/types';
+import { AztecAddress, Fr, Point } from '@aztec/circuits.js';
+import { Note, randomTxHash } from '@aztec/types';
 
 import { NoteDao } from './note_dao.js';
 
-export const randomNoteDao = (extendedNoteAttributes: Partial<ExtendedNote> = {}) => {
-  const extendedNote = randomExtendedNote(extendedNoteAttributes);
-  const nonce = Fr.random();
-  const innerNoteHash = Fr.random();
-  const siloedNullifier = Fr.random();
-  const index = BigInt(0);
-
-  return new NoteDao(extendedNote, nonce, innerNoteHash, siloedNullifier, index);
+export const randomNoteDao = ({
+  note = Note.random(),
+  contractAddress = AztecAddress.random(),
+  txHash = randomTxHash(),
+  storageSlot = Fr.random(),
+  nonce = Fr.random(),
+  innerNoteHash = Fr.random(),
+  siloedNullifier = Fr.random(),
+  index = BigInt(0),
+  publicKey = Point.random(),
+}: Partial<NoteDao> = {}) => {
+  return new NoteDao(
+    note,
+    contractAddress,
+    storageSlot,
+    txHash,
+    nonce,
+    innerNoteHash,
+    siloedNullifier,
+    index,
+    publicKey,
+  );
 };
 
 describe('Note DAO', () => {

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

@@ -1,14 +1,20 @@
-import { Fr } from '@aztec/circuits.js';
+import { AztecAddress, Fr, Point, PublicKey } from '@aztec/circuits.js';
 import { toBigIntBE, toBufferBE } from '@aztec/foundation/bigint-buffer';
-import { BufferReader, ExtendedNote } from '@aztec/types';
+import { BufferReader, Note, TxHash } from '@aztec/types';
 
 /**
  * A note with contextual data.
  */
 export class NoteDao {
   constructor(
-    /** The extended note. */
-    public extendedNote: ExtendedNote,
+    /** 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 nonce of the note. */
     public nonce: Fr,
     /**
@@ -20,27 +26,47 @@ export class NoteDao {
     public siloedNullifier: Fr,
     /** The location of the relevant note in the note hash tree. */
     public index: bigint,
+    /** The public key with which the note was encrypted. */
+    public publicKey: PublicKey,
   ) {}
 
   toBuffer(): Buffer {
     return Buffer.concat([
-      this.extendedNote.toBuffer(),
+      this.note.toBuffer(),
+      this.contractAddress.toBuffer(),
+      this.storageSlot.toBuffer(),
+      this.txHash.buffer,
       this.nonce.toBuffer(),
       this.innerNoteHash.toBuffer(),
       this.siloedNullifier.toBuffer(),
       toBufferBE(this.index, 32),
+      this.publicKey.toBuffer(),
     ]);
   }
   static fromBuffer(buffer: Buffer | BufferReader) {
     const reader = BufferReader.asReader(buffer);
 
-    const extendedNote = ExtendedNote.fromBuffer(reader);
+    const note = Note.fromBuffer(reader);
+    const contractAddress = AztecAddress.fromBuffer(reader);
+    const storageSlot = Fr.fromBuffer(reader);
+    const txHash = new TxHash(reader.readBytes(TxHash.SIZE));
     const nonce = Fr.fromBuffer(reader);
     const innerNoteHash = Fr.fromBuffer(reader);
     const siloedNullifier = Fr.fromBuffer(reader);
     const index = toBigIntBE(reader.readBytes(32));
+    const publicKey = Point.fromBuffer(reader);
 
-    return new this(extendedNote, nonce, innerNoteHash, siloedNullifier, index);
+    return new this(
+      note,
+      contractAddress,
+      storageSlot,
+      txHash,
+      nonce,
+      innerNoteHash,
+      siloedNullifier,
+      index,
+      publicKey,
+    );
   }
 
   toString() {
@@ -60,6 +86,6 @@ export class NoteDao {
     // TODO: update
     // 7 fields + 1 bigint + 1 buffer size (4 bytes) + 1 buffer
     const indexSize = Math.ceil(Math.log2(Number(this.index)));
-    return this.extendedNote.note.items.length * Fr.SIZE_IN_BYTES + 7 * Fr.SIZE_IN_BYTES + 4 + indexSize;
+    return this.note.items.length * Fr.SIZE_IN_BYTES + 7 * Fr.SIZE_IN_BYTES + 4 + indexSize;
   }
 }

yarn-project/pxe/src/note_processor/note_processor.test.ts

@@ -34,7 +34,6 @@ describe('Note Processor', () => {
   let addNotesSpy: any;
   let noteProcessor: NoteProcessor;
   let owner: KeyPair;
-  let ownerAddress: CompleteAddress;
   let keyStore: MockProxy<KeyStore>;
   let simulator: MockProxy<AcirSimulator>;
   const firstBlockNum = 123;
@@ -120,7 +119,6 @@ describe('Note Processor', () => {
     wasm = await CircuitsWasm.get();
     grumpkin = new Grumpkin(wasm);
     owner = ConstantKeyPair.random(grumpkin);
-    ownerAddress = await CompleteAddress.fromPrivateKeyAndPartialAddress(await owner.getPrivateKey(), Fr.random());
   });
 
   beforeEach(() => {
@@ -131,7 +129,14 @@ describe('Note Processor', () => {
     keyStore = mock<KeyStore>();
     simulator = mock<AcirSimulator>();
     keyStore.getAccountPrivateKey.mockResolvedValue(owner.getPrivateKey());
-    noteProcessor = new NoteProcessor(ownerAddress, keyStore, database, aztecNode, INITIAL_L2_BLOCK_NUM, simulator);
+    noteProcessor = new NoteProcessor(
+      owner.getPublicKey(),
+      keyStore,
+      database,
+      aztecNode,
+      INITIAL_L2_BLOCK_NUM,
+      simulator,
+    );
 
     simulator.computeNoteHashAndNullifier.mockImplementation((...args) =>
       Promise.resolve({
@@ -206,7 +211,7 @@ describe('Note Processor', () => {
     await noteProcessor.process(blockContexts, encryptedLogsArr);
 
     const addedNoteDaos: NoteDao[] = addNotesSpy.mock.calls[0][0];
-    expect(addedNoteDaos.map(dao => dao.extendedNote)).toEqual([
+    expect(addedNoteDaos.map(dao => dao)).toEqual([
       expect.objectContaining({ ...ownedL1NotePayloads[0] }),
       expect.objectContaining({ ...ownedL1NotePayloads[1] }),
       expect.objectContaining({ ...ownedL1NotePayloads[2] }),

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

@@ -3,6 +3,7 @@ import {
   CompleteAddress,
   MAX_NEW_COMMITMENTS_PER_TX,
   MAX_NEW_NULLIFIERS_PER_TX,
+  PublicKey,
 } from '@aztec/circuits.js';
 import { computeCommitmentNonce, siloNullifier } from '@aztec/circuits.js/abis';
 import { Grumpkin } from '@aztec/circuits.js/barretenberg';
@@ -48,7 +49,7 @@ export class NoteProcessor {
     /**
      * The public counterpart to the private key to be used in note decryption.
      */
-    public readonly account: CompleteAddress,
+    public readonly publicKey: PublicKey,
     private keyStore: KeyStore,
     private db: Database,
     private node: AztecNode,
@@ -112,7 +113,7 @@ export class NoteProcessor {
       // We are using set for `userPertainingTxIndices` to avoid duplicates. This would happen in case there were
       // multiple encrypted logs in a tx pertaining to a user.
       const noteDaos: NoteDao[] = [];
-      const privateKey = await this.keyStore.getAccountPrivateKey(this.account.publicKey);
+      const privateKey = await this.keyStore.getAccountPrivateKey(this.publicKey);
 
       // Iterate over all the encrypted logs and try decrypting them. If successful, store the note.
       for (let indexOfTxInABlock = 0; indexOfTxInABlock < txLogs.length; ++indexOfTxInABlock) {
@@ -147,17 +148,15 @@ export class NoteProcessor {
                 excludedIndices.add(commitmentIndex);
                 noteDaos.push(
                   new NoteDao(
-                    new ExtendedNote(
-                      payload.note,
-                      this.account.address,
-                      payload.contractAddress,
-                      payload.storageSlot,
-                      blockContext.getTxHash(indexOfTxInABlock),
-                    ),
+                    payload.note,
+                    payload.contractAddress,
+                    payload.storageSlot,
+                    blockContext.getTxHash(indexOfTxInABlock),
                     nonce,
                     innerNoteHash,
                     siloedNullifier,
                     index,
+                    this.publicKey,
                   ),
                 );
                 this.stats.decrypted++;
@@ -270,19 +269,19 @@ https://github.com/AztecProtocol/aztec-packages/issues/1641`;
       await this.db.addNotes(noteDaos);
       noteDaos.forEach(noteDao => {
         this.log(
-          `Added note for contract ${noteDao.extendedNote.contractAddress} at slot ${
-            noteDao.extendedNote.storageSlot
+          `Added note for contract ${noteDao.contractAddress} at slot ${
+            noteDao.storageSlot
           } with nullifier ${noteDao.siloedNullifier.toString()}`,
         );
       });
     }
 
     const newNullifiers: Fr[] = blocksAndNotes.flatMap(b => b.blockContext.block.newNullifiers);
-    const removedNotes = await this.db.removeNullifiedNotes(newNullifiers, this.account.address);
+    const removedNotes = await this.db.removeNullifiedNotes(newNullifiers, this.publicKey);
     removedNotes.forEach(noteDao => {
       this.log(
-        `Removed note for contract ${noteDao.extendedNote.contractAddress} at slot ${
-          noteDao.extendedNote.storageSlot
+        `Removed note for contract ${noteDao.contractAddress} at slot ${
+          noteDao.storageSlot
         } with nullifier ${noteDao.siloedNullifier.toString()}`,
       );
     });

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

@@ -124,8 +124,8 @@ export class PXEService implements PXE {
     const completeAddress = await CompleteAddress.fromPrivateKeyAndPartialAddress(privKey, partialAddress);
     const wasAdded = await this.db.addCompleteAddress(completeAddress);
     if (wasAdded) {
-      this.keyStore.addAccount(privKey);
-      this.synchronizer.addAccount(completeAddress, this.keyStore, this.config.l2StartingBlock);
+      const pubKey = this.keyStore.addAccount(privKey);
+      this.synchronizer.addAccount(pubKey, this.keyStore, this.config.l2StartingBlock);
       this.log.info(`Registered account ${completeAddress.address.toString()}`);
       this.log.debug(`Registered account\n ${completeAddress.toReadableString()}`);
     } else {
@@ -196,13 +196,29 @@ export class PXEService implements PXE {
 
   public async getNotes(filter: NoteFilter): Promise<ExtendedNote[]> {
     const noteDaos = await this.db.getNotes(filter);
-    return noteDaos.map(dao => dao.extendedNote);
+
+    // TODO(benesjan): Refactor --> This type conversion is ugly but I decided to keep it this way for now because
+    // key derivation will affect all this
+    const extendedNotes = noteDaos.map(async dao => {
+      let owner = filter.owner;
+      if (owner === undefined) {
+        const completeAddresses = (await this.db.getCompleteAddresses()).find(address =>
+          address.publicKey.equals(dao.publicKey),
+        );
+        if (completeAddresses === undefined) {
+          throw new Error(`Cannot find complete address for public key ${dao.publicKey.toString()}`);
+        }
+        owner = completeAddresses.address;
+      }
+      return new ExtendedNote(dao.note, owner, dao.contractAddress, dao.storageSlot, dao.txHash);
+    });
+    return Promise.all(extendedNotes);
   }
 
   public async addNote(note: ExtendedNote) {
-    const account = (await this.db.getCompleteAddress(note.owner)) ?? {};
-    if (!account) {
-      throw new Error('Unknown note owner: ' + note.owner.toString());
+    const { publicKey } = (await this.db.getCompleteAddress(note.owner)) ?? {};
+    if (!publicKey) {
+      throw new Error('Unknown account.');
     }
 
     const [nonce] = await this.getNoteNonces(note);
@@ -228,7 +244,19 @@ export class PXEService implements PXE {
       throw new Error('The note has been destroyed.');
     }
 
-    await this.db.addNote(new NoteDao(note, nonce, innerNoteHash, siloedNullifier, index));
+    await this.db.addNote(
+      new NoteDao(
+        note.note,
+        note.contractAddress,
+        note.storageSlot,
+        note.txHash,
+        nonce,
+        innerNoteHash,
+        siloedNullifier,
+        index,
+        publicKey,
+      ),
+    );
   }
 
   public async getNoteNonces(note: ExtendedNote): Promise<Fr[]> {