utils.ts

import { type AztecNodeConfig, type AztecNodeService } from '@aztec/aztec-node';
import { type AztecNode, BatchCall, INITIAL_L2_BLOCK_NUM, type SentTx, sleep } from '@aztec/aztec.js';
import { times } from '@aztec/foundation/collection';
import { randomInt } from '@aztec/foundation/crypto';
import { BenchmarkingContract } from '@aztec/noir-contracts.js/Benchmarking';
import { type PXEService, type PXEServiceConfig, createPXEService } from '@aztec/pxe';

import { mkdirpSync } from 'fs-extra';
import { globSync } from 'glob';
import { join } from 'path';

import { type EndToEndContext, setup } from '../fixtures/utils.js';

/**
 * Setup for benchmarks. Initializes a remote node with a single account and deploys a benchmark contract.
 */
export async function benchmarkSetup(opts: Partial<AztecNodeConfig>) {
  const context = await setup(1, { ...opts });
  const contract = await BenchmarkingContract.deploy(context.wallet).send().deployed();
  context.logger.info(`Deployed benchmarking contract at ${contract.address}`);
  const sequencer = (context.aztecNode as AztecNodeService).getSequencer()!;
  return { context, contract, sequencer };
}

/**
 * Creates and returns a directory with the current job name and a random number.
 * @param index - Index to merge into the dir path.
 * @returns A path to a created dir.
 */
export function makeDataDirectory(index: number) {
  const testName = expect.getState().currentTestName!.split(' ')[0].replaceAll('/', '_');
  const db = join('data', testName, index.toString(), `${randomInt(99)}`);
  mkdirpSync(db);
  return db;
}

/**
 * Returns the size in disk of a folder.
 * @param path - Path to the folder.
 * @returns Size in bytes.
 */
export function getFolderSize(path: string): number {
  return globSync('**', { stat: true, cwd: path, nodir: true, withFileTypes: true }).reduce(
    (accum, file) => accum + (file as any as { /** Size */ size: number }).size,
    0,
  );
}

/**
 * Returns a call to the benchmark contract. Each call has a private execution (account entrypoint),
 * a nested private call (create_note), a public call (increment_balance), and a nested public
 * call (broadcast). These include emitting one private note and one unencrypted log, two storage
 * reads and one write.
 * @param index - Index of the call within a block.
 * @param context - End to end context.
 * @param contract - Benchmarking contract.
 * @returns A BatchCall instance.
 */
export function makeCall(index: number, context: EndToEndContext, contract: BenchmarkingContract) {
  const owner = context.wallet.getAddress();
  const sender = owner;
  return new BatchCall(context.wallet, [
    contract.methods.create_note(owner, sender, index + 1).request(),
    contract.methods.increment_balance(owner, index + 1).request(),
  ]);
}

/**
 * Assembles and sends multiple transactions simultaneously to the node in context.
 * Each tx is the result of calling makeCall.
 * @param txCount - How many txs to send
 * @param context - End to end context.
 * @param contract - Target contract.
 * @returns Array of sent txs.
 */
export async function sendTxs(
  txCount: number,
  context: EndToEndContext,
  contract: BenchmarkingContract,
): Promise<SentTx[]> {
  const calls = times(txCount, index => makeCall(index, context, contract));
  const provenTxs = await Promise.all(calls.map(call => call.prove({ skipPublicSimulation: true })));
  const sentTxs = provenTxs.map(tx => tx.send());

  // Awaiting txHash waits until the aztec node has received the tx into its p2p pool
  await Promise.all(sentTxs.map(tx => tx.getTxHash()));
  await sleep(100);

  return sentTxs;
}

/**
 * Creates a new PXE
 * @param node - Node to connect the pxe to.
 * @param contract - Benchmark contract to add to the pxe.
 * @param startingBlock - First l2 block to process.
 * @returns The new PXE.
 */
export async function createNewPXE(
  node: AztecNode,
  contract: BenchmarkingContract,
  startingBlock: number = INITIAL_L2_BLOCK_NUM,
): Promise<PXEService> {
  const l1Contracts = await node.getL1ContractAddresses();
  const pxeConfig = {
    l2StartingBlock: startingBlock,
    l2BlockPollingIntervalMS: 100,
    dataDirectory: undefined,
    dataStoreMapSizeKB: 1024 * 1024,
    l1Contracts,
  } as PXEServiceConfig;
  const pxe = await createPXEService(node, pxeConfig);
  await pxe.registerContract(contract);
  return pxe;
}