fix(ssa): Constant fold Brillig calls using the SSA interpreter

compiler/noirc_evaluator/src/ssa/interpreter/mod.rs

@@ -1,4 +1,4 @@
-use std::{cmp::Ordering, io::Write};
+use std::{cmp::Ordering, collections::BTreeMap, io::Write};
 
 use super::{
     Ssa,
@@ -27,13 +27,13 @@ pub mod value;
 
 use value::Value;
 
-struct Interpreter<'ssa, W> {
+pub(crate) struct Interpreter<'ssa, W> {
     /// Contains each function called with `main` (or the first called function if
     /// the interpreter was manually invoked on a different function) at
     /// the front of the Vec.
     call_stack: Vec<CallContext>,
 
-    ssa: &'ssa Ssa,
+    functions: &'ssa BTreeMap<FunctionId, Function>,
 
     /// This variable can be modified by `enable_side_effects_if` instructions and is
     /// expected to have no effect if there are no such instructions or if the code
@@ -49,6 +49,8 @@ struct Interpreter<'ssa, W> {
 pub struct InterpreterOptions {
     /// If true, the interpreter will trace its execution.
     pub trace: bool,
+    /// If true, the interpreter treats all foreign function calls (e.g., `print`) as unknown
+    pub no_foreign_calls: bool,
 }
 
 struct CallContext {
@@ -104,7 +106,20 @@ impl Ssa {
 impl<'ssa, W: Write> Interpreter<'ssa, W> {
     fn new(ssa: &'ssa Ssa, options: InterpreterOptions, output: W) -> Self {
         let call_stack = vec![CallContext::global_context()];
-        Self { ssa, call_stack, side_effects_enabled: true, options, output }
+        Self { functions: &ssa.functions, call_stack, side_effects_enabled: true, options, output }
+    }
+
+    pub(crate) fn new_from_functions(
+        functions: &'ssa BTreeMap<FunctionId, Function>,
+        options: InterpreterOptions,
+        output: W,
+    ) -> Self {
+        let call_stack = vec![CallContext::global_context()];
+        Self { functions, call_stack, side_effects_enabled: true, options, output }
+    }
+
+    pub(crate) fn functions(&self) -> &BTreeMap<FunctionId, Function> {
+        self.functions
     }
 
     fn call_context(&self) -> &CallContext {
@@ -121,7 +136,7 @@ impl<'ssa, W: Write> Interpreter<'ssa, W> {
 
     fn try_current_function(&self) -> Option<&'ssa Function> {
         let current_function_id = self.call_context().called_function;
-        current_function_id.map(|current_function_id| &self.ssa.functions[&current_function_id])
+        current_function_id.map(|current_function_id| &self.functions[&current_function_id])
     }
 
     fn current_function(&self) -> &'ssa Function {
@@ -163,8 +178,9 @@ impl<'ssa, W: Write> Interpreter<'ssa, W> {
         Ok(())
     }
 
-    fn interpret_globals(&mut self) -> IResult<()> {
-        let globals = &self.ssa.main().dfg.globals;
+    pub(crate) fn interpret_globals(&mut self) -> IResult<()> {
+        let (_, function) = self.functions.first_key_value().unwrap();
+        let globals = &function.dfg.globals;
         for (global_id, global) in globals.values_iter() {
             let value = match global {
                 super::ir::value::Value::Instruction { instruction, .. } => {
@@ -189,10 +205,14 @@ impl<'ssa, W: Write> Interpreter<'ssa, W> {
         Ok(())
     }
 
-    fn call_function(&mut self, function_id: FunctionId, mut arguments: Vec<Value>) -> IResults {
+    pub(crate) fn call_function(
+        &mut self,
+        function_id: FunctionId,
+        mut arguments: Vec<Value>,
+    ) -> IResults {
         self.call_stack.push(CallContext::new(function_id));
 
-        let function = &self.ssa.functions[&function_id];
+        let function = &self.functions[&function_id];
         if self.options.trace {
             println!();
             println!("enter function {} ({})", function_id, function.name());
@@ -280,7 +300,7 @@ impl<'ssa, W: Write> Interpreter<'ssa, W> {
         if self.options.trace {
             if let Some(context) = self.call_stack.last() {
                 if let Some(function_id) = context.called_function {
-                    let function = &self.ssa.functions[&function_id];
+                    let function = &self.functions[&function_id];
                     println!("back in function {} ({})", function_id, function.name());
                 }
             }
@@ -741,7 +761,7 @@ impl<'ssa, W: Write> Interpreter<'ssa, W> {
                     // any shared mutable fields in our arguments since brillig should conceptually
                     // receive fresh array on each invocation.
                     if !self.in_unconstrained_context()
-                        && self.ssa.functions[&id].runtime().is_brillig()
+                        && self.functions[&id].runtime().is_brillig()
                     {
                         for argument in arguments.iter_mut() {
                             Self::reset_array_state(argument)?;
@@ -752,6 +772,9 @@ impl<'ssa, W: Write> Interpreter<'ssa, W> {
                 Value::Intrinsic(intrinsic) => {
                     self.call_intrinsic(intrinsic, argument_ids, results)?
                 }
+                Value::ForeignFunction(name) if self.options.no_foreign_calls => {
+                    return Err(InterpreterError::UnknownForeignFunctionCall { name });
+                }
                 Value::ForeignFunction(name) if name == "print" => self.call_print(arguments)?,
                 Value::ForeignFunction(name) => {
                     return Err(InterpreterError::UnknownForeignFunctionCall { name });
@@ -789,7 +812,7 @@ impl<'ssa, W: Write> Interpreter<'ssa, W> {
     /// Try to get a function's name or approximate it if it is not known
     fn try_get_function_name(&self, function: ValueId) -> String {
         match self.lookup(function) {
-            Ok(Value::Function(id)) => match self.ssa.functions.get(&id) {
+            Ok(Value::Function(id)) => match self.functions.get(&id) {
                 Some(function) => function.name().to_string(),
                 None => "unknown function".to_string(),
             },

compiler/noirc_evaluator/src/ssa/mod.rs

@@ -13,7 +13,7 @@ use std::{
 
 use crate::{
     acir::ssa::Artifacts,
-    brillig::{Brillig, BrilligOptions},
+    brillig::BrilligOptions,
     errors::{RuntimeError, SsaReport},
 };
 use acvm::{
@@ -102,9 +102,6 @@ pub struct SsaEvaluatorOptions {
 pub struct ArtifactsAndWarnings(pub Artifacts, pub Vec<SsaReport>);
 
 /// The default SSA optimization pipeline.
-///
-/// After these passes everything is ready for execution, which is
-/// something we take can advantage of in the [secondary_passes].
 pub fn primary_passes(options: &SsaEvaluatorOptions) -> Vec<SsaPass> {
     vec![
         SsaPass::new(Ssa::expand_signed_checks, "expand signed checks"),
@@ -188,14 +185,17 @@ pub fn primary_passes(options: &SsaEvaluatorOptions) -> Vec<SsaPass> {
             // Remove any potentially unnecessary duplication from the Brillig entry point analysis.
             .and_then(Ssa::remove_unreachable_functions),
         SsaPass::new(Ssa::remove_truncate_after_range_check, "Removing Truncate after RangeCheck"),
+        SsaPass::new(Ssa::checked_to_unchecked, "Checked to unchecked"),
+        SsaPass::new(Ssa::fold_constants_with_brillig, "Inlining Brillig Calls"),
+        SsaPass::new(Ssa::remove_unreachable_instructions, "Remove Unreachable Instructions")
+            .and_then(Ssa::remove_unreachable_functions),
         // This pass makes transformations specific to Brillig generation.
         // It must be the last pass to either alter or add new instructions before Brillig generation,
         // as other semantics in the compiler can potentially break (e.g. inserting instructions).
         SsaPass::new(Ssa::brillig_array_get_and_set, "Brillig Array Get and Set Optimizations"),
         SsaPass::new(Ssa::dead_instruction_elimination, "Dead Instruction Elimination")
             // A function can be potentially unreachable post-DIE if all calls to that function were removed.
             .and_then(Ssa::remove_unreachable_functions),
-        SsaPass::new(Ssa::checked_to_unchecked, "Checked to unchecked"),
         SsaPass::new_try(
             Ssa::verify_no_dynamic_indices_to_references,
             "Verifying no dynamic array indices to reference value elements",
@@ -209,20 +209,6 @@ pub fn primary_passes(options: &SsaEvaluatorOptions) -> Vec<SsaPass> {
     ]
 }
 
-/// The second SSA pipeline, in which we take the Brillig functions compiled after
-/// the primary pipeline, and execute the ones with all-constant arguments,
-/// to replace the calls with the return value.
-pub fn secondary_passes(brillig: &Brillig) -> Vec<SsaPass> {
-    vec![
-        SsaPass::new(move |ssa| ssa.fold_constants_with_brillig(brillig), "Inlining Brillig Calls"),
-        SsaPass::new(Ssa::remove_unreachable_instructions, "Remove Unreachable Instructions")
-            // It could happen that we inlined all calls to a given brillig function.
-            // In that case it's unused so we can remove it. This is what we check next.
-            .and_then(Ssa::remove_unreachable_functions),
-        SsaPass::new(Ssa::dead_instruction_elimination_acir, "Dead Instruction Elimination - ACIR"),
-    ]
-}
-
 /// For testing purposes we want a list of the minimum number of SSA passes that should
 /// return the same result as the full pipeline.
 ///
@@ -253,38 +239,24 @@ pub fn minimal_passes() -> Vec<SsaPass<'static>> {
 /// convert the final SSA into an ACIR program and return it.
 /// An ACIR program is made up of both ACIR functions
 /// and Brillig functions for unconstrained execution.
-///
-/// The `primary` SSA passes are applied on the initial SSA.
-/// Then we compile the Brillig functions, and use the output
-/// to run a `secondary` pass, which can use the Brillig
-/// artifacts to do constant folding.
-///
-/// See the [primary_passes] and [secondary_passes] for
-/// the default implementations.
-pub fn optimize_ssa_builder_into_acir<S>(
+pub fn optimize_ssa_builder_into_acir(
     builder: SsaBuilder,
     options: &SsaEvaluatorOptions,
-    primary: &[SsaPass],
-    secondary: S,
-) -> Result<ArtifactsAndWarnings, RuntimeError>
-where
-    S: for<'b> Fn(&'b Brillig) -> Vec<SsaPass<'b>>,
-{
+    passes: &[SsaPass],
+) -> Result<ArtifactsAndWarnings, RuntimeError> {
     let ssa_gen_span = span!(Level::TRACE, "ssa_generation");
     let ssa_gen_span_guard = ssa_gen_span.enter();
-    let builder = builder.with_skip_passes(options.skip_passes.clone()).run_passes(primary)?;
+    let builder = builder.with_skip_passes(options.skip_passes.clone()).run_passes(passes)?;
 
     drop(ssa_gen_span_guard);
 
-    let ssa = builder.ssa();
     let brillig = time("SSA to Brillig", options.print_codegen_timings, || {
-        ssa.to_brillig(&options.brillig_options)
+        builder.ssa().to_brillig(&options.brillig_options)
     });
 
     let ssa_gen_span_guard = ssa_gen_span.enter();
 
-    let mut ssa = builder.run_passes(&secondary(&brillig))?.finish();
-
+    let mut ssa = builder.finish();
     let mut ssa_level_warnings = vec![];
     if !options.skip_underconstrained_check {
         ssa_level_warnings.extend(time(
@@ -319,24 +291,12 @@ where
 /// convert the final SSA into an ACIR program and return it.
 /// An ACIR program is made up of both ACIR functions
 /// and Brillig functions for unconstrained execution.
-///
-/// The `primary` SSA passes are applied on the initial SSA.
-/// Then we compile the Brillig functions, and use the output
-/// to run a `secondary` pass, which can use the Brillig
-/// artifacts to do constant folding.
-///
-/// See the [primary_passes] and [secondary_passes] for
-/// the default implementations.
-pub fn optimize_into_acir<S>(
+pub fn optimize_into_acir(
     program: Program,
     options: &SsaEvaluatorOptions,
-    primary: &[SsaPass],
-    secondary: S,
+    passes: &[SsaPass],
     files: Option<&fm::FileManager>,
-) -> Result<ArtifactsAndWarnings, RuntimeError>
-where
-    S: for<'b> Fn(&'b Brillig) -> Vec<SsaPass<'b>>,
-{
+) -> Result<ArtifactsAndWarnings, RuntimeError> {
     let builder = SsaBuilder::from_program(
         program,
         options.ssa_logging.clone(),
@@ -345,7 +305,7 @@ where
         files,
     )?;
 
-    optimize_ssa_builder_into_acir(builder, options, primary, secondary)
+    optimize_ssa_builder_into_acir(builder, options, passes)
 }
 
 /// Compiles the [`Program`] into [`ACIR`][acvm::acir::circuit::Program].
@@ -357,7 +317,7 @@ pub fn create_program(
     options: &SsaEvaluatorOptions,
     files: Option<&fm::FileManager>,
 ) -> Result<SsaProgramArtifact, RuntimeError> {
-    create_program_with_passes(program, options, &primary_passes(options), secondary_passes, files)
+    create_program_with_passes(program, options, &primary_passes(options), files)
 }
 
 /// Compiles the [`Program`] into [`ACIR`][acvm::acir::circuit::Program] using the minimum amount of SSA passes.
@@ -377,30 +337,25 @@ pub fn create_program_with_minimal_passes(
             func.name
         );
     }
-    create_program_with_passes(program, options, &minimal_passes(), |_| vec![], Some(files))
+    create_program_with_passes(program, options, &minimal_passes(), Some(files))
 }
 
-/// Compiles the [`Program`] into [`ACIR`][acvm::acir::circuit::Program] by running it through
-/// `primary` and `secondary` SSA passes.
+/// Compiles the [`Program`] into [`ACIR`][acvm::acir::circuit::Program] by running it through SSA `passes`.
 #[tracing::instrument(level = "trace", skip_all)]
-pub fn create_program_with_passes<S>(
+pub fn create_program_with_passes(
     program: Program,
     options: &SsaEvaluatorOptions,
-    primary: &[SsaPass],
-    secondary: S,
+    passes: &[SsaPass],
     files: Option<&fm::FileManager>,
-) -> Result<SsaProgramArtifact, RuntimeError>
-where
-    S: for<'b> Fn(&'b Brillig) -> Vec<SsaPass<'b>>,
-{
+) -> Result<SsaProgramArtifact, RuntimeError> {
     let debug_variables = program.debug_variables.clone();
     let debug_types = program.debug_types.clone();
     let debug_functions = program.debug_functions.clone();
 
     let arg_size_and_visibilities: Vec<Vec<(u32, Visibility)>> =
         program.function_signatures.iter().map(resolve_function_signature).collect();
 
-    let artifacts = optimize_into_acir(program, options, primary, secondary, files)?;
+    let artifacts = optimize_into_acir(program, options, passes, files)?;
 
     Ok(combine_artifacts(
         artifacts,