feat: add aztec profile flamegraph command

yarn-project/aztec/scripts/aztec.sh

@@ -53,9 +53,13 @@ case $cmd in
 
     aztec start "$@"
     ;;
-  new|init|flamegraph)
+  new|init)
     $script_dir/${cmd}.sh "$@"
     ;;
+  flamegraph)
+    echo "Warning: 'aztec flamegraph' is deprecated. Use 'aztec profile flamegraph' instead." >&2
+    aztec profile flamegraph "$@"
+    ;;
   *)
     aztec $cmd "$@"
     ;;

yarn-project/aztec/src/cli/cmds/compile.ts

@@ -1,25 +1,11 @@
 import type { LogFn } from '@aztec/foundation/log';
 
-import { execFileSync, spawn } from 'child_process';
+import { execFileSync } from 'child_process';
 import type { Command } from 'commander';
 import { readFile, writeFile } from 'fs/promises';
 
 import { readArtifactFiles } from './utils/artifacts.js';
-
-/** Spawns a command with inherited stdio and rejects on non-zero exit. */
-function run(cmd: string, args: string[]): Promise<void> {
-  return new Promise((resolve, reject) => {
-    const child = spawn(cmd, args, { stdio: 'inherit' });
-    child.on('error', reject);
-    child.on('close', code => {
-      if (code !== 0) {
-        reject(new Error(`${cmd} exited with code ${code}`));
-      } else {
-        resolve();
-      }
-    });
-  });
-}
+import { run } from './utils/spawn.js';
 
 /** Returns paths to contract artifacts in the target directory. */
 async function collectContractArtifacts(): Promise<string[]> {

yarn-project/aztec/src/cli/cmds/profile.ts

@@ -2,6 +2,7 @@ import type { LogFn } from '@aztec/foundation/log';
 
 import type { Command } from 'commander';
 
+import { profileFlamegraph } from './profile_flamegraph.js';
 import { profileGates } from './profile_gates.js';
 
 export function injectProfileCommand(program: Command, log: LogFn): Command {
@@ -13,5 +14,12 @@ export function injectProfileCommand(program: Command, log: LogFn): Command {
     .description('Display gate counts for all compiled Aztec artifacts in a target directory.')
     .action((targetDir: string) => profileGates(targetDir, log));
 
+  profile
+    .command('flamegraph')
+    .argument('<contract-artifact>', 'Path to the compiled contract artifact JSON')
+    .argument('<function>', 'Name of the contract function to profile')
+    .description('Generate a gate count flamegraph SVG for a contract function.')
+    .action((artifactPath: string, functionName: string) => profileFlamegraph(artifactPath, functionName, log));
+
   return program;
 }

yarn-project/aztec/src/cli/cmds/profile_flamegraph.test.ts

@@ -0,0 +1,51 @@
+import { afterAll, beforeAll, describe, expect, it } from '@jest/globals';
+import { execFileSync } from 'child_process';
+import { existsSync, readFileSync, rmSync } from 'fs';
+import { dirname, join } from 'path';
+import { fileURLToPath } from 'url';
+
+const PACKAGE_ROOT = join(dirname(fileURLToPath(import.meta.url)), '../../..');
+const CLI = join(PACKAGE_ROOT, 'dest/bin/index.js');
+const WORKSPACE = join(PACKAGE_ROOT, 'test/mixed-workspace');
+const TARGET = join(WORKSPACE, 'target');
+const CONTRACT_ARTIFACT = join(TARGET, 'simple_contract-SimpleContract.json');
+
+describe('aztec profile flamegraph', () => {
+  const svgPath = join(TARGET, 'simple_contract-SimpleContract-private_function-flamegraph.svg');
+
+  beforeAll(() => {
+    rmSync(TARGET, { recursive: true, force: true });
+    runCompile();
+    runFlamegraph(CONTRACT_ARTIFACT, 'private_function');
+  }, 300_000);
+
+  afterAll(() => {
+    rmSync(TARGET, { recursive: true, force: true });
+  });
+
+  it('generates a valid flamegraph SVG', () => {
+    expect(existsSync(svgPath)).toBe(true);
+    const content = readFileSync(svgPath, 'utf-8');
+    expect(content).toContain('<svg');
+    expect(content).toContain('</svg>');
+  });
+});
+
+function runCompile() {
+  try {
+    execFileSync('node', [CLI, 'compile'], { cwd: WORKSPACE, stdio: 'pipe' });
+  } catch (e: any) {
+    throw new Error(`compile failed:\n${e.stderr?.toString() ?? e.message}`);
+  }
+}
+
+function runFlamegraph(artifactPath: string, functionName: string) {
+  try {
+    execFileSync('node', [CLI, 'profile', 'flamegraph', artifactPath, functionName], {
+      encoding: 'utf-8',
+      stdio: 'pipe',
+    });
+  } catch (e: any) {
+    throw new Error(`profile flamegraph failed:\n${e.stderr?.toString() ?? e.message}`);
+  }
+}

yarn-project/aztec/src/cli/cmds/profile_flamegraph.ts

@@ -0,0 +1,66 @@
+import type { LogFn } from '@aztec/foundation/log';
+
+import { mkdir, mkdtemp, readFile, rename, rm, writeFile } from 'fs/promises';
+import { tmpdir } from 'os';
+import { basename, dirname, join } from 'path';
+
+import { makeFunctionArtifact } from './profile_utils.js';
+import type { CompiledArtifact } from './utils/artifacts.js';
+import { run } from './utils/spawn.js';
+
+/** Generates a gate count flamegraph SVG for a single contract function. */
+export async function profileFlamegraph(artifactPath: string, functionName: string, log: LogFn): Promise<void> {
+  const raw = await readFile(artifactPath, 'utf-8');
+  const artifact: CompiledArtifact = JSON.parse(raw);
+
+  if (!Array.isArray(artifact.functions)) {
+    throw new Error(`${artifactPath} does not appear to be a contract artifact (no functions array)`);
+  }
+
+  const func = artifact.functions.find(f => f.name === functionName);
+  if (!func) {
+    const available = artifact.functions.map(f => f.name).join(', ');
+    throw new Error(`Function "${functionName}" not found in artifact. Available: ${available}`);
+  }
+  if (func.is_unconstrained) {
+    throw new Error(`Function "${functionName}" is unconstrained and cannot be profiled`);
+  }
+
+  const tmpDir = await mkdtemp(join(tmpdir(), 'aztec-flamegraph-'));
+
+  try {
+    const tmpArtifact = join(tmpDir, 'artifact.json');
+    await writeFile(tmpArtifact, makeFunctionArtifact(artifact, func));
+
+    const profiler = process.env.PROFILER_PATH ?? 'noir-profiler';
+    const bb = process.env.BB ?? 'bb';
+    const outputDir = join(tmpDir, 'output');
+    await mkdir(outputDir, { recursive: true });
+
+    await run(profiler, [
+      'gates',
+      '--artifact-path',
+      tmpArtifact,
+      '--backend-path',
+      bb,
+      '--backend-gates-command',
+      'gates',
+      '--output',
+      outputDir,
+      '--scheme',
+      'chonk',
+      '--include_gates_per_opcode',
+    ]);
+
+    // noir-profiler names the SVG using the internal function name which
+    // retains the __aztec_nr_internals__ prefix in the bytecode metadata.
+    const contractName = basename(artifactPath, '.json');
+    const srcSvg = join(outputDir, `__aztec_nr_internals__${functionName}_gates.svg`);
+    const destSvg = join(dirname(artifactPath), `${contractName}-${functionName}-flamegraph.svg`);
+    await rename(srcSvg, destSvg);
+
+    log(`Flamegraph written to ${destSvg}`);
+  } finally {
+    await rm(tmpDir, { recursive: true, force: true });
+  }
+}

yarn-project/aztec/src/cli/cmds/profile_utils.ts

@@ -44,7 +44,7 @@ export async function discoverArtifacts(
 }
 
 /** Extracts a contract function as a standalone program artifact JSON string. */
-function makeFunctionArtifact(artifact: CompiledArtifact, func: ContractFunction) {
+export function makeFunctionArtifact(artifact: CompiledArtifact, func: ContractFunction) {
   /* eslint-disable camelcase */
   return JSON.stringify({
     noir_version: artifact.noir_version,

yarn-project/aztec/src/cli/cmds/utils/spawn.ts

@@ -0,0 +1,16 @@
+import { spawn } from 'child_process';
+
+/** Spawns a command with inherited stdio and rejects on non-zero exit. */
+export function run(cmd: string, args: string[]): Promise<void> {
+  return new Promise((resolve, reject) => {
+    const child = spawn(cmd, args, { stdio: 'inherit' });
+    child.on('error', reject);
+    child.on('close', code => {
+      if (code !== 0) {
+        reject(new Error(`${cmd} exited with code ${code}`));
+      } else {
+        resolve();
+      }
+    });
+  });
+}