fix(defunctionalize): Higher order functions (HOF) dynamic dispatch and HOFs in arrays

compiler/noirc_evaluator/src/ssa/opt/flatten_cfg/branch_analysis.rs

@@ -104,7 +104,7 @@
             self.find_join_point(successors.next().unwrap())
         } else if successors.len() == 0 {
             unreachable!(
                 "return encountered before a join point was found. This can only happen if early-return was added to the language without implementing it by jmping to a join block first"
             )
         } else {
             unreachable!("A block can only have 0, 1, or 2 successors");
@@ -119,6 +119,7 @@
         function_builder::FunctionBuilder,
         ir::{cfg::ControlFlowGraph, map::Id, types::Type},
         opt::flatten_cfg::branch_analysis::find_branch_ends,
+        ssa_gen::Ssa,
     };
 
     #[test]
@@ -261,4 +262,61 @@
         let cfg = ControlFlowGraph::with_function(function);
         find_branch_ends(function, &cfg);
     }
+
+    #[test]
+    fn apply_function() {
+        // Make sure that our dynamic dispatch function created during defunctionalization
+        // passes branch analysis.
+        let src = "
+        acir(inline_always) fn apply f5 {
+          b0(v0: Field, v1: u32):
+            v4 = eq v0, Field 2
+            jmpif v4 then: b3, else: b2
+          b1(v2: u32):
+            return v2
+          b2():
+            v9 = eq v0, Field 3
+            jmpif v9 then: b6, else: b5
+          b3():
+            v6 = call f2(v1) -> u32
+            jmp b4(v6)
+          b4(v7: u32):
+            jmp b10(v7)
+          b5():
+            constrain v0 == Field 4
+            v15 = call f4(v1) -> u32
+            jmp b8(v15)
+          b6():
+            v11 = call f3(v1) -> u32
+            jmp b7(v11)
+          b7(v12: u32):
+            jmp b9(v12)
+          b8(v16: u32):
+            jmp b9(v16)
+          b9(v17: u32):
+            jmp b10(v17)
+          b10(v18: u32):
+            jmp b1(v18)
+        }
+        acir(inline) fn lambda f2 {
+          b0(v0: u32):
+            return v0
+        }
+        acir(inline) fn lambda f3 {
+          b0(v0: u32):
+            v2 = add v0, u32 1
+            return v2
+        }
+        acir(inline) fn lambda f4 {
+          b0(v0: u32):
+            v2 = add v0, u32 2
+            return v2
+        }
+        ";
+
+        let ssa = Ssa::from_str(src).unwrap();
+        let function = ssa.main();
+        let cfg = ControlFlowGraph::with_function(function);
+        find_branch_ends(function, &cfg);
+    }
 }

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

@@ -82,7 +82,12 @@ fn used_functions(func: &Function) -> BTreeSet<FunctionId> {
         for instruction_id in block.instructions() {
             let instruction = &func.dfg[*instruction_id];
 
-            if matches!(instruction, Instruction::Store { .. } | Instruction::Call { .. }) {
+            if matches!(
+                instruction,
+                Instruction::Store { .. }
+                    | Instruction::Call { .. }
+                    | Instruction::MakeArray { .. }
+            ) {
                 instruction.for_each_value(&mut find_functions);
             }
         }
@@ -177,4 +182,33 @@ mod tests {
         // It should not remove anything.
         assert_normalized_ssa_equals(ssa, src);
     }
+
+    #[test]
+    fn keep_functions_used_in_array() {
+        // f1 and f2 are used within an array. Thus, we do not want to remove them.
+        let src = r#"
+        acir(inline) fn main f0 {
+          b0(v0: u32):
+            v5 = make_array [f1, f2] : [function; 2]
+            v7 = lt v0, u32 4
+            constrain v7 == u1 1, "Index out of bounds"
+            v9 = array_get v5, index v0 -> function
+            call v9()
+            return
+        }
+        acir(inline) fn lambda f1 {
+          b0():
+            return
+        }
+        acir(inline) fn lambda f2 {
+          b0():
+            return
+        }
+        "#;
+
+        let ssa = Ssa::from_str(src).unwrap();
+        let ssa = ssa.remove_unreachable_functions();
+
+        assert_normalized_ssa_equals(ssa, src);
+    }
 }

test_programs/execution_success/lambda_from_array/Nargo.toml

@@ -0,0 +1,6 @@
+[package]
+name = "lambda_from_array"
+type = "bin"
+authors = [""]
+
+[dependencies]

test_programs/execution_success/lambda_from_array/Prover.toml

@@ -0,0 +1 @@
+x = 1

test_programs/execution_success/lambda_from_array/src/main.nr

@@ -0,0 +1,70 @@
+// Many parts of the code here are regressions from issue #5503 (https://github.com/noir-lang/noir/issues/5503)
+fn main(x: u32) {
+    lambdas_in_array_literal(x - 1);
+    lambdas_in_array_literal(x);
+    lambdas_in_array_literal(x + 2);
+    lambdas_in_array_literal(x + 1);
+
+    lambdas_in_slice_literal(x - 1);
+    lambdas_in_slice_literal(x);
+    lambdas_in_slice_literal(x + 1);
+    lambdas_in_slice_literal(x + 2);
+
+    functions_in_array_literal(x - 1);
+    functions_in_array_literal(x);
+    functions_in_slice_literal(x - 1);
+    functions_in_slice_literal(x);
+
+    let example_lambda: fn(u8) -> u8 = |x| x + 1;
+    let lambdas: [fn(u8) -> u8; 8] = [example_lambda; 8];
+    println(lambdas[0](5));
+    // Dynamic dispatch
+    println(lambdas[x - 1](5));
+
+    let lambdas: [fn(()) -> (); 1] = [|_: ()| {}];
+    lambdas[0](());
+    lambdas[x - 1](());
+
+    // Also check against slices
+    let lambdas: [fn(()) -> ()] = &[|_: ()| {}];
+    lambdas[0](());
+    lambdas[x - 1](());
+
+    // Still panics when there are no other lambdas
+    // This should fail either way as we are attempting to access an empty array at zero
+    // let lambdas: [fn(()) -> (); 0] = [];
+    // lambdas[0](());
+}
+
+fn lambdas_in_array_literal(x: u32) {
+    let xs = [|| println("hi"), || println("bye"), || println("wow"), || println("big")];
+    (xs[x])();
+}
+
+fn lambdas_in_slice_literal(x: u32) {
+    let xs = &[|| println("hi"), || println("bye"), || println("big"), || println("wow")];
+    (xs[x])();
+}
+
+fn functions_in_array_literal(x: u32) {
+    let xs = [foo, bar];
+    (xs[x])();
+}
+
+fn functions_in_slice_literal(x: u32) {
+    let xs = &[baz, qux];
+    (xs[x])();
+}
+
+fn foo() {
+    println("hi");
+}
+fn bar() {
+    println("bye");
+}
+fn baz() {
+    println("hi");
+}
+fn qux() {
+    println("bye");
+}

test_programs/execution_success/lambda_from_dynamic_if/Nargo.toml

@@ -0,0 +1,6 @@
+[package]
+name = "lambda_from_dynamic_if"
+type = "bin"
+authors = [""]
+
+[dependencies]

test_programs/execution_success/lambda_from_dynamic_if/Prover.toml

@@ -0,0 +1 @@
+x = 1

test_programs/execution_success/lambda_from_dynamic_if/src/main.nr

@@ -0,0 +1,31 @@
+fn main(x: u32) {
+    lambdas_no_inputs(x);
+    println(lambdas_with_input_and_return_values(x - 1));
+    println(lambdas_with_input_and_return_values(x));
+    println(lambdas_with_input_and_return_values(x + 1));
+}
+fn lambdas_no_inputs(x: u32) {
+    let f1 = || println("hi");
+    let f2 = || println("bye");
+    let f = if x == 0 {
+        f1
+    } else if x == 1 {
+        f2
+    } else {
+        panic(f"!")
+    };
+    f();
+}
+fn lambdas_with_input_and_return_values(x: u32) -> u32 {
+    let f1 = |value: u32| value;
+    let f2 = |value: u32| value + 1;
+    let f3 = |value: u32| value + 2;
+    let f = if x == 0 {
+        f1
+    } else if x == 1 {
+        f2
+    } else {
+        f3
+    };
+    f(x)
+}