feat(ir): split critical edges in ir

src/backend/lower.rs

@@ -694,71 +694,59 @@ where
                 let then_succ = inst.succ(self.ctx, 0);
                 let else_succ = inst.succ(self.ctx, 1);
 
+                let then_block = then_succ.block();
+                let else_block = else_succ.block();
+
                 let mblock_then = self.blocks[&then_succ.block()];
                 let mblock_else = self.blocks[&else_succ.block()];
 
-                // TODO: we can do better if there are no lost-copy problem between the two
-                // successors.
+                let this_block = self.curr_block.unwrap();
+
                 if then_succ.args().is_empty() {
                     // if there are no arugments to pass, we can just br, without creating a new
                     // block because no lost-copy problem will happen
                     S::gen_br(self, mcond, mblock_then);
-
-                    if else_succ.args().is_empty() {
-                        // if there are no arguments to pass, we can just br, without creating a new
-                        // block because no lost-copy problem will happen
-                        S::gen_jump(self, mblock_else);
-                    } else {
-                        // create a new block to jump
-                        let mblock_else_arg_passing = MBlock::new(
-                            &mut self.mctx,
-                            format!(".__machbb_{}", self.label_counter),
-                        );
-                        self.label_counter += 1;
-
-                        S::gen_jump(self, mblock_else_arg_passing);
-
-                        mblock_else.insert_before(self.mctx_mut(), mblock_else_arg_passing);
-
-                        self.curr_block = Some(mblock_else_arg_passing);
-
-                        self.lower_succ(else_succ);
-                        S::gen_jump(self, mblock_else);
-                    }
-                } else {
+                } else if then_block.preds(self.ctx).len() > 1 {
                     let mblock_then_arg_passing =
                         MBlock::new(&mut self.mctx, format!(".__machbb_{}", self.label_counter));
                     self.label_counter += 1;
 
-                    let this_mblock = self.curr_block.unwrap();
-                    this_mblock.insert_after(self.mctx_mut(), mblock_then_arg_passing);
-
+                    this_block.insert_after(self.mctx_mut(), mblock_then_arg_passing);
                     S::gen_br(self, mcond, mblock_then_arg_passing);
 
-                    if else_succ.args().is_empty() {
-                        // if there are no arguments to pass, we can just br, without creating a new
-                        // block because no lost-copy problem will happen
-                        S::gen_jump(self, mblock_else);
-                    } else {
-                        let mblock_else_arg_passing = MBlock::new(
-                            &mut self.mctx,
-                            format!(".__machbb_{}", self.label_counter),
-                        );
-                        self.label_counter += 1;
+                    self.curr_block = Some(mblock_then_arg_passing);
+                    self.lower_succ(then_succ);
+                    S::gen_jump(self, mblock_then);
+                } else {
+                    // pass arguments at the beginning of the destination block
+                    S::gen_br(self, mcond, mblock_then);
+                    self.curr_block = Some(mblock_then);
+                    self.lower_succ(then_succ);
+                }
 
-                        mblock_else.insert_before(self.mctx_mut(), mblock_else_arg_passing);
+                self.curr_block = Some(this_block);
 
-                        S::gen_jump(self, mblock_else_arg_passing);
+                if else_succ.args().is_empty() {
+                    // if there are no arguments to pass, we can just br, without creating a new
+                    // block because no lost-copy problem will happen
+                    S::gen_jump(self, mblock_else);
+                } else if else_block.preds(self.ctx).len() > 1 {
+                    let mblock_else_arg_passing =
+                        MBlock::new(&mut self.mctx, format!(".__machbb_{}", self.label_counter));
+                    self.label_counter += 1;
 
-                        self.curr_block = Some(mblock_else_arg_passing);
-                        self.lower_succ(else_succ);
-                        S::gen_jump(self, mblock_else);
-                    }
+                    S::gen_jump(self, mblock_else_arg_passing);
 
-                    self.curr_block = Some(mblock_then_arg_passing);
-                    self.lower_succ(then_succ);
+                    mblock_else.insert_before(self.mctx_mut(), mblock_else_arg_passing);
 
-                    S::gen_jump(self, mblock_then);
+                    self.curr_block = Some(mblock_else_arg_passing);
+                    self.lower_succ(else_succ);
+                    S::gen_jump(self, mblock_else);
+                } else {
+                    // pass arguments at the beginning of the destination block
+                    S::gen_jump(self, mblock_else);
+                    self.curr_block = Some(mblock_else);
+                    self.lower_succ(else_succ);
                 }
             }
             Ik::Call(symbol) => {

src/bin/compiler.rs

@@ -20,9 +20,11 @@ use orzcc::{
                 CfgCanonicalize,
                 CfgSimplify,
                 PHBlockLayout,
+                SplitCriticalEdge,
                 BLOCK_REORDER,
                 CFG_SIMPLIFY,
                 PH_BLOCK_LAYOUT,
+                SPLIT_CRITICAL_EDGE,
             },
             fold::{ConstantFolding, CONSTANT_FOLDING},
             gcm::{Gcm, GCM},
@@ -325,6 +327,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
             passman.run_transform(SIMPLE_DCE, &mut ir, 32);
             passman.run_transform(BRANCH_CONDITION_SINK, &mut ir, 1);
             passman.run_transform(BLOCK_REORDER, &mut ir, 1);
+            passman.run_transform(SPLIT_CRITICAL_EDGE, &mut ir, 1);
             passman.run_transform(PH_BLOCK_LAYOUT, &mut ir, 1);
         } else {
             passman.run_transform(LEGALIZE, &mut ir, 1);
@@ -416,6 +419,7 @@ fn register_passes(passman: &mut PassManager) {
 
     Legalize::register(passman);
     BlockReorder::register(passman);
+    SplitCriticalEdge::register(passman);
 }
 
 fn cli(passman: &mut PassManager) -> Command {

src/ir/inst.rs

@@ -1323,6 +1323,45 @@ impl Inst {
         }
     }
 
+    pub fn replace_single_succ_with_args(
+        self,
+        ctx: &mut Context,
+        old_idx: usize,
+        new: Block,
+        args: Vec<Value>
+    ) {
+        if !self.is_terminator(ctx) {
+            panic!("instruction is not a terminator");
+        }
+
+        let new_params = new.params(ctx).to_vec();
+        let successor = self.deref_mut(ctx).successors.get_mut(old_idx).unwrap();
+
+        let old_block = successor.block();
+        let set = successor.block.set_inner_if_eq(old_block, new);
+        assert!(set);
+
+        let mut args_to_drop = Vec::new();
+
+        for (_, arg) in successor.args.drain() {
+            args_to_drop.push(arg);
+        }
+
+        let mut new_args = FxHashMap::default();
+        for (param, arg) in new_params.iter().zip(args) {
+            new_args.insert(*param, Operand::new(ctx, arg, self));
+        }
+
+        self.deref_mut(ctx).successors.get_mut(old_idx).unwrap().args = new_args;
+
+        for arg in args_to_drop {
+            arg.drop(ctx);
+        }
+
+        old_block.remove_user(ctx, self);
+        new.add_user(ctx, self);
+    }
+
     /// Replace all the arguments passed to the `param` with the `arg`.
     pub fn replace_args(self, ctx: &mut Context, param: Value, arg: Value) {
         let mut operands_to_drop = Vec::new();

src/ir/passes/control_flow/mod.rs

@@ -2,8 +2,10 @@ mod canonicalize;
 mod ph_layout;
 mod reorder;
 mod simplify;
+mod split_critical_edge;
 
 pub use canonicalize::{CfgCanonicalize, CfgCanonicalizeError, CFG_CANONICALIZE};
 pub use ph_layout::{PHBlockLayout, PH_BLOCK_LAYOUT};
 pub use reorder::{BlockReorder, BLOCK_REORDER};
 pub use simplify::{CfgSimplify, CFG_SIMPLIFY};
+pub use split_critical_edge::{SplitCriticalEdge, SPLIT_CRITICAL_EDGE};

src/ir/passes/control_flow/split_critical_edge.rs

@@ -0,0 +1,98 @@
+use rustc_hash::FxHashMap;
+
+use super::CfgCanonicalize;
+use crate::{
+    collections::linked_list::{LinkedListContainerPtr, LinkedListNodePtr},
+    ir::{
+        passman::{GlobalPassMut, LocalPassMut, PassResult, TransformPass},
+        Block,
+        Context,
+        Func,
+        Inst,
+    },
+};
+
+pub const SPLIT_CRITICAL_EDGE: &str = "split-critical-edge";
+
+pub struct SplitCriticalEdge;
+
+impl LocalPassMut for SplitCriticalEdge {
+    type Output = ();
+
+    fn run(&mut self, ctx: &mut Context, func: Func) -> PassResult<(Self::Output, bool)> {
+        let mut cursor = func.cursor();
+
+        let mut critical_edges = FxHashMap::default();
+
+        while let Some(block) = cursor.next(ctx) {
+            let tail = block.tail(ctx).unwrap();
+
+            if tail.succs(ctx).len() <= 1 {
+                continue;
+            }
+
+            for (i, successor) in tail.succs(ctx).iter().enumerate() {
+                let succ = successor.block();
+                if succ.params(ctx).is_empty() {
+                    continue;
+                }
+                if succ.preds(ctx).len() > 1 {
+                    critical_edges.entry(tail).or_insert_with(Vec::new).push(i);
+                }
+            }
+        }
+
+        for (inst, edges) in critical_edges {
+            let curr_block = inst.container(ctx).unwrap();
+
+            let name = curr_block.name_or_alloc(ctx, "bb").clone();
+
+            for succ_idx in edges {
+                let succ = inst.succ(ctx, succ_idx).block();
+                let params = succ.params(ctx).to_vec();
+                let mut args = Vec::new();
+
+                for param in params {
+                    args.push(inst.succ(ctx, succ_idx).get_arg(param).unwrap());
+                }
+
+                let new_block = Block::new(ctx);
+
+                new_block.alloc_name(ctx, format!("{}.critical.{}_", name, succ_idx));
+
+                curr_block.insert_after(ctx, new_block);
+                inst.replace_single_succ_with_args(ctx, succ_idx, new_block, Vec::new());
+
+                let jump = Inst::jump(ctx, succ, args);
+                new_block.push_back(ctx, jump);
+            }
+        }
+
+        Ok(((), false)) // only run once.
+    }
+}
+
+impl GlobalPassMut for SplitCriticalEdge {
+    type Output = ();
+
+    fn run(&mut self, ctx: &mut Context) -> PassResult<(Self::Output, bool)> {
+        let mut changed = false;
+
+        for func in ctx.funcs() {
+            let ((), local_changed) = LocalPassMut::run(self, ctx, func)?;
+            changed |= local_changed;
+        }
+
+        Ok(((), changed))
+    }
+}
+
+impl TransformPass for SplitCriticalEdge {
+    fn register(passman: &mut crate::ir::passman::PassManager) {
+        passman.register_transform(
+            SPLIT_CRITICAL_EDGE,
+            SplitCriticalEdge,
+            vec![Box::new(CfgCanonicalize)],
+        );
+    }
+}