diff --git a/compiler/noirc_evaluator/src/ssa/opt/simplify_cfg.rs b/compiler/noirc_evaluator/src/ssa/opt/simplify_cfg.rs index 0a3b2a800ca..6887873dbc3 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/simplify_cfg.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/simplify_cfg.rs @@ -15,7 +15,9 @@ use acvm::acir::AcirField; use crate::ssa::{ ir::{ - basic_block::BasicBlockId, cfg::ControlFlowGraph, function::Function, + basic_block::BasicBlockId, + cfg::ControlFlowGraph, + function::{Function, RuntimeType}, instruction::TerminatorInstruction, }, ssa_gen::Ssa, @@ -30,7 +32,7 @@ impl Ssa { /// 5. Replacing any jmpifs with constant conditions with jmps. If this causes the block to have /// only 1 successor then (2) also will be applied. /// - /// Currently, 1 and 4 are unimplemented. + /// Currently, 1 is unimplemented. #[tracing::instrument(level = "trace", skip(self))] pub(crate) fn simplify_cfg(mut self) -> Self { for function in self.functions.values_mut() { @@ -72,6 +74,10 @@ fn simplify_function(function: &mut Function) { // optimizations performed after this point on the same block should check if // the inlining here was successful before continuing. try_inline_into_predecessor(function, &mut cfg, block, predecessor); + } else { + drop(predecessors); + + check_for_double_jmp(function, block, &mut cfg); } } } @@ -102,6 +108,71 @@ fn check_for_constant_jmpif( } } +/// Optimize a jmp to a block which immediately jmps elsewhere to just jmp to the second block. +fn check_for_double_jmp(function: &mut Function, block: BasicBlockId, cfg: &mut ControlFlowGraph) { + if matches!(function.runtime(), RuntimeType::Acir(_)) { + // We can't remove double jumps in ACIR functions as this interferes with the `flatten_cfg` pass. + return; + } + + if !function.dfg[block].instructions().is_empty() + || !function.dfg[block].parameters().is_empty() + { + return; + } + + let Some(TerminatorInstruction::Jmp { destination: final_destination, arguments, .. }) = + function.dfg[block].terminator() + else { + return; + }; + + if !arguments.is_empty() { + return; + } + + let final_destination = *final_destination; + + let predecessors: Vec<_> = cfg.predecessors(block).collect(); + for predecessor_block in predecessors { + let terminator_instruction = function.dfg[predecessor_block].take_terminator(); + let redirected_terminator_instruction = match terminator_instruction { + TerminatorInstruction::JmpIf { + condition, + then_destination, + else_destination, + call_stack, + } => { + let then_destination = + if then_destination == block { final_destination } else { then_destination }; + let else_destination = + if else_destination == block { final_destination } else { else_destination }; + TerminatorInstruction::JmpIf { + condition, + then_destination, + else_destination, + call_stack, + } + } + TerminatorInstruction::Jmp { destination, arguments, call_stack } => { + assert_eq!( + destination, block, + "ICE: predecessor block doesn't jump to current block" + ); + assert!(arguments.is_empty(), "ICE: predecessor jmp has arguments"); + TerminatorInstruction::Jmp { destination: final_destination, arguments, call_stack } + } + TerminatorInstruction::Return { .. } => { + unreachable!("ICE: predecessor block should not have return terminator instruction") + } + }; + + function.dfg[predecessor_block].set_terminator(redirected_terminator_instruction); + cfg.recompute_block(function, predecessor_block); + } + cfg.recompute_block(function, block); +} + /// If the given block has block parameters, replace them with the jump arguments from the predecessor. /// /// Currently, if this function is needed, `try_inline_into_predecessor` will also always apply,