feat(Misc): Multi wallet token simulator

yarn-project/end-to-end/src/e2e_blacklist_token_contract/access_control.test.ts

@@ -4,20 +4,18 @@ import { BlacklistTokenContractTest, Role } from './blacklist_token_contract_tes
 
 describe('e2e_blacklist_token_contract access control', () => {
   const t = new BlacklistTokenContractTest('access_control');
-  let { wallets } = t;
 
   beforeAll(async () => {
     await t.applyBaseSnapshots();
     await t.setup();
-    ({ wallets } = t);
   });
 
   afterAll(async () => {
     await t.teardown();
   });
 
   afterEach(async () => {
-    await t.tokenSim.check(wallets[0]);
+    await t.tokenSim.check();
   });
 
   it('grant mint permission to the admin', async () => {

yarn-project/end-to-end/src/e2e_blacklist_token_contract/blacklist_token_contract_test.ts

@@ -121,6 +121,7 @@ export class BlacklistTokenContractTest {
 
         this.tokenSim = new TokenSimulator(
           this.asset as unknown as TokenContract,
+          this.wallets[0],
           this.logger,
           this.accounts.map(a => a.address),
         );

yarn-project/end-to-end/src/e2e_blacklist_token_contract/burn.test.ts

@@ -21,7 +21,7 @@ describe('e2e_blacklist_token_contract burn', () => {
   });
 
   afterEach(async () => {
-    await t.tokenSim.check(wallets[0]);
+    await t.tokenSim.check();
   });
 
   describe('public', () => {

yarn-project/end-to-end/src/e2e_blacklist_token_contract/minting.test.ts

@@ -21,11 +21,11 @@ describe('e2e_blacklist_token_contract mint', () => {
   });
 
   beforeEach(async () => {
-    await t.tokenSim.check(wallets[0]);
+    await t.tokenSim.check();
   });
 
   afterEach(async () => {
-    await t.tokenSim.check(wallets[0]);
+    await t.tokenSim.check();
   });
 
   describe('Public', () => {

yarn-project/end-to-end/src/e2e_blacklist_token_contract/shielding.test.ts

@@ -20,7 +20,7 @@ describe('e2e_blacklist_token_contract shield + redeem_shield', () => {
   });
 
   afterEach(async () => {
-    await t.tokenSim.check(wallets[0]);
+    await t.tokenSim.check();
   });
 
   const secret = Fr.random();
@@ -38,7 +38,7 @@ describe('e2e_blacklist_token_contract shield + redeem_shield', () => {
     const receipt = await asset.methods.shield(wallets[0].getAddress(), amount, secretHash, 0).send().wait();
 
     tokenSim.shield(wallets[0].getAddress(), amount);
-    await t.tokenSim.check(wallets[0]);
+    await t.tokenSim.check();
 
     // Redeem it
     await t.addPendingShieldNoteToPXE(0, amount, secretHash, receipt.txHash);
@@ -60,7 +60,7 @@ describe('e2e_blacklist_token_contract shield + redeem_shield', () => {
     const receipt = await action.send().wait();
 
     tokenSim.shield(wallets[0].getAddress(), amount);
-    await t.tokenSim.check(wallets[0]);
+    await t.tokenSim.check();
 
     // Check that replaying the shield should fail!
     const txReplay = asset

yarn-project/end-to-end/src/e2e_blacklist_token_contract/transfer_private.test.ts

@@ -21,7 +21,7 @@ describe('e2e_blacklist_token_contract transfer private', () => {
   });
 
   afterEach(async () => {
-    await t.tokenSim.check(wallets[0]);
+    await t.tokenSim.check();
   });
 
   it('transfer less than balance', async () => {

yarn-project/end-to-end/src/e2e_blacklist_token_contract/transfer_public.test.ts

@@ -21,7 +21,7 @@ describe('e2e_blacklist_token_contract transfer public', () => {
   });
 
   afterEach(async () => {
-    await t.tokenSim.check(wallets[0]);
+    await t.tokenSim.check();
   });
 
   it('transfer less than balance', async () => {

yarn-project/end-to-end/src/e2e_blacklist_token_contract/unshielding.test.ts

@@ -21,7 +21,7 @@ describe('e2e_blacklist_token_contract unshielding', () => {
   });
 
   afterEach(async () => {
-    await t.tokenSim.check(wallets[0]);
+    await t.tokenSim.check();
   });
 
   it('on behalf of self', async () => {

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

@@ -74,15 +74,15 @@ describe('e2e_lending_contract', () => {
       lendingAccount,
       rate,
       lendingContract,
-      new TokenSimulator(collateralAsset, logger, [lendingContract.address, wallet.getAddress()]),
-      new TokenSimulator(stableCoin, logger, [lendingContract.address, wallet.getAddress()]),
+      new TokenSimulator(collateralAsset, wallet, logger, [lendingContract.address, wallet.getAddress()]),
+      new TokenSimulator(stableCoin, wallet, logger, [lendingContract.address, wallet.getAddress()]),
     );
   }, 300_000);
 
   afterAll(() => teardown());
 
   afterEach(async () => {
-    await lendingSim.check(wallet);
+    await lendingSim.check();
   });
 
   it('Mint assets for later usage', async () => {

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

@@ -7,22 +7,22 @@ const TIMEOUT = 1_800_000;
 
 describe('full_prover', () => {
   const t = new FullProverTest('full_prover');
-  let { provenAssets, accounts, tokenSim, logger, wallets } = t;
+  let { provenAssets, accounts, tokenSim, logger } = t;
 
   beforeAll(async () => {
     await t.applyBaseSnapshots();
     await t.applyMintSnapshot();
     await t.setup();
     await t.deployVerifier();
-    ({ provenAssets, accounts, tokenSim, logger, wallets } = t);
+    ({ provenAssets, accounts, tokenSim, logger } = t);
   });
 
   afterAll(async () => {
     await t.teardown();
   });
 
   afterEach(async () => {
-    await t.tokenSim.check(wallets[0]);
+    await t.tokenSim.check();
   });
 
   it(

yarn-project/end-to-end/src/e2e_prover/e2e_prover_test.ts

@@ -120,6 +120,7 @@ export class FullProverTest {
 
         this.tokenSim = new TokenSimulator(
           this.fakeProofsAsset,
+          this.wallets[0],
           this.logger,
           this.accounts.map(a => a.address),
         );

yarn-project/end-to-end/src/e2e_token_contract/access_control.test.ts

@@ -2,20 +2,18 @@ import { TokenContractTest } from './token_contract_test.js';
 
 describe('e2e_token_contract access control', () => {
   const t = new TokenContractTest('access_control');
-  let { wallets } = t;
 
   beforeAll(async () => {
     await t.applyBaseSnapshots();
     await t.setup();
-    ({ wallets } = t);
   });
 
   afterAll(async () => {
     await t.teardown();
   });
 
   afterEach(async () => {
-    await t.tokenSim.check(wallets[0]);
+    await t.tokenSim.check();
   });
 
   it('Set admin', async () => {

yarn-project/end-to-end/src/e2e_token_contract/burn.test.ts

@@ -20,7 +20,7 @@ describe('e2e_token_contract burn', () => {
   });
 
   afterEach(async () => {
-    await t.tokenSim.check(wallets[0]);
+    await t.tokenSim.check();
   });
 
   describe('public', () => {

yarn-project/end-to-end/src/e2e_token_contract/minting.test.ts

@@ -18,7 +18,7 @@ describe('e2e_token_contract minting', () => {
   });
 
   afterEach(async () => {
-    await t.tokenSim.check(wallets[0]);
+    await t.tokenSim.check();
   });
 
   describe('Public', () => {

yarn-project/end-to-end/src/e2e_token_contract/reading_constants.test.ts

@@ -17,13 +17,10 @@ const toString = ({ value }: { value: bigint }) => {
 describe('e2e_token_contract reading constants', () => {
   const t = new TokenContractTest('reading_constants');
   const { TOKEN_DECIMALS, TOKEN_NAME, TOKEN_SYMBOL } = TokenContractTest;
-  let { wallets } = t;
 
   beforeAll(async () => {
     await t.applyBaseSnapshots();
     await t.setup();
-
-    ({ wallets } = t);
   });
 
   afterAll(async () => {
@@ -33,7 +30,7 @@ describe('e2e_token_contract reading constants', () => {
   beforeEach(async () => {});
 
   afterEach(async () => {
-    await t.tokenSim.check(wallets[0]);
+    await t.tokenSim.check();
   });
 
   it('check name private', async () => {

yarn-project/end-to-end/src/e2e_token_contract/shielding.test.ts

@@ -23,7 +23,7 @@ describe('e2e_token_contract shield + redeem shield', () => {
   });
 
   afterEach(async () => {
-    await t.tokenSim.check(wallets[0]);
+    await t.tokenSim.check();
   });
 
   it('on behalf of self', async () => {
@@ -34,7 +34,7 @@ describe('e2e_token_contract shield + redeem shield', () => {
     const receipt = await asset.methods.shield(accounts[0].address, amount, secretHash, 0).send().wait();
 
     tokenSim.shield(accounts[0].address, amount);
-    await tokenSim.check(wallets[0]);
+    await tokenSim.check();
 
     // Redeem it
     await t.addPendingShieldNoteToPXE(0, amount, secretHash, receipt.txHash);
@@ -56,7 +56,7 @@ describe('e2e_token_contract shield + redeem shield', () => {
     const receipt = await action.send().wait();
 
     tokenSim.shield(accounts[0].address, amount);
-    await tokenSim.check(wallets[0]);
+    await tokenSim.check();
 
     // Check that replaying the shield should fail!
     const txReplay = asset.withWallet(wallets[1]).methods.shield(accounts[0].address, amount, secretHash, nonce).send();

yarn-project/end-to-end/src/e2e_token_contract/token_contract_test.ts

@@ -86,6 +86,7 @@ export class TokenContractTest {
 
         this.tokenSim = new TokenSimulator(
           this.asset,
+          this.wallets[0],
           this.logger,
           this.accounts.map(a => a.address),
         );

yarn-project/end-to-end/src/e2e_token_contract/transfer_private.test.ts

@@ -19,7 +19,7 @@ describe('e2e_token_contract transfer private', () => {
   });
 
   afterEach(async () => {
-    await t.tokenSim.check(wallets[0]);
+    await t.tokenSim.check();
   });
 
   it('transfer less than balance', async () => {

yarn-project/end-to-end/src/e2e_token_contract/transfer_public.test.ts

@@ -20,7 +20,7 @@ describe('e2e_token_contract transfer public', () => {
   });
 
   afterEach(async () => {
-    await t.tokenSim.check(wallets[0]);
+    await t.tokenSim.check();
   });
 
   it('transfer less than balance', async () => {

yarn-project/end-to-end/src/e2e_token_contract/unshielding.test.ts

@@ -20,7 +20,7 @@ describe('e2e_token_contract unshielding', () => {
   });
 
   afterEach(async () => {
-    await t.tokenSim.check(wallets[0]);
+    await t.tokenSim.check();
   });
 
   it('on behalf of self', async () => {

yarn-project/end-to-end/src/simulators/lending_simulator.ts

@@ -1,5 +1,5 @@
 // Convenience struct to hold an account's address and secret that can easily be passed around.
-import { type AztecAddress, type CheatCodes, Fr, type Wallet } from '@aztec/aztec.js';
+import { type AztecAddress, type CheatCodes, Fr } from '@aztec/aztec.js';
 import { pedersenHash } from '@aztec/foundation/crypto';
 import { type LendingContract } from '@aztec/noir-contracts.js/Lending';
 
@@ -154,10 +154,10 @@ export class LendingSimulator {
     this.mintedOutside += amount;
   }
 
-  async check(wallet: Wallet) {
+  async check() {
     // Run checks on both underlying assets
-    await this.collateralAsset.check(wallet);
-    await this.stableCoin.check(wallet);
+    await this.collateralAsset.check();
+    await this.stableCoin.check();
 
     // Check that total collateral equals total holdings by contract.
     const totalCollateral = Object.values(this.collateral).reduce((a, b) => new Fr(a.value + b.value), Fr.ZERO);

yarn-project/end-to-end/src/simulators/token_simulator.ts

@@ -9,7 +9,22 @@ export class TokenSimulator {
   private balancePublic: Map<string, bigint> = new Map();
   public totalSupply: bigint = 0n;
 
-  constructor(protected token: TokenContract, protected logger: DebugLogger, protected accounts: AztecAddress[]) {}
+  private lookupProvider: Map<string, Wallet> = new Map();
+
+  constructor(
+    protected token: TokenContract,
+    protected defaultWallet: Wallet,
+    protected logger: DebugLogger,
+    protected accounts: AztecAddress[],
+  ) {}
+
+  public addAccount(account: AztecAddress) {
+    this.accounts.push(account);
+  }
+
+  public setLookupProvider(account: AztecAddress, wallet: Wallet) {
+    this.lookupProvider.set(account.toString(), wallet);
+  }
 
   public mintPrivate(amount: bigint) {
     this.totalSupply += amount;
@@ -82,23 +97,61 @@ export class TokenSimulator {
     return this.balancesPrivate.get(address.toString()) || 0n;
   }
 
-  public async check(wallet: Wallet) {
+  async checkPublic() {
+    // public calls
     const calls = [this.token.methods.total_supply().request()];
     for (const address of this.accounts) {
       calls.push(this.token.methods.balance_of_public(address).request());
-      calls.push(this.token.methods.balance_of_private(address).request());
     }
 
-    const batchedCalls = chunk(calls, 4);
-
-    const results = (await Promise.all(batchedCalls.map(batch => new BatchCall(wallet, batch).simulate()))).flat();
-
+    const results = (
+      await Promise.all(chunk(calls, 4).map(batch => new BatchCall(this.defaultWallet, batch).simulate()))
+    ).flat();
     expect(results[0]).toEqual(this.totalSupply);
 
     // Check that all our balances match
     for (let i = 0; i < this.accounts.length; i++) {
-      expect(results[i * 2 + 1]).toEqual(this.balanceOfPublic(this.accounts[i]));
-      expect(results[i * 2 + 2]).toEqual(this.balanceOfPrivate(this.accounts[i]));
+      expect(results[i + 1]).toEqual(this.balanceOfPublic(this.accounts[i]));
+    }
+  }
+
+  async checkPrivate() {
+    // Private calls
+    const defaultLookups = [];
+    const nonDefaultLookups = [];
+
+    for (const address of this.accounts) {
+      if (this.lookupProvider.has(address.toString())) {
+        nonDefaultLookups.push(address);
+      } else {
+        defaultLookups.push(address);
+      }
+    }
+
+    const defaultCalls = [];
+    for (const address of defaultLookups) {
+      defaultCalls.push(this.token.methods.balance_of_private(address).request());
+    }
+    const results = (
+      await Promise.all(chunk(defaultCalls, 4).map(batch => new BatchCall(this.defaultWallet, batch).simulate()))
+    ).flat();
+    for (let i = 0; i < defaultLookups.length; i++) {
+      expect(results[i]).toEqual(this.balanceOfPrivate(defaultLookups[i]));
     }
+
+    // We are just running individual calls for the non-default lookups
+    // @todo We should also batch these
+    for (const address of nonDefaultLookups) {
+      const wallet = this.lookupProvider.get(address.toString());
+      const asset = wallet ? this.token.withWallet(wallet) : this.token;
+
+      const actualPrivateBalance = await asset.methods.balance_of_private({ address }).simulate();
+      expect(actualPrivateBalance).toEqual(this.balanceOfPrivate(address));
+    }
+  }
+
+  public async check() {
+    await this.checkPublic();
+    await this.checkPrivate();
   }
 }