feat: initial cheatcodes

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

@@ -0,0 +1,55 @@
+import { AztecAddress, CircuitsWasm, Fr } from '@aztec/circuits.js';
+import { pedersenPlookupCommitInputs } from '@aztec/circuits.js/barretenberg';
+import { AztecRPC } from '@aztec/types';
+
+/**
+ * A class that provides utility functions for interacting with the chain.
+ */
+export class CheatCodes {
+  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,
+  ) {}
+
+  static async create(aztecRpc: AztecRPC): Promise<CheatCodes> {
+    return new CheatCodes(aztecRpc, await CircuitsWasm.get());
+  }
+
+  /**
+   * 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 {
+    const mappingStorageSlot = typeof baseSlot === 'bigint' ? new Fr(baseSlot) : baseSlot;
+    const frKey = typeof key === 'bigint' ? new Fr(key) : key;
+
+    // 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,
+        [mappingStorageSlot, frKey].map(f => f.toBuffer()),
+      ),
+    );
+  }
+
+  /**
+   * 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 frSlot = typeof slot === 'bigint' ? new Fr(slot) : slot;
+    const publicStorage = await this.aztecRpc.getPublicStorageAt(who, frSlot);
+    return publicStorage ? Fr.fromBuffer(publicStorage) : Fr.ZERO;
+  }
+}

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 './CheatCodes.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();
@@ -33,6 +36,7 @@ describe('e2e_lending_contract', () => {
 
   beforeEach(async () => {
     ({ aztecNode, aztecRpcServer, wallet, accounts, logger } = await setup());
+    cc = await CheatCodes.create(aztecRpcServer);
   }, 100_000);
 
   afterEach(async () => {
@@ -45,24 +49,24 @@ 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 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.computeSlotInMap(1n, 0n);
+      storageValues['interestAccumulator'] = await cc.loadPublic(contract.address, baseSlot);
+      storageValues['last_updated_ts'] = await cc.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_collateral'] = await cc.loadPublic(contract.address, cc.computeSlotInMap(2n, accountKey));
+    storageValues['public_collateral'] = await cc.loadPublic(
+      contract.address,
+      cc.computeSlotInMap(2n, account.address.toField()),
+    );
+    storageValues['private_debt'] = await cc.loadPublic(contract.address, cc.computeSlotInMap(3n, accountKey));
+    storageValues['public_debt'] = await cc.loadPublic(
+      contract.address,
+      cc.computeSlotInMap(3n, account.address.toField()),
     );
-    storageValues['private_debt'] = await readValue(await calculateAztecStorageSlot(3n, accountKey));
-    storageValues['public_debt'] = await readValue(await calculateAztecStorageSlot(3n, account.address.toField()));
 
     return storageValues;
   };