From c908463ee10da16d9159182e851846b817af0dd1 Mon Sep 17 00:00:00 2001
From: Chris Fallin <chris@cfallin.org>
Date: Tue, 7 Feb 2023 12:38:02 -0800
Subject: [PATCH] Add remove-empty-blocks pass.

---
 src/ir/func.rs             |  17 +++++
 src/ir/module.rs           |   1 +
 src/passes.rs              |   3 +-
 src/passes/empty_blocks.rs | 128 +++++++++++++++++++++++++++++++++++++
 4 files changed, 148 insertions(+), 1 deletion(-)
 create mode 100644 src/passes/empty_blocks.rs

diff --git a/src/ir/func.rs b/src/ir/func.rs
index 13ced6c..183b854 100644
--- a/src/ir/func.rs
+++ b/src/ir/func.rs
@@ -94,6 +94,23 @@ impl FunctionBody {
         log::trace!("add_edge: from {} to {}", from, to);
     }
 
+    pub fn recompute_edges(&mut self) {
+        for block in self.blocks.values_mut() {
+            block.preds.clear();
+            block.succs.clear();
+            block.pos_in_succ_pred.clear();
+            block.pos_in_pred_succ.clear();
+        }
+
+        for block in 0..self.blocks.len() {
+            let block = Block::new(block);
+            let terminator = self.blocks[block].terminator.clone();
+            terminator.visit_successors(|succ| {
+                self.add_edge(block, succ);
+            });
+        }
+    }
+
     pub fn add_value(&mut self, value: ValueDef) -> Value {
         log::trace!("add_value: def {:?}", value);
         let value = self.values.push(value);
diff --git a/src/ir/module.rs b/src/ir/module.rs
index ad430b7..b72fbee 100644
--- a/src/ir/module.rs
+++ b/src/ir/module.rs
@@ -241,6 +241,7 @@ impl<'a> Module<'a> {
             let cfg = crate::cfg::CFGInfo::new(body);
             crate::passes::basic_opt::gvn(body, &cfg);
             crate::passes::resolve_aliases::run(body);
+            crate::passes::empty_blocks::run(body);
         });
     }
 
diff --git a/src/passes.rs b/src/passes.rs
index 96f923f..c001335 100644
--- a/src/passes.rs
+++ b/src/passes.rs
@@ -2,5 +2,6 @@
 
 pub mod basic_opt;
 pub mod dom_pass;
-pub mod resolve_aliases;
+pub mod empty_blocks;
 pub mod maxssa;
+pub mod resolve_aliases;
diff --git a/src/passes/empty_blocks.rs b/src/passes/empty_blocks.rs
new file mode 100644
index 0000000..1a7556b
--- /dev/null
+++ b/src/passes/empty_blocks.rs
@@ -0,0 +1,128 @@
+//! Pass to remove empty blocks.
+
+use crate::entity::EntityRef;
+use crate::ir::{Block, BlockTarget, FunctionBody, Terminator, Value, ValueDef};
+use std::borrow::Cow;
+use std::collections::HashSet;
+
+#[derive(Clone, Debug)]
+struct Forwarding {
+    to: Block,
+    args: Vec<ForwardingArg>,
+}
+
+#[derive(Clone, Copy, Debug)]
+enum ForwardingArg {
+    BlockParam(usize),
+    Value(Value),
+}
+
+impl Forwarding {
+    fn compose(a: &Forwarding, b: &Forwarding) -> Forwarding {
+        // `b` should be the target of `a.to`, but we can't assert
+        // that here. The composed target is thus `b.to`.
+        let to = b.to;
+
+        // For each arg in `b.args`, evaluate, replacing any
+        // `BlockParam` with the corresponding value from `a.args`.
+        let args = b
+            .args
+            .iter()
+            .map(|&arg| match arg {
+                ForwardingArg::BlockParam(idx) => a.args[idx].clone(),
+                ForwardingArg::Value(v) => ForwardingArg::Value(v),
+            })
+            .collect::<Vec<_>>();
+
+        Forwarding { to, args }
+    }
+}
+
+fn block_to_forwarding(body: &FunctionBody, block: Block) -> Option<Forwarding> {
+    // Must be empty except for terminator, and must have an
+    // unconditional-branch terminator.
+    if body.blocks[block].insts.len() > 0 {
+        return None;
+    }
+    let target = match &body.blocks[block].terminator {
+        &Terminator::Br { ref target } => target,
+        _ => return None,
+    };
+
+    // If conditions met, then gather ForwardingArgs.
+    let args = target
+        .args
+        .iter()
+        .map(|&arg| {
+            let arg = body.resolve_alias(arg);
+            match &body.values[arg] {
+                &ValueDef::BlockParam(param_block, index, _) if param_block == block => {
+                    ForwardingArg::BlockParam(index)
+                }
+                _ => ForwardingArg::Value(arg),
+            }
+        })
+        .collect::<Vec<_>>();
+
+    Some(Forwarding {
+        to: target.block,
+        args,
+    })
+}
+
+fn rewrite_target(forwardings: &[Option<Forwarding>], target: &BlockTarget) -> Option<BlockTarget> {
+    if !forwardings[target.block.index()].is_some() {
+        return None;
+    }
+
+    let mut forwarding = Cow::Borrowed(forwardings[target.block.index()].as_ref().unwrap());
+    let mut seen = HashSet::new();
+    while forwardings[forwarding.to.index()].is_some() && seen.insert(forwarding.to.index()) {
+        forwarding = Cow::Owned(Forwarding::compose(
+            &forwarding,
+            forwardings[forwarding.to.index()].as_ref().unwrap(),
+        ));
+    }
+
+    let args = forwarding
+        .args
+        .iter()
+        .map(|arg| match arg {
+            &ForwardingArg::Value(v) => v,
+            &ForwardingArg::BlockParam(idx) => target.args[idx],
+        })
+        .collect::<Vec<_>>();
+
+    Some(BlockTarget {
+        block: forwarding.to,
+        args,
+    })
+}
+
+pub fn run(body: &mut FunctionBody) {
+    // Identify empty blocks, and to where they should forward.
+    let forwardings = body
+        .blocks
+        .iter()
+        .map(|block| {
+            if block != body.entry {
+                block_to_forwarding(body, block)
+            } else {
+                None
+            }
+        })
+        .collect::<Vec<_>>();
+
+    // Rewrite every target according to a forwarding (or potentially
+    // a chain of composed forwardings).
+    for block_data in body.blocks.values_mut() {
+        block_data.terminator.update_targets(|target| {
+            if let Some(new_target) = rewrite_target(&forwardings[..], target) {
+                *target = new_target;
+            }
+        });
+    }
+
+    // Recompute preds/succs.
+    body.recompute_edges();
+}