From f391af2d61f4a38e02cb92c76fa4c2c148af3833 Mon Sep 17 00:00:00 2001
From: Tom French <15848336+TomAFrench@users.noreply.github.com>
Date: Thu, 5 Sep 2024 16:22:56 +0100
Subject: [PATCH] feat: remove blocks which consist of only a jump to another
 block (#5889)

# Description

## Problem\*

Resolves <!-- Link to GitHub Issue -->

## Summary\*

This PR removes blocks which are empty and just perform an unconditional
jump to another block. We instead just update any references to this
block with the block which it will eventually jump to.

I initially tried applying this optimization universally but ran into
troubles where ACIR functions would run into early returns due before
branches reached a common block again and so the cfg flattening pass
would fail. This optimisation then only applies to brillig functions
(although it wouldn't really affect ACIR anyway).

## Additional Context



## Documentation\*

Check one:
- [ ] No documentation needed.
- [ ] Documentation included in this PR.
- [ ] **[For Experimental Features]** Documentation to be submitted in a
separate PR.

# PR Checklist\*

- [ ] I have tested the changes locally.
- [ ] I have formatted the changes with [Prettier](https://prettier.io/)
and/or `cargo fmt` on default settings.
---
 .../src/ssa/opt/simplify_cfg.rs               | 75 ++++++++++++++++++-
 1 file changed, 73 insertions(+), 2 deletions(-)

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,