From 6794e3e9966d80dba25623c84698e816e75aa300 Mon Sep 17 00:00:00 2001
From: guipublic <guipublic@gmail.com>
Date: Tue, 28 Jan 2025 14:50:05 +0000
Subject: [PATCH 01/13] simplify simple conditionals for brillig

---
 compiler/noirc_evaluator/src/ssa.rs           |   1 +
 .../src/ssa/opt/basic_conditional.rs          | 280 ++++++++++++++++++
 .../src/ssa/opt/flatten_cfg.rs                |  76 +++--
 compiler/noirc_evaluator/src/ssa/opt/mod.rs   |   1 +
 4 files changed, 338 insertions(+), 20 deletions(-)
 create mode 100644 compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs

diff --git a/compiler/noirc_evaluator/src/ssa.rs b/compiler/noirc_evaluator/src/ssa.rs
index 4cefce1d647..82da6084760 100644
--- a/compiler/noirc_evaluator/src/ssa.rs
+++ b/compiler/noirc_evaluator/src/ssa.rs
@@ -190,6 +190,7 @@ fn optimize_all(builder: SsaBuilder, options: &SsaEvaluatorOptions) -> Result<Ss
         )
         .run_pass(Ssa::remove_if_else, "Remove IfElse")
         .run_pass(Ssa::fold_constants, "Constant Folding")
+        .run_pass(Ssa::flatten_basic_conditionals, "Simplify conditionals for unconstrained")
         .run_pass(Ssa::remove_enable_side_effects, "EnableSideEffectsIf removal")
         .run_pass(Ssa::fold_constants_using_constraints, "Constraint Folding")
         .run_pass(Ssa::make_constrain_not_equal_instructions, "Adding constrain not equal")
diff --git a/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs b/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
new file mode 100644
index 00000000000..35f765e8276
--- /dev/null
+++ b/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
@@ -0,0 +1,280 @@
+use std::collections::HashSet;
+
+use acvm::AcirField;
+use fxhash::FxHashMap;
+use iter_extended::vecmap;
+
+use crate::ssa::{
+    ir::{
+        basic_block::BasicBlockId,
+        cfg::ControlFlowGraph,
+        dfg::DataFlowGraph,
+        function::{Function, FunctionId, RuntimeType},
+        instruction::{BinaryOp, Instruction, TerminatorInstruction},
+    },
+    Ssa,
+};
+
+use super::flatten_cfg::Context;
+#[derive(Debug)]
+struct BasicConditional {
+    block_entry: BasicBlockId,
+    block_then: Option<BasicBlockId>,
+    block_else: Option<BasicBlockId>,
+    block_exit: BasicBlockId,
+}
+
+impl Ssa {
+    #[tracing::instrument(level = "trace", skip(self))]
+    pub(crate) fn flatten_basic_conditionals(mut self) -> Ssa {
+        // Retrieve the 'no_predicates' attribute of the functions in a map, to avoid problems with borrowing
+        let mut no_predicates = FxHashMap::default();
+        for function in self.functions.values() {
+            no_predicates.insert(function.id(), function.is_no_predicates());
+        }
+        for function in self.functions.values_mut() {
+            flatten_function(function, &mut no_predicates);
+        }
+        self
+    }
+}
+
+// Returns the blocks of the simple conditional sub-graph whose input block is the entry.
+// Returns None if the input block is not the entry block of a simple conditional.
+fn is_conditional(
+    block: BasicBlockId,
+    cfg: &ControlFlowGraph,
+    function: &Function,
+) -> Option<BasicConditional> {
+    // jump overhead is the cost for doing the conditional and jump around the blocks
+    // I use 10 as a rough estimate, real cost is less.
+    let jump_overhead = 10;
+    let mut successors = cfg.successors(block);
+    let mut result = None;
+    // a conditional must have 2 branches
+    if successors.len() == 2 {
+        let left = successors.next().unwrap();
+        let right = successors.next().unwrap();
+        let mut left_successors = cfg.successors(left);
+        let mut right_successors = cfg.successors(right);
+        let left_successors_len = left_successors.len();
+        let right_successors_len = right_successors.len();
+        let next_left = left_successors.next();
+        let next_right = right_successors.next();
+        if next_left == Some(block) || next_right == Some(block) {
+            // this is a loop, not a conditional
+            return None;
+        }
+        if left_successors_len == 1 && right_successors_len == 1 && next_left == next_right {
+            // The branches join on one block so it is a non-nested conditional
+            let cost_left = block_cost(left, &function.dfg);
+            let cost_right = block_cost(right, &function.dfg);
+            // For the flattening to be valuable, we compare the cost of the flattened code with the average cost of the 2 branches,
+            // including an overhead to take into account the jumps between the blocks.
+            let cost = cost_right.saturating_add(cost_left);
+            if cost < cost / 2 + jump_overhead {
+                result = Some(BasicConditional {
+                    block_entry: block,
+                    block_then: Some(left),
+                    block_else: Some(right),
+                    block_exit: next_left.unwrap(),
+                });
+            }
+        } else if left_successors_len == 1 && next_left == Some(right) {
+            // Left branch joins the right branch, it is a if/then statement with no else
+            // I am not sure whether this case can happen, but it is not difficult to handle it
+            let cost = block_cost(left, &function.dfg);
+            if cost < cost / 2 + jump_overhead {
+                result = Some(BasicConditional {
+                    block_entry: block,
+                    block_then: Some(left),
+                    block_else: None,
+                    block_exit: right,
+                });
+            }
+        } else if right_successors_len == 1 && next_right == Some(left) {
+            // Right branch joins the right branch, it is a if/else statement with no then
+            // I am not sure whether this case can happen, but it is not difficult to handle it
+            let cost = block_cost(right, &function.dfg);
+            if cost < cost / 2 + jump_overhead {
+                result = Some(BasicConditional {
+                    block_entry: block,
+                    block_then: None,
+                    block_else: Some(right),
+                    block_exit: right,
+                });
+            }
+        }
+    }
+    // A conditional exit would have exactly 2 predecessors
+    result.filter(|result| cfg.predecessors(result.block_exit).len() == 2)
+}
+
+/// Computes a cost estimate of a basic block
+/// returns u32::MAX if the block has side-effect instructions
+fn block_cost(block: BasicBlockId, dfg: &DataFlowGraph) -> u32 {
+    let mut cost: u32 = 0;
+
+    for instruction in dfg[block].instructions() {
+        let instruction_cost = match &dfg[*instruction] {
+            Instruction::Binary(binary) => {
+                match binary.operator {
+                    BinaryOp::Add { unchecked }
+                    | BinaryOp::Sub { unchecked }
+                    | BinaryOp::Mul { unchecked } => if unchecked { 1 } else { return u32::MAX },
+                    BinaryOp::Div
+                    | BinaryOp::Mod => return u32::MAX,
+                    BinaryOp::Eq => 2,
+                    BinaryOp::Lt => 5,
+                    BinaryOp::And
+                    | BinaryOp::Or
+                    | BinaryOp::Xor => 1,   //todo
+                    BinaryOp::Shl
+                    | BinaryOp::Shr => return u32::MAX,
+                }
+            },
+            Instruction::Cast(_, _) => 1,//TODO check if this instruction can fail
+            Instruction::Not(_) => 1,
+            Instruction::Truncate { .. } => 7,
+
+            Instruction::Constrain(_,_,_)
+            | Instruction::ConstrainNotEqual(_,_,_)
+            | Instruction::RangeCheck { .. }
+            | Instruction::Call { .. }  //TODO support calls with no-predicate set to true
+            |      Instruction::Load { .. }
+            | Instruction::Store { .. }
+            | Instruction::ArraySet { .. } => {return u32::MAX},
+
+            Instruction::ArrayGet { array, index  } => {
+                // check if index is in bound
+                let index = dfg.get_numeric_constant(*index);
+                let mut ok_bound = false;
+                if let Some(index) = index {
+                    let len = dfg.try_get_array_length(*array).unwrap();
+                    if index.to_u128() < len as u128 {
+                        ok_bound = true;
+                    }
+                }
+                if ok_bound {
+                    1
+                } else {
+                    return u32::MAX;
+                }
+            },
+
+            Instruction::Allocate => 0,
+            Instruction::EnableSideEffectsIf { .. } => 0,
+            Instruction::IncrementRc { .. } => 0,
+            Instruction::DecrementRc { .. } => 0,
+            Instruction::IfElse { .. } => 1,
+            Instruction::MakeArray { .. } => 0,
+            Instruction::Noop => 0,
+        };
+        cost += instruction_cost;
+    }
+    cost
+}
+
+/// Identifies all simple conditionals in the function and flattens them
+fn flatten_function(function: &mut Function, no_predicates: &mut FxHashMap<FunctionId, bool>) {
+    // This pass is dedicated to brillig functions
+    if !matches!(function.runtime(), RuntimeType::Brillig(_)) {
+        return;
+    }
+    let cfg = ControlFlowGraph::with_function(function);
+    let mut stack = vec![function.entry_block()];
+    let mut processed = HashSet::new();
+    let mut conditionals = Vec::new();
+
+    // 1. Process all blocks of the cfg, starting from the root and following the successors
+    while let Some(block) = stack.pop() {
+        // Avoid cycles
+        if processed.contains(&block) {
+            continue;
+        }
+        processed.insert(block);
+
+        // Identify the simple conditionals
+        if let Some(conditional) = is_conditional(block, &cfg, function) {
+            // no need to check the branches, process the join block directly
+            stack.push(conditional.block_exit);
+            conditionals.push(conditional);
+        } else {
+            stack.extend(cfg.successors(block));
+        }
+    }
+    dbg!(&conditionals.len());
+    // 2. Flatten all simple conditionals
+    // process in reverse order because the process of a conditional does not impact the previous ones
+    conditionals.reverse();
+    for simple_if in conditionals {
+        Context::flatten_single(&simple_if, function, no_predicates);
+    }
+}
+
+impl<'f> Context<'f> {
+    fn flatten_single(
+        conditional: &BasicConditional,
+        function: &mut Function,
+        no_predicates: &mut FxHashMap<FunctionId, bool>,
+    ) {
+        let mut branch_ends = FxHashMap::default();
+        branch_ends.insert(conditional.block_entry, conditional.block_exit);
+        let cfg = ControlFlowGraph::with_function(function);
+        let mut context = Context::new(function, cfg, branch_ends, conditional.block_entry);
+        context.no_predicate = true;
+
+        context.flatten_single_conditional(conditional, no_predicates);
+    }
+
+    fn flatten_single_conditional(
+        &mut self,
+        conditional: &BasicConditional,
+        no_predicates: &mut FxHashMap<FunctionId, bool>,
+    ) {
+        // Manually inline then, else and exit into the entry block
+        let queue = vec![];
+
+        self.inline_block(conditional.block_entry, no_predicates);
+        let _to_process = self.handle_terminator(conditional.block_entry, &queue);
+        if let Some(then) = conditional.block_then {
+            self.inline_block(then, no_predicates);
+            self.handle_terminator(then, &queue);
+        }
+        if let Some(else_b) = conditional.block_else {
+            self.inline_block(else_b, no_predicates);
+            self.handle_terminator(else_b, &queue);
+        }
+        self.inline_block(conditional.block_exit, no_predicates);
+        // Manually set the terminator of the entry block to the one of the exit block
+        let terminator =
+            self.inserter.function.dfg[conditional.block_exit].terminator().unwrap().clone();
+        let new_terminator = match terminator {
+            TerminatorInstruction::JmpIf {
+                condition,
+                then_destination,
+                else_destination,
+                call_stack,
+            } => {
+                let condition = self.inserter.resolve(condition);
+                TerminatorInstruction::JmpIf {
+                    condition,
+                    then_destination,
+                    else_destination,
+                    call_stack,
+                }
+            }
+            TerminatorInstruction::Jmp { destination, arguments, call_stack } => {
+                let arguments = vecmap(arguments, |value| self.inserter.resolve(value));
+                TerminatorInstruction::Jmp { destination, arguments, call_stack }
+            }
+            TerminatorInstruction::Return { return_values, call_stack } => {
+                let return_values = vecmap(return_values, |value| self.inserter.resolve(value));
+                TerminatorInstruction::Return { return_values, call_stack }
+            }
+        };
+        self.inserter.function.dfg.set_block_terminator(conditional.block_entry, new_terminator);
+
+        self.inserter.map_data_bus_in_place();
+    }
+}
diff --git a/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs b/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs
index 76f8495c009..d4e2968f113 100644
--- a/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs
+++ b/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs
@@ -176,12 +176,15 @@ impl Ssa {
     }
 }
 
-struct Context<'f> {
-    inserter: FunctionInserter<'f>,
+pub(crate) struct Context<'f> {
+    pub(crate) inserter: FunctionInserter<'f>,
 
     /// This ControlFlowGraph is the graph from before the function was modified by this flattening pass.
     cfg: ControlFlowGraph,
 
+    /// Target block of the flattening
+    target_block: BasicBlockId,
+
     /// Maps start of branch -> end of branch
     branch_ends: HashMap<BasicBlockId, BasicBlockId>,
 
@@ -213,6 +216,10 @@ struct Context<'f> {
     /// us from unnecessarily inserting extra instructions, and keeps ids unique which
     /// helps simplifications.
     not_instructions: HashMap<ValueId, ValueId>,
+
+    /// Flag to tell the context to not issue 'enable_side_effect' instructions during flattening.
+    /// This should set to true only by flatten_single(), when none instruction is known to fail.
+    pub(crate) no_predicate: bool,
 }
 
 #[derive(Clone)]
@@ -249,6 +256,7 @@ fn flatten_function_cfg(function: &mut Function, no_predicates: &HashMap<Functio
     }
     let cfg = ControlFlowGraph::with_function(function);
     let branch_ends = branch_analysis::find_branch_ends(function, &cfg);
+    let target_block = function.entry_block();
 
     let mut context = Context {
         inserter: FunctionInserter::new(function),
@@ -258,17 +266,38 @@ fn flatten_function_cfg(function: &mut Function, no_predicates: &HashMap<Functio
         arguments_stack: Vec::new(),
         local_allocations: HashSet::default(),
         not_instructions: HashMap::default(),
+        target_block,
+        no_predicate: false,
     };
     context.flatten(no_predicates);
 }
 
 impl<'f> Context<'f> {
-    fn flatten(&mut self, no_predicates: &HashMap<FunctionId, bool>) {
+    pub(crate) fn new(
+        function: &'f mut Function,
+        cfg: ControlFlowGraph,
+        branch_ends: HashMap<BasicBlockId, BasicBlockId>,
+        target_block: BasicBlockId,
+    ) -> Self {
+        Context {
+            inserter: FunctionInserter::new(function),
+            cfg,
+            branch_ends,
+            condition_stack: Vec::new(),
+            arguments_stack: Vec::new(),
+            local_allocations: HashSet::default(),
+            not_instructions: HashMap::default(),
+            target_block,
+            no_predicate: false,
+        }
+    }
+
+    pub(crate) fn flatten(&mut self, no_predicates: &HashMap<FunctionId, bool>) {
         // Flatten the CFG by inlining all instructions from the queued blocks
         // until all blocks have been flattened.
         // We follow the terminator of each block to determine which blocks to
         // process next
-        let mut queue = vec![self.inserter.function.entry_block()];
+        let mut queue = vec![self.target_block];
         while let Some(block) = queue.pop() {
             self.inline_block(block, no_predicates);
             let to_process = self.handle_terminator(block, &queue);
@@ -318,10 +347,14 @@ impl<'f> Context<'f> {
         result
     }
 
-    // Inline all instructions from the given block into the entry block, and track slice capacities
-    fn inline_block(&mut self, block: BasicBlockId, no_predicates: &HashMap<FunctionId, bool>) {
-        if self.inserter.function.entry_block() == block {
-            // we do not inline the entry block into itself
+    // Inline all instructions from the given block into the target block, and track slice capacities
+    pub(crate) fn inline_block(
+        &mut self,
+        block: BasicBlockId,
+        no_predicates: &HashMap<FunctionId, bool>,
+    ) {
+        if self.target_block == block {
+            // we do not inline the target block into itself
             // for the outer block before we start inlining
             return;
         }
@@ -354,7 +387,7 @@ impl<'f> Context<'f> {
     /// For a normal block, it would be its successor
     /// For blocks related to a conditional statement, we ensure to process
     /// the 'then-branch', then the 'else-branch' (if it exists), and finally the end block
-    fn handle_terminator(
+    pub(crate) fn handle_terminator(
         &mut self,
         block: BasicBlockId,
         work_list: &[BasicBlockId],
@@ -388,9 +421,9 @@ impl<'f> Context<'f> {
                 let return_values =
                     vecmap(return_values.clone(), |value| self.inserter.resolve(value));
                 let new_return = TerminatorInstruction::Return { return_values, call_stack };
-                let entry = self.inserter.function.entry_block();
+                let target = self.target_block;
 
-                self.inserter.function.dfg.set_block_terminator(entry, new_return);
+                self.inserter.function.dfg.set_block_terminator(target, new_return);
                 vec![]
             }
         }
@@ -544,7 +577,7 @@ impl<'f> Context<'f> {
         } else {
             self.inserter.function.dfg.make_constant(FieldElement::zero(), NumericType::bool())
         };
-        let block = self.inserter.function.entry_block();
+        let block = self.target_block;
 
         // Cannot include this in the previous vecmap since it requires exclusive access to self
         let args = vecmap(args, |(then_arg, else_arg)| {
@@ -568,11 +601,11 @@ impl<'f> Context<'f> {
         destination
     }
 
-    /// Insert a new instruction into the function's entry block.
+    /// Insert a new instruction into the target block.
     /// Unlike push_instruction, this function will not map any ValueIds.
     /// within the given instruction, nor will it modify self.values in any way.
     fn insert_instruction(&mut self, instruction: Instruction, call_stack: CallStackId) -> ValueId {
-        let block = self.inserter.function.entry_block();
+        let block = self.target_block;
         self.inserter
             .function
             .dfg
@@ -580,7 +613,7 @@ impl<'f> Context<'f> {
             .first()
     }
 
-    /// Inserts a new instruction into the function's entry block, using the given
+    /// Inserts a new instruction into the target block, using the given
     /// control type variables to specify result types if needed.
     /// Unlike push_instruction, this function will not map any ValueIds.
     /// within the given instruction, nor will it modify self.values in any way.
@@ -590,7 +623,7 @@ impl<'f> Context<'f> {
         ctrl_typevars: Option<Vec<Type>>,
         call_stack: CallStackId,
     ) -> InsertInstructionResult {
-        let block = self.inserter.function.entry_block();
+        let block = self.target_block;
         self.inserter.function.dfg.insert_instruction_and_results(
             instruction,
             block,
@@ -600,11 +633,14 @@ impl<'f> Context<'f> {
     }
 
     /// Checks the branch condition on the top of the stack and uses it to build and insert an
-    /// `EnableSideEffectsIf` instruction into the entry block.
+    /// `EnableSideEffectsIf` instruction into the target block.
     ///
     /// If the stack is empty, a "true" u1 constant is taken to be the active condition. This is
     /// necessary for re-enabling side-effects when re-emerging to a branch depth of 0.
     fn insert_current_side_effects_enabled(&mut self) {
+        if self.no_predicate {
+            return;
+        }
         let condition = match self.get_last_condition() {
             Some(cond) => cond,
             None => {
@@ -616,7 +652,7 @@ impl<'f> Context<'f> {
         self.insert_instruction_with_typevars(enable_side_effects, None, call_stack);
     }
 
-    /// Push the given instruction to the end of the entry block of the current function.
+    /// Push the given instruction to the end of the target block of the current function.
     ///
     /// Note that each ValueId of the instruction will be mapped via self.inserter.resolve.
     /// As a result, the instruction that will be pushed will actually be a new instruction
@@ -631,8 +667,8 @@ impl<'f> Context<'f> {
         let instruction = self.handle_instruction_side_effects(instruction, call_stack);
 
         let instruction_is_allocate = matches!(&instruction, Instruction::Allocate);
-        let entry = self.inserter.function.entry_block();
-        let results = self.inserter.push_instruction_value(instruction, id, entry, call_stack);
+        let results =
+            self.inserter.push_instruction_value(instruction, id, self.target_block, call_stack);
 
         // Remember an allocate was created local to this branch so that we do not try to merge store
         // values across branches for it later.
diff --git a/compiler/noirc_evaluator/src/ssa/opt/mod.rs b/compiler/noirc_evaluator/src/ssa/opt/mod.rs
index 44796e2531e..a2a8593658d 100644
--- a/compiler/noirc_evaluator/src/ssa/opt/mod.rs
+++ b/compiler/noirc_evaluator/src/ssa/opt/mod.rs
@@ -7,6 +7,7 @@
 mod array_set;
 mod as_slice_length;
 mod assert_constant;
+mod basic_conditional;
 mod constant_folding;
 mod defunctionalize;
 mod die;

From d3418150cf54b32f4dae975a33e3730b05cbfc70 Mon Sep 17 00:00:00 2001
From: guipublic <guipublic@gmail.com>
Date: Wed, 29 Jan 2025 09:03:14 +0000
Subject: [PATCH 02/13] fix the manual inline process

---
 .../src/ssa/opt/basic_conditional.rs          | 109 +++++++++++++-----
 1 file changed, 82 insertions(+), 27 deletions(-)

diff --git a/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs b/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
index 35f765e8276..cd0605a1154 100644
--- a/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
+++ b/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
@@ -16,7 +16,7 @@ use crate::ssa::{
 };
 
 use super::flatten_cfg::Context;
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 struct BasicConditional {
     block_entry: BasicBlockId,
     block_then: Option<BasicBlockId>,
@@ -73,36 +73,78 @@ fn is_conditional(
             // including an overhead to take into account the jumps between the blocks.
             let cost = cost_right.saturating_add(cost_left);
             if cost < cost / 2 + jump_overhead {
-                result = Some(BasicConditional {
-                    block_entry: block,
-                    block_then: Some(left),
-                    block_else: Some(right),
-                    block_exit: next_left.unwrap(),
-                });
+                if let Some(TerminatorInstruction::JmpIf {
+                    condition: _,
+                    then_destination,
+                    else_destination,
+                    call_stack: _,
+                }) = function.dfg[block].terminator()
+                {
+                    result = Some(BasicConditional {
+                        block_entry: block,
+                        block_then: Some(*then_destination),
+                        block_else: Some(*else_destination),
+                        block_exit: next_left.unwrap(),
+                    });
+                }
             }
         } else if left_successors_len == 1 && next_left == Some(right) {
             // Left branch joins the right branch, it is a if/then statement with no else
             // I am not sure whether this case can happen, but it is not difficult to handle it
             let cost = block_cost(left, &function.dfg);
             if cost < cost / 2 + jump_overhead {
-                result = Some(BasicConditional {
-                    block_entry: block,
-                    block_then: Some(left),
-                    block_else: None,
-                    block_exit: right,
-                });
+                if let Some(TerminatorInstruction::JmpIf {
+                    condition: _,
+                    then_destination,
+                    else_destination,
+                    call_stack: _,
+                }) = function.dfg[block].terminator()
+                {
+                    if left == *then_destination {
+                        result = Some(BasicConditional {
+                            block_entry: block,
+                            block_then: Some(left),
+                            block_else: None,
+                            block_exit: right,
+                        });
+                    } else if left == *else_destination {
+                        result = Some(BasicConditional {
+                            block_entry: block,
+                            block_then: None,
+                            block_else: Some(left),
+                            block_exit: right,
+                        });
+                    }
+                }
             }
         } else if right_successors_len == 1 && next_right == Some(left) {
             // Right branch joins the right branch, it is a if/else statement with no then
             // I am not sure whether this case can happen, but it is not difficult to handle it
             let cost = block_cost(right, &function.dfg);
             if cost < cost / 2 + jump_overhead {
-                result = Some(BasicConditional {
-                    block_entry: block,
-                    block_then: None,
-                    block_else: Some(right),
-                    block_exit: right,
-                });
+                if let Some(TerminatorInstruction::JmpIf {
+                    condition: _,
+                    then_destination,
+                    else_destination,
+                    call_stack: _,
+                }) = function.dfg[block].terminator()
+                {
+                    if right == *else_destination {
+                        result = Some(BasicConditional {
+                            block_entry: block,
+                            block_then: None,
+                            block_else: Some(right),
+                            block_exit: right,
+                        });
+                    } else if right == *then_destination {
+                        result = Some(BasicConditional {
+                            block_entry: block,
+                            block_then: Some(right),
+                            block_else: None,
+                            block_exit: right,
+                        });
+                    }
+                }
             }
         }
     }
@@ -232,19 +274,32 @@ impl<'f> Context<'f> {
         conditional: &BasicConditional,
         no_predicates: &mut FxHashMap<FunctionId, bool>,
     ) {
-        // Manually inline then, else and exit into the entry block
-        let queue = vec![];
-
+        // Manually inline 'then', 'else' and 'exit' into the entry block
+        let mut queue = vec![];
+        //1. process 'then' branch
         self.inline_block(conditional.block_entry, no_predicates);
-        let _to_process = self.handle_terminator(conditional.block_entry, &queue);
+        let to_process = self.handle_terminator(conditional.block_entry, &queue);
+        queue.extend(to_process);
         if let Some(then) = conditional.block_then {
+            assert_eq!(queue.pop(), conditional.block_then);
             self.inline_block(then, no_predicates);
-            self.handle_terminator(then, &queue);
+            let to_process = self.handle_terminator(then, &queue);
+
+            queue.extend(to_process);
         }
-        if let Some(else_b) = conditional.block_else {
-            self.inline_block(else_b, no_predicates);
-            self.handle_terminator(else_b, &queue);
+
+        //2. process 'else' branch, in case there is no 'then'
+        let next = queue.pop();
+        if next == conditional.block_else {
+            let next = next.unwrap();
+            self.inline_block(next, no_predicates);
+            let to_process = self.handle_terminator(next, &queue);
+            queue.extend(to_process);
+        } else {
+            assert_eq!(next, Some(conditional.block_exit));
         }
+
+        //3. process 'exit' block
         self.inline_block(conditional.block_exit, no_predicates);
         // Manually set the terminator of the entry block to the one of the exit block
         let terminator =

From 1757bc3da87527df50d2fdaf198d2351faa3f164 Mon Sep 17 00:00:00 2001
From: guipublic <guipublic@gmail.com>
Date: Wed, 29 Jan 2025 14:10:25 +0000
Subject: [PATCH 03/13] another fix to the manual inline

---
 .../noirc_evaluator/src/ssa/opt/basic_conditional.rs | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs b/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
index cd0605a1154..83912582028 100644
--- a/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
+++ b/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
@@ -285,7 +285,11 @@ impl<'f> Context<'f> {
             self.inline_block(then, no_predicates);
             let to_process = self.handle_terminator(then, &queue);
 
-            queue.extend(to_process);
+            for incoming_block in to_process {
+                if !queue.contains(&incoming_block) {
+                    queue.push(incoming_block);
+                }
+            }
         }
 
         //2. process 'else' branch, in case there is no 'then'
@@ -294,7 +298,11 @@ impl<'f> Context<'f> {
             let next = next.unwrap();
             self.inline_block(next, no_predicates);
             let to_process = self.handle_terminator(next, &queue);
-            queue.extend(to_process);
+            for incoming_block in to_process {
+                if !queue.contains(&incoming_block) {
+                    queue.push(incoming_block);
+                }
+            }
         } else {
             assert_eq!(next, Some(conditional.block_exit));
         }

From 69466844c8b73f754f36c735e6d56906ffee900b Mon Sep 17 00:00:00 2001
From: guipublic <guipublic@gmail.com>
Date: Wed, 29 Jan 2025 14:11:53 +0000
Subject: [PATCH 04/13] remove debug statement

---
 compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs b/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
index 83912582028..36e21281748 100644
--- a/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
+++ b/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
@@ -245,7 +245,7 @@ fn flatten_function(function: &mut Function, no_predicates: &mut FxHashMap<Funct
             stack.extend(cfg.successors(block));
         }
     }
-    dbg!(&conditionals.len());
+
     // 2. Flatten all simple conditionals
     // process in reverse order because the process of a conditional does not impact the previous ones
     conditionals.reverse();

From 39bc8e69ce151d057944c12819ba78a656c4b6d4 Mon Sep 17 00:00:00 2001
From: guipublic <guipublic@gmail.com>
Date: Wed, 29 Jan 2025 15:54:05 +0000
Subject: [PATCH 05/13] Add support for IfElse instruction in Brillig

---
 .../src/brillig/brillig_gen/brillig_block.rs  | 18 +++++++++++++--
 .../src/brillig/brillig_ir/debug_show.rs      | 18 +++++++++++++++
 .../src/brillig/brillig_ir/instructions.rs    | 22 +++++++++++++++++++
 compiler/noirc_evaluator/src/ssa/ir/dfg.rs    |  7 +++---
 4 files changed, 59 insertions(+), 6 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
index 97de1aea8c7..89c97f568ff 100644
--- a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs
+++ b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs
@@ -853,8 +853,22 @@ impl<'block, Registers: RegisterAllocator> BrilligBlock<'block, Registers> {
             Instruction::EnableSideEffectsIf { .. } => {
                 unreachable!("enable_side_effects not supported by brillig")
             }
-            Instruction::IfElse { .. } => {
-                unreachable!("IfElse instructions should not be possible in brillig")
+            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(
+                    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,
+                );
             }
             Instruction::MakeArray { elements: array, typ } => {
                 let value_id = dfg.instruction_results(instruction_id)[0];
diff --git a/compiler/noirc_evaluator/src/brillig/brillig_ir/debug_show.rs b/compiler/noirc_evaluator/src/brillig/brillig_ir/debug_show.rs
index ef1b5432128..797c8df430a 100644
--- a/compiler/noirc_evaluator/src/brillig/brillig_ir/debug_show.rs
+++ b/compiler/noirc_evaluator/src/brillig/brillig_ir/debug_show.rs
@@ -126,6 +126,24 @@ impl DebugShow {
         debug_println!(self.enable_debug_trace, "  MOV {}, {}", destination, source);
     }
 
+    /// Emits a `conditional mov` instruction.
+    pub(crate) fn conditional_mov_instruction(
+        &self,
+        destination: MemoryAddress,
+        source_a: MemoryAddress,
+        source_b: MemoryAddress,
+        condition: MemoryAddress,
+    ) {
+        debug_println!(
+            self.enable_debug_trace,
+            "  {} = CONDITIONAL MOV {} then {}, else {}",
+            destination,
+            condition,
+            source_a,
+            source_b
+        );
+    }
+
     /// Emits a `cast` instruction.
     pub(crate) fn cast_instruction(
         &self,
diff --git a/compiler/noirc_evaluator/src/brillig/brillig_ir/instructions.rs b/compiler/noirc_evaluator/src/brillig/brillig_ir/instructions.rs
index d67da423d44..793eaee8f8a 100644
--- a/compiler/noirc_evaluator/src/brillig/brillig_ir/instructions.rs
+++ b/compiler/noirc_evaluator/src/brillig/brillig_ir/instructions.rs
@@ -75,6 +75,28 @@ impl<F: AcirField + DebugToString, Registers: RegisterAllocator> BrilligContext<
         );
     }
 
+    /// Insert a conditional move instruction
+    pub(crate) fn conditional_move_instruction(
+        &mut self,
+        condition: SingleAddrVariable,
+        then_address: SingleAddrVariable,
+        else_address: SingleAddrVariable,
+        destination: SingleAddrVariable,
+    ) {
+        self.debug_show.conditional_mov_instruction(
+            destination.address,
+            then_address.address,
+            else_address.address,
+            condition.address,
+        );
+        self.push_opcode(BrilligOpcode::ConditionalMov {
+            destination: destination.address,
+            source_a: then_address.address,
+            source_b: else_address.address,
+            condition: condition.address,
+        });
+    }
+
     fn binary(
         &mut self,
         lhs: SingleAddrVariable,
diff --git a/compiler/noirc_evaluator/src/ssa/ir/dfg.rs b/compiler/noirc_evaluator/src/ssa/ir/dfg.rs
index 8e87db15caf..e02c205a103 100644
--- a/compiler/noirc_evaluator/src/ssa/ir/dfg.rs
+++ b/compiler/noirc_evaluator/src/ssa/ir/dfg.rs
@@ -232,10 +232,9 @@ impl DataFlowGraph {
                 instruction,
                 Instruction::IncrementRc { .. } | Instruction::DecrementRc { .. }
             ),
-            RuntimeType::Brillig(_) => !matches!(
-                instruction,
-                Instruction::EnableSideEffectsIf { .. } | Instruction::IfElse { .. }
-            ),
+            RuntimeType::Brillig(_) => {
+                !matches!(instruction, Instruction::EnableSideEffectsIf { .. })
+            }
         }
     }
 

From 1c91dbaa0f8f8865a623ff748e74729fb1eb4f82 Mon Sep 17 00:00:00 2001
From: guipublic <guipublic@gmail.com>
Date: Fri, 31 Jan 2025 12:00:52 +0000
Subject: [PATCH 06/13] map values in follow-up blocks

---
 .../src/ssa/ir/function_inserter.rs           | 20 +++++++++-
 .../src/ssa/opt/basic_conditional.rs          | 38 +++++++++++++++++--
 2 files changed, 54 insertions(+), 4 deletions(-)

diff --git a/compiler/noirc_evaluator/src/ssa/ir/function_inserter.rs b/compiler/noirc_evaluator/src/ssa/ir/function_inserter.rs
index 9e4557e06a6..cf5d20e8f22 100644
--- a/compiler/noirc_evaluator/src/ssa/ir/function_inserter.rs
+++ b/compiler/noirc_evaluator/src/ssa/ir/function_inserter.rs
@@ -7,7 +7,7 @@ use super::{
     call_stack::CallStackId,
     dfg::InsertInstructionResult,
     function::Function,
-    instruction::{Instruction, InstructionId},
+    instruction::{BinaryOp, Instruction, InstructionId},
     value::ValueId,
 };
 use fxhash::FxHashMap as HashMap;
@@ -79,6 +79,24 @@ impl<'f> FunctionInserter<'f> {
         (instruction, self.function.dfg.get_instruction_call_stack_id(id))
     }
 
+    /// Get an instruction, map all its values, and replace it with the resolved instruction.
+    pub(crate) fn map_instruction_in_place(&mut self, id: InstructionId, block: BasicBlockId) {
+        let mut instruction = self.function.dfg[id].clone();
+        instruction.map_values_mut(|id| self.resolve(id));
+        let call_stack = self.function.dfg.get_instruction_call_stack_id(id);
+        let old_results = self.function.dfg.instruction_results(id).to_vec();
+        let ctrl_typevars = instruction
+            .requires_ctrl_typevars()
+            .then(|| vecmap(old_results, |result| self.function.dfg.type_of_value(result)));
+
+        self.function.dfg.insert_instruction_and_results_if_simplified(
+            instruction,
+            block,
+            ctrl_typevars,
+            call_stack,
+            Some(id),
+        );
+    }
     /// Maps a terminator in place, replacing any ValueId in the terminator with the
     /// resolved version of that value id from this FunctionInserter's internal value mapping.
     pub(crate) fn map_terminator_in_place(&mut self, block: BasicBlockId) {
diff --git a/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs b/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
index 36e21281748..e445ad1b2b9 100644
--- a/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
+++ b/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
@@ -1,4 +1,4 @@
-use std::collections::HashSet;
+use std::collections::{vec_deque, HashSet, VecDeque};
 
 use acvm::AcirField;
 use fxhash::FxHashMap;
@@ -9,8 +9,10 @@ use crate::ssa::{
         basic_block::BasicBlockId,
         cfg::ControlFlowGraph,
         dfg::DataFlowGraph,
+        dom::DominatorTree,
         function::{Function, FunctionId, RuntimeType},
         instruction::{BinaryOp, Instruction, TerminatorInstruction},
+        post_order::PostOrder,
     },
     Ssa,
 };
@@ -185,7 +187,7 @@ fn block_cost(block: BasicBlockId, dfg: &DataFlowGraph) -> u32 {
             | Instruction::Call { .. }  //TODO support calls with no-predicate set to true
             |      Instruction::Load { .. }
             | Instruction::Store { .. }
-            | Instruction::ArraySet { .. } => {return u32::MAX},
+            | Instruction::ArraySet { .. } => return u32::MAX,
 
             Instruction::ArrayGet { array, index  } => {
                 // check if index is in bound
@@ -263,16 +265,19 @@ impl<'f> Context<'f> {
         let mut branch_ends = FxHashMap::default();
         branch_ends.insert(conditional.block_entry, conditional.block_exit);
         let cfg = ControlFlowGraph::with_function(function);
+        let post_order = PostOrder::with_function(function);
+        let mut dominator_tree = DominatorTree::with_cfg_and_post_order(&cfg, &post_order);
         let mut context = Context::new(function, cfg, branch_ends, conditional.block_entry);
         context.no_predicate = true;
 
-        context.flatten_single_conditional(conditional, no_predicates);
+        context.flatten_single_conditional(conditional, no_predicates, &mut dominator_tree);
     }
 
     fn flatten_single_conditional(
         &mut self,
         conditional: &BasicConditional,
         no_predicates: &mut FxHashMap<FunctionId, bool>,
+        dom: &mut DominatorTree,
     ) {
         // Manually inline 'then', 'else' and 'exit' into the entry block
         let mut queue = vec![];
@@ -312,6 +317,7 @@ impl<'f> Context<'f> {
         // Manually set the terminator of the entry block to the one of the exit block
         let terminator =
             self.inserter.function.dfg[conditional.block_exit].terminator().unwrap().clone();
+        let mut next_blocks = VecDeque::new();
         let new_terminator = match terminator {
             TerminatorInstruction::JmpIf {
                 condition,
@@ -320,6 +326,7 @@ impl<'f> Context<'f> {
                 call_stack,
             } => {
                 let condition = self.inserter.resolve(condition);
+                next_blocks.extend([then_destination, else_destination]);
                 TerminatorInstruction::JmpIf {
                     condition,
                     then_destination,
@@ -329,6 +336,7 @@ impl<'f> Context<'f> {
             }
             TerminatorInstruction::Jmp { destination, arguments, call_stack } => {
                 let arguments = vecmap(arguments, |value| self.inserter.resolve(value));
+                next_blocks.push_back(destination);
                 TerminatorInstruction::Jmp { destination, arguments, call_stack }
             }
             TerminatorInstruction::Return { return_values, call_stack } => {
@@ -339,5 +347,29 @@ impl<'f> Context<'f> {
         self.inserter.function.dfg.set_block_terminator(conditional.block_entry, new_terminator);
 
         self.inserter.map_data_bus_in_place();
+        // we need to map all the subsequent blocks in case some instructions are using the ones
+        // that have been optimised
+        while let Some(block) = next_blocks.pop_back() {
+            self.map_block(block);
+            let successors = self.inserter.function.dfg[block].successors();
+            for successor in successors {
+                // The function does not process the CFG from the root but from 'any' block (given as input),
+                // so we avoid looping backward instead of checking for already visited blocks
+                if !dom.dominates(successor, block) {
+                    // 'push front' ensures a block is mapped after his predecessors
+                    next_blocks.push_front(successor);
+                }
+            }
+        }
+    }
+
+    fn map_block(&mut self, block: BasicBlockId) {
+        // Map all instructions in the block
+        let instructions = self.inserter.function.dfg[block].instructions();
+        let instructions = instructions.iter().cloned().collect::<Vec<_>>();
+        for instruction in instructions {
+            self.inserter.map_instruction_in_place(instruction, block);
+        }
+        self.inserter.map_terminator_in_place(block);
     }
 }

From f24aac712c5a131592eb5e08d1df388a5c57c1ec Mon Sep 17 00:00:00 2001
From: guipublic <guipublic@gmail.com>
Date: Fri, 31 Jan 2025 12:13:34 +0000
Subject: [PATCH 07/13] fix clippy

---
 compiler/noirc_evaluator/src/ssa/ir/function_inserter.rs  | 2 +-
 compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs | 5 ++---
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/compiler/noirc_evaluator/src/ssa/ir/function_inserter.rs b/compiler/noirc_evaluator/src/ssa/ir/function_inserter.rs
index cf5d20e8f22..3eeab848c35 100644
--- a/compiler/noirc_evaluator/src/ssa/ir/function_inserter.rs
+++ b/compiler/noirc_evaluator/src/ssa/ir/function_inserter.rs
@@ -7,7 +7,7 @@ use super::{
     call_stack::CallStackId,
     dfg::InsertInstructionResult,
     function::Function,
-    instruction::{BinaryOp, Instruction, InstructionId},
+    instruction::{Instruction, InstructionId},
     value::ValueId,
 };
 use fxhash::FxHashMap as HashMap;
diff --git a/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs b/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
index e445ad1b2b9..550a7455dec 100644
--- a/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
+++ b/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
@@ -1,4 +1,4 @@
-use std::collections::{vec_deque, HashSet, VecDeque};
+use std::collections::{HashSet, VecDeque};
 
 use acvm::AcirField;
 use fxhash::FxHashMap;
@@ -365,8 +365,7 @@ impl<'f> Context<'f> {
 
     fn map_block(&mut self, block: BasicBlockId) {
         // Map all instructions in the block
-        let instructions = self.inserter.function.dfg[block].instructions();
-        let instructions = instructions.iter().cloned().collect::<Vec<_>>();
+        let instructions = self.inserter.function.dfg[block].instructions().to_vec();
         for instruction in instructions {
             self.inserter.map_instruction_in_place(instruction, block);
         }

From 77763b6ce939bbdc1bfc223bc7ef541b27727c1f Mon Sep 17 00:00:00 2001
From: guipublic <guipublic@gmail.com>
Date: Fri, 31 Jan 2025 14:27:03 +0000
Subject: [PATCH 08/13] Replace mapped instruction instead of inserting a new
 one

---
 compiler/noirc_evaluator/src/ssa/ir/dfg.rs       |  5 +++++
 .../src/ssa/ir/function_inserter.rs              | 16 ++--------------
 .../src/ssa/opt/basic_conditional.rs             |  7 +++++--
 3 files changed, 12 insertions(+), 16 deletions(-)

diff --git a/compiler/noirc_evaluator/src/ssa/ir/dfg.rs b/compiler/noirc_evaluator/src/ssa/ir/dfg.rs
index 6e9a7b47265..68bc15e10bd 100644
--- a/compiler/noirc_evaluator/src/ssa/ir/dfg.rs
+++ b/compiler/noirc_evaluator/src/ssa/ir/dfg.rs
@@ -367,6 +367,11 @@ impl DataFlowGraph {
         }
     }
 
+    /// Replace an existing instruction with a new one.
+    pub(crate) fn set_instruction(&mut self, id: InstructionId, instruction: Instruction) {
+        self.instructions[id] = instruction;
+    }
+
     /// Set the value of value_to_replace to refer to the value referred to by new_value.
     ///
     /// This is the preferred method to call for optimizations simplifying
diff --git a/compiler/noirc_evaluator/src/ssa/ir/function_inserter.rs b/compiler/noirc_evaluator/src/ssa/ir/function_inserter.rs
index 3eeab848c35..0ed33d018ff 100644
--- a/compiler/noirc_evaluator/src/ssa/ir/function_inserter.rs
+++ b/compiler/noirc_evaluator/src/ssa/ir/function_inserter.rs
@@ -80,22 +80,10 @@ impl<'f> FunctionInserter<'f> {
     }
 
     /// Get an instruction, map all its values, and replace it with the resolved instruction.
-    pub(crate) fn map_instruction_in_place(&mut self, id: InstructionId, block: BasicBlockId) {
+    pub(crate) fn map_instruction_in_place(&mut self, id: InstructionId) {
         let mut instruction = self.function.dfg[id].clone();
         instruction.map_values_mut(|id| self.resolve(id));
-        let call_stack = self.function.dfg.get_instruction_call_stack_id(id);
-        let old_results = self.function.dfg.instruction_results(id).to_vec();
-        let ctrl_typevars = instruction
-            .requires_ctrl_typevars()
-            .then(|| vecmap(old_results, |result| self.function.dfg.type_of_value(result)));
-
-        self.function.dfg.insert_instruction_and_results_if_simplified(
-            instruction,
-            block,
-            ctrl_typevars,
-            call_stack,
-            Some(id),
-        );
+        self.function.dfg.set_instruction(id, instruction);
     }
     /// Maps a terminator in place, replacing any ValueId in the terminator with the
     /// resolved version of that value id from this FunctionInserter's internal value mapping.
diff --git a/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs b/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
index 550a7455dec..23eaa9d120c 100644
--- a/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
+++ b/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
@@ -348,7 +348,10 @@ impl<'f> Context<'f> {
 
         self.inserter.map_data_bus_in_place();
         // we need to map all the subsequent blocks in case some instructions are using the ones
-        // that have been optimised
+        // that have been optimized
+        // Note that this is not efficient because flattening several blocks in the same function
+        // will lead to mapping the blocks at the end several times, while it can be done only once
+        // for all blocks flattening.
         while let Some(block) = next_blocks.pop_back() {
             self.map_block(block);
             let successors = self.inserter.function.dfg[block].successors();
@@ -367,7 +370,7 @@ impl<'f> Context<'f> {
         // Map all instructions in the block
         let instructions = self.inserter.function.dfg[block].instructions().to_vec();
         for instruction in instructions {
-            self.inserter.map_instruction_in_place(instruction, block);
+            self.inserter.map_instruction_in_place(instruction);
         }
         self.inserter.map_terminator_in_place(block);
     }

From 9fcf3482913ed6105406c4b7257d870ac2d793a9 Mon Sep 17 00:00:00 2001
From: guipublic <guipublic@gmail.com>
Date: Fri, 31 Jan 2025 16:35:10 +0000
Subject: [PATCH 09/13] Add support for arrays and vectors to IfElse in Brillig

---
 .../src/brillig/brillig_gen/brillig_block.rs  | 62 ++++++++++++++++---
 1 file changed, 52 insertions(+), 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
index 89c97f568ff..ae8d49ce0d8 100644
--- 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];

From 06d26db6d8ae1b85d6db0f1df10834442acf8c30 Mon Sep 17 00:00:00 2001
From: guipublic <guipublic@gmail.com>
Date: Mon, 3 Feb 2025 11:16:19 +0000
Subject: [PATCH 10/13] map blocks only once

---
 .../src/ssa/opt/basic_conditional.rs          | 53 +++++++++++--------
 .../src/ssa/opt/flatten_cfg.rs                |  2 +-
 2 files changed, 33 insertions(+), 22 deletions(-)

diff --git a/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs b/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
index 23eaa9d120c..f5f2a165976 100644
--- a/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
+++ b/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
@@ -249,38 +249,44 @@ fn flatten_function(function: &mut Function, no_predicates: &mut FxHashMap<Funct
     }
 
     // 2. Flatten all simple conditionals
-    // process in reverse order because the process of a conditional does not impact the previous ones
+    // process basic conditional in reverse order so that a conditional
+    // does not impact the previous ones
     conditionals.reverse();
-    for simple_if in conditionals {
-        Context::flatten_single(&simple_if, function, no_predicates);
-    }
+    Context::flatten_multiple(&conditionals, function, no_predicates);
 }
 
 impl<'f> Context<'f> {
-    fn flatten_single(
-        conditional: &BasicConditional,
+    fn flatten_multiple(
+        conditionals: &Vec<BasicConditional>,
         function: &mut Function,
         no_predicates: &mut FxHashMap<FunctionId, bool>,
     ) {
+        // 1. process each basic conditional
         let mut branch_ends = FxHashMap::default();
-        branch_ends.insert(conditional.block_entry, conditional.block_exit);
+        for conditional in conditionals {
+            branch_ends.insert(conditional.block_entry, conditional.block_exit);
+        }
         let cfg = ControlFlowGraph::with_function(function);
         let post_order = PostOrder::with_function(function);
         let mut dominator_tree = DominatorTree::with_cfg_and_post_order(&cfg, &post_order);
-        let mut context = Context::new(function, cfg, branch_ends, conditional.block_entry);
+        let cfg_root = function.entry_block();
+        let mut context = Context::new(function, cfg, branch_ends, cfg_root);
         context.no_predicate = true;
-
-        context.flatten_single_conditional(conditional, no_predicates, &mut dominator_tree);
+        for conditional in conditionals {
+            context.flatten_single_conditional(conditional, no_predicates);
+        }
+        // 2. re-map the full program for values that have been simplified.
+        context.map_cfg(cfg_root, &mut dominator_tree);
     }
 
     fn flatten_single_conditional(
         &mut self,
         conditional: &BasicConditional,
         no_predicates: &mut FxHashMap<FunctionId, bool>,
-        dom: &mut DominatorTree,
     ) {
         // Manually inline 'then', 'else' and 'exit' into the entry block
         let mut queue = vec![];
+        self.target_block = conditional.block_entry;
         //1. process 'then' branch
         self.inline_block(conditional.block_entry, no_predicates);
         let to_process = self.handle_terminator(conditional.block_entry, &queue);
@@ -347,11 +353,25 @@ impl<'f> Context<'f> {
         self.inserter.function.dfg.set_block_terminator(conditional.block_entry, new_terminator);
 
         self.inserter.map_data_bus_in_place();
+    }
+
+    fn map_block(&mut self, block: BasicBlockId) {
+        // Map all instructions in the block
+        let instructions = self.inserter.function.dfg[block].instructions().to_vec();
+        for instruction in instructions {
+            self.inserter.map_instruction_in_place(instruction);
+        }
+        self.inserter.map_terminator_in_place(block);
+    }
+
+    fn map_cfg(&mut self, entry_block: BasicBlockId, dom: &mut DominatorTree) {
         // we need to map all the subsequent blocks in case some instructions are using the ones
         // that have been optimized
         // Note that this is not efficient because flattening several blocks in the same function
         // will lead to mapping the blocks at the end several times, while it can be done only once
         // for all blocks flattening.
+        let mut next_blocks = VecDeque::new();
+        next_blocks.push_back(entry_block);
         while let Some(block) = next_blocks.pop_back() {
             self.map_block(block);
             let successors = self.inserter.function.dfg[block].successors();
@@ -359,19 +379,10 @@ impl<'f> Context<'f> {
                 // The function does not process the CFG from the root but from 'any' block (given as input),
                 // so we avoid looping backward instead of checking for already visited blocks
                 if !dom.dominates(successor, block) {
-                    // 'push front' ensures a block is mapped after his predecessors
+                    // 'push front' ensures a block is mapped after his predecessors are also mapped
                     next_blocks.push_front(successor);
                 }
             }
         }
     }
-
-    fn map_block(&mut self, block: BasicBlockId) {
-        // Map all instructions in the block
-        let instructions = self.inserter.function.dfg[block].instructions().to_vec();
-        for instruction in instructions {
-            self.inserter.map_instruction_in_place(instruction);
-        }
-        self.inserter.map_terminator_in_place(block);
-    }
 }
diff --git a/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs b/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs
index d4e2968f113..a20235f06bc 100644
--- a/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs
+++ b/compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs
@@ -183,7 +183,7 @@ pub(crate) struct Context<'f> {
     cfg: ControlFlowGraph,
 
     /// Target block of the flattening
-    target_block: BasicBlockId,
+    pub(crate) target_block: BasicBlockId,
 
     /// Maps start of branch -> end of branch
     branch_ends: HashMap<BasicBlockId, BasicBlockId>,

From 7a5ed52bbaf27df144f67fabc5e5423ad4bd3b14 Mon Sep 17 00:00:00 2001
From: guipublic <guipublic@gmail.com>
Date: Mon, 3 Feb 2025 11:56:08 +0000
Subject: [PATCH 11/13] map blocks using post order

---
 .../src/ssa/opt/basic_conditional.rs          | 32 ++++---------------
 1 file changed, 6 insertions(+), 26 deletions(-)

diff --git a/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs b/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
index f5f2a165976..22923181c35 100644
--- a/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
+++ b/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
@@ -9,7 +9,6 @@ use crate::ssa::{
         basic_block::BasicBlockId,
         cfg::ControlFlowGraph,
         dfg::DataFlowGraph,
-        dom::DominatorTree,
         function::{Function, FunctionId, RuntimeType},
         instruction::{BinaryOp, Instruction, TerminatorInstruction},
         post_order::PostOrder,
@@ -268,15 +267,18 @@ impl<'f> Context<'f> {
         }
         let cfg = ControlFlowGraph::with_function(function);
         let post_order = PostOrder::with_function(function);
-        let mut dominator_tree = DominatorTree::with_cfg_and_post_order(&cfg, &post_order);
         let cfg_root = function.entry_block();
         let mut context = Context::new(function, cfg, branch_ends, cfg_root);
         context.no_predicate = true;
         for conditional in conditionals {
             context.flatten_single_conditional(conditional, no_predicates);
         }
-        // 2. re-map the full program for values that have been simplified.
-        context.map_cfg(cfg_root, &mut dominator_tree);
+        // 2. re-map the full program for values that may been simplified.
+        if !conditionals.is_empty() {
+            for block in post_order.as_slice() {
+                context.map_block(*block);
+            }
+        }
     }
 
     fn flatten_single_conditional(
@@ -363,26 +365,4 @@ impl<'f> Context<'f> {
         }
         self.inserter.map_terminator_in_place(block);
     }
-
-    fn map_cfg(&mut self, entry_block: BasicBlockId, dom: &mut DominatorTree) {
-        // we need to map all the subsequent blocks in case some instructions are using the ones
-        // that have been optimized
-        // Note that this is not efficient because flattening several blocks in the same function
-        // will lead to mapping the blocks at the end several times, while it can be done only once
-        // for all blocks flattening.
-        let mut next_blocks = VecDeque::new();
-        next_blocks.push_back(entry_block);
-        while let Some(block) = next_blocks.pop_back() {
-            self.map_block(block);
-            let successors = self.inserter.function.dfg[block].successors();
-            for successor in successors {
-                // The function does not process the CFG from the root but from 'any' block (given as input),
-                // so we avoid looping backward instead of checking for already visited blocks
-                if !dom.dominates(successor, block) {
-                    // 'push front' ensures a block is mapped after his predecessors are also mapped
-                    next_blocks.push_front(successor);
-                }
-            }
-        }
-    }
 }

From 90aa52c6906cda622ea6d59ab4594b38183fd4ce Mon Sep 17 00:00:00 2001
From: guipublic <guipublic@gmail.com>
Date: Mon, 3 Feb 2025 17:56:22 +0000
Subject: [PATCH 12/13] Use one flattening context per conditional

---
 .../src/ssa/ir/function_inserter.rs           | 18 +++++
 .../src/ssa/opt/basic_conditional.rs          | 77 +++++++++++--------
 2 files changed, 62 insertions(+), 33 deletions(-)

diff --git a/compiler/noirc_evaluator/src/ssa/ir/function_inserter.rs b/compiler/noirc_evaluator/src/ssa/ir/function_inserter.rs
index 0ed33d018ff..553f96a651b 100644
--- a/compiler/noirc_evaluator/src/ssa/ir/function_inserter.rs
+++ b/compiler/noirc_evaluator/src/ssa/ir/function_inserter.rs
@@ -257,4 +257,22 @@ impl<'f> FunctionInserter<'f> {
             self.values.entry(*param).or_insert(*new_param);
         }
     }
+
+    /// Merge the internal mapping into the given mapping
+    /// The merge is guaranteed to be coherent because ambiguous cases are prevented
+    pub(crate) fn extract_mapping(&self, mapping: &mut HashMap<ValueId, ValueId>) {
+        for (k, v) in &self.values {
+            if mapping.contains_key(k) {
+                unreachable!("cannot merge key");
+            }
+            if mapping.contains_key(v) {
+                unreachable!("cannot merge value");
+            }
+            mapping.insert(*k, *v);
+        }
+    }
+
+    pub(crate) fn set_mapping(&mut self, mapping: HashMap<ValueId, ValueId>) {
+        self.values = mapping;
+    }
 }
diff --git a/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs b/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
index 22923181c35..35dad3067f5 100644
--- a/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
+++ b/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
@@ -10,8 +10,10 @@ use crate::ssa::{
         cfg::ControlFlowGraph,
         dfg::DataFlowGraph,
         function::{Function, FunctionId, RuntimeType},
+        function_inserter::FunctionInserter,
         instruction::{BinaryOp, Instruction, TerminatorInstruction},
         post_order::PostOrder,
+        value::ValueId,
     },
     Ssa,
 };
@@ -248,47 +250,23 @@ fn flatten_function(function: &mut Function, no_predicates: &mut FxHashMap<Funct
     }
 
     // 2. Flatten all simple conditionals
-    // process basic conditional in reverse order so that a conditional
-    // does not impact the previous ones
+    // process basic conditionals in reverse order so that
+    // a conditional does not impact the previous ones
     conditionals.reverse();
     Context::flatten_multiple(&conditionals, function, no_predicates);
 }
 
 impl<'f> Context<'f> {
-    fn flatten_multiple(
-        conditionals: &Vec<BasicConditional>,
-        function: &mut Function,
-        no_predicates: &mut FxHashMap<FunctionId, bool>,
-    ) {
-        // 1. process each basic conditional
-        let mut branch_ends = FxHashMap::default();
-        for conditional in conditionals {
-            branch_ends.insert(conditional.block_entry, conditional.block_exit);
-        }
-        let cfg = ControlFlowGraph::with_function(function);
-        let post_order = PostOrder::with_function(function);
-        let cfg_root = function.entry_block();
-        let mut context = Context::new(function, cfg, branch_ends, cfg_root);
-        context.no_predicate = true;
-        for conditional in conditionals {
-            context.flatten_single_conditional(conditional, no_predicates);
-        }
-        // 2. re-map the full program for values that may been simplified.
-        if !conditionals.is_empty() {
-            for block in post_order.as_slice() {
-                context.map_block(*block);
-            }
-        }
-    }
-
     fn flatten_single_conditional(
         &mut self,
         conditional: &BasicConditional,
         no_predicates: &mut FxHashMap<FunctionId, bool>,
     ) {
         // Manually inline 'then', 'else' and 'exit' into the entry block
+        //0. initialize the context for flattening a 'single conditional'
         let mut queue = vec![];
         self.target_block = conditional.block_entry;
+        self.no_predicate = true;
         //1. process 'then' branch
         self.inline_block(conditional.block_entry, no_predicates);
         let to_process = self.handle_terminator(conditional.block_entry, &queue);
@@ -353,16 +331,49 @@ impl<'f> Context<'f> {
             }
         };
         self.inserter.function.dfg.set_block_terminator(conditional.block_entry, new_terminator);
-
         self.inserter.map_data_bus_in_place();
     }
 
-    fn map_block(&mut self, block: BasicBlockId) {
+    fn map_block_with_mapping(
+        mapping: FxHashMap<ValueId, ValueId>,
+        func: &mut Function,
+        block: BasicBlockId,
+    ) {
         // Map all instructions in the block
-        let instructions = self.inserter.function.dfg[block].instructions().to_vec();
+        let mut inserter = FunctionInserter::new(func);
+        inserter.set_mapping(mapping);
+        let instructions = inserter.function.dfg[block].instructions().to_vec();
         for instruction in instructions {
-            self.inserter.map_instruction_in_place(instruction);
+            inserter.map_instruction_in_place(instruction);
+        }
+        inserter.map_terminator_in_place(block);
+    }
+
+    fn flatten_multiple(
+        conditionals: &Vec<BasicConditional>,
+        function: &mut Function,
+        no_predicates: &mut FxHashMap<FunctionId, bool>,
+    ) {
+        // 1. process each basic conditional, using a new context per conditional
+        let post_order = PostOrder::with_function(function);
+
+        let mut mapping = FxHashMap::default();
+        for conditional in conditionals {
+            let cfg = ControlFlowGraph::with_function(function);
+            let cfg_root = function.entry_block();
+            let mut branch_ends = FxHashMap::default();
+            branch_ends.insert(conditional.block_entry, conditional.block_exit);
+            let mut context = Context::new(function, cfg, branch_ends, cfg_root);
+            context.flatten_single_conditional(conditional, no_predicates);
+            // extract the mapping into 'mapping
+            context.inserter.extract_mapping(&mut mapping);
+        }
+        // 2. re-map the full program for values that may been simplified.
+        if !mapping.is_empty() {
+            let po = post_order.as_slice();
+            for block in po {
+                Context::map_block_with_mapping(mapping.clone(), function, *block);
+            }
         }
-        self.inserter.map_terminator_in_place(block);
     }
 }

From 1db3d6105c2613116fad8d38066f9a71a6d7fae7 Mon Sep 17 00:00:00 2001
From: guipublic <guipublic@gmail.com>
Date: Wed, 5 Feb 2025 17:40:48 +0000
Subject: [PATCH 13/13] code review

---
 .../src/ssa/ir/function_inserter.rs           |   1 +
 .../src/ssa/opt/basic_conditional.rs          | 250 +++++++++---------
 2 files changed, 125 insertions(+), 126 deletions(-)

diff --git a/compiler/noirc_evaluator/src/ssa/ir/function_inserter.rs b/compiler/noirc_evaluator/src/ssa/ir/function_inserter.rs
index 553f96a651b..13b5ead5eb6 100644
--- a/compiler/noirc_evaluator/src/ssa/ir/function_inserter.rs
+++ b/compiler/noirc_evaluator/src/ssa/ir/function_inserter.rs
@@ -85,6 +85,7 @@ impl<'f> FunctionInserter<'f> {
         instruction.map_values_mut(|id| self.resolve(id));
         self.function.dfg.set_instruction(id, instruction);
     }
+
     /// Maps a terminator in place, replacing any ValueId in the terminator with the
     /// resolved version of that value id from this FunctionInserter's internal value mapping.
     pub(crate) fn map_terminator_in_place(&mut self, block: BasicBlockId) {
diff --git a/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs b/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
index 35dad3067f5..a08ac1fbf0d 100644
--- a/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
+++ b/compiler/noirc_evaluator/src/ssa/opt/basic_conditional.rs
@@ -1,7 +1,7 @@
 use std::collections::{HashSet, VecDeque};
 
 use acvm::AcirField;
-use fxhash::FxHashMap;
+use fxhash::FxHashMap as HashMap;
 use iter_extended::vecmap;
 
 use crate::ssa::{
@@ -9,7 +9,7 @@ use crate::ssa::{
         basic_block::BasicBlockId,
         cfg::ControlFlowGraph,
         dfg::DataFlowGraph,
-        function::{Function, FunctionId, RuntimeType},
+        function::{Function, FunctionId},
         function_inserter::FunctionInserter,
         instruction::{BinaryOp, Instruction, TerminatorInstruction},
         post_order::PostOrder,
@@ -31,7 +31,7 @@ impl Ssa {
     #[tracing::instrument(level = "trace", skip(self))]
     pub(crate) fn flatten_basic_conditionals(mut self) -> Ssa {
         // Retrieve the 'no_predicates' attribute of the functions in a map, to avoid problems with borrowing
-        let mut no_predicates = FxHashMap::default();
+        let mut no_predicates = HashMap::default();
         for function in self.functions.values() {
             no_predicates.insert(function.id(), function.is_no_predicates());
         }
@@ -50,104 +50,102 @@ fn is_conditional(
     function: &Function,
 ) -> Option<BasicConditional> {
     // jump overhead is the cost for doing the conditional and jump around the blocks
-    // I use 10 as a rough estimate, real cost is less.
+    // We use 10 as a rough estimate, the real cost is less.
     let jump_overhead = 10;
     let mut successors = cfg.successors(block);
     let mut result = None;
     // a conditional must have 2 branches
-    if successors.len() == 2 {
-        let left = successors.next().unwrap();
-        let right = successors.next().unwrap();
-        let mut left_successors = cfg.successors(left);
-        let mut right_successors = cfg.successors(right);
-        let left_successors_len = left_successors.len();
-        let right_successors_len = right_successors.len();
-        let next_left = left_successors.next();
-        let next_right = right_successors.next();
-        if next_left == Some(block) || next_right == Some(block) {
-            // this is a loop, not a conditional
-            return None;
-        }
-        if left_successors_len == 1 && right_successors_len == 1 && next_left == next_right {
-            // The branches join on one block so it is a non-nested conditional
-            let cost_left = block_cost(left, &function.dfg);
-            let cost_right = block_cost(right, &function.dfg);
-            // For the flattening to be valuable, we compare the cost of the flattened code with the average cost of the 2 branches,
-            // including an overhead to take into account the jumps between the blocks.
-            let cost = cost_right.saturating_add(cost_left);
-            if cost < cost / 2 + jump_overhead {
-                if let Some(TerminatorInstruction::JmpIf {
-                    condition: _,
-                    then_destination,
-                    else_destination,
-                    call_stack: _,
-                }) = function.dfg[block].terminator()
-                {
-                    result = Some(BasicConditional {
-                        block_entry: block,
-                        block_then: Some(*then_destination),
-                        block_else: Some(*else_destination),
-                        block_exit: next_left.unwrap(),
-                    });
-                }
+    if successors.len() != 2 {
+        return None;
+    }
+    let left = successors.next().unwrap();
+    let right = successors.next().unwrap();
+    let mut left_successors = cfg.successors(left);
+    let mut right_successors = cfg.successors(right);
+    let left_successors_len = left_successors.len();
+    let right_successors_len = right_successors.len();
+    let next_left = left_successors.next();
+    let next_right = right_successors.next();
+    if next_left == Some(block) || next_right == Some(block) {
+        // this is a loop, not a conditional
+        return None;
+    }
+    if left_successors_len == 1 && right_successors_len == 1 && next_left == next_right {
+        // The branches join on one block so it is a non-nested conditional
+        let cost_left = block_cost(left, &function.dfg);
+        let cost_right = block_cost(right, &function.dfg);
+        // For the flattening to be valuable, we compare the cost of the flattened code with the average cost of the 2 branches,
+        // including an overhead to take into account the jumps between the blocks.
+        let cost = cost_right.saturating_add(cost_left);
+        if cost < cost / 2 + jump_overhead {
+            if let Some(TerminatorInstruction::JmpIf {
+                condition: _,
+                then_destination,
+                else_destination,
+                call_stack: _,
+            }) = function.dfg[block].terminator()
+            {
+                result = Some(BasicConditional {
+                    block_entry: block,
+                    block_then: Some(*then_destination),
+                    block_else: Some(*else_destination),
+                    block_exit: next_left.unwrap(),
+                });
             }
-        } else if left_successors_len == 1 && next_left == Some(right) {
-            // Left branch joins the right branch, it is a if/then statement with no else
-            // I am not sure whether this case can happen, but it is not difficult to handle it
-            let cost = block_cost(left, &function.dfg);
-            if cost < cost / 2 + jump_overhead {
-                if let Some(TerminatorInstruction::JmpIf {
-                    condition: _,
-                    then_destination,
-                    else_destination,
-                    call_stack: _,
-                }) = function.dfg[block].terminator()
-                {
-                    if left == *then_destination {
-                        result = Some(BasicConditional {
-                            block_entry: block,
-                            block_then: Some(left),
-                            block_else: None,
-                            block_exit: right,
-                        });
-                    } else if left == *else_destination {
-                        result = Some(BasicConditional {
-                            block_entry: block,
-                            block_then: None,
-                            block_else: Some(left),
-                            block_exit: right,
-                        });
-                    }
-                }
+        }
+    } else if left_successors_len == 1 && next_left == Some(right) {
+        // Left branch joins the right branch, it is a if/then statement with no else
+        // I am not sure whether this case can happen, but it is not difficult to handle it
+        let cost = block_cost(left, &function.dfg);
+        if cost < cost / 2 + jump_overhead {
+            if let Some(TerminatorInstruction::JmpIf {
+                condition: _,
+                then_destination,
+                else_destination,
+                call_stack: _,
+            }) = function.dfg[block].terminator()
+            {
+                let (block_then, block_else) = if left == *then_destination {
+                    (Some(left), None)
+                } else if left == *else_destination {
+                    (None, Some(left))
+                } else {
+                    return None;
+                };
+
+                result = Some(BasicConditional {
+                    block_entry: block,
+                    block_then,
+                    block_else,
+                    block_exit: right,
+                });
             }
-        } else if right_successors_len == 1 && next_right == Some(left) {
-            // Right branch joins the right branch, it is a if/else statement with no then
-            // I am not sure whether this case can happen, but it is not difficult to handle it
-            let cost = block_cost(right, &function.dfg);
-            if cost < cost / 2 + jump_overhead {
-                if let Some(TerminatorInstruction::JmpIf {
-                    condition: _,
-                    then_destination,
-                    else_destination,
-                    call_stack: _,
-                }) = function.dfg[block].terminator()
-                {
-                    if right == *else_destination {
-                        result = Some(BasicConditional {
-                            block_entry: block,
-                            block_then: None,
-                            block_else: Some(right),
-                            block_exit: right,
-                        });
-                    } else if right == *then_destination {
-                        result = Some(BasicConditional {
-                            block_entry: block,
-                            block_then: Some(right),
-                            block_else: None,
-                            block_exit: right,
-                        });
-                    }
-                }
+        }
+    } else if right_successors_len == 1 && next_right == Some(left) {
+        // Right branch joins the right branch, it is a if/else statement with no then
+        // I am not sure whether this case can happen, but it is not difficult to handle it
+        let cost = block_cost(right, &function.dfg);
+        if cost < cost / 2 + jump_overhead {
+            if let Some(TerminatorInstruction::JmpIf {
+                condition: _,
+                then_destination,
+                else_destination,
+                call_stack: _,
+            }) = function.dfg[block].terminator()
+            {
+                let (block_then, block_else) = if right == *then_destination {
+                    (Some(right), None)
+                } else if right == *else_destination {
+                    (None, Some(right))
+                } else {
+                    return None;
+                };
+                result = Some(BasicConditional {
+                    block_entry: block,
+                    block_then,
+                    block_else,
+                    block_exit: right,
+                });
             }
         }
     }
@@ -157,63 +155,64 @@ fn is_conditional(
 
 /// Computes a cost estimate of a basic block
 /// returns u32::MAX if the block has side-effect instructions
+/// WARNING: the estimates are estimate of the runtime cost of each instructions,
+/// 1 being the cost of the simplest instruction. These numbers can be improved.
 fn block_cost(block: BasicBlockId, dfg: &DataFlowGraph) -> u32 {
     let mut cost: u32 = 0;
-
     for instruction in dfg[block].instructions() {
         let instruction_cost = match &dfg[*instruction] {
             Instruction::Binary(binary) => {
                 match binary.operator {
                     BinaryOp::Add { unchecked }
                     | BinaryOp::Sub { unchecked }
-                    | BinaryOp::Mul { unchecked } => if unchecked { 1 } else { return u32::MAX },
+                    | BinaryOp::Mul { unchecked } => if unchecked { 3 } else { return u32::MAX },
                     BinaryOp::Div
                     | BinaryOp::Mod => return u32::MAX,
-                    BinaryOp::Eq => 2,
+                    BinaryOp::Eq => 1,
                     BinaryOp::Lt => 5,
                     BinaryOp::And
                     | BinaryOp::Or
-                    | BinaryOp::Xor => 1,   //todo
+                    | BinaryOp::Xor => 1,
                     BinaryOp::Shl
                     | BinaryOp::Shr => return u32::MAX,
                 }
             },
-            Instruction::Cast(_, _) => 1,//TODO check if this instruction can fail
+            // A Cast can be either simplified, or lead to a truncate
+            Instruction::Cast(_, _) => 3,
             Instruction::Not(_) => 1,
             Instruction::Truncate { .. } => 7,
 
             Instruction::Constrain(_,_,_)
             | Instruction::ConstrainNotEqual(_,_,_)
             | Instruction::RangeCheck { .. }
-            | Instruction::Call { .. }  //TODO support calls with no-predicate set to true
+            // Calls with no-predicate set to true could be supported, but
+            // they are likely to be too costly anyways. Simple calls would
+            // have been inlined already.
+            | Instruction::Call { .. }
             |      Instruction::Load { .. }
             | Instruction::Store { .. }
             | Instruction::ArraySet { .. } => return u32::MAX,
 
             Instruction::ArrayGet { array, index  } => {
+                // A get can fail because of out-of-bound index
+                let mut get_cost = u32::MAX;
                 // check if index is in bound
-                let index = dfg.get_numeric_constant(*index);
-                let mut ok_bound = false;
-                if let Some(index) = index {
-                    let len = dfg.try_get_array_length(*array).unwrap();
+                if let (Some(index), Some(len)) = (dfg.get_numeric_constant(*index), dfg.try_get_array_length(*array)) {
+                    // The index is in-bounds
                     if index.to_u128() < len as u128 {
-                        ok_bound = true;
+                        get_cost = 1;
                     }
                 }
-                if ok_bound {
-                    1
-                } else {
-                    return u32::MAX;
-                }
+                get_cost
             },
 
-            Instruction::Allocate => 0,
-            Instruction::EnableSideEffectsIf { .. } => 0,
-            Instruction::IncrementRc { .. } => 0,
-            Instruction::DecrementRc { .. } => 0,
+            Instruction::Allocate
+            | Instruction::EnableSideEffectsIf { .. }
+            | Instruction::IncrementRc { .. }
+            | Instruction::DecrementRc { .. }
+            | Instruction::MakeArray { .. }
+            | Instruction::Noop => 0,
             Instruction::IfElse { .. } => 1,
-            Instruction::MakeArray { .. } => 0,
-            Instruction::Noop => 0,
         };
         cost += instruction_cost;
     }
@@ -221,9 +220,9 @@ fn block_cost(block: BasicBlockId, dfg: &DataFlowGraph) -> u32 {
 }
 
 /// Identifies all simple conditionals in the function and flattens them
-fn flatten_function(function: &mut Function, no_predicates: &mut FxHashMap<FunctionId, bool>) {
+fn flatten_function(function: &mut Function, no_predicates: &mut HashMap<FunctionId, bool>) {
     // This pass is dedicated to brillig functions
-    if !matches!(function.runtime(), RuntimeType::Brillig(_)) {
+    if !function.runtime().is_brillig() {
         return;
     }
     let cfg = ControlFlowGraph::with_function(function);
@@ -260,7 +259,7 @@ impl<'f> Context<'f> {
     fn flatten_single_conditional(
         &mut self,
         conditional: &BasicConditional,
-        no_predicates: &mut FxHashMap<FunctionId, bool>,
+        no_predicates: &mut HashMap<FunctionId, bool>,
     ) {
         // Manually inline 'then', 'else' and 'exit' into the entry block
         //0. initialize the context for flattening a 'single conditional'
@@ -335,7 +334,7 @@ impl<'f> Context<'f> {
     }
 
     fn map_block_with_mapping(
-        mapping: FxHashMap<ValueId, ValueId>,
+        mapping: HashMap<ValueId, ValueId>,
         func: &mut Function,
         block: BasicBlockId,
     ) {
@@ -352,16 +351,16 @@ impl<'f> Context<'f> {
     fn flatten_multiple(
         conditionals: &Vec<BasicConditional>,
         function: &mut Function,
-        no_predicates: &mut FxHashMap<FunctionId, bool>,
+        no_predicates: &mut HashMap<FunctionId, bool>,
     ) {
         // 1. process each basic conditional, using a new context per conditional
         let post_order = PostOrder::with_function(function);
 
-        let mut mapping = FxHashMap::default();
+        let mut mapping = HashMap::default();
         for conditional in conditionals {
             let cfg = ControlFlowGraph::with_function(function);
             let cfg_root = function.entry_block();
-            let mut branch_ends = FxHashMap::default();
+            let mut branch_ends = HashMap::default();
             branch_ends.insert(conditional.block_entry, conditional.block_exit);
             let mut context = Context::new(function, cfg, branch_ends, cfg_root);
             context.flatten_single_conditional(conditional, no_predicates);
@@ -370,8 +369,7 @@ impl<'f> Context<'f> {
         }
         // 2. re-map the full program for values that may been simplified.
         if !mapping.is_empty() {
-            let po = post_order.as_slice();
-            for block in po {
+            for block in post_order.as_slice() {
                 Context::map_block_with_mapping(mapping.clone(), function, *block);
             }
         }