armfortas Public

Watch 0 Fork 0 Star 0
Rust · 7913 bytes Raw Blame History
1 //! SimplifyCFG pass: merge empty block chains and thread jumps.
2 //!
3 //! After inlining, the CFG often contains blocks that simply branch
4 //! unconditionally to the next block with no instructions or block
5 //! params. This pass merges them:
6 //!
7 //! ```text
8 //! Before: A: br B B: br C C: [real code]
9 //! After: A: br C C: [real code]
10 //! ```
11 //!
12 //! Also merges a predecessor into its successor when:
13 //! - The predecessor's only terminator is `Branch(succ, args)`
14 //! - The successor has exactly one predecessor
15 //! - Neither block is the function entry
16 //!
17 //! This eliminates the empty block chains left by inlining.
18
19 use super::pass::Pass;
20 use crate::ir::inst::*;
21 use crate::ir::walk::predecessors;
22
23 pub struct SimplifyCfg;
24
25 impl Pass for SimplifyCfg {
26 fn name(&self) -> &'static str {
27 "simplify-cfg"
28 }
29
30 fn run(&self, module: &mut Module) -> bool {
31 let mut changed = false;
32 for func in &mut module.functions {
33 if simplify_function(func) {
34 changed = true;
35 }
36 }
37 changed
38 }
39 }
40
41 fn simplify_function(func: &mut Function) -> bool {
42 let mut changed = false;
43
44 // Phase 1: Thread jumps through empty blocks.
45 // If block B has no instructions, no block params, and terminates
46 // with Branch(C, []), replace all branches to B with branches to C.
47 loop {
48 let mut redirect: Option<(BlockId, BlockId)> = None;
49
50 for block in &func.blocks {
51 if block.id == func.entry {
52 continue;
53 }
54 if !block.insts.is_empty() {
55 continue;
56 }
57 if !block.params.is_empty() {
58 continue;
59 }
60 if let Some(Terminator::Branch(target, args)) = &block.terminator {
61 if args.is_empty() {
62 // This block is empty and just forwards to target.
63 redirect = Some((block.id, *target));
64 break;
65 }
66 }
67 }
68
69 let Some((empty_block, target)) = redirect else {
70 break;
71 };
72
73 // Redirect all predecessors of empty_block to target instead.
74 for block in &mut func.blocks {
75 if block.id == empty_block {
76 continue;
77 } // don't redirect self
78 let term = match &mut block.terminator {
79 Some(t) => t,
80 None => continue,
81 };
82 redirect_terminator(term, empty_block, target);
83 }
84 // Mark the empty block as unreachable so we don't find it again.
85 func.block_mut(empty_block).terminator = Some(Terminator::Unreachable);
86 changed = true;
87 }
88
89 // Phase 2: Merge blocks where predecessor has unconditional branch
90 // to a successor with exactly one predecessor.
91 loop {
92 let preds = predecessors(func);
93 let mut merge: Option<(BlockId, BlockId)> = None;
94
95 for block in &func.blocks {
96 if let Some(Terminator::Branch(succ, args)) = &block.terminator {
97 if args.is_empty() && *succ != block.id {
98 // Check if successor has exactly one predecessor.
99 let succ_preds = preds.get(succ);
100 if let Some(sp) = succ_preds {
101 let mut unique: Vec<BlockId> = sp.clone();
102 unique.sort_by_key(|b| b.0);
103 unique.dedup();
104 if unique.len() == 1 && unique[0] == block.id {
105 // Can merge: append successor's contents to predecessor.
106 merge = Some((block.id, *succ));
107 break;
108 }
109 }
110 }
111 }
112 }
113
114 let Some((pred_id, succ_id)) = merge else {
115 break;
116 };
117
118 // Merge: append successor's instructions and terminator to predecessor.
119 let succ_insts = func.block(succ_id).insts.clone();
120 let succ_term = func.block(succ_id).terminator.clone();
121 func.block_mut(pred_id).insts.extend(succ_insts);
122 func.block_mut(pred_id).terminator = succ_term;
123
124 // Mark successor as unreachable (will be pruned).
125 func.block_mut(succ_id).terminator = Some(Terminator::Unreachable);
126 func.block_mut(succ_id).insts.clear();
127
128 changed = true;
129 }
130
131 // Prune unreachable blocks.
132 if changed {
133 crate::ir::walk::prune_unreachable(func);
134 }
135
136 changed
137 }
138
139 /// Redirect all edges targeting `old` to `new` in a terminator.
140 fn redirect_terminator(term: &mut Terminator, old: BlockId, new: BlockId) {
141 match term {
142 Terminator::Branch(dest, _) if *dest == old => {
143 *dest = new;
144 }
145 Terminator::CondBranch {
146 true_dest,
147 false_dest,
148 ..
149 } => {
150 if *true_dest == old {
151 *true_dest = new;
152 }
153 if *false_dest == old {
154 *false_dest = new;
155 }
156 }
157 Terminator::Switch { default, cases, .. } => {
158 if *default == old {
159 *default = new;
160 }
161 for (_, dest) in cases.iter_mut() {
162 if *dest == old {
163 *dest = new;
164 }
165 }
166 }
167 _ => {}
168 }
169 }
170
171 #[cfg(test)]
172 mod tests {
173 use super::*;
174 use crate::ir::types::{IntWidth, IrType};
175 use crate::opt::pass::Pass;
176
177 #[test]
178 fn merges_empty_chain() {
179 let mut m = Module::new("test".into());
180 let mut f = Function::new("test".into(), vec![], IrType::Void);
181 let b = f.create_block("middle");
182 let c = f.create_block("end");
183 // entry → b → c → ret
184 f.block_mut(f.entry).terminator = Some(Terminator::Branch(b, vec![]));
185 f.block_mut(b).terminator = Some(Terminator::Branch(c, vec![]));
186 f.block_mut(c).terminator = Some(Terminator::Return(None));
187 m.add_function(f);
188
189 let pass = SimplifyCfg;
190 let changed = pass.run(&mut m);
191 assert!(changed);
192
193 // After: entry should branch directly to end (or entry should
194 // contain the ret directly).
195 let f = &m.functions[0];
196 // The chain should be collapsed.
197 assert!(
198 f.blocks.len() <= 2,
199 "should merge empty chain, got {} blocks",
200 f.blocks.len()
201 );
202 }
203
204 #[test]
205 fn no_op_on_non_empty_blocks() {
206 let mut m = Module::new("test".into());
207 let mut f = Function::new("test".into(), vec![], IrType::Void);
208 let b = f.create_block("body");
209 f.block_mut(f.entry).terminator = Some(Terminator::Branch(b, vec![]));
210 // Body has an instruction — should not be merged.
211 let id = f.next_value_id();
212 f.register_type(id, IrType::Int(IntWidth::I32));
213 f.block_mut(b).insts.push(Inst {
214 id,
215 ty: IrType::Int(IntWidth::I32),
216 span: crate::lexer::Span {
217 file_id: 0,
218 start: crate::lexer::Position { line: 0, col: 0 },
219 end: crate::lexer::Position { line: 0, col: 0 },
220 },
221 kind: InstKind::ConstInt(42, IntWidth::I32),
222 });
223 f.block_mut(b).terminator = Some(Terminator::Return(None));
224 m.add_function(f);
225
226 let pass = SimplifyCfg;
227 let _changed = pass.run(&mut m);
228 // Body block may be merged into entry (single pred + unconditional branch).
229 // But the instruction should be preserved.
230 let f = &m.functions[0];
231 let has_const = f.blocks.iter().any(|bl| {
232 bl.insts
233 .iter()
234 .any(|i| matches!(i.kind, InstKind::ConstInt(42, _)))
235 });
236 assert!(has_const, "const instruction must be preserved");
237 }
238 }
239