chore(profiler): Add option to only get the total sample count for the execution-opcodes command

tooling/profiler/src/cli/execution_flamegraph_cmd.rs

@@ -30,13 +30,22 @@ pub(crate) struct ExecutionFlamegraphCommand {
 
     /// The output folder for the flamegraph svg files
     #[clap(long, short)]
-    output: PathBuf,
+    output: Option<PathBuf>,
 
     /// Use pedantic ACVM solving, i.e. double-check some black-box function
     /// assumptions when solving.
     /// This is disabled by default.
     #[clap(long, default_value = "false")]
     pedantic_solving: bool,
+
+    /// A single number representing the total opcodes executed.
+    /// Outputs to stdout and skips generating a flamegraph.
+    #[clap(long, default_value = "false")]
+    sample_count: bool,
+
+    /// Enables additional logging
+    #[clap(long, default_value = "false")]
+    verbose: bool,
 }
 
 pub(crate) fn run(args: ExecutionFlamegraphCommand) -> eyre::Result<()> {
@@ -46,27 +55,38 @@ pub(crate) fn run(args: ExecutionFlamegraphCommand) -> eyre::Result<()> {
         &InfernoFlamegraphGenerator { count_name: "samples".to_string() },
         &args.output,
         args.pedantic_solving,
+        args.sample_count,
+        args.verbose,
     )
 }
 
 fn run_with_generator(
     artifact_path: &Path,
     prover_toml_path: &Path,
     flamegraph_generator: &impl FlamegraphGenerator,
-    output_path: &Path,
+    output_path: &Option<PathBuf>,
     pedantic_solving: bool,
+    print_sample_count: bool,
+    verbose: bool,
 ) -> eyre::Result<()> {
     let program =
         read_program_from_file(artifact_path).context("Error reading program from file")?;
 
     ensure_brillig_entry_point(&program)?;
 
+    if !print_sample_count && output_path.is_none() {
+        return report_error("Missing --output <OUTPUT> argument for when building a flamegraph")
+            .map_err(Into::into);
+    }
+
     let (inputs_map, _) =
         read_inputs_from_file(&prover_toml_path.with_extension("toml"), &program.abi)?;
 
     let initial_witness = program.abi.encode(&inputs_map, None)?;
 
-    println!("Executing...");
+    if verbose {
+        println!("Executing...");
+    }
 
     let solved_witness_stack_err = nargo::ops::execute_program_with_profiling(
         &program.bytecode,
@@ -94,9 +114,21 @@ fn run_with_generator(
         }
     };
 
-    println!("Executed");
+    if verbose {
+        println!("Executed");
+    }
+
+    if print_sample_count {
+        println!("{}", profiling_samples.len());
+        return Ok(());
+    }
 
-    println!("Collecting {} samples", profiling_samples.len());
+    // We place this logging output before the transforming and collection of the samples.
+    // This is done because large traces can take some time, and can make it look
+    // as if the profiler has stalled.
+    if verbose {
+        println!("Generating flamegraph for {} samples...", profiling_samples.len());
+    }
 
     let profiling_samples: Vec<BrilligExecutionSample> = profiling_samples
         .iter_mut()
@@ -118,24 +150,28 @@ fn run_with_generator(
         })
         .collect();
 
-    let debug_artifact: DebugArtifact = program.into();
-
-    println!("Generating flamegraph with {} samples", profiling_samples.len());
+    let output_path =
+        output_path.as_ref().expect("Should have already checked for the output path");
 
+    let debug_artifact: DebugArtifact = program.into();
     flamegraph_generator.generate_flamegraph(
         profiling_samples,
         &debug_artifact.debug_symbols[0],
         &debug_artifact,
         artifact_path.to_str().unwrap(),
         "main",
-        &Path::new(&output_path).join(Path::new(&format!("{}_brillig_trace.svg", "main"))),
+        &Path::new(output_path).join(Path::new(&format!("{}_brillig_trace.svg", "main"))),
     )?;
 
+    if verbose {
+        println!("Generated flamegraph");
+    }
+
     Ok(())
 }
 
 fn ensure_brillig_entry_point(artifact: &ProgramArtifact) -> Result<(), CliError> {
-    let err_msg = "Command only supports fully unconstrained Noir programs e.g. `unconstrained fn main() { .. }".to_owned();
+    let err_msg = "Command only supports fully unconstrained Noir programs e.g. `unconstrained fn main() { .. }";
     let program = &artifact.bytecode;
     if program.functions.len() != 1 || program.unconstrained_functions.len() != 1 {
         return report_error(err_msg);
@@ -152,3 +188,86 @@ fn ensure_brillig_entry_point(artifact: &ProgramArtifact) -> Result<(), CliError
 
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use acir::circuit::{Circuit, Program, brillig::BrilligBytecode};
+    use color_eyre::eyre;
+    use fm::codespan_files::Files;
+    use noirc_artifacts::program::ProgramArtifact;
+    use noirc_driver::CrateName;
+    use noirc_errors::debug_info::{DebugInfo, ProgramDebugInfo};
+    use std::{collections::BTreeMap, path::Path, str::FromStr};
+
+    use crate::flamegraph::Sample;
+
+    #[derive(Default)]
+    struct TestFlamegraphGenerator {}
+
+    impl super::FlamegraphGenerator for TestFlamegraphGenerator {
+        fn generate_flamegraph<'files, S: Sample>(
+            &self,
+            _samples: Vec<S>,
+            _debug_symbols: &DebugInfo,
+            _files: &'files impl Files<'files, FileId = fm::FileId>,
+            _artifact_name: &str,
+            _function_name: &str,
+            output_path: &Path,
+        ) -> eyre::Result<()> {
+            let output_file = std::fs::File::create(output_path).unwrap();
+            std::io::Write::write_all(&mut std::io::BufWriter::new(output_file), b"success")
+                .unwrap();
+
+            Ok(())
+        }
+    }
+
+    #[test]
+    fn error_reporter_smoke_test() {
+        // This test purposefully uses an artifact that does not represent a Brillig entry point.
+        // The goal is to see that our program fails gracefully and does not panic.
+        let temp_dir = tempfile::tempdir().unwrap();
+
+        let prover_toml_path = temp_dir.path().join("Prover.toml");
+
+        let artifact = ProgramArtifact {
+            noir_version: "0.0.0".to_string(),
+            hash: 27,
+            abi: noirc_abi::Abi::default(),
+            bytecode: Program {
+                functions: vec![Circuit::default()],
+                unconstrained_functions: vec![
+                    BrilligBytecode::default(),
+                    BrilligBytecode::default(),
+                ],
+            },
+            debug_symbols: ProgramDebugInfo { debug_infos: vec![DebugInfo::default()] },
+            file_map: BTreeMap::default(),
+            names: vec!["main".to_string()],
+            brillig_names: Vec::new(),
+        };
+
+        // Write the artifact to a file
+        let artifact_path = noir_artifact_cli::fs::artifact::save_program_to_file(
+            &artifact,
+            &CrateName::from_str("test").unwrap(),
+            temp_dir.path(),
+        )
+        .unwrap();
+
+        let flamegraph_generator = TestFlamegraphGenerator::default();
+
+        assert!(
+            super::run_with_generator(
+                &artifact_path,
+                &prover_toml_path,
+                &flamegraph_generator,
+                &Some(temp_dir.into_path()),
+                false,
+                false,
+                false
+            )
+            .is_err()
+        );
+    }
+}

tooling/profiler/src/errors.rs

@@ -1,5 +1,5 @@
-use fm::FileMap;
-use noirc_errors::{CustomDiagnostic, Location};
+use fm::{FileId, FileMap};
+use noirc_errors::CustomDiagnostic;
 use thiserror::Error;
 
 #[derive(Debug, Error)]
@@ -9,8 +9,8 @@ pub(crate) enum CliError {
 }
 
 /// Report an error from the CLI that is not reliant on a stack trace.
-pub(crate) fn report_error(message: String) -> Result<(), CliError> {
-    let error = CustomDiagnostic::simple_error(message.clone(), String::new(), Location::dummy());
+pub(crate) fn report_error(message: &str) -> Result<(), CliError> {
+    let error = CustomDiagnostic::from_message(message, FileId::dummy());
     noirc_errors::reporter::report(&FileMap::default(), &error, false);
     Err(CliError::Generic)
 }