test: persistence uses TokenContract (AztecProtocol#3930)

This PR extends the end to end test suite added in AztecProtocol#3911 to validate
that the system recovers correctly if the chain advances while its
offline. It also changes the test contract from `EasyPrivateToken` to
`TokenContract` in order to mix in public state as well as private
state.

yarn-project/aztec.js/src/account_manager/index.ts

@@ -1,4 +1,4 @@
-import { PublicKey, getContractDeploymentInfo } from '@aztec/circuits.js';
+import { EthAddress, PublicKey, getContractDeploymentInfo } from '@aztec/circuits.js';
 import { Fr } from '@aztec/foundation/fields';
 import { CompleteAddress, GrumpkinPrivateKey, PXE } from '@aztec/types';
 
@@ -92,6 +92,15 @@ export class AccountManager {
    */
   public async register(opts: WaitOpts = DefaultWaitOpts): Promise<AccountWalletWithPrivateKey> {
     const address = await this.#register();
+
+    await this.pxe.addContracts([
+      {
+        artifact: this.accountContract.getContractArtifact(),
+        completeAddress: address,
+        portalContract: EthAddress.ZERO,
+      },
+    ]);
+
     await waitForAccountSynch(this.pxe, address, opts);
     return this.getWallet();
   }

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

@@ -1,8 +1,15 @@
 import { getUnsafeSchnorrAccount, getUnsafeSchnorrWallet } from '@aztec/accounts/single_key';
-import { AccountWallet, waitForAccountSynch } from '@aztec/aztec.js';
+import {
+  AccountWallet,
+  ExtendedNote,
+  Note,
+  TxHash,
+  computeMessageSecretHash,
+  waitForAccountSynch,
+} from '@aztec/aztec.js';
 import { CompleteAddress, EthAddress, Fq, Fr } from '@aztec/circuits.js';
 import { DeployL1Contracts } from '@aztec/ethereum';
-import { EasyPrivateTokenContract } from '@aztec/noir-contracts/EasyPrivateToken';
+import { TokenContract } from '@aztec/noir-contracts/Token';
 
 import { mkdtemp } from 'fs/promises';
 import { tmpdir } from 'os';
@@ -14,13 +21,14 @@ describe('Aztec persistence', () => {
   /**
    * These tests check that the Aztec Node and PXE can be shutdown and restarted without losing data.
    *
-   * There are four scenarios to check:
+   * There are five scenarios to check:
    * 1. Node and PXE are started with an existing databases
    * 2. PXE is started with an existing database and connects to a Node with an empty database
    * 3. PXE is started with an empty database and connects to a Node with an existing database
    * 4. PXE is started with an empty database and connects to a Node with an empty database
+   * 5. Node and PXE are started with existing databases, but the chain has advanced since they were shutdown
    *
-   * All four scenarios use the same L1 state, which is deployed in the `beforeAll` hook.
+   * All five scenarios use the same L1 state, which is deployed in the `beforeAll` hook.
    */
 
   // the test contract and account deploying it
@@ -48,12 +56,30 @@ describe('Aztec persistence', () => {
     const ownerWallet = await getUnsafeSchnorrAccount(initialContext.pxe, ownerPrivateKey, Fr.ZERO).waitDeploy();
     ownerAddress = ownerWallet.getCompleteAddress();
 
-    const deployer = EasyPrivateTokenContract.deploy(ownerWallet, 1000n, ownerWallet.getAddress());
+    const deployer = TokenContract.deploy(ownerWallet, ownerWallet.getAddress(), 'Test token', 'TEST', 2);
     await deployer.simulate({});
 
     const contract = await deployer.send().deployed();
     contractAddress = contract.completeAddress;
 
+    const secret = Fr.random();
+
+    const mintTx = contract.methods.mint_private(1000n, computeMessageSecretHash(secret));
+    await mintTx.simulate();
+    const mintTxReceipt = await mintTx.send().wait();
+
+    await addPendingShieldNoteToPXE(
+      ownerWallet,
+      contractAddress,
+      1000n,
+      computeMessageSecretHash(secret),
+      mintTxReceipt.txHash,
+    );
+
+    const redeemTx = contract.methods.redeem_shield(ownerAddress.address, 1000n, secret);
+    await redeemTx.simulate();
+    await redeemTx.send().wait();
+
     await initialContext.teardown();
   }, 100_000);
 
@@ -72,37 +98,57 @@ describe('Aztec persistence', () => {
     ],
   ])('%s', (_, contextSetup, timeout) => {
     let ownerWallet: AccountWallet;
-    let contract: EasyPrivateTokenContract;
+    let contract: TokenContract;
 
     beforeEach(async () => {
       context = await contextSetup();
       ownerWallet = await getUnsafeSchnorrWallet(context.pxe, ownerAddress.address, ownerPrivateKey);
-      contract = await EasyPrivateTokenContract.at(contractAddress.address, ownerWallet);
+      contract = await TokenContract.at(contractAddress.address, ownerWallet);
     }, timeout);
 
     afterEach(async () => {
       await context.teardown();
     });
 
-    it('correctly restores balances', async () => {
+    it('correctly restores private notes', async () => {
       // test for >0 instead of exact value so test isn't dependent on run order
-      await expect(contract.methods.getBalance(ownerWallet.getAddress()).view()).resolves.toBeGreaterThan(0n);
+      await expect(contract.methods.balance_of_private(ownerWallet.getAddress()).view()).resolves.toBeGreaterThan(0n);
+    });
+
+    it('correctly restores public storage', async () => {
+      await expect(contract.methods.total_supply().view()).resolves.toBeGreaterThan(0n);
     });
 
     it('tracks new notes for the owner', async () => {
-      const balance = await contract.methods.getBalance(ownerWallet.getAddress()).view();
-      await contract.methods.mint(1000n, ownerWallet.getAddress()).send().wait();
-      await expect(contract.methods.getBalance(ownerWallet.getAddress()).view()).resolves.toEqual(balance + 1000n);
+      const balance = await contract.methods.balance_of_private(ownerWallet.getAddress()).view();
+
+      const secret = Fr.random();
+      const mintTxReceipt = await contract.methods.mint_private(1000n, computeMessageSecretHash(secret)).send().wait();
+      await addPendingShieldNoteToPXE(
+        ownerWallet,
+        contractAddress,
+        1000n,
+        computeMessageSecretHash(secret),
+        mintTxReceipt.txHash,
+      );
+
+      await contract.methods.redeem_shield(ownerWallet.getAddress(), 1000n, secret).send().wait();
+
+      await expect(contract.methods.balance_of_private(ownerWallet.getAddress()).view()).resolves.toEqual(
+        balance + 1000n,
+      );
     });
 
-    it('allows transfers of tokens from owner', async () => {
+    it('allows spending of private notes', async () => {
       const otherWallet = await getUnsafeSchnorrAccount(context.pxe, Fq.random(), Fr.ZERO).waitDeploy();
 
-      const initialOwnerBalance = await contract.methods.getBalance(ownerWallet.getAddress()).view();
-      await contract.methods.transfer(500n, ownerWallet.getAddress(), otherWallet.getAddress()).send().wait();
+      const initialOwnerBalance = await contract.methods.balance_of_private(ownerWallet.getAddress()).view();
+
+      await contract.methods.transfer(ownerWallet.getAddress(), otherWallet.getAddress(), 500n, Fr.ZERO).send().wait();
+
       const [ownerBalance, targetBalance] = await Promise.all([
-        contract.methods.getBalance(ownerWallet.getAddress()).view(),
-        contract.methods.getBalance(otherWallet.getAddress()).view(),
+        contract.methods.balance_of_private(ownerWallet.getAddress()).view(),
+        contract.methods.balance_of_private(otherWallet.getAddress()).view(),
       ]);
 
       expect(ownerBalance).toEqual(initialOwnerBalance - 500n);
@@ -143,42 +189,158 @@ describe('Aztec persistence', () => {
       await context.pxe.registerRecipient(ownerAddress);
 
       const wallet = await getUnsafeSchnorrAccount(context.pxe, Fq.random(), Fr.ZERO).waitDeploy();
-      const contract = await EasyPrivateTokenContract.at(contractAddress.address, wallet);
-      await expect(contract.methods.getBalance(ownerAddress.address).view()).rejects.toThrowError(/Unknown contract/);
+      const contract = await TokenContract.at(contractAddress.address, wallet);
+      await expect(contract.methods.balance_of_private(ownerAddress.address).view()).rejects.toThrowError(
+        /Unknown contract/,
+      );
     });
 
-    it("pxe does not have owner's notes", async () => {
+    it("pxe does not have owner's private notes", async () => {
       await context.pxe.addContracts([
         {
-          artifact: EasyPrivateTokenContract.artifact,
+          artifact: TokenContract.artifact,
           completeAddress: contractAddress,
           portalContract: EthAddress.ZERO,
         },
       ]);
       await context.pxe.registerRecipient(ownerAddress);
 
       const wallet = await getUnsafeSchnorrAccount(context.pxe, Fq.random(), Fr.ZERO).waitDeploy();
-      const contract = await EasyPrivateTokenContract.at(contractAddress.address, wallet);
-      await expect(contract.methods.getBalance(ownerAddress.address).view()).resolves.toEqual(0n);
+      const contract = await TokenContract.at(contractAddress.address, wallet);
+      await expect(contract.methods.balance_of_private(ownerAddress.address).view()).resolves.toEqual(0n);
+    });
+
+    it('has access to public storage', async () => {
+      await context.pxe.addContracts([
+        {
+          artifact: TokenContract.artifact,
+          completeAddress: contractAddress,
+          portalContract: EthAddress.ZERO,
+        },
+      ]);
+
+      const wallet = await getUnsafeSchnorrAccount(context.pxe, Fq.random(), Fr.ZERO).waitDeploy();
+      const contract = await TokenContract.at(contractAddress.address, wallet);
+
+      await expect(contract.methods.total_supply().view()).resolves.toBeGreaterThan(0n);
     });
 
     it('pxe restores notes after registering the owner', async () => {
       await context.pxe.addContracts([
         {
-          artifact: EasyPrivateTokenContract.artifact,
+          artifact: TokenContract.artifact,
           completeAddress: contractAddress,
           portalContract: EthAddress.ZERO,
         },
       ]);
 
-      await context.pxe.registerAccount(ownerPrivateKey, ownerAddress.partialAddress);
-      const ownerWallet = await getUnsafeSchnorrAccount(context.pxe, ownerPrivateKey, ownerAddress).getWallet();
-      const contract = await EasyPrivateTokenContract.at(contractAddress.address, ownerWallet);
+      const ownerAccount = getUnsafeSchnorrAccount(context.pxe, ownerPrivateKey, ownerAddress);
+      await ownerAccount.register();
+      const ownerWallet = await ownerAccount.getWallet();
+      const contract = await TokenContract.at(contractAddress.address, ownerWallet);
 
       await waitForAccountSynch(context.pxe, ownerAddress, { interval: 1, timeout: 10 });
 
       // check that notes total more than 0 so that this test isn't dependent on run order
-      await expect(contract.methods.getBalance(ownerAddress.address).view()).resolves.toBeGreaterThan(0n);
+      await expect(contract.methods.balance_of_private(ownerAddress.address).view()).resolves.toBeGreaterThan(0n);
+    });
+  });
+
+  describe('when starting Node and PXE with existing databases, but chain has advanced since they were shutdown', () => {
+    let secret: Fr;
+    let mintTxHash: TxHash;
+    let mintAmount: bigint;
+    let revealedAmount: bigint;
+
+    // The test system is shutdown. Its state is saved to disk
+    // Start a temporary node and PXE, synch it and add the contract and account to it.
+    // Perform some actions with these temporary components to advance the chain
+    // Then shutdown the temporary components and restart the original components
+    // They should sync up from where they left off and be able to see the actions performed by the temporary node & PXE.
+    beforeAll(async () => {
+      const temporaryContext = await setup(0, { deployL1ContractsValues }, {});
+
+      await temporaryContext.pxe.addContracts([
+        {
+          artifact: TokenContract.artifact,
+          completeAddress: contractAddress,
+          portalContract: EthAddress.ZERO,
+        },
+      ]);
+
+      const ownerAccount = getUnsafeSchnorrAccount(temporaryContext.pxe, ownerPrivateKey, ownerAddress);
+      await ownerAccount.register();
+      const ownerWallet = await ownerAccount.getWallet();
+
+      const contract = await TokenContract.at(contractAddress.address, ownerWallet);
+
+      // mint some tokens with a secret we know and redeem later on a separate PXE
+      secret = Fr.random();
+      mintAmount = 1000n;
+      const mintTxReceipt = await contract.methods
+        .mint_private(mintAmount, computeMessageSecretHash(secret))
+        .send()
+        .wait();
+      mintTxHash = mintTxReceipt.txHash;
+
+      // publicly reveal that I have 1000 tokens
+      revealedAmount = 1000n;
+      await contract.methods.unshield(ownerAddress, ownerAddress, revealedAmount, 0).send().wait();
+
+      // shut everything down
+      await temporaryContext.teardown();
+    }, 100_000);
+
+    let ownerWallet: AccountWallet;
+    let contract: TokenContract;
+
+    beforeEach(async () => {
+      context = await setup(0, { dataDirectory, deployL1ContractsValues }, { dataDirectory });
+      ownerWallet = await getUnsafeSchnorrWallet(context.pxe, ownerAddress.address, ownerPrivateKey);
+      contract = await TokenContract.at(contractAddress.address, ownerWallet);
+
+      await waitForAccountSynch(context.pxe, ownerAddress, { interval: 0.1, timeout: 5 });
+    }, 5000);
+
+    afterEach(async () => {
+      await context.teardown();
+    });
+
+    it("restores owner's public balance", async () => {
+      await expect(contract.methods.balance_of_public(ownerAddress.address).view()).resolves.toEqual(revealedAmount);
+    });
+
+    it('allows consuming transparent note created on another PXE', async () => {
+      // this was created in the temporary PXE in `beforeAll`
+      await addPendingShieldNoteToPXE(
+        ownerWallet,
+        contractAddress,
+        mintAmount,
+        computeMessageSecretHash(secret),
+        mintTxHash,
+      );
+
+      const balanceBeforeRedeem = await contract.methods.balance_of_private(ownerWallet.getAddress()).view();
+
+      await contract.methods.redeem_shield(ownerWallet.getAddress(), mintAmount, secret).send().wait();
+      const balanceAfterRedeem = await contract.methods.balance_of_private(ownerWallet.getAddress()).view();
+
+      expect(balanceAfterRedeem).toEqual(balanceBeforeRedeem + mintAmount);
     });
   });
 });
+
+async function addPendingShieldNoteToPXE(
+  wallet: AccountWallet,
+  asset: CompleteAddress,
+  amount: bigint,
+  secretHash: Fr,
+  txHash: TxHash,
+) {
+  // The storage slot of `pending_shields` is 5.
+  // TODO AlexG, this feels brittle
+  const storageSlot = new Fr(5);
+  const note = new Note([new Fr(amount), secretHash]);
+  const extendedNote = new ExtendedNote(note, wallet.getAddress(), asset.address, storageSlot, txHash);
+  await wallet.addNote(extendedNote);
+}