feat: Initial cheatcode loadPublic

yarn-project/end-to-end/src/cheat_codes.ts

@@ -0,0 +1,170 @@
+import { AztecAddress, CircuitsWasm, Fr } from '@aztec/circuits.js';
+import { pedersenPlookupCommitInputs } from '@aztec/circuits.js/barretenberg';
+import { createDebugLogger } from '@aztec/foundation/log';
+import { AztecRPC } from '@aztec/types';
+
+const toFr = (value: Fr | bigint): Fr => {
+  return typeof value === 'bigint' ? new Fr(value) : value;
+};
+
+/**
+ * A class that provides utility functions for interacting with the chain.
+ */
+export class CheatCodes {
+  constructor(
+    /**
+     * The L1 cheat codes.
+     */
+    public l1: L1CheatCodes,
+    /**
+     * The L2 cheat codes.
+     */
+    public l2: L2CheatCodes,
+  ) {}
+
+  static async create(rpcUrl: string, aztecRpc: AztecRPC): Promise<CheatCodes> {
+    const l1CheatCodes = new L1CheatCodes(rpcUrl);
+    const l2CheatCodes = new L2CheatCodes(aztecRpc, await CircuitsWasm.get(), l1CheatCodes);
+    return new CheatCodes(l1CheatCodes, l2CheatCodes);
+  }
+}
+
+/**
+ * A class that provides utility functions for interacting with the L1 chain.
+ */
+class L1CheatCodes {
+  constructor(
+    /**
+     * The RPC client to use for interacting with the chain
+     */
+    public rpcUrl: string,
+    /**
+     * The logger to use for the l1 cheatcodes
+     */
+    public logger = createDebugLogger('aztec:cheat_codes:l1'),
+  ) {}
+
+  async rpcCall(method: string, params: any[]) {
+    const paramsString = JSON.stringify(params);
+    const content = {
+      body: `{"jsonrpc":"2.0", "method": "${method}", "params": ${paramsString}, "id": 1}`,
+      method: 'POST',
+      headers: { 'Content-Type': 'application/json' },
+    };
+    return await (await fetch(this.rpcUrl, content)).json();
+  }
+
+  /**
+   * Get the current blocknumber
+   * @returns The current block number
+   */
+  public async blockNumber(): Promise<number> {
+    const res = await this.rpcCall('eth_blockNumber', []);
+    return parseInt(res.result, 16);
+  }
+
+  /**
+   * Get the current chainId
+   * @returns The current chainId
+   */
+  public async chainId(): Promise<number> {
+    const res = await this.rpcCall('eth_chainId', []);
+    return parseInt(res.result, 16);
+  }
+
+  /**
+   * Get the current timestamp
+   * @returns The current timestamp
+   */
+  public async timestamp(): Promise<number> {
+    const res = await this.rpcCall('eth_getBlockByNumber', ['latest', true]);
+    return parseInt(res.result.timestamp, 16);
+  }
+
+  /**
+   * Get the current chainId
+   * @param numberOfBlocks - The number of blocks to mine
+   * @returns The current chainId
+   */
+  public async mine(numberOfBlocks = 1): Promise<void> {
+    const res = await this.rpcCall('anvil_mine', [numberOfBlocks]);
+    if (res.error) throw new Error(`Error mining: ${res.error.message}`);
+    this.logger(`Mined ${numberOfBlocks} blocks`);
+  }
+
+  /**
+   * Set the next block timestamp
+   * @param timestamp - The timestamp to set the next block to
+   */
+  public async setNextBlockTimestamp(timestamp: number): Promise<void> {
+    const res = await this.rpcCall('anvil_setNextBlockTimestamp', [timestamp]);
+    if (res.error) throw new Error(`Error setting next block timestamp: ${res.error.message}`);
+    this.logger(`Set next block timestamp to ${timestamp}`);
+  }
+
+  // Good basis for the remaining functions:
+  // https://github.com/foundry-rs/foundry/blob/master/anvil/core/src/eth/mod.rs
+}
+
+/**
+ * A class that provides utility functions for interacting with the L2 chain.
+ */
+class L2CheatCodes {
+  constructor(
+    /**
+     * The RPC client to use for interacting with the chain
+     */
+    public aztecRpc: AztecRPC,
+    /**
+     * The circuits wasm module used for pedersen hashing
+     */
+    public wasm: CircuitsWasm,
+    /**
+     * The L1 cheat codes.
+     */
+    public l1: L1CheatCodes,
+    /**
+     * The logger to use for the l2 cheatcodes
+     */
+    public logger = createDebugLogger('aztec:cheat_codes:l2'),
+  ) {}
+
+  /**
+   * Computes the slot value for a given map and key.
+   * @param baseSlot - The base slot of the map (specified in noir contract)
+   * @param key - The key to lookup in the map
+   * @returns The storage slot of the value in the map
+   */
+  public computeSlotInMap(baseSlot: Fr | bigint, key: Fr | bigint): Fr {
+    // Based on `at` function in
+    // aztec3-packages/yarn-project/noir-contracts/src/contracts/noir-aztec/src/state_vars/map.nr
+    return Fr.fromBuffer(
+      pedersenPlookupCommitInputs(
+        this.wasm,
+        [toFr(baseSlot), toFr(key)].map(f => f.toBuffer()),
+      ),
+    );
+  }
+
+  /**
+   * Get the current blocknumber
+   * @returns The current block number
+   */
+  public async blockNumber(): Promise<number> {
+    return await this.aztecRpc.getBlockNum();
+  }
+
+  /**
+   * Loads the value stored at the given slot in the public storage of the given contract.
+   * @param who - The address of the contract
+   * @param slot - The storage slot to lookup
+   * @returns The value stored at the given slot
+   */
+  public async loadPublic(who: AztecAddress, slot: Fr | bigint): Promise<Fr> {
+    const storageValue = await this.aztecRpc.getPublicStorageAt(who, toFr(slot));
+    if (storageValue === undefined) {
+      throw new Error(`Storage slot ${slot} not found`);
+    }
+    return Fr.fromBuffer(storageValue);
+  }
+}

yarn-project/end-to-end/src/cross_chain/test_harness.ts

@@ -12,7 +12,8 @@ import { AztecRPC, TxStatus } from '@aztec/types';
 
 import { Chain, HttpTransport, PublicClient, getContract } from 'viem';
 
-import { deployAndInitializeNonNativeL2TokenContracts, expectAztecStorageSlot } from '../utils.js';
+import { CheatCodes } from '../cheat_codes.js';
+import { deployAndInitializeNonNativeL2TokenContracts } from '../utils.js';
 
 /**
  * A Class for testing cross chain interactions, contains common interactions
@@ -27,6 +28,7 @@ export class CrossChainTestHarness {
     accounts: AztecAddress[],
     wallet: Wallet,
     logger: DebugLogger,
+    cheatCodes: CheatCodes,
   ): Promise<CrossChainTestHarness> {
     const walletClient = deployL1ContractsValues.walletClient;
     const publicClient = deployL1ContractsValues.publicClient;
@@ -60,6 +62,7 @@ export class CrossChainTestHarness {
     return new CrossChainTestHarness(
       aztecNode,
       aztecRpcServer,
+      cheatCodes,
       accounts,
       logger,
       l2Contract,
@@ -80,6 +83,8 @@ export class CrossChainTestHarness {
     public aztecNode: AztecNodeService | undefined,
     /** AztecRpcServer. */
     public aztecRpcServer: AztecRPC,
+    /** CheatCodes. */
+    public cc: CheatCodes,
     /** Accounts. */
     public accounts: AztecAddress[],
     /** Logger. */
@@ -199,14 +204,11 @@ export class CrossChainTestHarness {
   }
 
   async expectPublicBalanceOnL2(owner: AztecAddress, expectedBalance: bigint, publicBalanceSlot: bigint) {
-    await expectAztecStorageSlot(
-      this.logger,
-      this.aztecRpcServer,
-      this.l2Contract,
-      publicBalanceSlot,
-      owner.toField(),
-      expectedBalance,
+    const balance = await this.cc.l2.loadPublic(
+      this.l2Contract.address,
+      this.cc.l2.computeSlotInMap(publicBalanceSlot, owner.toField()),
     );
+    expect(balance.value).toBe(expectedBalance);
   }
 
   async checkEntryIsNotInOutbox(withdrawAmount: bigint, callerOnL1: EthAddress = EthAddress.ZERO): Promise<Fr> {

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

@@ -0,0 +1,53 @@
+import { AztecNodeService } from '@aztec/aztec-node';
+import { AztecRPCServer } from '@aztec/aztec-rpc';
+import { AztecRPC } from '@aztec/types';
+
+import { CheatCodes } from './cheat_codes.js';
+import { setup } from './utils.js';
+
+describe('e2e_cheat_codes', () => {
+  let aztecNode: AztecNodeService | undefined;
+  let aztecRpcServer: AztecRPC;
+
+  let cc: CheatCodes;
+
+  beforeAll(async () => {
+    ({ aztecNode, aztecRpcServer, cheatCodes: cc } = await setup());
+  }, 100_000);
+
+  afterAll(async () => {
+    await aztecNode?.stop();
+    if (aztecRpcServer instanceof AztecRPCServer) {
+      await aztecRpcServer?.stop();
+    }
+  });
+
+  describe('L1 only', () => {
+    describe('mine', () => {
+      it(`mine block`, async () => {
+        const blockNumber = await cc.l1.blockNumber();
+        await cc.l1.mine();
+        expect(await cc.l1.blockNumber()).toBe(blockNumber + 1);
+      });
+
+      it.each([10, 42, 99])(`mine blocks`, async increment => {
+        const blockNumber = await cc.l1.blockNumber();
+        await cc.l1.mine(increment);
+        expect(await cc.l1.blockNumber()).toBe(blockNumber + increment);
+      });
+    });
+
+    it.each([100, 42, 99])('setNextBlockTimestamp', async increment => {
+      const blockNumber = await cc.l1.blockNumber();
+      const timestamp = await cc.l1.timestamp();
+      await cc.l1.setNextBlockTimestamp(timestamp + increment);
+
+      expect(await cc.l1.timestamp()).toBe(timestamp);
+
+      await cc.l1.mine();
+
+      expect(await cc.l1.blockNumber()).toBe(blockNumber + 1);
+      expect(await cc.l1.timestamp()).toBe(timestamp + increment);
+    });
+  });
+});

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

@@ -33,6 +33,7 @@ describe('e2e_cross_chain_messaging', () => {
       accounts,
       wallet,
       logger: logger_,
+      cheatCodes,
     } = await setup(2);
     crossChainTestHarness = await CrossChainTestHarness.new(
       initialBalance,
@@ -42,6 +43,7 @@ describe('e2e_cross_chain_messaging', () => {
       accounts,
       wallet,
       logger_,
+      cheatCodes,
     );
 
     l2Contract = crossChainTestHarness.l2Contract;

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

@@ -7,7 +7,8 @@ import { DebugLogger } from '@aztec/foundation/log';
 import { LendingContract } from '@aztec/noir-contracts/types';
 import { AztecRPC, TxStatus } from '@aztec/types';
 
-import { calculateAztecStorageSlot, setup } from './utils.js';
+import { CheatCodes } from './cheat_codes.js';
+import { setup } from './utils.js';
 
 describe('e2e_lending_contract', () => {
   let aztecNode: AztecNodeService | undefined;
@@ -18,6 +19,8 @@ describe('e2e_lending_contract', () => {
 
   let contract: Contract;
 
+  let cc: CheatCodes;
+
   const deployContract = async () => {
     logger(`Deploying L2 public contract...`);
     const tx = LendingContract.deploy(aztecRpcServer).send();
@@ -32,7 +35,7 @@ describe('e2e_lending_contract', () => {
   };
 
   beforeEach(async () => {
-    ({ aztecNode, aztecRpcServer, wallet, accounts, logger } = await setup());
+    ({ aztecNode, aztecRpcServer, wallet, accounts, logger, cheatCodes: cc } = await setup());
   }, 100_000);
 
   afterEach(async () => {
@@ -44,25 +47,23 @@ describe('e2e_lending_contract', () => {
 
   // Fetch a storage snapshot from the contract that we can use to compare between transitions.
   const getStorageSnapshot = async (contract: Contract, aztecNode: AztecRPC, account: Account) => {
-    const storageValues: { [key: string]: any } = {};
-
-    const readValue = async (slot: Fr) =>
-      Fr.fromBuffer((await aztecNode.getPublicStorageAt(contract.address, slot)) ?? Buffer.alloc(0));
+    const loadPublicStorageInMap = async (slot: Fr | bigint, key: Fr | bigint) => {
+      return await cc.l2.loadPublic(contract.address, cc.l2.computeSlotInMap(slot, key));
+    };
 
+    const storageValues: { [key: string]: any } = {};
     {
-      const baseSlot = await calculateAztecStorageSlot(1n, Fr.ZERO);
-      storageValues['interestAccumulator'] = await readValue(baseSlot);
-      storageValues['last_updated_ts'] = await readValue(new Fr(baseSlot.value + 1n));
+      const baseSlot = cc.l2.computeSlotInMap(1n, 0n);
+      storageValues['interestAccumulator'] = await cc.l2.loadPublic(contract.address, baseSlot);
+      storageValues['last_updated_ts'] = await cc.l2.loadPublic(contract.address, baseSlot.value + 1n);
     }
 
     const accountKey = await account.key();
 
-    storageValues['private_collateral'] = await readValue(await calculateAztecStorageSlot(2n, accountKey));
-    storageValues['public_collateral'] = await readValue(
-      await calculateAztecStorageSlot(2n, account.address.toField()),
-    );
-    storageValues['private_debt'] = await readValue(await calculateAztecStorageSlot(3n, accountKey));
-    storageValues['public_debt'] = await readValue(await calculateAztecStorageSlot(3n, account.address.toField()));
+    storageValues['private_collateral'] = await loadPublicStorageInMap(2n, accountKey);
+    storageValues['public_collateral'] = await loadPublicStorageInMap(2n, account.address.toField());
+    storageValues['private_debt'] = await loadPublicStorageInMap(3n, accountKey);
+    storageValues['public_debt'] = await loadPublicStorageInMap(3n, account.address.toField());
 
     return storageValues;
   };

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

@@ -33,6 +33,7 @@ describe('e2e_public_cross_chain_messaging', () => {
       accounts,
       wallet,
       logger: logger_,
+      cheatCodes,
     } = await setup(2);
     crossChainTestHarness = await CrossChainTestHarness.new(
       initialBalance,
@@ -42,6 +43,7 @@ describe('e2e_public_cross_chain_messaging', () => {
       accounts,
       wallet,
       logger_,
+      cheatCodes,
     );
 
     l2Contract = crossChainTestHarness.l2Contract;

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

@@ -29,6 +29,7 @@ describe('e2e_public_to_private_messaging', () => {
       accounts,
       wallet,
       logger: logger_,
+      cheatCodes,
     } = await setup(2);
     crossChainTestHarness = await CrossChainTestHarness.new(
       initialBalance,
@@ -38,6 +39,7 @@ describe('e2e_public_to_private_messaging', () => {
       accounts,
       wallet,
       logger_,
+      cheatCodes,
     );
 
     ethAccount = crossChainTestHarness.ethAccount;