1use std::collections::HashMap;
4
5use crate::ast::{Block, Expr, FunctionDecl, Literal, Program, Stmt};
6use crate::enhanced_error::CompilerError;
7
8use super::advanced::{AdvancedOptimizer, CSEOptimizer, FunctionInliner, LoopOptimizer};
9use super::modular_pipeline::{
10 OptimizationContext, OptimizationLevel, OptimizationPass, OptimizationResult,
11 OptimizationStats,
12};
13use super::Optimizer;
14
15fn optimize_expr_in_place(expr: &mut Expr, optimizer: &Optimizer) {
16 *expr = optimizer.optimize_expr(expr);
17}
18
19fn fold_block(block: &mut Block, optimizer: &Optimizer) -> usize {
20 let mut modifications = 0;
21
22 for stmt in &mut block.statements {
23 if optimize_stmt_tree(stmt, optimizer) {
24 modifications += 1;
25 }
26 }
27
28 for proc in &mut block.procedures {
29 modifications += fold_block(&mut proc.block, optimizer);
30 }
31 for func in &mut block.functions {
32 modifications += fold_block(&mut func.block, optimizer);
33 }
34
35 modifications
36}
37
38fn optimize_stmt_tree(stmt: &mut Stmt, optimizer: &Optimizer) -> bool {
39 let mut changed = false;
40
41 match stmt {
42 Stmt::Assignment { value, .. } => {
43 let before = format!("{value:?}");
44 optimize_expr_in_place(value, optimizer);
45 if format!("{value:?}") != before {
46 changed = true;
47 }
48 }
49 Stmt::If {
50 condition,
51 then_branch,
52 else_branch,
53 } => {
54 optimize_expr_in_place(condition, optimizer);
55 if let Expr::Literal(Literal::Boolean(val)) = &*condition {
56 if *val {
57 if then_branch.len() == 1 {
58 *stmt = then_branch[0].clone();
59 } else {
60 *stmt = Stmt::Block(Block::with_statements(then_branch.clone()));
61 }
62 } else if let Some(else_stmts) = else_branch {
63 if else_stmts.len() == 1 {
64 *stmt = else_stmts[0].clone();
65 } else {
66 *stmt = Stmt::Block(Block::with_statements(else_stmts.clone()));
67 }
68 } else {
69 return false;
70 }
71 return true;
72 }
73 for s in then_branch {
74 changed |= optimize_stmt_tree(s, optimizer);
75 }
76 if let Some(else_stmts) = else_branch {
77 for s in else_stmts {
78 changed |= optimize_stmt_tree(s, optimizer);
79 }
80 }
81 }
82 Stmt::While { condition, body } => {
83 optimize_expr_in_place(condition, optimizer);
84 if matches!(condition, Expr::Literal(Literal::Boolean(false))) {
85 return false;
86 }
87 for s in body {
88 changed |= optimize_stmt_tree(s, optimizer);
89 }
90 }
91 Stmt::For { start, end, body, .. } => {
92 optimize_expr_in_place(start, optimizer);
93 optimize_expr_in_place(end, optimizer);
94 for s in body {
95 changed |= optimize_stmt_tree(s, optimizer);
96 }
97 }
98 Stmt::Repeat {
99 body,
100 until_condition,
101 } => {
102 for s in body {
103 changed |= optimize_stmt_tree(s, optimizer);
104 }
105 optimize_expr_in_place(until_condition, optimizer);
106 }
107 Stmt::Case {
108 expression,
109 branches,
110 else_branch,
111 } => {
112 optimize_expr_in_place(expression, optimizer);
113 for branch in branches {
114 if let Some(guard) = &mut branch.guard {
115 optimize_expr_in_place(guard, optimizer);
116 }
117 for s in &mut branch.body {
118 changed |= optimize_stmt_tree(s, optimizer);
119 }
120 }
121 if let Some(else_stmts) = else_branch {
122 for s in else_stmts {
123 changed |= optimize_stmt_tree(s, optimizer);
124 }
125 }
126 }
127 Stmt::Block(b) => {
128 changed |= fold_block(b, optimizer) > 0;
129 }
130 Stmt::ProcedureCall { arguments, .. } => {
131 for arg in arguments {
132 optimize_expr_in_place(arg, optimizer);
133 }
134 }
135 _ => {}
136 }
137
138 changed
139}
140
141fn apply_advanced_expr(block: &mut Block, advanced: &mut AdvancedOptimizer) -> usize {
142 let mut modifications = 0;
143
144 fn walk_expr(expr: &mut Expr, advanced: &mut AdvancedOptimizer, mods: &mut usize) {
145 match expr {
146 Expr::BinaryOp { left, right, .. } => {
147 walk_expr(left, advanced, mods);
148 walk_expr(right, advanced, mods);
149 }
150 Expr::UnaryOp { operand, .. } => walk_expr(operand, advanced, mods),
151 Expr::FunctionCall { arguments, .. } => {
152 for arg in arguments {
153 walk_expr(arg, advanced, mods);
154 }
155 }
156 _ => {}
157 }
158 let before = format!("{expr:?}");
159 *expr = advanced.optimize_expr(expr);
160 if format!("{expr:?}") != before {
161 *mods += 1;
162 }
163 }
164
165 fn walk_stmt(stmt: &mut Stmt, advanced: &mut AdvancedOptimizer, mods: &mut usize) {
166 match stmt {
167 Stmt::Assignment { value, .. } => walk_expr(value, advanced, mods),
168 Stmt::If {
169 condition,
170 then_branch,
171 else_branch,
172 } => {
173 walk_expr(condition, advanced, mods);
174 for s in then_branch {
175 walk_stmt(s, advanced, mods);
176 }
177 if let Some(else_stmts) = else_branch {
178 for s in else_stmts {
179 walk_stmt(s, advanced, mods);
180 }
181 }
182 }
183 Stmt::While { condition, body } => {
184 walk_expr(condition, advanced, mods);
185 for s in body {
186 walk_stmt(s, advanced, mods);
187 }
188 }
189 Stmt::For { start, end, body, .. } => {
190 walk_expr(start, advanced, mods);
191 walk_expr(end, advanced, mods);
192 for s in body {
193 walk_stmt(s, advanced, mods);
194 }
195 }
196 Stmt::Repeat {
197 body,
198 until_condition,
199 } => {
200 for s in body {
201 walk_stmt(s, advanced, mods);
202 }
203 walk_expr(until_condition, advanced, mods);
204 }
205 Stmt::Block(b) => walk_block(b, advanced, mods),
206 _ => {}
207 }
208 let before = format!("{stmt:?}");
209 *stmt = advanced.optimize_stmt(stmt);
210 if format!("{stmt:?}") != before {
211 *mods += 1;
212 }
213 }
214
215 fn walk_block(block: &mut Block, advanced: &mut AdvancedOptimizer, mods: &mut usize) {
216 for stmt in &mut block.statements {
217 walk_stmt(stmt, advanced, mods);
218 }
219 for proc in &mut block.procedures {
220 walk_block(&mut proc.block, advanced, mods);
221 }
222 for func in &mut block.functions {
223 walk_block(&mut func.block, advanced, mods);
224 }
225 }
226
227 walk_block(block, advanced, &mut modifications);
228 modifications
229}
230
231fn inline_functions_in_block(block: &mut Block, inliner: &mut FunctionInliner) -> usize {
232 let functions: HashMap<String, FunctionDecl> = block
233 .functions
234 .iter()
235 .map(|f| (f.name.to_lowercase(), f.clone()))
236 .collect();
237
238 let mut modifications = 0;
239
240 fn inline_stmt(stmt: &mut Stmt, functions: &HashMap<String, FunctionDecl>, inliner: &mut FunctionInliner, mods: &mut usize) {
241 match stmt {
242 Stmt::Assignment {
243 target,
244 value: Expr::FunctionCall { name, arguments },
245 } if functions.contains_key(&name.to_lowercase()) => {
246 if let Some(func) = functions.get(&name.to_lowercase()) {
247 if inliner.should_inline(func) {
248 let inlined = inliner.inline_call(func, arguments);
249 if let Expr::Variable(var) = target {
250 if let Some(last) = inlined.last() {
251 if let Stmt::Assignment { value, .. } = last {
252 *stmt = Stmt::Assignment {
253 target: Expr::Variable(var.clone()),
254 value: value.clone(),
255 };
256 *mods += 1;
257 }
258 }
259 }
260 }
261 }
262 }
263 Stmt::Block(b) => inline_block(b, functions, inliner, mods),
264 Stmt::If {
265 then_branch,
266 else_branch,
267 ..
268 } => {
269 for s in then_branch {
270 inline_stmt(s, functions, inliner, mods);
271 }
272 if let Some(else_stmts) = else_branch {
273 for s in else_stmts {
274 inline_stmt(s, functions, inliner, mods);
275 }
276 }
277 }
278 Stmt::While { body, .. } | Stmt::For { body, .. } => {
279 for s in body {
280 inline_stmt(s, functions, inliner, mods);
281 }
282 }
283 _ => {}
284 }
285 }
286
287 fn inline_block(block: &mut Block, functions: &HashMap<String, FunctionDecl>, inliner: &mut FunctionInliner, mods: &mut usize) {
288 for stmt in &mut block.statements {
289 inline_stmt(stmt, functions, inliner, mods);
290 }
291 for proc in &mut block.procedures {
292 inline_block(&mut proc.block, functions, inliner, mods);
293 }
294 for func in &mut block.functions {
295 inline_block(&mut func.block, functions, inliner, mods);
296 }
297 }
298
299 inline_block(block, &functions, inliner, &mut modifications);
300 modifications
301}
302
303macro_rules! define_pass {
304 ($struct_name:ident, $pass_name:literal, $level:expr, $desc:literal) => {
305 pub struct $struct_name;
306
307 impl $struct_name {
308 pub fn new() -> Self {
309 Self
310 }
311
312 fn pass_name(&self) -> &str {
313 $pass_name
314 }
315
316 fn pass_description(&self) -> &str {
317 $desc
318 }
319
320 fn pass_level(&self) -> OptimizationLevel {
321 $level
322 }
323
324 fn should_run(&self, context: &OptimizationContext) -> bool {
325 context.level.includes($level)
326 }
327 }
328 };
329}
330
331define_pass!(
332 ConstantFoldingPass,
333 "constant_folding",
334 OptimizationLevel::Basic,
335 "Fold constant expressions at compile time"
336);
337
338impl OptimizationPass for ConstantFoldingPass {
339 fn name(&self) -> &str {
340 self.pass_name()
341 }
342
343 fn description(&self) -> &str {
344 self.pass_description()
345 }
346
347 fn optimization_level(&self) -> OptimizationLevel {
348 self.pass_level()
349 }
350
351 fn should_apply(&self, context: &OptimizationContext) -> bool {
352 self.should_run(context)
353 }
354
355 fn dependencies(&self) -> Vec<String> {
356 Vec::new()
357 }
358
359 fn optimize(&mut self, ast: &mut Program) -> Result<OptimizationResult, CompilerError> {
360 let optimizer = Optimizer::new();
361 let modifications = fold_block(&mut ast.block, &optimizer);
362 let mut stats = OptimizationStats::new();
363 stats.expressions_simplified = modifications;
364 Ok(OptimizationResult::new(modifications, 0, stats))
365 }
366}
367
368define_pass!(
369 DeadCodeEliminationPass,
370 "dead_code_elimination",
371 OptimizationLevel::Standard,
372 "Remove unreachable code and uncalled routines"
373);
374
375impl OptimizationPass for DeadCodeEliminationPass {
376 fn name(&self) -> &str {
377 self.pass_name()
378 }
379
380 fn description(&self) -> &str {
381 self.pass_description()
382 }
383
384 fn optimization_level(&self) -> OptimizationLevel {
385 self.pass_level()
386 }
387
388 fn should_apply(&self, context: &OptimizationContext) -> bool {
389 self.should_run(context)
390 }
391
392 fn dependencies(&self) -> Vec<String> {
393 Vec::new()
394 }
395
396 fn optimize(&mut self, ast: &mut Program) -> Result<OptimizationResult, CompilerError> {
397 let before_procs = ast.block.procedures.len();
398 let before_funcs = ast.block.functions.len();
399 Optimizer::eliminate_dead_procedures_and_functions(&mut ast.block);
400
401 let optimizer = Optimizer::new();
402 ast.block.statements = ast
403 .block
404 .statements
405 .iter()
406 .filter_map(|stmt| optimizer.optimize_stmt(stmt))
407 .collect();
408
409 let removed = (before_procs - ast.block.procedures.len())
410 + (before_funcs - ast.block.functions.len());
411 let mut stats = OptimizationStats::new();
412 stats.dead_code_removed = removed;
413 Ok(OptimizationResult::new(removed, 0, stats))
414 }
415}
416
417define_pass!(
418 CommonSubexpressionEliminationPass,
419 "cse",
420 OptimizationLevel::Standard,
421 "Eliminate duplicate subexpressions"
422);
423
424impl OptimizationPass for CommonSubexpressionEliminationPass {
425 fn name(&self) -> &str {
426 self.pass_name()
427 }
428
429 fn description(&self) -> &str {
430 self.pass_description()
431 }
432
433 fn optimization_level(&self) -> OptimizationLevel {
434 self.pass_level()
435 }
436
437 fn should_apply(&self, context: &OptimizationContext) -> bool {
438 self.should_run(context)
439 }
440
441 fn dependencies(&self) -> Vec<String> {
442 Vec::new()
443 }
444
445 fn optimize(&mut self, ast: &mut Program) -> Result<OptimizationResult, CompilerError> {
446 let mut cse = CSEOptimizer::new();
447 let mut modifications = 0;
448
449 fn walk_exprs(block: &mut Block, cse: &mut CSEOptimizer, mods: &mut usize) {
450 fn walk(expr: &mut Expr, cse: &mut CSEOptimizer, mods: &mut usize) {
451 if let Expr::BinaryOp { left, right, .. } = expr {
452 walk(left, cse, mods);
453 walk(right, cse, mods);
454 let (_optimized, temps) = cse.optimize_expr(expr);
455 if !temps.is_empty() {
456 *mods += temps.len();
457 }
458 }
459 }
460 for stmt in &mut block.statements {
461 if let Stmt::Assignment { value, .. } = stmt {
462 walk(value, cse, mods);
463 }
464 }
465 for func in &mut block.functions {
466 walk_exprs(&mut func.block, cse, mods);
467 }
468 }
469
470 walk_exprs(&mut ast.block, &mut cse, &mut modifications);
471 let mut stats = OptimizationStats::new();
472 stats.expressions_simplified = modifications;
473 Ok(OptimizationResult::new(modifications, 0, stats))
474 }
475}
476
477define_pass!(
478 FunctionInliningPass,
479 "function_inlining",
480 OptimizationLevel::Aggressive,
481 "Inline small functions"
482);
483
484impl OptimizationPass for FunctionInliningPass {
485 fn name(&self) -> &str {
486 self.pass_name()
487 }
488
489 fn description(&self) -> &str {
490 self.pass_description()
491 }
492
493 fn optimization_level(&self) -> OptimizationLevel {
494 self.pass_level()
495 }
496
497 fn should_apply(&self, context: &OptimizationContext) -> bool {
498 self.should_run(context)
499 }
500
501 fn dependencies(&self) -> Vec<String> {
502 Vec::new()
503 }
504
505 fn optimize(&mut self, ast: &mut Program) -> Result<OptimizationResult, CompilerError> {
506 let mut inliner = FunctionInliner::new(10);
507 let modifications = inline_functions_in_block(&mut ast.block, &mut inliner);
508 let mut stats = OptimizationStats::new();
509 stats.functions_inlined = inliner.inlined_count();
510 Ok(OptimizationResult::new(modifications, 0, stats))
511 }
512}
513
514define_pass!(
515 LoopOptimizationPass,
516 "loop_optimization",
517 OptimizationLevel::Standard,
518 "Optimize loops via unrolling and strength reduction"
519);
520
521impl OptimizationPass for LoopOptimizationPass {
522 fn name(&self) -> &str {
523 self.pass_name()
524 }
525
526 fn description(&self) -> &str {
527 self.pass_description()
528 }
529
530 fn optimization_level(&self) -> OptimizationLevel {
531 self.pass_level()
532 }
533
534 fn should_apply(&self, context: &OptimizationContext) -> bool {
535 self.should_run(context)
536 }
537
538 fn dependencies(&self) -> Vec<String> {
539 Vec::new()
540 }
541
542 fn optimize(&mut self, ast: &mut Program) -> Result<OptimizationResult, CompilerError> {
543 let loop_opt = LoopOptimizer::new(4);
544 let mut advanced = AdvancedOptimizer::new();
545 let mut modifications = 0;
546
547 fn walk_loops(block: &mut Block, loop_opt: &LoopOptimizer, mods: &mut usize) {
548 block.statements = block
549 .statements
550 .iter()
551 .map(|stmt| {
552 let optimized = loop_opt.optimize_loop(stmt);
553 if format!("{optimized:?}") != format!("{stmt:?}") {
554 *mods += 1;
555 }
556 optimized
557 })
558 .collect();
559 for func in &mut block.functions {
560 walk_loops(&mut func.block, loop_opt, mods);
561 }
562 }
563
564 walk_loops(&mut ast.block, &loop_opt, &mut modifications);
565 modifications += apply_advanced_expr(&mut ast.block, &mut advanced);
566
567 let mut stats = OptimizationStats::new();
568 stats.loops_unrolled = modifications;
569 Ok(OptimizationResult::new(modifications, 0, stats))
570 }
571}
572
573define_pass!(
574 RegisterAllocationPass,
575 "register_allocation",
576 OptimizationLevel::Aggressive,
577 "Analyze register pressure for generated code"
578);
579
580impl OptimizationPass for RegisterAllocationPass {
581 fn name(&self) -> &str {
582 self.pass_name()
583 }
584
585 fn description(&self) -> &str {
586 self.pass_description()
587 }
588
589 fn optimization_level(&self) -> OptimizationLevel {
590 self.pass_level()
591 }
592
593 fn should_apply(&self, context: &OptimizationContext) -> bool {
594 self.should_run(context)
595 }
596
597 fn dependencies(&self) -> Vec<String> {
598 Vec::new()
599 }
600
601 fn optimize(&mut self, ast: &mut Program) -> Result<OptimizationResult, CompilerError> {
602 use crate::register_allocator::{LiveRangeAnalyzer, RegisterAllocator};
603
604 let mut analyzer = LiveRangeAnalyzer::new();
605 for var in &ast.block.vars {
606 analyzer.define_variable(&var.name);
607 }
608 for stmt in &ast.block.statements {
609 if let Stmt::Assignment {
610 target: Expr::Variable(name),
611 ..
612 } = stmt
613 {
614 analyzer.use_variable(name);
615 }
616 analyzer.next_instruction();
617 }
618
619 let live_ranges = analyzer.get_live_ranges();
620 let mut allocator = RegisterAllocator::new();
621 if !live_ranges.is_empty() {
622 allocator.build_interference_graph(live_ranges);
623 allocator.allocate().map_err(|e| CompilerError::OptimizerError {
624 location: None,
625 message: e.to_string(),
626 })?;
627 }
628
629 let modifications = allocator.get_allocations().len();
630 let mut stats = OptimizationStats::new();
631 stats.memory_usage_optimized = modifications;
632 Ok(OptimizationResult::new(modifications, 0, stats))
633 }
634}