diff --git a/src/ir/cfg.rs b/src/ir/cfg.rs new file mode 100644 index 0000000..e8df204 --- /dev/null +++ b/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; + +#[derive(Clone)] +pub struct BasicBlock { + pub id: usize, + pub instrs: Vec, +} + +pub struct ControlFlowGraph { + pub func_name: String, + pub blocks: Vec, + 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, 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::>(); + + 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 { + self.blocks + .into_iter() + .flat_map(|block| block.instrs) + .collect() + } +} + +pub fn optimize_ir(ir: Vec) -> Vec { + 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) -> Vec { + 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) -> (Vec, Vec) { + 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 { + 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::>(); + 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 { + 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 { + 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 { + 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) -> Option { + (index + 1..block_count).find(|candidate| !remove.contains(candidate)) +} + +fn first_label(block: &BasicBlock) -> Option { + match block.instrs.first() { + Some(IRInstr::Label(label)) => Some(*label), + _ => None, + } +} + +fn label_only_block(block: &BasicBlock) -> Option { + 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) { + 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 { + let mut current = label; + let mut visited = HashSet::new(); + while let Some(next) = redirect.get(¤t) { + if !visited.insert(current) { + break; + } + current = *next; + } + current +} + +fn remove_unused_labels(instrs: &mut Vec) { + 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) { + 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::>() + .join("\\l"); + format!("{}::B{}\\l{}\\l", func_name, block.id, instrs) +} diff --git a/src/ir/mod.rs b/src/ir/mod.rs index a7c4177..190e1a1 100644 --- a/src/ir/mod.rs +++ b/src/ir/mod.rs @@ -1,3 +1,4 @@ pub mod generator; pub mod types; -pub mod err; \ No newline at end of file +pub mod err; +pub mod cfg; diff --git a/src/ir/types.rs b/src/ir/types.rs index e6d401c..903ddaf 100644 --- a/src/ir/types.rs +++ b/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, Vec), @@ -169,6 +171,7 @@ impl From 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 } diff --git a/src/main.rs b/src/main.rs index 2352495..9379d1e 100644 --- a/src/main.rs +++ b/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, + /// Output CFG graph in Graphviz dot format + #[arg(long = "cfg-dot")] + cfg_dot: Option, /// 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::>().join("\n")) {