Add support for arrays and vectors to IfElse in Brillig

noir-lang · Jan 31, 2025 · 9fcf348 · 9fcf348
1 parent 56ef385
commit 9fcf348
Showing 1 changed file with 52 additions and 10 deletions.
diff --git a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs
@@ -1,6 +1,6 @@
 use crate::brillig::brillig_ir::artifact::Label;
 use crate::brillig::brillig_ir::brillig_variable::{
-    type_to_heap_value_type, BrilligArray, BrilligVariable, SingleAddrVariable,
+    type_to_heap_value_type, BrilligArray, BrilligVariable, BrilligVector, SingleAddrVariable,
 };
 
 use crate::brillig::brillig_ir::registers::RegisterAllocator;
@@ -855,20 +855,62 @@ impl<'block, Registers: RegisterAllocator> BrilligBlock<'block, Registers> {
             }
             Instruction::IfElse { then_condition, then_value, else_condition: _, else_value } => {
                 let then_condition = self.convert_ssa_single_addr_value(*then_condition, dfg);
-                let then_value = self.convert_ssa_single_addr_value(*then_value, dfg);
-                let else_value = self.convert_ssa_single_addr_value(*else_value, dfg);
-                let result_register = self.variables.define_single_addr_variable(
+                let then_value = self.convert_ssa_value(*then_value, dfg);
+                let else_value = self.convert_ssa_value(*else_value, dfg);
+                let result = self.variables.define_variable(
                     self.function_context,
                     self.brillig_context,
                     dfg.instruction_results(instruction_id)[0],
                     dfg,
                 );
-                self.brillig_context.conditional_move_instruction(
-                    then_condition,
-                    then_value,
-                    else_value,
-                    result_register,
-                );
+                match (then_value, else_value) {
+                    (
+                        BrilligVariable::SingleAddr(then_address),
+                        BrilligVariable::SingleAddr(else_address),
+                    ) => {
+                        self.brillig_context.conditional_move_instruction(
+                            then_condition,
+                            then_address,
+                            else_address,
+                            result.extract_single_addr(),
+                        );
+                    }
+                    (
+                        BrilligVariable::BrilligArray(then_array),
+                        BrilligVariable::BrilligArray(else_array),
+                    ) => {
+                        // Pointer to the array which result from the if-else
+                        let pointer = self.brillig_context.allocate_register();
+                        self.brillig_context.conditional_move_instruction(
+                            then_condition,
+                            SingleAddrVariable::new_usize(then_array.pointer),
+                            SingleAddrVariable::new_usize(else_array.pointer),
+                            SingleAddrVariable::new_usize(pointer),
+                        );
+                        let if_else_array = BrilligArray { pointer, size: then_array.size };
+                        // Copy the if-else array to the result
+                        self.brillig_context
+                            .call_array_copy_procedure(if_else_array, result.extract_array());
+                    }
+                    (
+                        BrilligVariable::BrilligVector(then_vector),
+                        BrilligVariable::BrilligVector(else_vector),
+                    ) => {
+                        // Pointer to the vector which result from the if-else
+                        let pointer = self.brillig_context.allocate_register();
+                        self.brillig_context.conditional_move_instruction(
+                            then_condition,
+                            SingleAddrVariable::new_usize(then_vector.pointer),
+                            SingleAddrVariable::new_usize(else_vector.pointer),
+                            SingleAddrVariable::new_usize(pointer),
+                        );
+                        let if_else_vector = BrilligVector { pointer };
+                        // Copy the if-else vector to the result
+                        self.brillig_context
+                            .call_vector_copy_procedure(if_else_vector, result.extract_vector());
+                    }
+                    _ => unreachable!("ICE - then and else values must have the same type"),
+                }
             }
             Instruction::MakeArray { elements: array, typ } => {
                 let value_id = dfg.instruction_results(instruction_id)[0];