feat(profiler): Add support for brillig functions in opcodes-flamegra…

…ph (#7698)

The opcodes-flamegraph command of the profiler now also profiles brillig
opcode generation.

Sample output:

![0_brillig_opcodes](https://github.com/user-attachments/assets/52bc6413-f3e0-4972-b2fb-156f908cc8b1)

noir/noir-repo/tooling/profiler/src/cli/gates_flamegraph_cmd.rs

@@ -1,13 +1,15 @@
 use std::path::{Path, PathBuf};
 
+use acir::circuit::OpcodeLocation;
 use clap::Args;
 use color_eyre::eyre::{self, Context};
 
 use noirc_artifacts::debug::DebugArtifact;
 
-use crate::flamegraph::{FlamegraphGenerator, InfernoFlamegraphGenerator};
+use crate::flamegraph::{FlamegraphGenerator, InfernoFlamegraphGenerator, Sample};
 use crate::fs::read_program_from_file;
 use crate::gates_provider::{BackendGatesProvider, GatesProvider};
+use crate::opcode_formatter::AcirOrBrilligOpcode;
 
 #[derive(Debug, Clone, Args)]
 pub(crate) struct GatesFlamegraphCommand {
@@ -76,9 +78,20 @@ fn run_with_provider<Provider: GatesProvider, Generator: FlamegraphGenerator>(
             func_gates.circuit_size
         );
 
+        let samples = func_gates
+            .gates_per_opcode
+            .into_iter()
+            .zip(bytecode.opcodes)
+            .enumerate()
+            .map(|(index, (gates, opcode))| Sample {
+                opcode: AcirOrBrilligOpcode::Acir(opcode),
+                call_stack: vec![OpcodeLocation::Acir(index)],
+                count: gates,
+            })
+            .collect();
+
         flamegraph_generator.generate_flamegraph(
-            func_gates.gates_per_opcode,
-            bytecode.opcodes,
+            samples,
             &debug_artifact.debug_symbols[func_idx],
             &debug_artifact,
             artifact_path.to_str().unwrap(),
@@ -92,7 +105,10 @@ fn run_with_provider<Provider: GatesProvider, Generator: FlamegraphGenerator>(
 
 #[cfg(test)]
 mod tests {
-    use acir::circuit::{Circuit, Opcode, Program};
+    use acir::{
+        circuit::{Circuit, Program},
+        AcirField,
+    };
     use color_eyre::eyre::{self};
     use fm::codespan_files::Files;
     use noirc_artifacts::program::ProgramArtifact;
@@ -102,7 +118,10 @@ mod tests {
         path::{Path, PathBuf},
     };
 
-    use crate::gates_provider::{BackendGatesReport, BackendGatesResponse, GatesProvider};
+    use crate::{
+        flamegraph::Sample,
+        gates_provider::{BackendGatesReport, BackendGatesResponse, GatesProvider},
+    };
 
     struct TestGateProvider {
         mock_responses: HashMap<PathBuf, BackendGatesResponse>,
@@ -123,10 +142,9 @@ mod tests {
     struct TestFlamegraphGenerator {}
 
     impl super::FlamegraphGenerator for TestFlamegraphGenerator {
-        fn generate_flamegraph<'files, F>(
+        fn generate_flamegraph<'files, F: AcirField>(
             &self,
-            _samples_per_opcode: Vec<usize>,
-            _opcodes: Vec<Opcode<F>>,
+            _samples: Vec<Sample<F>>,
             _debug_symbols: &DebugInfo,
             _files: &'files impl Files<'files, FileId = fm::FileId>,
             _artifact_name: &str,

noir/noir-repo/tooling/profiler/src/cli/opcodes_flamegraph_cmd.rs

@@ -1,12 +1,14 @@
 use std::path::{Path, PathBuf};
 
+use acir::circuit::{Circuit, Opcode, OpcodeLocation};
 use clap::Args;
 use color_eyre::eyre::{self, Context};
 
 use noirc_artifacts::debug::DebugArtifact;
 
-use crate::flamegraph::{FlamegraphGenerator, InfernoFlamegraphGenerator};
+use crate::flamegraph::{FlamegraphGenerator, InfernoFlamegraphGenerator, Sample};
 use crate::fs::read_program_from_file;
+use crate::opcode_formatter::AcirOrBrilligOpcode;
 
 #[derive(Debug, Clone, Args)]
 pub(crate) struct OpcodesFlamegraphCommand {
@@ -17,20 +19,26 @@ pub(crate) struct OpcodesFlamegraphCommand {
     /// The output folder for the flamegraph svg files
     #[clap(long, short)]
     output: String,
+
+    /// Wether to skip brillig functions
+    #[clap(long, short, action)]
+    skip_brillig: bool,
 }
 
 pub(crate) fn run(args: OpcodesFlamegraphCommand) -> eyre::Result<()> {
     run_with_generator(
         &PathBuf::from(args.artifact_path),
         &InfernoFlamegraphGenerator { count_name: "opcodes".to_string() },
         &PathBuf::from(args.output),
+        args.skip_brillig,
     )
 }
 
 fn run_with_generator<Generator: FlamegraphGenerator>(
     artifact_path: &Path,
     flamegraph_generator: &Generator,
     output_path: &Path,
+    skip_brillig: bool,
 ) -> eyre::Result<()> {
     let mut program =
         read_program_from_file(artifact_path).context("Error reading program from file")?;
@@ -42,41 +50,113 @@ fn run_with_generator<Generator: FlamegraphGenerator>(
     let debug_artifact: DebugArtifact = program.into();
 
     for (func_idx, (func_name, bytecode)) in
-        function_names.into_iter().zip(bytecode.functions).enumerate()
+        function_names.into_iter().zip(bytecode.functions.iter()).enumerate()
     {
-        println!("Opcode count: {}", bytecode.opcodes.len());
+        println!("Opcode count for {}: {}", func_name, bytecode.opcodes.len());
+
+        let samples = bytecode
+            .opcodes
+            .iter()
+            .enumerate()
+            .map(|(index, opcode)| Sample {
+                opcode: AcirOrBrilligOpcode::Acir(opcode.clone()),
+                call_stack: vec![OpcodeLocation::Acir(index)],
+                count: 1,
+            })
+            .collect();
 
         flamegraph_generator.generate_flamegraph(
-            bytecode.opcodes.iter().map(|_op| 1).collect(),
-            bytecode.opcodes,
+            samples,
             &debug_artifact.debug_symbols[func_idx],
             &debug_artifact,
             artifact_path.to_str().unwrap(),
             &func_name,
-            &Path::new(&output_path).join(Path::new(&format!("{}_opcodes.svg", &func_name))),
+            &Path::new(&output_path).join(Path::new(&format!("{}_acir_opcodes.svg", &func_name))),
         )?;
     }
 
+    if !skip_brillig {
+        for (brillig_fn_index, brillig_bytecode) in
+            bytecode.unconstrained_functions.into_iter().enumerate()
+        {
+            let acir_location = locate_brillig_call(brillig_fn_index, &bytecode.functions);
+            let Some((acir_fn_index, acir_opcode_index)) = acir_location else {
+                continue;
+            };
+
+            println!(
+                "Opcode count for brillig_{}: {}",
+                brillig_fn_index,
+                brillig_bytecode.bytecode.len()
+            );
+
+            let samples = brillig_bytecode
+                .bytecode
+                .into_iter()
+                .enumerate()
+                .map(|(brillig_index, opcode)| Sample {
+                    opcode: AcirOrBrilligOpcode::Brillig(opcode),
+                    call_stack: vec![OpcodeLocation::Brillig {
+                        acir_index: acir_opcode_index,
+                        brillig_index,
+                    }],
+                    count: 1,
+                })
+                .collect();
+
+            flamegraph_generator.generate_flamegraph(
+                samples,
+                &debug_artifact.debug_symbols[acir_fn_index],
+                &debug_artifact,
+                artifact_path.to_str().unwrap(),
+                &format!("brillig_{}", brillig_fn_index),
+                &Path::new(&output_path)
+                    .join(Path::new(&format!("{}_brillig_opcodes.svg", &brillig_fn_index))),
+            )?;
+        }
+    }
+
     Ok(())
 }
 
+fn locate_brillig_call<F>(
+    brillig_fn_index: usize,
+    acir_functions: &[Circuit<F>],
+) -> Option<(usize, usize)> {
+    for (acir_fn_index, acir_fn) in acir_functions.iter().enumerate() {
+        for (acir_opcode_index, acir_opcode) in acir_fn.opcodes.iter().enumerate() {
+            match acir_opcode {
+                Opcode::BrilligCall { id, .. } if (*id) as usize == brillig_fn_index => {
+                    return Some((acir_fn_index, acir_opcode_index))
+                }
+                _ => {}
+            }
+        }
+    }
+    None
+}
+
 #[cfg(test)]
 mod tests {
-    use acir::circuit::{Circuit, Opcode, Program};
+    use acir::{
+        circuit::{brillig::BrilligBytecode, Circuit, Opcode, Program},
+        AcirField, FieldElement,
+    };
     use color_eyre::eyre::{self};
     use fm::codespan_files::Files;
     use noirc_artifacts::program::ProgramArtifact;
     use noirc_errors::debug_info::{DebugInfo, ProgramDebugInfo};
     use std::{collections::BTreeMap, path::Path};
 
+    use crate::flamegraph::Sample;
+
     #[derive(Default)]
     struct TestFlamegraphGenerator {}
 
     impl super::FlamegraphGenerator for TestFlamegraphGenerator {
-        fn generate_flamegraph<'files, F>(
+        fn generate_flamegraph<'files, F: AcirField>(
             &self,
-            _samples_per_opcode: Vec<usize>,
-            _opcodes: Vec<Opcode<F>>,
+            _samples: Vec<Sample<F>>,
             _debug_symbols: &DebugInfo,
             _files: &'files impl Files<'files, FileId = fm::FileId>,
             _artifact_name: &str,
@@ -113,11 +193,50 @@ mod tests {
 
         let flamegraph_generator = TestFlamegraphGenerator::default();
 
-        super::run_with_generator(&artifact_path, &flamegraph_generator, temp_dir.path())
+        super::run_with_generator(&artifact_path, &flamegraph_generator, temp_dir.path(), true)
             .expect("should run without errors");
 
         // Check that the output file was written to
-        let output_file = temp_dir.path().join("main_opcodes.svg");
+        let output_file = temp_dir.path().join("main_acir_opcodes.svg");
+        assert!(output_file.exists());
+    }
+
+    #[test]
+    fn brillig_test() {
+        let temp_dir = tempfile::tempdir().unwrap();
+
+        let artifact_path = temp_dir.path().join("test.json");
+
+        let acir: Vec<Opcode<FieldElement>> =
+            vec![Opcode::BrilligCall { id: 0, inputs: vec![], outputs: vec![], predicate: None }];
+
+        let artifact = ProgramArtifact {
+            noir_version: "0.0.0".to_string(),
+            hash: 27,
+            abi: noirc_abi::Abi::default(),
+            bytecode: Program {
+                functions: vec![Circuit { opcodes: acir, ..Circuit::default() }],
+                unconstrained_functions: vec![BrilligBytecode::default()],
+            },
+            debug_symbols: ProgramDebugInfo { debug_infos: vec![DebugInfo::default()] },
+            file_map: BTreeMap::default(),
+            names: vec!["main".to_string()],
+        };
+
+        // Write the artifact to a file
+        let artifact_file = std::fs::File::create(&artifact_path).unwrap();
+        serde_json::to_writer(artifact_file, &artifact).unwrap();
+
+        let flamegraph_generator = TestFlamegraphGenerator::default();
+
+        super::run_with_generator(&artifact_path, &flamegraph_generator, temp_dir.path(), false)
+            .expect("should run without errors");
+
+        // Check that the output files ware written to
+        let output_file = temp_dir.path().join("main_acir_opcodes.svg");
+        assert!(output_file.exists());
+
+        let output_file = temp_dir.path().join("0_brillig_opcodes.svg");
         assert!(output_file.exists());
     }
 }