feat(ssa): Pass to preprocess functions

compiler/noirc_driver/src/lib.rs

@@ -141,6 +141,10 @@ pub struct CompileOptions {
     #[arg(long)]
     pub skip_brillig_constraints_check: bool,
 
+    /// Flag to turn off preprocessing functions during SSA passes.
+    #[arg(long)]
+    pub skip_preprocess_fns: bool,
+
     /// Setting to decide on an inlining strategy for Brillig functions.
     /// A more aggressive inliner should generate larger programs but more optimized
     /// A less aggressive inliner should generate smaller programs
@@ -679,6 +683,7 @@ pub fn compile_no_check(
         emit_ssa: if options.emit_ssa { Some(context.package_build_path.clone()) } else { None },
         skip_underconstrained_check: options.skip_underconstrained_check,
         skip_brillig_constraints_check: options.skip_brillig_constraints_check,
+        skip_preprocess_fns: options.skip_preprocess_fns,
         inliner_aggressiveness: options.inliner_aggressiveness,
         max_bytecode_increase_percent: options.max_bytecode_increase_percent,
     };

compiler/noirc_evaluator/src/ssa.rs

@@ -71,6 +71,9 @@ pub struct SsaEvaluatorOptions {
     /// Skip the missing Brillig call constraints check
     pub skip_brillig_constraints_check: bool,
 
+    /// Skip preprocessing functions.
+    pub skip_preprocess_fns: bool,
+
     /// The higher the value, the more inlined Brillig functions will be.
     pub inliner_aggressiveness: i64,
 
@@ -150,15 +153,24 @@ pub(crate) fn optimize_into_acir(
 /// Run all SSA passes.
 fn optimize_all(builder: SsaBuilder, options: &SsaEvaluatorOptions) -> Result<Ssa, RuntimeError> {
     Ok(builder
-        .run_pass(Ssa::remove_unreachable_functions, "Removing Unreachable Functions")
+        .run_pass(Ssa::remove_unreachable_functions, "Removing Unreachable Functions (1st)")
         .run_pass(Ssa::defunctionalize, "Defunctionalization")
         .run_pass(Ssa::remove_paired_rc, "Removing Paired rc_inc & rc_decs")
+        .run_pass(
+            |ssa| {
+                if options.skip_preprocess_fns {
+                    return ssa;
+                }
+                ssa.preprocess_functions(options.inliner_aggressiveness)
+            },
+            "Preprocessing Functions",
+        )
         .run_pass(|ssa| ssa.inline_functions(options.inliner_aggressiveness), "Inlining (1st)")
         // Run mem2reg with the CFG separated into blocks
         .run_pass(Ssa::mem2reg, "Mem2Reg (1st)")
         .run_pass(Ssa::simplify_cfg, "Simplifying (1st)")
         .run_pass(Ssa::as_slice_optimization, "`as_slice` optimization")
-        .run_pass(Ssa::remove_unreachable_functions, "Removing Unreachable Functions")
+        .run_pass(Ssa::remove_unreachable_functions, "Removing Unreachable Functions (2nd)")
         .try_run_pass(
             Ssa::evaluate_static_assert_and_assert_constant,
             "`static_assert` and `assert_constant`",
@@ -188,7 +200,7 @@ fn optimize_all(builder: SsaBuilder, options: &SsaEvaluatorOptions) -> Result<Ss
         .run_pass(Ssa::fold_constants_using_constraints, "Constraint Folding")
         .run_pass(Ssa::make_constrain_not_equal_instructions, "Adding constrain not equal")
         .run_pass(Ssa::dead_instruction_elimination, "Dead Instruction Elimination (1st)")
-        .run_pass(Ssa::simplify_cfg, "Simplifying:")
+        .run_pass(Ssa::simplify_cfg, "Simplifying (3rd):")
         .run_pass(Ssa::array_set_optimization, "Array Set Optimizations")
         .finish())
 }

compiler/noirc_evaluator/src/ssa/ir/function.rs

@@ -12,7 +12,7 @@ use super::map::Id;
 use super::types::Type;
 use super::value::ValueId;
 
-#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash, Serialize, Deserialize)]
+#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash, Serialize, Deserialize, PartialOrd, Ord)]
 pub(crate) enum RuntimeType {
     // A noir function, to be compiled in ACIR and executed by ACVM
     Acir(InlineType),

compiler/noirc_evaluator/src/ssa/opt/defunctionalize.rs

@@ -13,7 +13,7 @@ use crate::ssa::{
     function_builder::FunctionBuilder,
     ir::{
         basic_block::BasicBlockId,
-        function::{Function, FunctionId, Signature},
+        function::{Function, FunctionId, RuntimeType, Signature},
         instruction::{BinaryOp, Instruction},
         types::{NumericType, Type},
         value::{Value, ValueId},
@@ -43,12 +43,15 @@ struct ApplyFunction {
     dispatches_to_multiple_functions: bool,
 }
 
+type Variants = BTreeMap<(Signature, RuntimeType), Vec<FunctionId>>;
+type ApplyFunctions = HashMap<(Signature, RuntimeType), ApplyFunction>;
+
 /// Performs defunctionalization on all functions
 /// This is done by changing all functions as value to be a number (FieldElement)
 /// And creating apply functions that dispatch to the correct target by runtime comparisons with constants
 #[derive(Debug, Clone)]
 struct DefunctionalizationContext {
-    apply_functions: HashMap<Signature, ApplyFunction>,
+    apply_functions: ApplyFunctions,
 }
 
 impl Ssa {
@@ -104,7 +107,7 @@ impl DefunctionalizationContext {
                         };
 
                         // Find the correct apply function
-                        let apply_function = self.get_apply_function(&signature);
+                        let apply_function = self.get_apply_function(signature, func.runtime());
 
                         // Replace the instruction with a call to apply
                         let apply_function_value_id = func.dfg.import_function(apply_function.id);
@@ -152,19 +155,21 @@ impl DefunctionalizationContext {
     }
 
     /// Returns the apply function for the given signature
-    fn get_apply_function(&self, signature: &Signature) -> ApplyFunction {
-        *self.apply_functions.get(signature).expect("Could not find apply function")
+    fn get_apply_function(&self, signature: Signature, runtime: RuntimeType) -> ApplyFunction {
+        *self.apply_functions.get(&(signature, runtime)).expect("Could not find apply function")
     }
 }
 
 /// Collects all functions used as values that can be called by their signatures
-fn find_variants(ssa: &Ssa) -> BTreeMap<Signature, Vec<FunctionId>> {
-    let mut dynamic_dispatches: BTreeSet<Signature> = BTreeSet::new();
+fn find_variants(ssa: &Ssa) -> Variants {
+    let mut dynamic_dispatches: BTreeSet<(Signature, RuntimeType)> = BTreeSet::new();
     let mut functions_as_values: BTreeSet<FunctionId> = BTreeSet::new();
 
     for function in ssa.functions.values() {
         functions_as_values.extend(find_functions_as_values(function));
-        dynamic_dispatches.extend(find_dynamic_dispatches(function));
+        dynamic_dispatches.extend(
+            find_dynamic_dispatches(function).into_iter().map(|sig| (sig, function.runtime())),
+        );
     }
 
     let mut signature_to_functions_as_value: BTreeMap<Signature, Vec<FunctionId>> = BTreeMap::new();
@@ -174,16 +179,12 @@ fn find_variants(ssa: &Ssa) -> BTreeMap<Signature, Vec<FunctionId>> {
         signature_to_functions_as_value.entry(signature).or_default().push(function_id);
     }
 
-    let mut variants = BTreeMap::new();
+    let mut variants: Variants = BTreeMap::new();
 
-    for dispatch_signature in dynamic_dispatches {
-        let mut target_fns = vec![];
-        for (target_signature, functions) in &signature_to_functions_as_value {
-            if &dispatch_signature == target_signature {
-                target_fns.extend(functions);
-            }
-        }
-        variants.insert(dispatch_signature, target_fns);
+    for (dispatch_signature, caller_runtime) in dynamic_dispatches {
+        let target_fns =
+            signature_to_functions_as_value.get(&dispatch_signature).cloned().unwrap_or_default();
+        variants.insert((dispatch_signature, caller_runtime), target_fns);
     }
 
     variants
@@ -247,22 +248,23 @@ fn find_dynamic_dispatches(func: &Function) -> BTreeSet<Signature> {
 
 fn create_apply_functions(
     ssa: &mut Ssa,
-    variants_map: BTreeMap<Signature, Vec<FunctionId>>,
-) -> HashMap<Signature, ApplyFunction> {
+    variants_map: BTreeMap<(Signature, RuntimeType), Vec<FunctionId>>,
+) -> ApplyFunctions {
     let mut apply_functions = HashMap::default();
-    for (signature, variants) in variants_map.into_iter() {
+    for ((signature, runtime), variants) in variants_map.into_iter() {
         assert!(
             !variants.is_empty(),
             "ICE: at least one variant should exist for a dynamic call {signature:?}"
         );
         let dispatches_to_multiple_functions = variants.len() > 1;
 
         let id = if dispatches_to_multiple_functions {
-            create_apply_function(ssa, signature.clone(), variants)
+            create_apply_function(ssa, signature.clone(), runtime, variants)
         } else {
             variants[0]
         };
-        apply_functions.insert(signature, ApplyFunction { id, dispatches_to_multiple_functions });
+        apply_functions
+            .insert((signature, runtime), ApplyFunction { id, dispatches_to_multiple_functions });
     }
     apply_functions
 }
@@ -275,13 +277,15 @@ fn function_id_to_field(function_id: FunctionId) -> FieldElement {
 fn create_apply_function(
     ssa: &mut Ssa,
     signature: Signature,
+    runtime: RuntimeType,
     function_ids: Vec<FunctionId>,
 ) -> FunctionId {
     assert!(!function_ids.is_empty());
     let globals = ssa.functions[&function_ids[0]].dfg.globals.clone();
     ssa.add_fn(|id| {
         let mut function_builder = FunctionBuilder::new("apply".to_string(), id);
         function_builder.set_globals(globals);
+        function_builder.set_runtime(runtime);
         let target_id = function_builder.add_parameter(Type::field());
         let params_ids = vecmap(signature.params, |typ| function_builder.add_parameter(typ));
 
@@ -339,22 +343,21 @@ fn create_apply_function(
     })
 }
 
-/// Crates a return block, if no previous return exists, it will create a final return
-/// Else, it will create a bypass return block that points to the previous return block
+/// If no previous return target exists, it will create a final return,
+/// otherwise returns the existing return block to jump to.
 fn build_return_block(
     builder: &mut FunctionBuilder,
     previous_block: BasicBlockId,
     passed_types: &[Type],
     target: Option<BasicBlockId>,
 ) -> BasicBlockId {
+    if let Some(return_block) = target {
+        return return_block;
+    }
     let return_block = builder.insert_block();
     builder.switch_to_block(return_block);
-
     let params = vecmap(passed_types, |typ| builder.add_block_parameter(return_block, typ.clone()));
-    match target {
-        None => builder.terminate_with_return(params),
-        Some(target) => builder.terminate_with_jmp(target, params),
-    }
+    builder.terminate_with_return(params);
     builder.switch_to_block(previous_block);
     return_block
 }

compiler/noirc_evaluator/src/ssa/opt/hint.rs

@@ -19,6 +19,7 @@ mod tests {
             emit_ssa: None,
             skip_underconstrained_check: true,
             skip_brillig_constraints_check: true,
+            skip_preprocess_fns: true,
             inliner_aggressiveness: 0,
             max_bytecode_increase_percent: None,
         };