fix(experimental): Fix execution of match expressions with multiple branches (#7570)

Co-authored-by: Maxim Vezenov <mvezenov@gmail.com>

Author: jfecher

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

@@ -24,6 +24,7 @@ use self::{
     value::{Tree, Values},
 };
 
+use super::ir::basic_block::BasicBlockId;
 use super::ir::dfg::GlobalsGraph;
 use super::ir::instruction::ErrorType;
 use super::ir::types::NumericType;
@@ -767,7 +768,15 @@ impl FunctionContext<'_> {
         let tag = self.enum_tag(&variable);
         let tag_type = self.builder.type_of_value(tag).unwrap_numeric();
 
-        let end_block = self.builder.insert_block();
+        let make_end_block = |this: &mut Self| -> (BasicBlockId, Values) {
+            let block = this.builder.insert_block();
+            let results = Self::map_type(&match_expr.typ, |typ| {
+                this.builder.add_block_parameter(block, typ).into()
+            });
+            (block, results)
+        };
+
+        let (end_block, end_results) = make_end_block(self);
 
         // Optimization: if there is no default case we can jump directly to the last case
         // when finished with the previous case instead of using a jmpif with an unreachable
@@ -778,6 +787,8 @@ impl FunctionContext<'_> {
             match_expr.cases.len() - 1
         };
 
+        let mut blocks_to_merge = Vec::with_capacity(last_case);
+
         for i in 0..last_case {
             let case = &match_expr.cases[i];
             let variant_tag = self.variant_index_value(&case.constructor, tag_type)?;
@@ -790,28 +801,70 @@ impl FunctionContext<'_> {
             self.builder.switch_to_block(case_block);
             self.bind_case_arguments(variable.clone(), case);
             let results = self.codegen_expression(&case.branch)?.into_value_list(self);
-            self.builder.terminate_with_jmp(end_block, results);
+
+            // Each branch will jump to a different end block for now. We have to merge them all
+            // later since SSA doesn't support more than two blocks jumping to the same end block.
+            let local_end_block = make_end_block(self);
+            self.builder.terminate_with_jmp(local_end_block.0, results);
+            blocks_to_merge.push(local_end_block);
 
             self.builder.switch_to_block(else_block);
         }
 
+        let (last_local_end_block, last_results) = make_end_block(self);
+        blocks_to_merge.push((last_local_end_block, last_results));
+
         if let Some(branch) = &match_expr.default_case {
             let results = self.codegen_expression(branch)?.into_value_list(self);
-            self.builder.terminate_with_jmp(end_block, results);
+            self.builder.terminate_with_jmp(last_local_end_block, results);
         } else {
             // If there is no default case, assume we saved the last case from the
             // last_case optimization above
             let case = match_expr.cases.last().unwrap();
             self.bind_case_arguments(variable, case);
             let results = self.codegen_expression(&case.branch)?.into_value_list(self);
-            self.builder.terminate_with_jmp(end_block, results);
+            self.builder.terminate_with_jmp(last_local_end_block, results);
+        }
+
+        // Merge blocks as last-in first-out:
+        //
+        // local_end_block0-----------------------------------------\
+        //                                                           end block
+        //                                                          /
+        // local_end_block1---------------------\                  /
+        //                                       new merge block2-/
+        // local_end_block2--\                  /
+        //                    new merge block1-/
+        // local_end_block3--/
+        //
+        // This is necessary since SSA panics during flattening if we immediately
+        // try to jump directly to end block instead: https://github.com/noir-lang/noir/issues/7323.
+        //
+        // It'd also be more efficient to merge them tournament-bracket style but that
+        // also leads to panics during flattening for similar reasons.
+        while let Some((block, results)) = blocks_to_merge.pop() {
+            self.builder.switch_to_block(block);
+
+            if let Some((block2, results2)) = blocks_to_merge.pop() {
+                // Merge two blocks in the queue together
+                let (new_merge, new_merge_results) = make_end_block(self);
+                blocks_to_merge.push((new_merge, new_merge_results));
+
+                let results = results.into_value_list(self);
+                self.builder.terminate_with_jmp(new_merge, results);
+
+                self.builder.switch_to_block(block2);
+                let results2 = results2.into_value_list(self);
+                self.builder.terminate_with_jmp(new_merge, results2);
+            } else {
+                // Finally done, jump to the end
+                let results = results.into_value_list(self);
+                self.builder.terminate_with_jmp(end_block, results);
+            }
         }
 
         self.builder.switch_to_block(end_block);
-        let result = Self::map_type(&match_expr.typ, |typ| {
-            self.builder.add_block_parameter(end_block, typ).into()
-        });
-        Ok(result)
+        Ok(end_results)
     }
 
     fn variant_index_value(

test_programs/execution_success/regression_7323/Nargo.toml

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

test_programs/execution_success/regression_7323/Prover.toml

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

test_programs/execution_success/regression_7323/src/main.nr

@@ -0,0 +1,11 @@
+// This program lead to panics previously due to the compiler lowering it to multiple blocks
+// which all jumped to the same end block. It runs now due to the compiler lowering to the
+// equivalent of a nested series of if-else instead.
+fn main(x: Field) {
+    match x {
+        1 => panic(f"Branch 1 should not be taken"),
+        2 => panic(f"Branch 2 should not be taken"),
+        3 => panic(f"Branch 3 should not be taken"),
+        _ => (),
+    }
+}

test_programs/rebuild.sh

@@ -47,7 +47,8 @@ process_dir() {
 
 export -f process_dir
 
-excluded_dirs=("workspace" "workspace_default_member")
+# Reactive `regression_7323` once enums are ready
+excluded_dirs=("workspace" "workspace_default_member" "regression_7323")
 current_dir=$(pwd)
 base_path="$current_dir/execution_success"
 dirs_to_process=()

tooling/debugger/ignored-tests.txt

@@ -5,5 +5,6 @@ macros
 reference_counts
 references
 regression_4709
+regression_7323
 reference_only_used_as_alias
 brillig_rc_regression_6123

tooling/nargo_cli/src/cli/compile_cmd.rs

@@ -325,6 +325,7 @@ mod tests {
     use nargo::ops::compile_program;
     use nargo_toml::PackageSelection;
     use noirc_driver::{CompileOptions, CrateName};
+    use noirc_frontend::elaborator::UnstableFeature;
     use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
 
     use crate::cli::test_cmd::formatters::diagnostic_to_string;
@@ -403,12 +404,17 @@ mod tests {
             let binary_packages = workspace.into_iter().filter(|package| package.is_binary());
 
             for package in binary_packages {
+                let options = CompileOptions {
+                    unstable_features: vec![UnstableFeature::Enums],
+                    ..Default::default()
+                };
+
                 let (program_0, _warnings) = compile_program(
                     &file_manager,
                     &parsed_files,
                     workspace,
                     package,
-                    &CompileOptions::default(),
+                    &options,
                     None,
                 )
                 .unwrap_or_else(|err| {