bring back flamegraph

Author: charlielye

noir-projects/noir-contracts/scripts/flamegraph.sh

@@ -45,7 +45,7 @@ case $cmd in
 
     aztec start "$@"
     ;;
-  compile|new|init)
+  compile|new|init|flamegraph)
     $script_dir/${cmd}.sh "$@"
     ;;
   fmt|check|lsp)
@@ -55,12 +55,4 @@ case $cmd in
   *)
     aztec $cmd "$@"
     ;;
-  # flamegraph)
-  #   docker run -it \
-  #     --entrypoint /usr/src/noir-projects/noir-contracts/scripts/flamegraph.sh \
-  #     --env SERVE=${SERVE:-0} \
-  #     $([ "${SERVE:-0}" == "1" ] && echo "-p 8000:8000" || echo "") \
-  #     -v $(realpath $(dirname $2))/:/tmp \
-  #     $DOCKER_REPO:$VERSION /tmp/$(basename $2) $3
-  #   ;;
 esac

yarn-project/aztec/scripts/aztec.sh

@@ -1,26 +1,21 @@
-const path = require("path");
-const fs = require("fs").promises;
+#!/usr/bin/env node
+import fs from 'fs/promises';
+import path from 'path';
 
 // Simple script to extract a contract function as a separate Noir artifact.
 // We need to use this since the transpiling that we do on public functions make the contract artifacts
 // unreadable by noir tooling, since they are no longer following the noir artifact format.
 async function main() {
   let [contractArtifactPath, functionName] = process.argv.slice(2);
   if (!contractArtifactPath || !functionName) {
-    console.log(
-      "Usage: node extractFunctionAsNoirArtifact.js <contractArtifactPath> <functionName>"
-    );
+    console.log('Usage: node extractFunctionAsNoirArtifact.js <contractArtifactPath> <functionName>');
     return;
   }
 
-  const contractArtifact = JSON.parse(
-    await fs.readFile(contractArtifactPath, "utf8")
-  );
-  const func = contractArtifact.functions.find((f) => f.name === functionName);
+  const contractArtifact = JSON.parse(await fs.readFile(contractArtifactPath, 'utf8'));
+  const func = contractArtifact.functions.find(f => f.name === functionName);
   if (!func) {
-    console.error(
-      `Function ${functionName} not found in ${contractArtifactPath}`
-    );
+    console.error(`Function ${functionName} not found in ${contractArtifactPath}`);
     return;
   }
 
@@ -39,17 +34,14 @@ async function main() {
   };
 
   const outputDir = path.dirname(contractArtifactPath);
-  const outputName =
-    path.basename(contractArtifactPath, ".json") + `-${functionName}.json`;
+  const outputName = path.basename(contractArtifactPath, '.json') + `-${functionName}.json`;
 
   const outPath = path.join(outputDir, outputName);
 
-  console.log(`Writing to ${outPath}`);
-
   await fs.writeFile(outPath, JSON.stringify(artifact, null, 2));
 }
 
-main().catch((err) => {
+main().catch(err => {
   console.error(err);
   process.exit(1);
 });

…scripts/extractFunctionAsNoirArtifact.js …roject/aztec/scripts/extract_function.jsnoir-projects/noir-contracts/scripts/extractFunctionAsNoirArtifact.js renamed to yarn-project/aztec/scripts/extract_function.js noir-projects/noir-contracts/scripts/extractFunctionAsNoirArtifact.js renamed to yarn-project/aztec/scripts/extract_function.js

@@ -0,0 +1,59 @@
+#!/usr/bin/env bash
+set -eu
+
+# If first arg is -h or --help, print usage.
+if [ $# -lt 2 ] || [ "$1" == "-h" ] || [ "$1" == "--help" ]; then
+    cat << 'EOF'
+Aztec Flamegraph - Generate a gate count flamegraph for an aztec contract function.
+
+Usage: aztec flamegraph <contract_artifact> <function>
+
+Options:
+  -h, --help     Print help
+
+Will output an svg at <artifact_path>/<contract>-<function>-flamegraph.svg.
+You can open it in your browser to view it.
+
+EOF
+    exit 0
+fi
+
+cleanup() {
+  set +e
+  if [ -f "$function_artifact" ]; then
+    rm -f "$function_artifact"
+  fi
+}
+
+trap cleanup EXIT
+
+# Get the directory of the script
+script_dir=$(realpath $(dirname $0))
+
+PROFILER=${PROFILER_PATH:-noir-profiler}
+BB=${BB:-bb}
+
+# first console arg is contract name in camel case or path to contract artifact
+contract=$1
+
+# second console arg is the contract function
+function=$2
+
+if [ ! -f "$contract" ]; then
+  echo "Error: Contract artifact not found at: $contract"
+  exit 1
+fi
+artifact_path=$contract
+function_artifact="${artifact_path%%.json}-${function}.json"
+output_dir=$(dirname "$artifact_path")
+
+# Extract artifact for the specific function.
+node $script_dir/extract_function.js "$artifact_path" $function
+
+# Generate the flamegraph
+$PROFILER gates --artifact-path "$function_artifact" --backend-path "$BB" --backend-gates-command "gates" --output "$output_dir" --scheme chonk --include_gates_per_opcode
+
+# Save as $artifact_name-$function-flamegraph.svg
+output_file="${function_artifact%%.json}-flamegraph.svg"
+mv "$output_dir/__aztec_nr_internals__${function}_gates.svg" "$output_file"
+echo "Flamegraph generated at: $output_file"