diff --git a/compiler/noirc_evaluator/src/ssa/opt/defunctionalize.rs b/compiler/noirc_evaluator/src/ssa/opt/defunctionalize.rs index 110b029a6ae..0e368ee3ab5 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/defunctionalize.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/defunctionalize.rs @@ -2127,22 +2127,22 @@ mod tests { brillig(inline_always) fn apply f6 { b0(v0: Field): v2 = eq v0, Field 1 - jmpif v2 then: b3, else: b2 + jmpif v2 then: b2, else: b1 b1(): - return - b2(): v5 = eq v0, Field 2 - jmpif v5 then: b5, else: b4 - b3(): + jmpif v5 then: b4, else: b3 + b2(): call f1() - jmp b1() - b4(): + jmp b5() + b3(): constrain v0 == Field 5 call f5() - jmp b1() - b5(): + jmp b5() + b4(): call f2() - jmp b1() + jmp b5() + b5(): + return } "); } diff --git a/compiler/noirc_evaluator/src/ssa/opt/simplify_cfg.rs b/compiler/noirc_evaluator/src/ssa/opt/simplify_cfg.rs index b73fcb34a8e..4bbefb42c5b 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/simplify_cfg.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/simplify_cfg.rs @@ -60,7 +60,7 @@ impl Function { // First perform any simplifications on the current block, this ensures that we add the proper successors // to the stack. - simplify_current_block(self, block, &mut cfg); + simplify_current_block(self, block, &mut cfg, &mut values_to_replace); if visited.insert(block) { stack.extend(self.dfg[block].successors().filter(|block| !visited.contains(block))); @@ -72,16 +72,7 @@ impl Function { predecessors.next().expect("Already checked length of predecessors"); drop(predecessors); - // If the block has only 1 predecessor, we can safely remove its block parameters - remove_block_parameters(self, block, predecessor, &mut values_to_replace); - - // Note: this function relies on `remove_block_parameters` being called first. - // Otherwise the inlined block will refer to parameters that no longer exist. - // - // If successful, `block` will be empty and unreachable after this call, so any - // optimizations performed after this point on the same block should check if - // the inlining here was successful before continuing. - try_inline_into_predecessor(self, &mut cfg, block, predecessor); + try_inline_successor(self, &mut cfg, predecessor, &mut values_to_replace); } else { drop(predecessors); @@ -110,16 +101,16 @@ fn simplify_current_block( function: &mut Function, block: BasicBlockId, cfg: &mut ControlFlowGraph, + values_to_replace: &mut ValueMapping, ) { // These functions return `true` if they successfully simplified the CFG for the current block. - let mut simplified = false; - - simplified |= check_for_negated_jmpif_condition(function, block, cfg); - simplified |= check_for_constant_jmpif(function, block, cfg); - simplified |= check_for_converging_jmpif(function, block, cfg); + let mut simplified = true; - if simplified { - simplify_current_block(function, block, cfg); + while simplified { + simplified = check_for_negated_jmpif_condition(function, block, cfg) + | check_for_constant_jmpif(function, block, cfg) + | check_for_converging_jmpif(function, block, cfg) + | try_inline_successor(function, cfg, block, values_to_replace); } } @@ -391,6 +382,40 @@ fn remove_block_parameters( } } +/// Try to inline a block into its predecessor, returning true if successful. +/// +/// This will only occur if the predecessor's only successor is the given block. +/// It is also expected that the given block's only predecessor is the given one. +fn try_inline_successor( + function: &mut Function, + cfg: &mut ControlFlowGraph, + block: BasicBlockId, + values_to_replace: &mut ValueMapping, +) -> bool { + if let Some(TerminatorInstruction::Jmp { destination, .. }) = function.dfg[block].terminator() { + let destination = *destination; + let predecessors = cfg.predecessors(destination); + if predecessors.len() == 1 { + drop(predecessors); + + // If the block has only 1 predecessor, we can safely remove its block parameters + remove_block_parameters(function, destination, block, values_to_replace); + + // Note: this function relies on `remove_block_parameters` being called first. + // Otherwise the inlined block will refer to parameters that no longer exist. + // + // If successful, `block` will be empty and unreachable after this call, so any + // optimizations performed after this point on the same block should check if + // the inlining here was successful before continuing. + try_inline_into_predecessor(function, cfg, destination, block) + } else { + false + } + } else { + false + } +} + /// Try to inline a block into its predecessor, returning true if successful. /// /// This will only occur if the predecessor's only successor is the given block. @@ -647,11 +672,11 @@ mod test { brillig(inline) predicate_pure fn main f0 { b0(v0: i16): v2 = lt i16 1, v0 - jmpif v2 then: b2, else: b1 + jmpif v2 then: b1, else: b2 b1(): - return u32 2 + jmp b1() b2(): - jmp b2() + return u32 2 } "); } @@ -916,6 +941,30 @@ mod test { "); } + #[test] + fn double_jmp_with_args_blocks() { + let src = " + brillig(inline) fn test f0 { + b0(v0: Field): + jmp b1(v0, Field 2) + b1(v1: Field, v2: Field): + jmp b2(v1) + b2(v3: Field): + return v3 + } + "; + + let ssa = Ssa::from_str(src).unwrap(); + let ssa = ssa.simplify_cfg(); + + assert_ssa_snapshot!(ssa, @r" + brillig(inline) fn test f0 { + b0(v0: Field): + return v0 + } + "); + } + #[test] fn deep_jmp_empty_blocks() { let src = "