feat(cfg): Support cfg analyse

src/ir/cfg.rs

@@ -0,0 +1,446 @@
+use std::collections::{BTreeMap, BTreeSet, HashSet};
+
+use petgraph::dot::Dot;
+use petgraph::graph::Graph;
+
+use crate::ir::types::IRInstr;
+
+pub type CfgGraph = Graph<String, String>;
+
+#[derive(Clone)]
+pub struct BasicBlock {
+    pub id: usize,
+    pub instrs: Vec<IRInstr>,
+}
+
+pub struct ControlFlowGraph {
+    pub func_name: String,
+    pub blocks: Vec<BasicBlock>,
+    pub edges: Vec<(usize, usize, String)>,
+}
+
+#[derive(Clone, Copy, PartialEq, Eq)]
+enum EdgeKind {
+    Branch,
+    Fallthrough,
+}
+
+impl ControlFlowGraph {
+    pub fn from_body(func_name: impl Into<String>, body: &[IRInstr]) -> Self {
+        let blocks = split_basic_blocks(body);
+        let edges = collect_edges(&blocks)
+            .into_iter()
+            .map(|(from, to, label, _)| (from, to, label))
+            .collect();
+        Self {
+            func_name: func_name.into(),
+            blocks,
+            edges,
+        }
+    }
+
+    fn add_to_graph(&self, graph: &mut CfgGraph) {
+        let mut nodes = Vec::new();
+        for block in &self.blocks {
+            nodes.push(graph.add_node(block_dot_label(&self.func_name, block)));
+        }
+        for (from, to, label) in &self.edges {
+            if let (Some(from), Some(to)) = (nodes.get(*from), nodes.get(*to)) {
+                graph.add_edge(*from, *to, label.clone());
+            }
+        }
+    }
+
+    fn refresh_edges(&mut self) {
+        self.edges = collect_edges(&self.blocks)
+            .into_iter()
+            .map(|(from, to, label, _)| (from, to, label))
+            .collect();
+    }
+
+    fn remove_unreachable(&mut self) -> bool {
+        if self.blocks.is_empty() {
+            return false;
+        }
+
+        let edges = collect_edges(&self.blocks);
+        let mut adjacency = vec![Vec::new(); self.blocks.len()];
+        for (from, to, _, _) in edges {
+            adjacency[from].push(to);
+        }
+
+        let mut reachable = vec![false; self.blocks.len()];
+        let mut stack = vec![0];
+        while let Some(index) = stack.pop() {
+            if reachable[index] {
+                continue;
+            }
+            reachable[index] = true;
+            for next in &adjacency[index] {
+                stack.push(*next);
+            }
+        }
+
+        let old_len = self.blocks.len();
+        self.blocks = self
+            .blocks
+            .drain(..)
+            .enumerate()
+            .filter_map(|(index, block)| reachable[index].then_some(block))
+            .collect();
+        let changed = self.blocks.len() != old_len;
+        if changed {
+            self.refresh_edges();
+        }
+        changed
+    }
+
+    fn remove_trivial_blocks(&mut self) -> bool {
+        if self.blocks.is_empty() {
+            return false;
+        }
+
+        let edges = collect_edges(&self.blocks);
+        let explicit_pred_count = explicit_predecessor_counts(self.blocks.len(), &edges);
+        let mut redirect = BTreeMap::new();
+        let mut remove = HashSet::new();
+        let mut fallthrough_gotos = Vec::new();
+
+        for (index, block) in self.blocks.iter().enumerate() {
+            if let Some((label, target)) = jump_only_block(block) {
+                let target = resolve_label(target, &redirect);
+                redirect.insert(label, target);
+                for (from, to, _, kind) in &edges {
+                    if *to == index && *kind == EdgeKind::Fallthrough {
+                        fallthrough_gotos.push((*from, target));
+                    }
+                }
+                remove.insert(index);
+                continue;
+            }
+
+            if let Some(label) = label_only_block(block) {
+                if let Some(next_index) = next_kept_index(index, self.blocks.len(), &remove) {
+                    if let Some(next_label) = first_label(&self.blocks[next_index]) {
+                        redirect.insert(label, resolve_label(next_label, &redirect));
+                        remove.insert(index);
+                    } else if explicit_pred_count[index] == 0 {
+                        remove.insert(index);
+                    }
+                } else if explicit_pred_count[index] == 0 {
+                    remove.insert(index);
+                }
+            }
+        }
+
+        if redirect.is_empty() && remove.is_empty() && fallthrough_gotos.is_empty() {
+            return false;
+        }
+
+        for (from, target) in fallthrough_gotos {
+            if !remove.contains(&from) && !ends_with_terminal(&self.blocks[from]) {
+                self.blocks[from].instrs.push(IRInstr::Goto(target));
+            }
+        }
+
+        for block in &mut self.blocks {
+            rewrite_targets(&mut block.instrs, &redirect);
+        }
+
+        self.blocks = self
+            .blocks
+            .drain(..)
+            .enumerate()
+            .filter_map(|(index, block)| (!remove.contains(&index)).then_some(block))
+            .collect();
+        self.refresh_edges();
+        true
+    }
+
+    fn merge_linear_blocks(&mut self) -> bool {
+        let edges = collect_edges(&self.blocks);
+        let pred_count = predecessor_counts(self.blocks.len(), &edges);
+
+        let mut index = 0;
+        while index + 1 < self.blocks.len() {
+            let successors = edges
+                .iter()
+                .filter_map(|(from, to, _, _)| (*from == index).then_some(*to))
+                .collect::<BTreeSet<_>>();
+
+            if successors.len() == 1 {
+                let successor = *successors.iter().next().unwrap();
+                if successor == index + 1 && pred_count[successor] == 1 {
+                    if matches!(self.blocks[index].instrs.last(), Some(IRInstr::Goto(_))) {
+                        self.blocks[index].instrs.pop();
+                    }
+                    let successor_block = self.blocks.remove(successor);
+                    self.blocks[index].instrs.extend(successor_block.instrs);
+                    self.refresh_edges();
+                    return true;
+                }
+            }
+            index += 1;
+        }
+
+        false
+    }
+
+    fn into_instrs(self) -> Vec<IRInstr> {
+        self.blocks
+            .into_iter()
+            .flat_map(|block| block.instrs)
+            .collect()
+    }
+}
+
+pub fn optimize_ir(ir: Vec<IRInstr>) -> Vec<IRInstr> {
+    ir.into_iter()
+        .map(|instr| match instr {
+            IRInstr::DefineFunc(func, args, body) => {
+                let func_name = func.name.clone();
+                IRInstr::DefineFunc(func, args, optimize_body(&func_name, body))
+            }
+            other => other,
+        })
+        .collect()
+}
+
+pub fn ir_to_cfg_dot(ir: &[IRInstr]) -> String {
+    let mut graph = Graph::new();
+    for instr in ir {
+        if let IRInstr::DefineFunc(func, _, body) = instr {
+            ControlFlowGraph::from_body(func.name.clone(), body).add_to_graph(&mut graph);
+        }
+    }
+    format!("{}", Dot::new(&graph))
+}
+
+pub fn optimize_body(func_name: &str, body: Vec<IRInstr>) -> Vec<IRInstr> {
+    let (mut prologue, cfg_body) = split_function_prologue(body);
+    if cfg_body.is_empty() {
+        return prologue;
+    }
+
+    let mut cfg = ControlFlowGraph::from_body(func_name, &cfg_body);
+    cfg.remove_unreachable();
+    loop {
+        let mut changed = false;
+        changed |= cfg.remove_trivial_blocks();
+        changed |= cfg.remove_unreachable();
+        changed |= cfg.merge_linear_blocks();
+        if !changed {
+            break;
+        }
+    }
+
+    let mut body = cfg.into_instrs();
+    remove_redundant_fallthrough_gotos(&mut body);
+    remove_unused_labels(&mut body);
+    prologue.extend(body);
+    prologue
+}
+
+fn split_function_prologue(mut body: Vec<IRInstr>) -> (Vec<IRInstr>, Vec<IRInstr>) {
+    if let Some(entry_index) = body.iter().position(|instr| matches!(instr, IRInstr::Entry)) {
+        let cfg_body = body.split_off(entry_index + 1);
+        (body, cfg_body)
+    } else {
+        (Vec::new(), body)
+    }
+}
+
+fn split_basic_blocks(instrs: &[IRInstr]) -> Vec<BasicBlock> {
+    if instrs.is_empty() {
+        return Vec::new();
+    }
+
+    let mut leaders = BTreeSet::new();
+    leaders.insert(0);
+    for (index, instr) in instrs.iter().enumerate() {
+        if matches!(instr, IRInstr::Label(_)) {
+            leaders.insert(index);
+        }
+        if is_terminal(instr) && index + 1 < instrs.len() {
+            leaders.insert(index + 1);
+        }
+    }
+
+    let leader_list = leaders.into_iter().collect::<Vec<_>>();
+    let mut blocks = Vec::new();
+    for (block_id, start) in leader_list.iter().enumerate() {
+        let end = leader_list
+            .get(block_id + 1)
+            .copied()
+            .unwrap_or(instrs.len());
+        blocks.push(BasicBlock {
+            id: block_id,
+            instrs: instrs[*start..end].to_vec(),
+        });
+    }
+    blocks
+}
+
+fn collect_edges(blocks: &[BasicBlock]) -> Vec<(usize, usize, String, EdgeKind)> {
+    let label_to_block = label_to_block(blocks);
+    let mut edges = Vec::new();
+
+    for (index, block) in blocks.iter().enumerate() {
+        match block.instrs.last() {
+            Some(IRInstr::Goto(label)) => {
+                if let Some(target) = label_to_block.get(label) {
+                    edges.push((index, *target, "br".to_string(), EdgeKind::Branch));
+                }
+            }
+            Some(IRInstr::CondGoto(_, true_label, false_label)) => {
+                if let Some(target) = label_to_block.get(true_label) {
+                    edges.push((index, *target, "true".to_string(), EdgeKind::Branch));
+                }
+                if let Some(target) = label_to_block.get(false_label) {
+                    edges.push((index, *target, "false".to_string(), EdgeKind::Branch));
+                }
+            }
+            Some(IRInstr::Exit(_)) => {}
+            Some(_) | None => {
+                if index + 1 < blocks.len() {
+                    edges.push((index, index + 1, "fallthrough".to_string(), EdgeKind::Fallthrough));
+                }
+            }
+        }
+    }
+
+    edges
+}
+
+fn label_to_block(blocks: &[BasicBlock]) -> BTreeMap<usize, usize> {
+    let mut map = BTreeMap::new();
+    for (index, block) in blocks.iter().enumerate() {
+        for instr in &block.instrs {
+            if let IRInstr::Label(label) = instr {
+                map.insert(*label, index);
+            }
+        }
+    }
+    map
+}
+
+fn predecessor_counts(block_count: usize, edges: &[(usize, usize, String, EdgeKind)]) -> Vec<usize> {
+    let mut counts = vec![0; block_count];
+    for (_, to, _, _) in edges {
+        counts[*to] += 1;
+    }
+    counts
+}
+
+fn explicit_predecessor_counts(block_count: usize, edges: &[(usize, usize, String, EdgeKind)]) -> Vec<usize> {
+    let mut counts = vec![0; block_count];
+    for (_, to, _, kind) in edges {
+        if *kind == EdgeKind::Branch {
+            counts[*to] += 1;
+        }
+    }
+    counts
+}
+
+fn next_kept_index(index: usize, block_count: usize, remove: &HashSet<usize>) -> Option<usize> {
+    (index + 1..block_count).find(|candidate| !remove.contains(candidate))
+}
+
+fn first_label(block: &BasicBlock) -> Option<usize> {
+    match block.instrs.first() {
+        Some(IRInstr::Label(label)) => Some(*label),
+        _ => None,
+    }
+}
+
+fn label_only_block(block: &BasicBlock) -> Option<usize> {
+    match block.instrs.as_slice() {
+        [IRInstr::Label(label)] => Some(*label),
+        _ => None,
+    }
+}
+
+fn jump_only_block(block: &BasicBlock) -> Option<(usize, usize)> {
+    match block.instrs.as_slice() {
+        [IRInstr::Label(label), IRInstr::Goto(target)] => Some((*label, *target)),
+        _ => None,
+    }
+}
+
+fn is_terminal(instr: &IRInstr) -> bool {
+    matches!(instr, IRInstr::Goto(_) | IRInstr::CondGoto(_, _, _) | IRInstr::Exit(_))
+}
+
+fn ends_with_terminal(block: &BasicBlock) -> bool {
+    block.instrs.last().is_some_and(is_terminal)
+}
+
+fn rewrite_targets(instrs: &mut [IRInstr], redirect: &BTreeMap<usize, usize>) {
+    for instr in instrs {
+        match instr {
+            IRInstr::Goto(label) => *label = resolve_label(*label, redirect),
+            IRInstr::CondGoto(_, true_label, false_label) => {
+                *true_label = resolve_label(*true_label, redirect);
+                *false_label = resolve_label(*false_label, redirect);
+            }
+            _ => {}
+        }
+    }
+}
+
+fn resolve_label(label: usize, redirect: &BTreeMap<usize, usize>) -> usize {
+    let mut current = label;
+    let mut visited = HashSet::new();
+    while let Some(next) = redirect.get(&current) {
+        if !visited.insert(current) {
+            break;
+        }
+        current = *next;
+    }
+    current
+}
+
+fn remove_unused_labels(instrs: &mut Vec<IRInstr>) {
+    let mut used = BTreeSet::new();
+    for instr in instrs.iter() {
+        match instr {
+            IRInstr::Goto(label) => {
+                used.insert(*label);
+            }
+            IRInstr::CondGoto(_, true_label, false_label) => {
+                used.insert(*true_label);
+                used.insert(*false_label);
+            }
+            _ => {}
+        }
+    }
+    instrs.retain(|instr| match instr {
+        IRInstr::Label(label) => used.contains(label),
+        _ => true,
+    });
+}
+
+fn remove_redundant_fallthrough_gotos(instrs: &mut Vec<IRInstr>) {
+    let mut index = 0;
+    while index + 1 < instrs.len() {
+        let remove = match (&instrs[index], &instrs[index + 1]) {
+            (IRInstr::Goto(target), IRInstr::Label(label)) => target == label,
+            _ => false,
+        };
+        if remove {
+            instrs.remove(index);
+        } else {
+            index += 1;
+        }
+    }
+}
+
+fn block_dot_label(func_name: &str, block: &BasicBlock) -> String {
+    let instrs = block
+        .instrs
+        .iter()
+        .map(|instr| instr.to_string())
+        .collect::<Vec<_>>()
+        .join("\\l");
+    format!("{}::B{}\\l{}\\l", func_name, block.id, instrs)
+}

src/ir/mod.rs

@@ -1,3 +1,4 @@
 pub mod generator;
 pub mod types;
 pub mod err;
+pub mod cfg;

src/ir/types.rs

@@ -4,6 +4,7 @@ use crate::ast::types::Type as AstType;
 use crate::ast::types::BinaryOp as AstBinaryOp;
 use crate::ir::err::IRError;
 
+#[derive(Clone)]
 pub enum VariableOrIntLit {
     Var(Variable),
     IntLit(i32),
@@ -17,6 +18,7 @@ impl Display for VariableOrIntLit {
         }
     }
 }
+#[derive(Clone)]
 pub enum IRInstr {
     Declare(Variable),
     DefineFunc(Function, Vec<Variable>, Vec<IRInstr>),
@@ -169,6 +171,7 @@ impl From<AstType> for IRType {
         }
     }
 }
+#[derive(Clone)]
 pub enum MoveRValue {
     Var(Variable),
     ConstInt(i32),
@@ -276,6 +279,7 @@ impl Function {
         format!("{} @{}({})", self.return_type, self.name, params_str)
     }
 }
+#[derive(Clone)]
 pub enum BinaryOp {
     Add,
     Sub,
@@ -283,6 +287,7 @@ pub enum BinaryOp {
     Div,
     Mod,
 }
+#[derive(Clone)]
 pub enum CmpOp {
     Le,
     Lt,
@@ -291,6 +296,7 @@ pub enum CmpOp {
     Ne,
     Eq,
 }
+#[derive(Clone)]
 pub enum UnaryOp {
     Neg
 }

src/main.rs

@@ -11,7 +11,7 @@ use std::{fs::File, io::BufRead};
 
 use clap::Parser as ArgParser;
 
-use crate::{ast::parser::Parser, ir::generator::Generator, lexer::lexer::Lexer, sema::analyzer::Analyzer};
+use crate::{ast::parser::Parser, ir::{cfg, generator::Generator}, lexer::lexer::Lexer, sema::analyzer::Analyzer};
 use crate::backend::generator::Generator as ASMGerenerator;
 /// Simple minic compiler built by Rust
 #[derive(ArgParser, Debug)]
@@ -33,6 +33,9 @@ struct Args {
     /// Output file for the generated code (default: stdout)
     #[arg(short = 'o', long = "output")]
     output: Option<String>,
+    /// Output CFG graph in Graphviz dot format
+    #[arg(long = "cfg-dot")]
+    cfg_dot: Option<String>,
     /// Source file to compile
     source: String,
 }
@@ -85,7 +88,13 @@ fn main() {
         analyzer.get_diagnostics().print(&format!("{}", source_path.display()), &full_text);
     }
     let mut generator = Generator::new(&analyzer);
-    let ir = generator.emit(hir);
+    let ir = cfg::optimize_ir(generator.emit(hir));
+    if let Some(cfg_dot_path) = &args.cfg_dot {
+        match std::fs::write(cfg_dot_path, cfg::ir_to_cfg_dot(&ir)) {
+            Ok(_) => println!("CFG graph written to {}", cfg_dot_path),
+            Err(e) => eprintln!("Failed to write CFG graph to {}: {}", cfg_dot_path, e),
+        }
+    }
     if args.output_ir {
         if let Some(output_path) = args.output {
             match std::fs::write(&output_path, ir.iter().map(|instr| instr.to_string()).collect::<Vec<_>>().join("\n")) {