diff --git a/rust_compiler/libs/integration_tests/src/snapshots/integration_tests__tests__leaf_function_no_stack_frame.snap b/rust_compiler/libs/integration_tests/src/snapshots/integration_tests__tests__leaf_function_no_stack_frame.snap index 8037ef9..3362dcd 100644 --- a/rust_compiler/libs/integration_tests/src/snapshots/integration_tests__tests__leaf_function_no_stack_frame.snap +++ b/rust_compiler/libs/integration_tests/src/snapshots/integration_tests__tests__leaf_function_no_stack_frame.snap @@ -1,5 +1,6 @@ --- source: libs/integration_tests/src/lib.rs +assertion_line: 144 expression: output --- ## Unoptimized Output @@ -20,6 +21,10 @@ j ra j main pop r8 +push sp +push ra add r1 r8 1 move r8 r1 +pop ra +pop sp j ra diff --git a/rust_compiler/libs/integration_tests/src/snapshots/integration_tests__tests__peephole_comparison_fusion.snap b/rust_compiler/libs/integration_tests/src/snapshots/integration_tests__tests__peephole_comparison_fusion.snap index b376dca..880034c 100644 --- a/rust_compiler/libs/integration_tests/src/snapshots/integration_tests__tests__peephole_comparison_fusion.snap +++ b/rust_compiler/libs/integration_tests/src/snapshots/integration_tests__tests__peephole_comparison_fusion.snap @@ -1,5 +1,6 @@ --- source: libs/integration_tests/src/lib.rs +assertion_line: 116 expression: output --- ## Unoptimized Output @@ -24,6 +25,10 @@ j ra j main pop r8 pop r9 -ble r9 r8 5 +push sp +push ra +ble r9 r8 7 move r10 1 +pop ra +pop sp j ra diff --git a/rust_compiler/libs/integration_tests/src/snapshots/integration_tests__tests__simple_leaf_function.snap b/rust_compiler/libs/integration_tests/src/snapshots/integration_tests__tests__simple_leaf_function.snap index 2baef8d..f093e06 100644 --- a/rust_compiler/libs/integration_tests/src/snapshots/integration_tests__tests__simple_leaf_function.snap +++ b/rust_compiler/libs/integration_tests/src/snapshots/integration_tests__tests__simple_leaf_function.snap @@ -1,5 +1,6 @@ --- source: libs/integration_tests/src/lib.rs +assertion_line: 60 expression: output --- ## Unoptimized Output @@ -17,5 +18,9 @@ j ra ## Optimized Output j main +push sp +push ra move r8 10 +pop ra +pop sp j ra diff --git a/rust_compiler/libs/optimizer/src/leaf_function_optimization.rs b/rust_compiler/libs/optimizer/src/leaf_function_optimization.rs index 8447e95..d65f10a 100644 --- a/rust_compiler/libs/optimizer/src/leaf_function_optimization.rs +++ b/rust_compiler/libs/optimizer/src/leaf_function_optimization.rs @@ -1,150 +1,41 @@ use crate::leaf_function::find_leaf_functions; -use il::{Instruction, InstructionNode, Operand}; -use rust_decimal::Decimal; -use std::collections::{HashMap, HashSet}; - -/// Helper: Check if a function body contains unsafe stack manipulation. -fn function_has_complex_stack_ops( - instructions: &[InstructionNode], - start_idx: usize, - end_idx: usize, -) -> bool { - for instruction in instructions.iter().take(end_idx).skip(start_idx) { - match instruction.instruction { - Instruction::Push(_) | Instruction::Pop(_) => return true, - Instruction::Add(Operand::StackPointer, _, _) - | Instruction::Sub(Operand::StackPointer, _, _) - | Instruction::Mul(Operand::StackPointer, _, _) - | Instruction::Div(Operand::StackPointer, _, _) - | Instruction::Move(Operand::StackPointer, _) => return true, - _ => {} - } - } - false -} +use il::InstructionNode; /// Pass: Leaf Function Optimization /// If a function makes no calls (is a leaf), it doesn't need to save/restore `ra`. +/// +/// NOTE: This optimization is DISABLED due to correctness issues. +/// The optimization was designed for a specific calling convention (GET/PUT for RA) +/// but the compiler generates POP ra for return address restoration. Without proper +/// tracking of both conventions and validation of balanced push/pop pairs, this +/// optimization corrupts the stack frame by: +/// +/// 1. Removing `push ra` but not `pop ra`, leaving unbalanced push/pop pairs +/// 2. Not accounting for parameter pops that occur before `push sp` +/// 3. Assuming all RA restoration uses GET instruction, but code uses POP +/// +/// Example of broken optimization: +/// ``` +/// Unoptimized: Optimized (BROKEN): +/// compare: pop r8 +/// pop r8 pop r9 +/// pop r9 ble r9 r8 5 +/// push sp move r10 1 +/// push ra j ra +/// sgt r1 r9 r8 ^ Missing stack frame! +/// ... +/// pop ra +/// pop sp +/// j ra +/// ``` +/// +/// Future work: Fix by handling both POP and GET calling conventions, validating +/// balanced push/pop pairs, and accounting for parameter pops. pub fn optimize_leaf_functions<'a>( input: Vec>, ) -> (Vec>, bool) { - let leaves = find_leaf_functions(&input); - if leaves.is_empty() { - return (input, false); - } - - let mut changed = false; - let mut to_remove = HashSet::new(); - let mut func_restore_indices = HashMap::new(); - let mut func_ra_offsets = HashMap::new(); - let mut current_function: Option = None; - let mut function_start_indices = HashMap::new(); - - // First scan: Identify instructions to remove and capture RA offsets - for (i, node) in input.iter().enumerate() { - match &node.instruction { - Instruction::LabelDef(label) if !label.starts_with("__internal_L") => { - current_function = Some(label.to_string()); - function_start_indices.insert(label.to_string(), i); - } - Instruction::Push(Operand::ReturnAddress) => { - if let Some(func) = ¤t_function - && leaves.contains(func) - { - to_remove.insert(i); - } - } - Instruction::Get(Operand::ReturnAddress, _, Operand::Register(_)) => { - if let Some(func) = ¤t_function - && leaves.contains(func) - { - to_remove.insert(i); - func_restore_indices.insert(func.clone(), i); - - // Look back for the address calc: `sub r0 sp OFFSET` - if i > 0 - && let Instruction::Sub(_, Operand::StackPointer, Operand::Number(n)) = - &input[i - 1].instruction - { - func_ra_offsets.insert(func.clone(), *n); - to_remove.insert(i - 1); - } - } - } - _ => {} - } - } - - // Safety Check: Verify functions don't have complex stack ops - let mut safe_functions = HashSet::new(); - - for (func, start_idx) in &function_start_indices { - if let Some(restore_idx) = func_restore_indices.get(func) { - let check_start = if to_remove.contains(&(start_idx + 1)) { - start_idx + 2 - } else { - start_idx + 1 - }; - - if !function_has_complex_stack_ops(&input, check_start, *restore_idx) { - safe_functions.insert(func.clone()); - changed = true; - } - } - } - - if !changed { - return (input, false); - } - - // Second scan: Rebuild with adjustments - let mut output = Vec::with_capacity(input.len()); - let mut processing_function: Option = None; - - for (i, mut node) in input.into_iter().enumerate() { - if to_remove.contains(&i) - && let Some(func) = &processing_function - && safe_functions.contains(func) - { - continue; - } - - if let Instruction::LabelDef(l) = &node.instruction - && !l.starts_with("__internal_L") - { - processing_function = Some(l.to_string()); - } - - // Apply Stack Adjustments - if let Some(func) = &processing_function - && safe_functions.contains(func) - && let Some(ra_offset) = func_ra_offsets.get(func) - { - // Stack Cleanup Adjustment - if let Instruction::Sub( - Operand::StackPointer, - Operand::StackPointer, - Operand::Number(n), - ) = &mut node.instruction - { - let new_n = *n - Decimal::from(1); - if new_n.is_zero() { - continue; - } - *n = new_n; - } - - // Stack Variable Offset Adjustment - if let Instruction::Sub(_, Operand::StackPointer, Operand::Number(n)) = - &mut node.instruction - && *n > *ra_offset - { - *n -= Decimal::from(1); - } - } - - output.push(node); - } - - (output, true) + // Optimization disabled - returns input unchanged + #[allow(unused)] + let _leaves = find_leaf_functions(&input); + (input, false) } diff --git a/rust_compiler/libs/optimizer/src/peephole_optimization.rs b/rust_compiler/libs/optimizer/src/peephole_optimization.rs index f4ecc59..eb78676 100644 --- a/rust_compiler/libs/optimizer/src/peephole_optimization.rs +++ b/rust_compiler/libs/optimizer/src/peephole_optimization.rs @@ -358,11 +358,12 @@ fn find_matching_ra_pop<'a>( return Some((idx, &instructions[1..idx])); } - // Stop searching if we hit a jump (different control flow) - // Labels are OK - they're just markers + // Stop searching if we hit a jump (different control flow) or a function label + // Labels are OK - they're just markers EXCEPT for user-defined function labels + // which indicate a function boundary if matches!( node.instruction, - Instruction::Jump(_) | Instruction::JumpRelative(_) + Instruction::Jump(_) | Instruction::JumpRelative(_) | Instruction::LabelDef(_) ) { return None; }