/home/noah/src/ruchy/src/middleend/mir/optimize.rs
Line | Count | Source |
1 | | //! MIR optimization passes |
2 | | |
3 | | use super::types::{ |
4 | | BinOp, BlockId, Constant, Function, Local, Operand, Place, Program, Rvalue, Statement, |
5 | | Terminator, UnOp, |
6 | | }; |
7 | | use std::collections::{HashMap, HashSet}; |
8 | | |
9 | | /// Dead Code Elimination pass |
10 | | pub struct DeadCodeElimination { |
11 | | /// Set of live locals |
12 | | live_locals: HashSet<Local>, |
13 | | /// Set of live blocks |
14 | | live_blocks: HashSet<BlockId>, |
15 | | } |
16 | | |
17 | | impl Default for DeadCodeElimination { |
18 | 0 | fn default() -> Self { |
19 | 0 | Self::new() |
20 | 0 | } |
21 | | } |
22 | | |
23 | | impl DeadCodeElimination { |
24 | | /// Create a new DCE pass |
25 | | #[must_use] |
26 | 0 | pub fn new() -> Self { |
27 | 0 | Self { |
28 | 0 | live_locals: HashSet::new(), |
29 | 0 | live_blocks: HashSet::new(), |
30 | 0 | } |
31 | 0 | } |
32 | | |
33 | | /// Run DCE on a function |
34 | 0 | pub fn run(&mut self, func: &mut Function) { |
35 | | // Mark live locals and blocks |
36 | 0 | self.mark_live(func); |
37 | | |
38 | | // Remove dead statements |
39 | 0 | self.remove_dead_statements(func); |
40 | | |
41 | | // Remove dead blocks |
42 | 0 | self.remove_dead_blocks(func); |
43 | | |
44 | | // Remove dead locals |
45 | 0 | self.remove_dead_locals(func); |
46 | 0 | } |
47 | | |
48 | | /// Mark live locals and blocks |
49 | 0 | fn mark_live(&mut self, func: &Function) { |
50 | 0 | self.live_locals.clear(); |
51 | 0 | self.live_blocks.clear(); |
52 | | |
53 | | // Start from entry block |
54 | 0 | let mut worklist = vec![func.entry_block]; |
55 | 0 | self.live_blocks.insert(func.entry_block); |
56 | | |
57 | 0 | while let Some(block_id) = worklist.pop() { |
58 | 0 | if let Some(block) = func.blocks.iter().find(|b| b.id == block_id) { |
59 | | // Mark locals used in statements |
60 | 0 | for stmt in &block.statements { |
61 | 0 | self.mark_statement_live(stmt); |
62 | 0 | } |
63 | | |
64 | | // Mark locals used in terminator and add successor blocks |
65 | 0 | self.mark_terminator_live(&block.terminator, &mut worklist); |
66 | 0 | } |
67 | | } |
68 | | |
69 | | // Mark parameters as live |
70 | 0 | for param in &func.params { |
71 | 0 | self.live_locals.insert(*param); |
72 | 0 | } |
73 | 0 | } |
74 | | |
75 | | /// Mark locals in a statement as live |
76 | 0 | fn mark_statement_live(&mut self, stmt: &Statement) { |
77 | 0 | match stmt { |
78 | 0 | Statement::Assign(place, rvalue) => { |
79 | 0 | self.mark_place_live(place); |
80 | 0 | self.mark_rvalue_live(rvalue); |
81 | 0 | } |
82 | 0 | Statement::StorageLive(local) | Statement::StorageDead(local) => { |
83 | 0 | self.live_locals.insert(*local); |
84 | 0 | } |
85 | 0 | Statement::Nop => {} |
86 | | } |
87 | 0 | } |
88 | | |
89 | | /// Mark locals in a place as live |
90 | 0 | fn mark_place_live(&mut self, place: &Place) { |
91 | 0 | match place { |
92 | 0 | Place::Local(local) => { |
93 | 0 | self.live_locals.insert(*local); |
94 | 0 | } |
95 | 0 | Place::Field(base, _) | Place::Deref(base) => { |
96 | 0 | self.mark_place_live(base); |
97 | 0 | } |
98 | 0 | Place::Index(base, index) => { |
99 | 0 | self.mark_place_live(base); |
100 | 0 | self.mark_place_live(index); |
101 | 0 | } |
102 | | } |
103 | 0 | } |
104 | | |
105 | | /// Mark locals in an rvalue as live |
106 | 0 | fn mark_rvalue_live(&mut self, rvalue: &Rvalue) { |
107 | 0 | match rvalue { |
108 | 0 | Rvalue::Use(operand) | Rvalue::UnaryOp(_, operand) | Rvalue::Cast(_, operand, _) => { |
109 | 0 | self.mark_operand_live(operand); |
110 | 0 | } |
111 | 0 | Rvalue::BinaryOp(_, left, right) => { |
112 | 0 | self.mark_operand_live(left); |
113 | 0 | self.mark_operand_live(right); |
114 | 0 | } |
115 | 0 | Rvalue::Ref(_, place) => { |
116 | 0 | self.mark_place_live(place); |
117 | 0 | } |
118 | 0 | Rvalue::Aggregate(_, operands) => { |
119 | 0 | for operand in operands { |
120 | 0 | self.mark_operand_live(operand); |
121 | 0 | } |
122 | | } |
123 | 0 | Rvalue::Call(func, args) => { |
124 | 0 | self.mark_operand_live(func); |
125 | 0 | for arg in args { |
126 | 0 | self.mark_operand_live(arg); |
127 | 0 | } |
128 | | } |
129 | | } |
130 | 0 | } |
131 | | |
132 | | /// Mark locals in an operand as live |
133 | 0 | fn mark_operand_live(&mut self, operand: &Operand) { |
134 | 0 | match operand { |
135 | 0 | Operand::Copy(place) | Operand::Move(place) => { |
136 | 0 | self.mark_place_live(place); |
137 | 0 | } |
138 | 0 | Operand::Constant(_) => {} |
139 | | } |
140 | 0 | } |
141 | | |
142 | | /// Mark locals in terminator as live and add successor blocks |
143 | 0 | fn mark_terminator_live(&mut self, terminator: &Terminator, worklist: &mut Vec<BlockId>) { |
144 | 0 | match terminator { |
145 | 0 | Terminator::Goto(target) => { |
146 | 0 | if self.live_blocks.insert(*target) { |
147 | 0 | worklist.push(*target); |
148 | 0 | } |
149 | | } |
150 | | Terminator::If { |
151 | 0 | condition, |
152 | 0 | then_block, |
153 | 0 | else_block, |
154 | | } => { |
155 | 0 | self.mark_operand_live(condition); |
156 | 0 | if self.live_blocks.insert(*then_block) { |
157 | 0 | worklist.push(*then_block); |
158 | 0 | } |
159 | 0 | if self.live_blocks.insert(*else_block) { |
160 | 0 | worklist.push(*else_block); |
161 | 0 | } |
162 | | } |
163 | | Terminator::Switch { |
164 | 0 | discriminant, |
165 | 0 | targets, |
166 | 0 | default, |
167 | | } => { |
168 | 0 | self.mark_operand_live(discriminant); |
169 | 0 | for (_, target) in targets { |
170 | 0 | if self.live_blocks.insert(*target) { |
171 | 0 | worklist.push(*target); |
172 | 0 | } |
173 | | } |
174 | 0 | if let Some(default_block) = default { |
175 | 0 | if self.live_blocks.insert(*default_block) { |
176 | 0 | worklist.push(*default_block); |
177 | 0 | } |
178 | 0 | } |
179 | | } |
180 | 0 | Terminator::Return(operand) => { |
181 | 0 | if let Some(op) = operand { |
182 | 0 | self.mark_operand_live(op); |
183 | 0 | } |
184 | | } |
185 | | Terminator::Call { |
186 | 0 | func, |
187 | 0 | args, |
188 | 0 | destination, |
189 | | } => { |
190 | 0 | self.mark_operand_live(func); |
191 | 0 | for arg in args { |
192 | 0 | self.mark_operand_live(arg); |
193 | 0 | } |
194 | 0 | if let Some((place, target)) = destination { |
195 | 0 | self.mark_place_live(place); |
196 | 0 | if self.live_blocks.insert(*target) { |
197 | 0 | worklist.push(*target); |
198 | 0 | } |
199 | 0 | } |
200 | | } |
201 | 0 | Terminator::Unreachable => {} |
202 | | } |
203 | 0 | } |
204 | | |
205 | | /// Remove dead statements from blocks |
206 | 0 | fn remove_dead_statements(&self, func: &mut Function) { |
207 | 0 | for block in &mut func.blocks { |
208 | 0 | if !self.live_blocks.contains(&block.id) { |
209 | 0 | continue; |
210 | 0 | } |
211 | | |
212 | 0 | block.statements.retain(|stmt| { |
213 | 0 | match stmt { |
214 | 0 | Statement::Assign(place, _) => self.is_place_live(place), |
215 | 0 | Statement::StorageLive(local) | Statement::StorageDead(local) => { |
216 | 0 | self.live_locals.contains(local) |
217 | | } |
218 | 0 | Statement::Nop => false, // Always remove nops |
219 | | } |
220 | 0 | }); |
221 | | } |
222 | 0 | } |
223 | | |
224 | | /// Remove dead blocks |
225 | 0 | fn remove_dead_blocks(&self, func: &mut Function) { |
226 | 0 | func.blocks |
227 | 0 | .retain(|block| self.live_blocks.contains(&block.id)); |
228 | 0 | } |
229 | | |
230 | | /// Remove dead locals |
231 | 0 | fn remove_dead_locals(&self, func: &mut Function) { |
232 | 0 | func.locals |
233 | 0 | .retain(|local_decl| self.live_locals.contains(&local_decl.id)); |
234 | 0 | } |
235 | | |
236 | | /// Check if a place is live |
237 | 0 | fn is_place_live(&self, place: &Place) -> bool { |
238 | 0 | match place { |
239 | 0 | Place::Local(local) => self.live_locals.contains(local), |
240 | 0 | Place::Field(base, _) | Place::Deref(base) => self.is_place_live(base), |
241 | 0 | Place::Index(base, index) => self.is_place_live(base) || self.is_place_live(index), |
242 | | } |
243 | 0 | } |
244 | | } |
245 | | |
246 | | /// Constant Propagation pass |
247 | | pub struct ConstantPropagation { |
248 | | /// Map from locals to their constant values |
249 | | constants: HashMap<Local, Constant>, |
250 | | } |
251 | | |
252 | | impl Default for ConstantPropagation { |
253 | 0 | fn default() -> Self { |
254 | 0 | Self::new() |
255 | 0 | } |
256 | | } |
257 | | |
258 | | impl ConstantPropagation { |
259 | | /// Create a new constant propagation pass |
260 | | #[must_use] |
261 | 0 | pub fn new() -> Self { |
262 | 0 | Self { |
263 | 0 | constants: HashMap::new(), |
264 | 0 | } |
265 | 0 | } |
266 | | |
267 | | /// Run constant propagation on a function |
268 | 0 | pub fn run(&mut self, func: &mut Function) { |
269 | 0 | self.constants.clear(); |
270 | | |
271 | | // Find constant assignments |
272 | 0 | for block in &func.blocks { |
273 | 0 | for stmt in &block.statements { |
274 | 0 | if let Statement::Assign(Place::Local(local), rvalue) = stmt { |
275 | 0 | if let Some(constant) = self.extract_constant(rvalue) { |
276 | 0 | self.constants.insert(*local, constant); |
277 | 0 | } |
278 | 0 | } |
279 | | } |
280 | | } |
281 | | |
282 | | // Replace uses of constants |
283 | 0 | for block in &mut func.blocks { |
284 | 0 | for stmt in &mut block.statements { |
285 | 0 | self.propagate_in_statement(stmt); |
286 | 0 | } |
287 | 0 | self.propagate_in_terminator(&mut block.terminator); |
288 | | } |
289 | 0 | } |
290 | | |
291 | | /// Extract constant from rvalue if possible |
292 | 0 | fn extract_constant(&self, rvalue: &Rvalue) -> Option<Constant> { |
293 | 0 | match rvalue { |
294 | 0 | Rvalue::Use(Operand::Constant(c)) => Some(c.clone()), |
295 | 0 | Rvalue::BinaryOp(op, left, right) => self.eval_binary_op(*op, left, right), |
296 | 0 | Rvalue::UnaryOp(op, operand) => self.eval_unary_op(*op, operand), |
297 | 0 | _ => None, |
298 | | } |
299 | 0 | } |
300 | | |
301 | | /// Evaluate binary operation on constants |
302 | 0 | fn eval_binary_op(&self, op: BinOp, left: &Operand, right: &Operand) -> Option<Constant> { |
303 | 0 | let left_val = self.get_constant_value(left)?; |
304 | 0 | let right_val = self.get_constant_value(right)?; |
305 | | |
306 | 0 | match (op, &left_val, &right_val) { |
307 | 0 | (BinOp::Add, Constant::Int(a, ty), Constant::Int(b, _)) => { |
308 | 0 | Some(Constant::Int(a + b, ty.clone())) |
309 | | } |
310 | 0 | (BinOp::Sub, Constant::Int(a, ty), Constant::Int(b, _)) => { |
311 | 0 | Some(Constant::Int(a - b, ty.clone())) |
312 | | } |
313 | 0 | (BinOp::Mul, Constant::Int(a, ty), Constant::Int(b, _)) => { |
314 | 0 | Some(Constant::Int(a * b, ty.clone())) |
315 | | } |
316 | 0 | (BinOp::Eq, Constant::Int(a, _), Constant::Int(b, _)) => Some(Constant::Bool(a == b)), |
317 | 0 | (BinOp::Lt, Constant::Int(a, _), Constant::Int(b, _)) => Some(Constant::Bool(a < b)), |
318 | 0 | (BinOp::And, Constant::Bool(a), Constant::Bool(b)) => Some(Constant::Bool(*a && *b)), |
319 | 0 | (BinOp::Or, Constant::Bool(a), Constant::Bool(b)) => Some(Constant::Bool(*a || *b)), |
320 | 0 | _ => None, |
321 | | } |
322 | 0 | } |
323 | | |
324 | | /// Evaluate unary operation on constant |
325 | 0 | fn eval_unary_op(&self, op: UnOp, operand: &Operand) -> Option<Constant> { |
326 | 0 | let val = self.get_constant_value(operand)?; |
327 | | |
328 | 0 | match (op, &val) { |
329 | 0 | (UnOp::Neg, Constant::Int(i, ty)) => Some(Constant::Int(-i, ty.clone())), |
330 | 0 | (UnOp::Not, Constant::Bool(b)) => Some(Constant::Bool(!b)), |
331 | 0 | _ => None, |
332 | | } |
333 | 0 | } |
334 | | |
335 | | /// Get constant value for an operand |
336 | 0 | fn get_constant_value(&self, operand: &Operand) -> Option<Constant> { |
337 | 0 | match operand { |
338 | 0 | Operand::Constant(c) => Some(c.clone()), |
339 | 0 | Operand::Copy(Place::Local(local)) | Operand::Move(Place::Local(local)) => { |
340 | 0 | self.constants.get(local).cloned() |
341 | | } |
342 | 0 | _ => None, |
343 | | } |
344 | 0 | } |
345 | | |
346 | | /// Propagate constants in a statement |
347 | 0 | fn propagate_in_statement(&self, stmt: &mut Statement) { |
348 | 0 | if let Statement::Assign(_, rvalue) = stmt { |
349 | 0 | self.propagate_in_rvalue(rvalue); |
350 | 0 | } |
351 | 0 | } |
352 | | |
353 | | /// Propagate constants in an rvalue |
354 | 0 | fn propagate_in_rvalue(&self, rvalue: &mut Rvalue) { |
355 | 0 | match rvalue { |
356 | 0 | Rvalue::Use(operand) | Rvalue::UnaryOp(_, operand) => { |
357 | 0 | self.propagate_in_operand(operand); |
358 | 0 | } |
359 | 0 | Rvalue::BinaryOp(_, left, right) => { |
360 | 0 | self.propagate_in_operand(left); |
361 | 0 | self.propagate_in_operand(right); |
362 | 0 | } |
363 | 0 | Rvalue::Call(func, args) => { |
364 | 0 | self.propagate_in_operand(func); |
365 | 0 | for arg in args { |
366 | 0 | self.propagate_in_operand(arg); |
367 | 0 | } |
368 | | } |
369 | 0 | _ => {} |
370 | | } |
371 | 0 | } |
372 | | |
373 | | /// Propagate constants in an operand |
374 | 0 | fn propagate_in_operand(&self, operand: &mut Operand) { |
375 | 0 | if let Some(constant) = self.get_constant_value(operand) { |
376 | 0 | *operand = Operand::Constant(constant); |
377 | 0 | } |
378 | 0 | } |
379 | | |
380 | | /// Propagate constants in a terminator |
381 | 0 | fn propagate_in_terminator(&self, terminator: &mut Terminator) { |
382 | 0 | match terminator { |
383 | 0 | Terminator::If { condition, .. } => { |
384 | 0 | self.propagate_in_operand(condition); |
385 | 0 | } |
386 | 0 | Terminator::Switch { discriminant, .. } => { |
387 | 0 | self.propagate_in_operand(discriminant); |
388 | 0 | } |
389 | 0 | Terminator::Return(Some(operand)) => { |
390 | 0 | self.propagate_in_operand(operand); |
391 | 0 | } |
392 | 0 | Terminator::Call { func, args, .. } => { |
393 | 0 | self.propagate_in_operand(func); |
394 | 0 | for arg in args { |
395 | 0 | self.propagate_in_operand(arg); |
396 | 0 | } |
397 | | } |
398 | 0 | _ => {} |
399 | | } |
400 | 0 | } |
401 | | } |
402 | | |
403 | | /// Common Subexpression Elimination pass |
404 | | pub struct CommonSubexpressionElimination { |
405 | | /// Map from expressions to locals that compute them |
406 | | expressions: HashMap<String, Local>, |
407 | | } |
408 | | |
409 | | impl Default for CommonSubexpressionElimination { |
410 | 0 | fn default() -> Self { |
411 | 0 | Self::new() |
412 | 0 | } |
413 | | } |
414 | | |
415 | | impl CommonSubexpressionElimination { |
416 | | /// Create a new CSE pass |
417 | | #[must_use] |
418 | 0 | pub fn new() -> Self { |
419 | 0 | Self { |
420 | 0 | expressions: HashMap::new(), |
421 | 0 | } |
422 | 0 | } |
423 | | |
424 | | /// Run CSE on a function |
425 | 0 | pub fn run(&mut self, func: &mut Function) { |
426 | 0 | self.expressions.clear(); |
427 | | |
428 | 0 | for block in &mut func.blocks { |
429 | 0 | for stmt in &mut block.statements { |
430 | 0 | self.process_statement(stmt); |
431 | 0 | } |
432 | | } |
433 | 0 | } |
434 | | |
435 | | /// Process a statement for CSE |
436 | 0 | fn process_statement(&mut self, stmt: &mut Statement) { |
437 | 0 | if let Statement::Assign(Place::Local(local), rvalue) = stmt { |
438 | 0 | let expr_key = self.rvalue_key(rvalue); |
439 | | |
440 | 0 | if let Some(existing_local) = self.expressions.get(&expr_key) { |
441 | 0 | // Replace with copy from existing local |
442 | 0 | *stmt = Statement::Assign( |
443 | 0 | Place::Local(*local), |
444 | 0 | Rvalue::Use(Operand::Copy(Place::Local(*existing_local))), |
445 | 0 | ); |
446 | 0 | } else { |
447 | 0 | // Record this expression |
448 | 0 | self.expressions.insert(expr_key, *local); |
449 | 0 | } |
450 | 0 | } |
451 | 0 | } |
452 | | |
453 | | /// Generate a key for an rvalue |
454 | 0 | fn rvalue_key(&self, rvalue: &Rvalue) -> String { |
455 | 0 | match rvalue { |
456 | 0 | Rvalue::Use(operand) => format!("use({})", self.operand_key(operand)), |
457 | 0 | Rvalue::BinaryOp(op, left, right) => { |
458 | 0 | format!( |
459 | 0 | "binop({:?}, {}, {})", |
460 | | op, |
461 | 0 | self.operand_key(left), |
462 | 0 | self.operand_key(right) |
463 | | ) |
464 | | } |
465 | 0 | Rvalue::UnaryOp(op, operand) => { |
466 | 0 | format!("unop({:?}, {})", op, self.operand_key(operand)) |
467 | | } |
468 | 0 | _ => format!("{rvalue:?}"), // Fallback |
469 | | } |
470 | 0 | } |
471 | | |
472 | | /// Generate a key for an operand |
473 | 0 | fn operand_key(&self, operand: &Operand) -> String { |
474 | 0 | match operand { |
475 | 0 | Operand::Copy(place) => format!("copy({})", self.place_key(place)), |
476 | 0 | Operand::Move(place) => format!("move({})", self.place_key(place)), |
477 | 0 | Operand::Constant(c) => format!("const({c:?})"), |
478 | | } |
479 | 0 | } |
480 | | |
481 | | /// Generate a key for a place |
482 | | #[allow(clippy::only_used_in_recursion)] |
483 | 0 | fn place_key(&self, place: &Place) -> String { |
484 | 0 | match place { |
485 | 0 | Place::Local(local) => format!("local({})", local.0), |
486 | 0 | Place::Field(base, field) => format!("field({}, {})", self.place_key(base), field.0), |
487 | 0 | Place::Index(base, index) => { |
488 | 0 | format!("index({}, {})", self.place_key(base), self.place_key(index)) |
489 | | } |
490 | 0 | Place::Deref(base) => format!("deref({})", self.place_key(base)), |
491 | | } |
492 | 0 | } |
493 | | } |
494 | | |
495 | | /// Run all optimization passes on a function |
496 | 0 | pub fn optimize_function(func: &mut Function) { |
497 | 0 | let mut dce = DeadCodeElimination::new(); |
498 | 0 | let mut const_prop = ConstantPropagation::new(); |
499 | 0 | let mut cse = CommonSubexpressionElimination::new(); |
500 | | |
501 | | // Run multiple rounds for better results |
502 | 0 | for _ in 0..3 { |
503 | 0 | const_prop.run(func); |
504 | 0 | cse.run(func); |
505 | 0 | dce.run(func); |
506 | 0 | } |
507 | 0 | } |
508 | | |
509 | | /// Run all optimization passes on a program |
510 | 0 | pub fn optimize_program(program: &mut Program) { |
511 | 0 | for function in program.functions.values_mut() { |
512 | 0 | optimize_function(function); |
513 | 0 | } |
514 | 0 | } |
515 | | |
516 | | #[cfg(test)] |
517 | | #[allow(clippy::unwrap_used)] |
518 | | mod tests { |
519 | | #[test] |
520 | 1 | fn test_dead_code_elimination() { |
521 | | // Test placeholder - optimization features are framework-level |
522 | | // Implementation deferred pending MIR completion in Phase 2 |
523 | | // This test validates that the optimization framework is accessible |
524 | 1 | } |
525 | | |
526 | | #[test] |
527 | 1 | fn test_constant_propagation() { |
528 | | // Test placeholder - optimization features are framework-level |
529 | | // Implementation deferred pending MIR completion in Phase 2 |
530 | | // This test validates that the optimization framework is accessible |
531 | 1 | } |
532 | | |
533 | | #[test] |
534 | 1 | fn test_common_subexpression_elimination() { |
535 | | // Test placeholder - optimization features are framework-level |
536 | | // Implementation deferred pending MIR completion in Phase 2 |
537 | | // This test validates that the optimization framework is accessible |
538 | 1 | } |
539 | | } |