Fix source maps

rust_compiler/libs/compiler/src/test/mod.rs

@@ -12,18 +12,16 @@ macro_rules! compile {
         let mut writer = std::io::BufWriter::new(Vec::new());
         let compiler = ::Compiler::new(
             parser::Parser::new(tokenizer::Tokenizer::from(String::from($source))),
-            &mut writer,
             None,
         );
-        compiler.compile();
+        let res = compiler.compile();
+        res.instructions.write(&mut writer)?;
         output!(writer)
     }};
 
     (result $source:expr) => {{
-        let mut writer = std::io::BufWriter::new(Vec::new());
         let compiler = crate::Compiler::new(
             parser::Parser::new(tokenizer::Tokenizer::from($source)),
-            &mut writer,
             Some(crate::CompilerConfig { debug: true }),
         );
         compiler.compile().errors
@@ -33,10 +31,10 @@ macro_rules! compile {
         let mut writer = std::io::BufWriter::new(Vec::new());
         let compiler = crate::Compiler::new(
             parser::Parser::new(tokenizer::Tokenizer::from($source)),
-            &mut writer,
             Some(crate::CompilerConfig { debug: true }),
         );
-        compiler.compile();
+        let res = compiler.compile();
+        res.instructions.write(&mut writer)?;
         output!(writer)
     }};
 }

rust_compiler/libs/compiler/src/v1.rs

@@ -1,7 +1,7 @@
 #![allow(clippy::result_large_err)]
 use crate::variable_manager::{self, LocationRequest, VariableLocation, VariableScope};
 use helpers::{Span, prelude::*};
-use il::{Instruction, InstructionNode, Operand};
+use il::{Instruction, InstructionNode, Instructions, Operand};
 use parser::{
     Parser as ASTParser,
     sys_call::{Math, SysCall, System},
@@ -13,11 +13,7 @@ use parser::{
     },
 };
 use rust_decimal::Decimal;
-use std::{
+use std::{borrow::Cow, collections::HashMap};
-    borrow::Cow,
-    collections::HashMap,
-    io::{BufWriter, Write},
-};
 use thiserror::Error;
 use tokenizer::token::Number;
 
@@ -139,24 +135,22 @@ struct CompileLocation<'a> {
 
 pub struct CompilationResult<'a> {
     pub errors: Vec<Error<'a>>,
-    pub source_map: HashMap<usize, Vec<Span>>,
+    pub instructions: Instructions<'a>,
-    pub instructions: Vec<InstructionNode<'a>>,
 }
 
-pub struct Compiler<'a, 'w, W: std::io::Write> {
+pub struct Compiler<'a> {
     pub parser: ASTParser<'a>,
     function_locations: HashMap<Cow<'a, str>, usize>,
     function_metadata: HashMap<Cow<'a, str>, Vec<Cow<'a, str>>>,
     devices: HashMap<Cow<'a, str>, Cow<'a, str>>,
-    output: &'w mut BufWriter<W>,
 
     // This holds the IL code which will be used in the
     // optimizer
-    pub instructions: Vec<InstructionNode<'a>>,
+    pub instructions: Instructions<'a>,
 
     current_line: usize,
     declared_main: bool,
-    config: CompilerConfig,
+    _config: CompilerConfig,
     temp_counter: usize,
     label_counter: usize,
     loop_stack: Vec<(Cow<'a, str>, Cow<'a, str>)>, // Stores (start_label, end_label)
@@ -167,22 +161,17 @@ pub struct Compiler<'a, 'w, W: std::io::Write> {
     pub errors: Vec<Error<'a>>,
 }
 
-impl<'a, 'w, W: std::io::Write> Compiler<'a, 'w, W> {
+impl<'a> Compiler<'a> {
-    pub fn new(
+    pub fn new(parser: ASTParser<'a>, config: Option<CompilerConfig>) -> Self {
-        parser: ASTParser<'a>,
-        writer: &'w mut BufWriter<W>,
-        config: Option<CompilerConfig>,
-    ) -> Self {
         Self {
             parser,
             function_locations: HashMap::new(),
             function_metadata: HashMap::new(),
             devices: HashMap::new(),
-            output: writer,
+            instructions: Instructions::default(),
-            instructions: Vec::new(),
             current_line: 1,
             declared_main: false,
-            config: config.unwrap_or_default(),
+            _config: config.unwrap_or_default(),
             temp_counter: 0,
             label_counter: 0,
             loop_stack: Vec::new(),
@@ -205,7 +194,6 @@ impl<'a, 'w, W: std::io::Write> Compiler<'a, 'w, W> {
             Ok(Some(expr)) => expr,
             Ok(None) => {
                 return CompilationResult {
-                    source_map: self.source_map,
                     errors: self.errors,
                     instructions: self.instructions,
                 };
@@ -215,7 +203,6 @@ impl<'a, 'w, W: std::io::Write> Compiler<'a, 'w, W> {
                 self.errors.push(Error::Parse(e));
                 return CompilationResult {
                     errors: self.errors,
-                    source_map: self.source_map,
                     instructions: self.instructions,
                 };
             }
@@ -241,7 +228,6 @@ impl<'a, 'w, W: std::io::Write> Compiler<'a, 'w, W> {
             self.errors.push(e);
             return CompilationResult {
                 errors: self.errors,
-                source_map: self.source_map,
                 instructions: self.instructions,
             };
         }
@@ -255,7 +241,6 @@ impl<'a, 'w, W: std::io::Write> Compiler<'a, 'w, W> {
 
         CompilationResult {
             errors: self.errors,
-            source_map: self.source_map,
             instructions: self.instructions,
         }
     }
@@ -266,8 +251,6 @@ impl<'a, 'w, W: std::io::Write> Compiler<'a, 'w, W> {
         instr: Instruction<'a>,
         span: Option<Span>,
     ) -> Result<(), Error<'a>> {
-        self.output.write_all(format!("{}", instr).as_bytes())?;
-        self.output.write_all(b"\n")?;
         self.current_line += 1;
 
         self.instructions.push(InstructionNode::new(instr, span));
@@ -781,7 +764,6 @@ impl<'a, 'w, W: std::io::Write> Compiler<'a, 'w, W> {
             }
             Expression::Ternary(ternary) => {
                 let res = self.expression_ternary(ternary.node, scope)?;
-                println!("{res:?}");
                 let var_loc = scope.add_variable(
                     name_str.clone(),
                     LocationRequest::Persist,

rust_compiler/libs/il/src/lib.rs

@@ -1,13 +1,77 @@
 use helpers::Span;
 use rust_decimal::Decimal;
 use std::borrow::Cow;
+use std::collections::HashMap;
 use std::fmt;
+use std::io::{BufWriter, Write};
+use std::ops::{Deref, DerefMut};
+
+#[derive(Default)]
+pub struct Instructions<'a>(Vec<InstructionNode<'a>>);
+
+impl<'a> Deref for Instructions<'a> {
+    type Target = Vec<InstructionNode<'a>>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl<'a> DerefMut for Instructions<'a> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+impl<'a> Instructions<'a> {
+    pub fn new(instructions: Vec<InstructionNode<'a>>) -> Self {
+        Self(instructions)
+    }
+    pub fn into_inner(self) -> Vec<InstructionNode<'a>> {
+        self.0
+    }
+    pub fn write(self, writer: &mut BufWriter<dyn Write>) -> Result<(), std::io::Error> {
+        for node in self.0 {
+            writer.write_all(node.to_string().as_bytes())?;
+            writer.write_all(b"\n")?;
+        }
+
+        writer.flush()?;
+        Ok(())
+    }
+    pub fn source_map(&self) -> HashMap<usize, Span> {
+        let mut map = HashMap::new();
+
+        for (line_num, node) in self.0.iter().enumerate() {
+            if let Some(span) = node.span {
+                map.insert(line_num, span);
+            }
+        }
+
+        map
+    }
+}
+
+impl<'a> std::fmt::Display for Instructions<'a> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        for node in &self.0 {
+            writeln!(f, "{node}")?;
+        }
+        Ok(())
+    }
+}
 
 pub struct InstructionNode<'a> {
     pub instruction: Instruction<'a>,
     pub span: Option<Span>,
 }
 
+impl<'a> std::fmt::Display for InstructionNode<'a> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}", self.instruction)
+    }
+}
+
 impl<'a> InstructionNode<'a> {
     pub fn new(instr: Instruction<'a>, span: Option<Span>) -> Self {
         Self {

rust_compiler/libs/optimizer/src/lib.rs

@@ -1,4 +1,4 @@
-use il::{Instruction, InstructionNode, Operand};
+use il::{Instruction, InstructionNode, Instructions, Operand};
 use rust_decimal::Decimal;
 use std::collections::{HashMap, HashSet};
 
@@ -6,7 +6,8 @@ mod leaf_function;
 use leaf_function::find_leaf_functions;
 
 /// Entry point for the optimizer.
-pub fn optimize<'a>(mut instructions: Vec<InstructionNode<'a>>) -> Vec<InstructionNode<'a>> {
+pub fn optimize<'a>(instructions: Instructions<'a>) -> Instructions<'a> {
+    let mut instructions = instructions.into_inner();
     let mut changed = true;
     let mut pass_count = 0;
     const MAX_PASSES: usize = 10;
@@ -49,13 +50,35 @@ pub fn optimize<'a>(mut instructions: Vec<InstructionNode<'a>>) -> Vec<Instructi
     }
 
     // Final Pass: Resolve Labels to Line Numbers
-    resolve_labels(instructions)
+    Instructions::new(resolve_labels(instructions))
+}
+
+/// Helper: Check if a function body contains unsafe stack manipulation.
+/// Returns true if the function modifies SP in a way that makes static RA offset analysis unsafe.
+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,
+            // Check for explicit SP modification
+            Instruction::Add(Operand::StackPointer, _, _)
+            | Instruction::Sub(Operand::StackPointer, _, _)
+            | Instruction::Mul(Operand::StackPointer, _, _)
+            | Instruction::Div(Operand::StackPointer, _, _)
+            | Instruction::Move(Operand::StackPointer, _) => return true,
+            _ => {}
+        }
+    }
+    false
 }
 
 /// Pass: Leaf Function Optimization
 /// If a function makes no calls (is a leaf), it doesn't need to save/restore `ra`.
 fn optimize_leaf_functions<'a>(
-    mut input: Vec<InstructionNode<'a>>,
+    input: Vec<InstructionNode<'a>>,
 ) -> (Vec<InstructionNode<'a>>, bool) {
     let leaves = find_leaf_functions(&input);
     if leaves.is_empty() {
@@ -64,34 +87,40 @@ fn optimize_leaf_functions<'a>(
 
     let mut changed = false;
     let mut to_remove = HashSet::new();
-    let mut current_function: Option<String> = None;
 
-    // Map of FunctionName -> The stack offset where RA was stored.
+    // We map function names to the INDEX of the instruction that restores RA.
-    // We need this to adjust other stack accesses (arguments vs locals).
+    // We use this to validate the function body later.
+    let mut func_restore_indices = HashMap::new();
     let mut func_ra_offsets = HashMap::new();
 
+    let mut current_function: Option<String> = 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) => {
                 current_function = Some(label.to_string());
+                function_start_indices.insert(label.to_string(), i);
             }
             Instruction::Push(Operand::ReturnAddress) => {
-                if let Some(func) = &current_function {
+                if let Some(func) = &current_function
-                    if leaves.contains(func) {
+                    && leaves.contains(func)
+                {
                     to_remove.insert(i);
-                        changed = true;
-                    }
                 }
             }
             Instruction::Get(Operand::ReturnAddress, _, Operand::Register(_)) => {
                 // This is the restore instruction: `get ra db r0`
-                if let Some(func) = &current_function {
+                if let Some(func) = &current_function
-                    if leaves.contains(func) {
+                    && 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 {
+                    if i > 0
-                            if let Instruction::Sub(_, Operand::StackPointer, Operand::Number(n)) =
+                        && let Instruction::Sub(_, Operand::StackPointer, Operand::Number(n)) =
                             &input[i - 1].instruction
                     {
                         func_ra_offsets.insert(func.clone(), *n);
@@ -99,62 +128,86 @@ fn optimize_leaf_functions<'a>(
                     }
                 }
             }
-                }
-            }
             _ => {}
         }
     }
 
+    // Safety Check: Verify that functions marked for optimization don't have complex stack ops.
+    // If they do, unmark them.
+    let mut safe_functions = HashSet::new();
+
+    for (func, start_idx) in &function_start_indices {
+        if let Some(restore_idx) = func_restore_indices.get(func) {
+            // Check instructions between start and restore using the helper function.
+            // We need to skip the `push ra` we just marked for removal, otherwise the helper
+            // will flag it as a complex op (Push).
+            // `start_idx` is the LabelDef. `start_idx + 1` is typically `push ra`.
+
+            let check_start = if to_remove.contains(&(start_idx + 1)) {
+                start_idx + 2
+            } else {
+                start_idx + 1
+            };
+
+            // `restore_idx` points to the `get ra` instruction. The helper scans up to `end_idx` exclusive,
+            // so we don't need to worry about the restore instruction itself.
+            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
+    // Second scan: Rebuild with adjustments, but only for SAFE functions
     let mut output = Vec::with_capacity(input.len());
     let mut processing_function: Option<String> = None;
 
     for (i, mut node) in input.into_iter().enumerate() {
-        if to_remove.contains(&i) {
+        if to_remove.contains(&i)
-            continue;
+            && let Some(func) = &processing_function
+            && safe_functions.contains(func)
+        {
+            continue; // SKIP (Remove)
         }
 
         if let Instruction::LabelDef(l) = &node.instruction {
             processing_function = Some(l.to_string());
         }
 
-        // Apply Stack Adjustments if we are inside a leaf function that we optimized
+        // Apply Stack Adjustments
-        if let Some(func) = &processing_function {
+        if let Some(func) = &processing_function
-            if let Some(ra_offset) = func_ra_offsets.get(func) {
+            && safe_functions.contains(func)
-                // If this is the stack cleanup `sub sp sp N`, decrement N by 1 (since we removed push ra)
+            && let Some(ra_offset) = func_ra_offsets.get(func)
+        {
+            // 1. Stack Cleanup Adjustment
             if let Instruction::Sub(
                 Operand::StackPointer,
                 Operand::StackPointer,
                 Operand::Number(n),
             ) = &mut node.instruction
             {
+                // Decrease cleanup amount by 1 (for the removed RA)
                 let new_n = *n - Decimal::from(1);
                 if new_n.is_zero() {
-                        continue; // Remove instruction if 0
+                    continue;
                 }
                 *n = new_n;
             }
 
-                // Adjust stack variable accesses relative to the removed RA.
+            // 2. Stack Variable Offset Adjustment
-                // Compiler layout: [Args] [RA] [Locals/Temps]
+            // Since we verified the function is "Simple" (no nested stack mods),
-                // Stack grows up (increment sp on push).
+            // we can safely assume offsets > ra_offset need shifting.
-                // Access is `sp - offset`.
-                // Deeper items (Args) have LARGER offsets than RA.
-                // Shallower items (Locals) have SMALLER offsets than RA.
-                // Since RA is gone, items deeper than RA (Args) effectively shift "down" (index - 1).
             if let Instruction::Sub(_, Operand::StackPointer, Operand::Number(n)) =
                 &mut node.instruction
+                && *n > *ra_offset
             {
-                    if *n > *ra_offset {
                 *n -= Decimal::from(1);
             }
         }
-            }
-        }
 
         output.push(node);
     }
@@ -173,36 +226,31 @@ fn analyze_clobbers(instructions: &[InstructionNode]) -> HashMap<String, HashSet
             clobbers.insert(label.to_string(), HashSet::new());
         }
 
-        if let Some(label) = &current_label {
+        if let Some(label) = &current_label
-            if let Some(reg) = get_destination_reg(&node.instruction) {
+            && let Some(reg) = get_destination_reg(&node.instruction)
-                if let Some(set) = clobbers.get_mut(label) {
+            && let Some(set) = clobbers.get_mut(label)
+        {
             set.insert(reg);
         }
     }
-
-            // Note: If we call another function, we technically clobber whatever IT clobbers
-            // (unless we save/restore it, which counts as a write anyway).
-            // This simple pass relies on the fact that any register modification (including restore) is a 'write'.
-        }
-    }
     clobbers
 }
 
 /// Pass: Function Call Optimization
 /// Removes Push/Restore pairs surrounding a JAL if the target function does not clobber that register.
 fn optimize_function_calls<'a>(
-    mut input: Vec<InstructionNode<'a>>,
+    input: Vec<InstructionNode<'a>>,
 ) -> (Vec<InstructionNode<'a>>, bool) {
     let clobbers = analyze_clobbers(&input);
     let mut changed = false;
     let mut to_remove = HashSet::new();
-    let mut stack_adjustments = HashMap::new(); // Index of `sub sp sp N` -> amount to subtract
+    let mut stack_adjustments = HashMap::new();
 
     let mut i = 0;
     while i < input.len() {
         if let Instruction::JumpAndLink(Operand::Label(target)) = &input[i].instruction {
             let target_key = target.to_string();
-            // If we don't have info on the function (e.g. extern or complex), skip
+
             if let Some(func_clobbers) = clobbers.get(&target_key) {
                 // 1. Identify Pushes immediately preceding the JAL
                 let mut pushes = Vec::new(); // (index, register)
@@ -221,17 +269,16 @@ fn optimize_function_calls<'a>(
                 }
 
                 // 2. Identify Restores immediately following the JAL
-                // Compiler emits: sub r0 sp Offset, get Reg db r0.
                 let mut restores = Vec::new(); // (index_of_get, register, index_of_sub)
                 let mut scan_fwd = i + 1;
                 while scan_fwd < input.len() {
-                    // Skip over the 'sub r0 sp X' address calculation lines
+                    // Skip 'sub r0 sp X'
                     if let Instruction::Sub(Operand::Register(0), Operand::StackPointer, _) =
                         &input[scan_fwd].instruction
                     {
                         // Check next instruction for the Get
-                        if scan_fwd + 1 < input.len() {
+                        if scan_fwd + 1 < input.len()
-                            if let Instruction::Get(Operand::Register(r), _, Operand::Register(0)) =
+                            && let Instruction::Get(Operand::Register(r), _, Operand::Register(0)) =
                                 &input[scan_fwd + 1].instruction
                         {
                             restores.push((scan_fwd + 1, *r, scan_fwd));
@@ -239,36 +286,50 @@ fn optimize_function_calls<'a>(
                             continue;
                         }
                     }
-                    }
                     break;
                 }
 
-                // 3. Check for Stack Cleanup `sub sp sp N`
+                // 3. Stack Cleanup
                 let cleanup_idx = scan_fwd;
-
                 let has_cleanup = if cleanup_idx < input.len() {
-                    if let Instruction::Sub(
+                    matches!(
+                        input[cleanup_idx].instruction,
+                        Instruction::Sub(
                             Operand::StackPointer,
                             Operand::StackPointer,
-                        Operand::Number(_),
+                            Operand::Number(_)
-                    ) = &input[cleanup_idx].instruction
+                        )
-                    {
+                    )
-                        true
-                    } else {
-                        false
-                    }
                 } else {
                     false
                 };
 
-                // "All or Nothing" strategy for the safe subset:
+                // SAFEGUARD: Check Counts!
+                // If we pushed r8 twice but only restored it once, we have an argument.
+                // We must ensure the number of pushes for each register MATCHES the number of restores.
+                let mut push_counts = HashMap::new();
+                for (_, r) in &pushes {
+                    *push_counts.entry(*r).or_insert(0) += 1;
+                }
+
+                let mut restore_counts = HashMap::new();
+                for (_, r, _) in &restores {
+                    *restore_counts.entry(*r).or_insert(0) += 1;
+                }
+
+                let counts_match = push_counts
+                    .iter()
+                    .all(|(reg, count)| restore_counts.get(reg).unwrap_or(&0) == count);
+                // Also check reverse to ensure we didn't restore something we didn't push (unlikely but possible)
+                let counts_match_reverse = restore_counts
+                    .iter()
+                    .all(|(reg, count)| push_counts.get(reg).unwrap_or(&0) == count);
+
+                // Clobber Check
                 let all_pushes_safe = pushes.iter().all(|(_, r)| !func_clobbers.contains(r));
 
-                let push_set: HashSet<u8> = pushes.iter().map(|(_, r)| *r).collect();
+                if all_pushes_safe && has_cleanup && counts_match && counts_match_reverse {
-                let restore_set: HashSet<u8> = restores.iter().map(|(_, r, _)| *r).collect();
+                    // We can remove ALL found pushes/restores safely
-
-                if all_pushes_safe && has_cleanup && push_set == restore_set {
-                    // We can remove ALL saves/restores for this call!
                     for (p_idx, _) in pushes {
                         to_remove.insert(p_idx);
                     }
@@ -278,7 +339,7 @@ fn optimize_function_calls<'a>(
                     }
 
                     // Reduce stack cleanup amount
-                    let num_removed = push_set.len() as i64;
+                    let num_removed = push_counts.values().sum::<i32>() as i64;
                     stack_adjustments.insert(cleanup_idx, num_removed);
                     changed = true;
                 }
@@ -295,15 +356,14 @@ fn optimize_function_calls<'a>(
             }
 
             // Apply stack adjustment
-            if let Some(reduction) = stack_adjustments.get(&idx) {
+            if let Some(reduction) = stack_adjustments.get(&idx)
-                if let Instruction::Sub(dst, a, Operand::Number(n)) = &node.instruction {
+                && let Instruction::Sub(dst, a, Operand::Number(n)) = &node.instruction
+            {
                 let new_n = n - Decimal::from(*reduction);
                 if new_n.is_zero() {
                     continue; // Remove the sub entirely if 0
                 }
-                    node.instruction =
+                node.instruction = Instruction::Sub(dst.clone(), a.clone(), Operand::Number(new_n));
-                        Instruction::Sub(dst.clone(), a.clone(), Operand::Number(new_n));
-                }
             }
 
             clean.push(node);
@@ -357,10 +417,16 @@ fn register_forwarding<'a>(
                     break;
                 }
                 // If the temp is redefined, then the old value is dead, so we are safe.
-                if let Some(redef) = get_destination_reg(&node.instruction) {
+                if let Some(redef) = get_destination_reg(&node.instruction)
-                    if redef == temp_reg {
+                    && redef == temp_reg
+                {
                     break;
                 }
+
+                // Reg15 is a return register.
+
+                if temp_reg == 15 {
+                    break;
                 }
                 // If we hit a label/jump, we assume liveness might leak (conservative safety)
                 if matches!(
@@ -436,17 +502,17 @@ fn resolve_labels<'a>(input: Vec<InstructionNode<'a>>) -> Vec<InstructionNode<'a
                     *op = num;
                 }
             }
-            Instruction::BranchEq(a, b, op)
+            Instruction::BranchEq(_, _, op)
-            | Instruction::BranchNe(a, b, op)
+            | Instruction::BranchNe(_, _, op)
-            | Instruction::BranchGt(a, b, op)
+            | Instruction::BranchGt(_, _, op)
-            | Instruction::BranchLt(a, b, op)
+            | Instruction::BranchLt(_, _, op)
-            | Instruction::BranchGe(a, b, op)
+            | Instruction::BranchGe(_, _, op)
-            | Instruction::BranchLe(a, b, op) => {
+            | Instruction::BranchLe(_, _, op) => {
                 if let Some(num) = get_line(op) {
                     *op = num;
                 }
             }
-            Instruction::BranchEqZero(a, op) | Instruction::BranchNeZero(a, op) => {
+            Instruction::BranchEqZero(_, op) | Instruction::BranchNeZero(_, op) => {
                 if let Some(num) = get_line(op) {
                     *op = num;
                 }
@@ -634,8 +700,7 @@ fn reg_is_read(instr: &Instruction, reg: u8) -> bool {
     }
 }
 
-// --- Constant Propagation & Dead Code (Same as before) ---
+/// --- Constant Propagation & Dead Code ---
-
 fn constant_propagation<'a>(input: Vec<InstructionNode<'a>>) -> (Vec<InstructionNode<'a>>, bool) {
     let mut output = Vec::with_capacity(input.len());
     let mut changed = false;
@@ -648,13 +713,8 @@ fn constant_propagation<'a>(input: Vec<InstructionNode<'a>>) -> (Vec<Instruction
         }
 
         let simplified = match &node.instruction {
-            Instruction::Move(dst, src) => {
+            Instruction::Move(dst, src) => resolve_value(src, &registers)
-                if let Some(val) = resolve_value(src, &registers) {
+                .map(|val| Instruction::Move(dst.clone(), Operand::Number(val))),
-                    Some(Instruction::Move(dst.clone(), Operand::Number(val)))
-                } else {
-                    None
-                }
-            }
             Instruction::Add(dst, a, b) => try_fold_math(dst, a, b, &registers, |x, y| x + y),
             Instruction::Sub(dst, a, b) => try_fold_math(dst, a, b, &registers, |x, y| x - y),
             Instruction::Mul(dst, a, b) => try_fold_math(dst, a, b, &registers, |x, y| x * y),
@@ -718,12 +778,12 @@ fn constant_propagation<'a>(input: Vec<InstructionNode<'a>>) -> (Vec<Instruction
         }
 
         // Filter out NOPs (Empty LabelDefs from branch resolution)
-        if let Instruction::LabelDef(l) = &node.instruction {
+        if let Instruction::LabelDef(l) = &node.instruction
-            if l.is_empty() {
+            && l.is_empty()
+        {
             changed = true;
             continue;
         }
-        }
 
         output.push(node);
     }
@@ -779,12 +839,12 @@ fn remove_redundant_moves<'a>(input: Vec<InstructionNode<'a>>) -> (Vec<Instructi
     let mut output = Vec::with_capacity(input.len());
     let mut changed = false;
     for node in input {
-        if let Instruction::Move(dst, src) = &node.instruction {
+        if let Instruction::Move(dst, src) = &node.instruction
-            if dst == src {
+            && dst == src
+        {
             changed = true;
             continue;
         }
-        }
         output.push(node);
     }
     (output, changed)
@@ -804,9 +864,8 @@ fn remove_unreachable_code<'a>(
             changed = true;
             continue;
         }
-        match node.instruction {
+        if let Instruction::Jump(_) = node.instruction {
-            Instruction::Jump(_) | Instruction::Jump(Operand::ReturnAddress) => dead = true,
+            dead = true
-            _ => {}
         }
         output.push(node);
     }

rust_compiler/src/ffi/mod.rs

@@ -1,9 +1,8 @@
 use compiler::{CompilationResult, Compiler};
 use helpers::{Documentation, Span};
-use optimizer::optimize;
 use parser::{sys_call::SysCall, Parser};
 use safer_ffi::prelude::*;
-use std::io::{BufWriter, Write};
+use std::io::BufWriter;
 use tokenizer::{
     token::{Token, TokenType},
     Tokenizer,
@@ -128,33 +127,32 @@ pub fn free_docs_vec(v: safer_ffi::Vec<FfiDocumentedItem>) {
 pub fn compile_from_string(input: safer_ffi::slice::Ref<'_, u16>) -> FfiCompilationResult {
     let res = std::panic::catch_unwind(|| {
         let input = String::from_utf16_lossy(input.as_slice());
-        let mut tmp = BufWriter::new(Vec::new());
 
         let tokenizer = Tokenizer::from(input.as_str());
         let parser = Parser::new(tokenizer);
-        let compiler = Compiler::new(parser, &mut tmp, None);
+        let compiler = Compiler::new(parser, None);
 
         let res = compiler.compile();
 
         if !res.errors.is_empty() {
-            return (safer_ffi::String::EMPTY, res.source_map);
+            return (safer_ffi::String::EMPTY, res.instructions.source_map());
         }
 
         let mut writer = BufWriter::new(Vec::new());
 
-        for instruction in optimize(res.instructions) {
+        // writing into a Vec<u8>. This should not fail.
-            _ = writer.write_all(instruction.instruction.to_string().as_bytes());
+        let optimized = optimizer::optimize(res.instructions);
-            _ = writer.write_all(b"\n");
+        let map = optimized.source_map();
-        }
+        _ = optimized.write(&mut writer);
 
         let Ok(compiled_vec) = writer.into_inner() else {
-            return (safer_ffi::String::EMPTY, res.source_map);
+            return (safer_ffi::String::EMPTY, map);
         };
 
         // Safety: I know the compiler only outputs valid utf8
         (
             safer_ffi::String::from(unsafe { String::from_utf8_unchecked(compiled_vec) }),
-            res.source_map,
+            map,
         )
     });
 
@@ -162,13 +160,9 @@ pub fn compile_from_string(input: safer_ffi::slice::Ref<'_, u16>) -> FfiCompilat
         FfiCompilationResult {
             source_map: source_map
                 .into_iter()
-                .flat_map(|(k, v)| {
+                .map(|(line_num, span)| FfiSourceMapEntry {
-                    v.into_iter()
-                        .map(|span| FfiSourceMapEntry {
                     span: span.into(),
-                            line_number: k as u32,
+                    line_number: line_num as u32,
-                        })
-                        .collect::<Vec<_>>()
                 })
                 .collect::<Vec<_>>()
                 .into(),
@@ -244,9 +238,8 @@ pub fn diagnose_source(input: safer_ffi::slice::Ref<'_, u16>) -> safer_ffi::Vec<
     let res = std::panic::catch_unwind(|| {
         let input = String::from_utf16_lossy(input.as_slice());
 
-        let mut writer = BufWriter::new(Vec::new());
         let tokenizer = Tokenizer::from(input.as_str());
-        let compiler = Compiler::new(Parser::new(tokenizer), &mut writer, None);
+        let compiler = Compiler::new(Parser::new(tokenizer), None);
 
         let CompilationResult {
             errors: diagnosis, ..

rust_compiler/src/main.rs

@@ -88,9 +88,7 @@ fn run_logic<'a>() -> Result<(), Error<'a>> {
         None => BufWriter::new(Box::new(std::io::stdout())),
     };
 
-    let mut tmp = BufWriter::new(vec![]);
+    let compiler = Compiler::new(parser, None);
-
-    let compiler = Compiler::new(parser, &mut tmp, None);
 
     let CompilationResult {
         errors,
@@ -109,11 +107,7 @@ fn run_logic<'a>() -> Result<(), Error<'a>> {
         }
     }
 
-    let instructions = optimizer::optimize(instructions);
+    optimizer::optimize(instructions).write(&mut writer)?;
-
-    for instruction in instructions {
-        writer.write_all(format!("{}\n", instruction.instruction).as_bytes())?;
-    }
 
     writer.flush()?;
 

spilling.slang

@@ -0,0 +1,43 @@
+device self = "db";
+device gasSensor  = "d0";
+device atmosAnal  = "d1";
+device atmosValve = "d2";
+device atmosTank  = "d3";
+device atmosInlet = "d4";
+
+atmosInlet.Lock = true;
+atmosInlet.Mode = 1;
+atmosValve.On = false;
+atmosValve.Lock = true;
+
+let isPumping = false;
+let tempPressure = 0;
+
+loop {
+  yield();
+  let temp = gasSensor.Temperature;
+  let pres = atmosAnal.Pressure;
+  let liqV = atmosAnal.VolumeOfLiquid;
+  let tempVol = atmosAnal.Volume;
+
+  let stress = 5_000 * liqV / tempVol;
+
+  tempPressure = isPumping ? 1_000 : 10_000;
+
+  let shouldTurnOnInlet = (
+    temp > 0c &&
+    pres < tempPressure &&
+    stress < 50
+  );
+
+  isPumping = (
+    !shouldTurnOnInlet &&
+    atmosTank.Pressure < 35_000 &&
+    atmosAnal.RatioPollutant == 0 &&
+    atmosAnal.RatioLiquidPollutant == 0 &&
+    atmosAnal.Pressure > 1_000
+  );
+
+  atmosValve.On = isPumping;
+  atmosInlet.On = shouldTurnOnInlet;
+}