feat: add aztec flamegraph command

aztec-nargo/Dockerfile

@@ -11,7 +11,7 @@ FROM --platform=linux/amd64 aztecprotocol/barretenberg-x86_64-linux-clang as bar
 FROM ubuntu:noble
 # Install Tini as nargo doesn't handle signals properly.
 # Install git as nargo needs it to clone.
-RUN apt-get update && apt-get install -y git tini jq curl && rm -rf /var/lib/apt/lists/* && apt-get clean
+RUN apt-get update && apt-get install -y git tini jq curl nodejs npm && rm -rf /var/lib/apt/lists/* && apt-get clean
 
 # Copy binaries to /usr/bin
 COPY --from=built-noir /usr/src/noir/noir-repo/target/release/nargo /usr/bin/nargo
@@ -21,4 +21,9 @@ COPY --from=barretenberg /usr/src/barretenberg/cpp/build/bin/bb /usr/bin/bb
 # Copy in script that calls both binaries
 COPY ./aztec-nargo/compile_then_postprocess.sh /usr/src/aztec-nargo/compile_then_postprocess.sh
 
+# Profiler
+COPY --from=built-noir /usr/src/noir/noir-repo/target/release/noir-profiler /usr/bin/noir-profiler
+COPY ./noir-projects/noir-contracts/scripts/flamegraph.sh /usr/bin/flamegraph.sh
+COPY ./noir-projects/noir-contracts/scripts/extractFunctionAsNoirArtifact.js /usr/bin/extractFunctionAsNoirArtifact.js
+
 ENTRYPOINT ["/usr/bin/tini", "--", "/usr/src/aztec-nargo/compile_then_postprocess.sh"]

aztec-nargo/Earthfile

@@ -4,7 +4,7 @@ run:
   FROM ubuntu:noble
   # Install Tini as nargo doesn't handle signals properly.
   # Install git as nargo needs it to clone.
-  RUN apt-get update && apt-get install -y git tini jq curl && rm -rf /var/lib/apt/lists/* && apt-get clean
+  RUN apt-get update && apt-get install -y git tini jq curl nodejs npm && rm -rf /var/lib/apt/lists/* && apt-get clean
 
   # Copy binaries to /usr/bin
   COPY ../+bootstrap/usr/src/noir/noir-repo/target/release/nargo /usr/bin/nargo
@@ -14,6 +14,11 @@ run:
   # Copy in script that calls both binaries
   COPY ./compile_then_postprocess.sh /usr/bin/compile_then_postprocess.sh
 
+  # Profiler
+  COPY ../+bootstrap/usr/src/noir/noir-repo/target/release/noir-profiler /usr/bin/noir-profiler
+  COPY ../+bootstrap/usr/src/noir-projects/noir-contracts/scripts/flamegraph.sh /usr/bin/flamegraph.sh
+  COPY ../+bootstrap/usr/src/noir-projects/noir-contracts/scripts/extractFunctionAsNoirArtifact.js /usr/bin/extractFunctionAsNoirArtifact.js
+
   ENV PATH "/usr/bin:${PATH}"
   ENTRYPOINT ["/usr/bin/tini", "--", "/usr/bin/compile_then_postprocess.sh"]
   SAVE IMAGE aztecprotocol/aztec-nargo

aztec-up/bin/aztec

@@ -72,6 +72,15 @@ elif [ "${1:-}" == "start" ]; then
   readarray -t ENV_VARS_TO_INJECT <"$LOCAL_ENV_VAR_FILE"
   export ENV_VARS_TO_INJECT="${ENV_VARS_TO_INJECT[*]}"
   ENV_VARS_TO_INJECT="${ENV_VARS_TO_INJECT[*]}" INHERIT_USER=0 $(dirname $0)/.aztec-run aztecprotocol/aztec "$@"
+elif [ "${1:-}" == "flamegraph" ]; then
+  docker run -it \
+    --entrypoint /usr/bin/flamegraph.sh \
+    --env PROFILER_PATH=/usr/bin/noir-profiler \
+    --env BACKEND_PATH=/usr/bin/bb \
+    --env SERVE=${SERVE:-0} \
+    $([ "${SERVE:-0}" == "1" ] && echo "-p 8000:8000" || echo "") \
+    -v $(realpath $(dirname $2))/:/tmp \
+    aztecprotocol/aztec-nargo:$VERSION /tmp/$(basename $2) $3
 else
   ENV_VARS_TO_INJECT="SECRET_KEY" SKIP_PORT_ASSIGNMENT=1 $(dirname $0)/.aztec-run aztecprotocol/aztec "$@"
 fi

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

@@ -3,35 +3,38 @@ set -eu
 
 # Function to clean up and exit
 cleanup_and_exit() {
-    echo "Cleaning up..."
-    rm -f "$SCRIPT_DIR/../target/$FUNCTION_ARTIFACT"
-    exit 0
+  echo "Cleaning up..."
+  rm -f "$FUNCTION_ARTIFACT"
+  exit 0
 }
 
 # Trap SIGINT (Ctrl+C) and call cleanup_and_exit
 trap cleanup_and_exit SIGINT
 
 # If first  arg is -h or --help, print usage
 if [ "$1" == "-h" ] || [ "$1" == "--help" ]; then
-    echo "Usage: $0 <contract> <function>"
-    echo "e.g.: $0 Token transfer"
-    echo "Generates a flamegraph for the given contract and function"
-    exit 0
+  echo "Generates a flamegraph for the given contract and function"
+  echo "Usage: $0 <contract_artifact | contract_name> <function>"
+  echo "e.g.: $0 ./target/voting_contract_Voting.json vote"
+  echo "e.g.: $0 Token transfer"
+  exit 0
 fi
 
 # Get the directory of the script
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 
-PROFILER="$SCRIPT_DIR/../../../noir/noir-repo/target/release/noir-profiler"
+PROFILER=${PROFILER_PATH:-"$SCRIPT_DIR/../../../noir/noir-repo/target/release/noir-profiler"}
+BACKEND_PATH=${BACKEND_PATH:-"$SCRIPT_DIR/../../../barretenberg/cpp/build/bin/bb"}
+SERVE=${SERVE:-"1"}
 
 if [ ! -f $PROFILER ]; then
-    echo "Profiler not found, building profiler"
-    cd "$SCRIPT_DIR/../../../noir/noir-repo/tooling/profiler"
-    cargo build --release
-    cd "$SCRIPT_DIR"
+  echo "Profiler not found, building profiler"
+  cd "$SCRIPT_DIR/../../../noir/noir-repo/tooling/profiler"
+  cargo build --release
+  cd "$SCRIPT_DIR"
 fi
 
-# first console arg is contract name in camel case (e.g. TokenBridge)
+# first console arg is contract name in camel case or path to contract artifact
 CONTRACT=$1
 
 # second console arg is the contract function
@@ -49,25 +52,48 @@ function sed_wrapper() {
   fi
 }
 
-# convert contract name to following format: token_bridge_contract-TokenBridge.json
-ARTIFACT=$(echo "$CONTRACT" | sed_wrapper -r 's/^([A-Z])/\L\1/; s/([a-z0-9])([A-Z])/\1_\L\2/g')
-ARTIFACT=$(echo "$ARTIFACT" | tr '[:upper:]' '[:lower:]')
-ARTIFACT_NAME="${ARTIFACT}_contract-${CONTRACT}"
-
-# Extract artifact for the specific function
-node "$SCRIPT_DIR/extractFunctionAsNoirArtifact.js" "$SCRIPT_DIR/../target/${ARTIFACT_NAME}.json" $FUNCTION
+if [[ "$CONTRACT" == *.json ]]; then
+  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"
+else
+  # convert contract name to following format: token_bridge_contract-TokenBridge.json
+  ARTIFACT=$(echo "$CONTRACT" | sed_wrapper -r 's/^([A-Z])/\L\1/; s/([a-z0-9])([A-Z])/\1_\L\2/g')
+  ARTIFACT=$(echo "$ARTIFACT" | tr '[:upper:]' '[:lower:]')
+  ARTIFACT_NAME="${ARTIFACT}_contract-${CONTRACT}"
+  ARTIFACT_PATH="$SCRIPT_DIR/../target/${ARTIFACT_NAME}.json"
+  FUNCTION_ARTIFACT="$SCRIPT_DIR/../target/${ARTIFACT_NAME}-${FUNCTION}.json"
+fi
 
-FUNCTION_ARTIFACT="${ARTIFACT_NAME}-${FUNCTION}.json"
+# Extract artifact for the specific function (will save to $FUNCTION_ARTIFACT)
+node "$SCRIPT_DIR/extractFunctionAsNoirArtifact.js" "$ARTIFACT_PATH" $FUNCTION
 
-# We create dest directory and use it as an output for the generated main.svg file
-mkdir -p "$SCRIPT_DIR/../dest"
+if [ "$SERVE" == "true" ]; then
+  # We create dest directory and use it as an output for the generated main.svg file
+  OUTPUT_DIR="$SCRIPT_DIR/../dest"
+  mkdir -p "$OUTPUT_DIR"
+else
+  # Save the flamegraph to the same directory as the artifact
+  OUTPUT_DIR=$(dirname "$ARTIFACT_PATH")
+fi
 
 # At last, generate the flamegraph
-$PROFILER gates --artifact-path "$SCRIPT_DIR/../target/$FUNCTION_ARTIFACT" --backend-path "$SCRIPT_DIR/../../../barretenberg/cpp/build/bin/bb"  --backend-gates-command "gates_for_ivc" --output "$SCRIPT_DIR/../dest"
-
-# serve the file over http
-echo "Serving flamegraph at http://0.0.0.0:8000/main::gates.svg"
-python3 -m http.server --directory "$SCRIPT_DIR/../dest" 8000
+$PROFILER gates --artifact-path "$FUNCTION_ARTIFACT" --backend-path "$BACKEND_PATH" --backend-gates-command "gates_for_ivc" --output "$OUTPUT_DIR"
+
+# save as $ARTIFACT_NAME-${FUNCTION}-flamegraph.svg
+OUTPUT_FILE="${OUTPUT_DIR}/$(basename ${ARTIFACT_PATH%%.json})-${FUNCTION}-flamegraph.svg"
+mv "$OUTPUT_DIR/main::gates.svg" "$OUTPUT_FILE"
+
+if [ "$SERVE" == "1" ]; then
+  # serve the file over http
+  echo -e "\nServing flamegraph at http://0.0.0.0:8000/$(basename $OUTPUT_FILE)\n"
+  npx -y http-server --silent -p 8000 "$OUTPUT_DIR"
+else
+  echo -e "\nFlamegraph $(basename $OUTPUT_FILE) saved to artifacts directory. You can open it in your browser.\n"
+fi
 
 # Clean up before exiting
-cleanup_and_exit
+cleanup_and_exit

noir/bootstrap.sh

@@ -12,7 +12,7 @@ function build {
     export COMMIT_HASH="$(echo "$hash" | sed 's/-.*//g')"
     denoise ./scripts/bootstrap_native.sh
     denoise ./scripts/bootstrap_packages.sh
-    cache_upload noir-$hash.tar.gz noir-repo/target/release/nargo noir-repo/target/release/acvm packages
+    cache_upload noir-$hash.tar.gz noir-repo/target/release/nargo noir-repo/target/release/acvm noir-repo/target/release/noir-profiler packages
   fi
   github_endgroup
 }