feat: Apply optimizations to unconstrained code

crates/noirc_evaluator/src/ssa.rs

@@ -21,7 +21,7 @@ use noirc_abi::Abi;
 
 use noirc_frontend::monomorphization::ast::Program;
 
-use self::{abi_gen::gen_abi, acir_gen::GeneratedAcir, ir::function::RuntimeType, ssa_gen::Ssa};
+use self::{abi_gen::gen_abi, acir_gen::GeneratedAcir, ssa_gen::Ssa};
 
 mod abi_gen;
 mod acir_gen;
@@ -39,38 +39,25 @@ pub(crate) fn optimize_into_acir(
     print_brillig_trace: bool,
 ) -> Result<GeneratedAcir, RuntimeError> {
     let abi_distinctness = program.return_distinctness;
-    let mut ssa = ssa_gen::generate_ssa(program)
-        .print(print_ssa_passes, "Initial SSA:")
-        .defunctionalize()
-        .print(print_ssa_passes, "After Defunctionalization:");
+    let ssa = SsaBuilder::new(program, print_ssa_passes)
+        .run_pass(Ssa::defunctionalize, "After Defunctionalization:")
+        .run_pass(Ssa::inline_functions, "After Inlining:")
+        // Run mem2reg with the CFG separated into blocks
+        .run_pass(Ssa::mem2reg, "After Mem2Reg:")
+        .try_run_pass(Ssa::evaluate_assert_constant, "After Assert Constant:")?
+        .try_run_pass(Ssa::unroll_loops, "After Unrolling:")?
+        .run_pass(Ssa::simplify_cfg, "After Simplifying:")
+        // Run mem2reg before flattening to handle any promotion
+        // of values that can be accessed after loop unrolling
+        .run_pass(Ssa::mem2reg, "After Mem2Reg:")
+        .run_pass(Ssa::flatten_cfg, "After Flattening:")
+        // Run mem2reg once more with the flattened CFG to catch any remaining loads/stores
+        .run_pass(Ssa::mem2reg, "After Mem2Reg:")
+        .run_pass(Ssa::fold_constants, "After Constant Folding:")
+        .run_pass(Ssa::dead_instruction_elimination, "After Dead Instruction Elimination:")
+        .finish();
 
     let brillig = ssa.to_brillig(print_brillig_trace);
-    if let RuntimeType::Acir = ssa.main().runtime() {
-        ssa = ssa
-            .inline_functions()
-            .print(print_ssa_passes, "After Inlining:")
-            // Run mem2reg with the CFG separated into blocks
-            .mem2reg()
-            .print(print_ssa_passes, "After Mem2Reg:")
-            .evaluate_assert_constant()?
-            .unroll_loops()?
-            .print(print_ssa_passes, "After Unrolling:")
-            .simplify_cfg()
-            .print(print_ssa_passes, "After Simplifying:")
-            // Run mem2reg before flattening to handle any promotion
-            // of values that can be accessed after loop unrolling
-            .mem2reg()
-            .print(print_ssa_passes, "After Mem2Reg:")
-            .flatten_cfg()
-            .print(print_ssa_passes, "After Flattening:")
-            // Run mem2reg once more with the flattened CFG to catch any remaining loads/stores
-            .mem2reg()
-            .print(print_ssa_passes, "After Mem2Reg:")
-            .fold_constants()
-            .print(print_ssa_passes, "After Constant Folding:")
-            .dead_instruction_elimination()
-            .print(print_ssa_passes, "After Dead Instruction Elimination:");
-    }
     let last_array_uses = ssa.find_last_array_uses();
     ssa.into_acir(brillig, abi_distinctness, &last_array_uses)
 }
@@ -124,10 +111,40 @@ pub fn create_circuit(
     Ok((circuit, debug_info, abi))
 }
 
-impl Ssa {
-    fn print(self, print_ssa_passes: bool, msg: &str) -> Ssa {
-        if print_ssa_passes {
-            println!("{msg}\n{self}");
+// This is just a convenience object to bundle the ssa with `print_ssa_passes` for debug printing.
+struct SsaBuilder {
+    ssa: Ssa,
+    print_ssa_passes: bool,
+}
+
+impl SsaBuilder {
+    fn new(program: Program, print_ssa_passes: bool) -> SsaBuilder {
+        SsaBuilder { print_ssa_passes, ssa: ssa_gen::generate_ssa(program) }.print("Initial SSA:")
+    }
+
+    fn finish(self) -> Ssa {
+        self.ssa
+    }
+
+    /// Runs the given SSA pass and prints the SSA afterward if `print_ssa_passes` is true.
+    fn run_pass(mut self, pass: fn(Ssa) -> Ssa, msg: &str) -> Self {
+        self.ssa = pass(self.ssa);
+        self.print(msg)
+    }
+
+    /// The same as `run_pass` but for passes that may fail
+    fn try_run_pass(
+        mut self,
+        pass: fn(Ssa) -> Result<Ssa, RuntimeError>,
+        msg: &str,
+    ) -> Result<Self, RuntimeError> {
+        self.ssa = pass(self.ssa)?;
+        Ok(self.print(msg))
+    }
+
+    fn print(self, msg: &str) -> Self {
+        if self.print_ssa_passes {
+            println!("{msg}\n{}", self.ssa);
         }
         self
     }

crates/noirc_evaluator/src/ssa/opt/flatten_cfg.rs

@@ -206,12 +206,9 @@ struct Branch {
 }
 
 fn flatten_function_cfg(function: &mut Function) {
-    // TODO This pass will run forever on a brillig function.
-    // TODO In particular, analyze will check if the predecessors
-    // TODO have been processed and push the block to the back of the queue
-    // TODO This loops forever, if the predecessors are not then processed
-    // TODO Because it will visit the same block again, pop it out of the queue
-    // TODO then back into the queue again.
+    // This pass may run forever on a brillig function.
+    // Analyze will check if the predecessors have been processed and push the block to the back of
+    // the queue. This loops forever if there are still any loops present in the program.
     if let crate::ssa::ir::function::RuntimeType::Brillig = function.runtime() {
         return;
     }

crates/noirc_evaluator/src/ssa/opt/inlining.rs

@@ -2,9 +2,9 @@
 //! The purpose of this pass is to inline the instructions of each function call
 //! within the function caller. If all function calls are known, there will only
 //! be a single function remaining when the pass finishes.
-use std::collections::{HashMap, HashSet};
+use std::collections::{BTreeSet, HashMap, HashSet};
 
-use iter_extended::vecmap;
+use iter_extended::{btree_map, vecmap};
 
 use crate::ssa::{
     ir::{
@@ -35,8 +35,13 @@ impl Ssa {
     /// changes. This is because if the function's id later becomes known by a later
     /// pass, we would need to re-run all of inlining anyway to inline it, so we might
     /// as well save the work for later instead of performing it twice.
-    pub(crate) fn inline_functions(self) -> Ssa {
-        InlineContext::new(&self).inline_all(self)
+    pub(crate) fn inline_functions(mut self) -> Ssa {
+        self.functions = btree_map(get_entry_point_functions(&self), |entry_point| {
+            let new_function = InlineContext::new(&self, entry_point).inline_all(&self);
+            (entry_point, new_function)
+        });
+
+        self
     }
 }
 
@@ -51,10 +56,8 @@ struct InlineContext {
 
     call_stack: CallStack,
 
-    /// True if we failed to inline at least one call. If this is still false when finishing
-    /// inlining we can remove all other functions from the resulting Ssa struct and keep only
-    /// the function that was inlined into.
-    failed_to_inline_a_call: bool,
+    // The FunctionId of the entry point function we're inlining into in the old, unmodified Ssa.
+    entry_point: FunctionId,
 }
 
 /// The per-function inlining context contains information that is only valid for one function.
@@ -87,8 +90,22 @@ struct PerFunctionContext<'function> {
     /// Maps InstructionIds from the function being inlined to the function being inlined into.
     instructions: HashMap<InstructionId, InstructionId>,
 
-    /// True if we're currently working on the main function.
-    inlining_main: bool,
+    /// True if we're currently working on the entry point function.
+    inlining_entry: bool,
+}
+
+/// The entry point functions are each function we should inline into - and each function that
+/// should be left in the final program. This is usually just `main` but also includes any
+/// brillig functions used.
+fn get_entry_point_functions(ssa: &Ssa) -> BTreeSet<FunctionId> {
+    let functions = ssa.functions.iter();
+    let mut entry_points = functions
+        .filter(|(_, function)| function.runtime() == RuntimeType::Brillig)
+        .map(|(id, _)| *id)
+        .collect::<BTreeSet<_>>();
+
+    entry_points.insert(ssa.main_id);
+    entry_points
 }
 
 impl InlineContext {
@@ -97,22 +114,18 @@ impl InlineContext {
     /// The function being inlined into will always be the main function, although it is
     /// actually a copy that is created in case the original main is still needed from a function
     /// that could not be inlined calling it.
-    fn new(ssa: &Ssa) -> InlineContext {
-        let main_name = ssa.main().name().to_owned();
-        let builder = FunctionBuilder::new(main_name, ssa.next_id.next(), RuntimeType::Acir);
-        Self {
-            builder,
-            recursion_level: 0,
-            call_stack: CallStack::new(),
-            failed_to_inline_a_call: false,
-        }
+    fn new(ssa: &Ssa, entry_point: FunctionId) -> InlineContext {
+        let source = &ssa.functions[&entry_point];
+        let builder = FunctionBuilder::new(source.name().to_owned(), entry_point, source.runtime());
+        Self { builder, recursion_level: 0, entry_point, call_stack: CallStack::new() }
     }
 
-    /// Start inlining the main function and all functions reachable from it.
-    fn inline_all(mut self, ssa: Ssa) -> Ssa {
-        let main = ssa.main();
-        let mut context = PerFunctionContext::new(&mut self, main);
-        context.inlining_main = true;
+    /// Start inlining the entry point function and all functions reachable from it.
+    fn inline_all(mut self, ssa: &Ssa) -> Function {
+        let entry_point = &ssa.functions[&self.entry_point];
+
+        let mut context = PerFunctionContext::new(&mut self, entry_point);
+        context.inlining_entry = true;
 
         // The main block is already inserted so we have to add it to context.blocks and add
         // its parameters here. Failing to do so would cause context.translate_block() to add
@@ -127,8 +140,12 @@ impl InlineContext {
         }
 
         context.blocks.insert(context.source_function.entry_block(), entry_block);
-        context.inline_blocks(&ssa);
-        self.finish(ssa)
+        context.inline_blocks(ssa);
+
+        // Finally, we should have 1 function left representing the inlined version of the target function.
+        let mut new_ssa = self.builder.finish();
+        assert_eq!(new_ssa.functions.len(), 1);
+        new_ssa.functions.pop_first().unwrap().1
     }
 
     /// Inlines a function into the current function and returns the translated return values
@@ -161,26 +178,6 @@ impl InlineContext {
         self.recursion_level -= 1;
         return_values
     }
-
-    /// Finish inlining and return the new Ssa struct with the inlined version of main.
-    /// If any functions failed to inline, they are not removed from the final Ssa struct.
-    fn finish(self, mut ssa: Ssa) -> Ssa {
-        let mut new_ssa = self.builder.finish();
-        assert_eq!(new_ssa.functions.len(), 1);
-
-        // If we failed to inline any call, any function may still be reachable so we
-        // don't remove any from the final program. We could be more precise here and
-        // do a reachability analysis but it should be fine to keep the extra functions
-        // around longer if they are not called.
-        if self.failed_to_inline_a_call {
-            let new_main = new_ssa.functions.pop_first().unwrap().1;
-            ssa.main_id = new_main.id();
-            ssa.functions.insert(new_main.id(), new_main);
-            ssa
-        } else {
-            new_ssa
-        }
-    }
 }
 
 impl<'function> PerFunctionContext<'function> {
@@ -195,7 +192,7 @@ impl<'function> PerFunctionContext<'function> {
             blocks: HashMap::new(),
             instructions: HashMap::new(),
             values: HashMap::new(),
-            inlining_main: false,
+            inlining_entry: false,
         }
     }
 
@@ -279,10 +276,7 @@ impl<'function> PerFunctionContext<'function> {
             // don't correspond to actual functions in the SSA program that would
             // need to be removed afterward.
             Value::Intrinsic(_) => None,
-            _ => {
-                self.context.failed_to_inline_a_call = true;
-                None
-            }
+            _ => None,
         }
     }
 
@@ -355,10 +349,7 @@ impl<'function> PerFunctionContext<'function> {
                 Instruction::Call { func, arguments } => match self.get_function(*func) {
                     Some(function) => match ssa.functions[&function].runtime() {
                         RuntimeType::Acir => self.inline_function(ssa, *id, function, arguments),
-                        RuntimeType::Brillig => {
-                            self.context.failed_to_inline_a_call = true;
-                            self.push_instruction(*id);
-                        }
+                        RuntimeType::Brillig => self.push_instruction(*id),
                     },
                     None => self.push_instruction(*id),
                 },
@@ -494,7 +485,7 @@ impl<'function> PerFunctionContext<'function> {
             }
             TerminatorInstruction::Return { return_values } => {
                 let return_values = vecmap(return_values, |value| self.translate_value(*value));
-                if self.inlining_main {
+                if self.inlining_entry {
                     self.context.builder.terminate_with_return(return_values.clone());
                 }
                 // Note that `translate_block` would take us back to the point at which the