More optimizations

rust_compiler/libs/optimizer/src/dead_store_elimination.rs

@@ -32,11 +32,29 @@ pub fn dead_store_elimination<'a>(
             last_write.insert(dest_reg, i);
         }
 
-        // On labels/jumps, conservatively clear tracking (value might be used elsewhere)
+        // Before clearing on labels/calls, check if current tracked writes are dead
         if matches!(
             node.instruction,
-            Instruction::LabelDef(_) | Instruction::Jump(_) | Instruction::JumpAndLink(_)
+            Instruction::LabelDef(_) | Instruction::JumpAndLink(_)
         ) {
+            // Check all currently tracked writes to see if they're dead
+            for (&reg, &idx) in &last_write {
+                // Don't remove writes to r15 (return register)
+                if reg == 15 {
+                    continue;
+                }
+
+                // Check if this write was used between write and now
+                let was_used = input[idx + 1..i]
+                    .iter()
+                    .any(|n| reg_is_read_or_affects_control(&n.instruction, reg));
+
+                if !was_used && !to_remove.contains(&idx) {
+                    to_remove.push(idx);
+                    changed = true;
+                }
+            }
+
             last_write.clear();
         }
     }
@@ -59,29 +77,12 @@ pub fn dead_store_elimination<'a>(
     }
 }
 
-/// Simplified check: Does this instruction read the register or affect control flow?
+/// Simplified check: Does this instruction read the register?
 fn reg_is_read_or_affects_control(instr: &Instruction, reg: u8) -> bool {
     use crate::helpers::reg_is_read;
 
     // If it reads the register, it's used
-    if reg_is_read(instr, reg) {
-        return true;
-    }
-
-    // Conservatively assume register might be used if there's control flow
-    matches!(
-        instr,
-        Instruction::Jump(_)
-            | Instruction::JumpAndLink(_)
-            | Instruction::BranchEq(_, _, _)
-            | Instruction::BranchNe(_, _, _)
-            | Instruction::BranchGt(_, _, _)
-            | Instruction::BranchLt(_, _, _)
-            | Instruction::BranchGe(_, _, _)
-            | Instruction::BranchLe(_, _, _)
-            | Instruction::BranchEqZero(_, _)
-            | Instruction::BranchNeZero(_, _)
-    )
+    reg_is_read(instr, reg)
 }
 
 #[cfg(test)]

rust_compiler/libs/optimizer/src/peephole_optimization.rs

@@ -38,6 +38,8 @@ pub fn peephole_optimization<'a>(
                                 // Safe to remove all four: push sp, push ra, pop ra, pop sp
                                 // Also need to adjust stack pointer offsets in between by -2
                                 let absolute_sp_pop = absolute_ra_pop + 1;
+                                // Clear output since we're going to reprocess the entire input
+                                output.clear();
                                 for (idx, node) in input.iter().enumerate() {
                                     if idx == i
                                         || idx == i + 1
@@ -83,6 +85,8 @@ pub fn peephole_optimization<'a>(
                     // Safe to remove both push and pop
                     // Also need to adjust stack pointer offsets in between
                     let absolute_pop_idx = i + pop_idx;
+                    // Clear output since we're going to reprocess the entire input
+                    output.clear();
                     for (idx, node) in input.iter().enumerate() {
                         if idx == i || idx == absolute_pop_idx {
                             // Skip the push and pop