Working logic and if / else / else if statements

libs/compiler/src/test/branching.rs

@@ -0,0 +1,158 @@
+use crate::compile;
+use indoc::indoc;
+use pretty_assertions::assert_eq;
+
+#[test]
+fn test_if_statement() -> anyhow::Result<()> {
+    let compiled = compile! {
+        debug
+        "
+        let a = 10;
+        if (a > 5) {
+            a = 20;
+        }
+        "
+    };
+
+    assert_eq!(
+        compiled,
+        indoc! {
+            "
+            j main
+            main:
+            move r8 10 #a
+            sgt r1 r8 5
+            beq r1 0 L1
+            move r8 20 #a
+            L1:
+            "
+        }
+    );
+
+    Ok(())
+}
+
+#[test]
+fn test_if_else_statement() -> anyhow::Result<()> {
+    let compiled = compile! {
+        debug
+        "
+        let a = 0;
+        if (10 > 5) {
+            a = 1;
+        } else {
+            a = 2;
+        }
+        "
+    };
+
+    assert_eq!(
+        compiled,
+        indoc! {
+            "
+            j main
+            main:
+            move r8 0 #a
+            sgt r1 10 5
+            beq r1 0 L2
+            move r8 1 #a
+            j L1
+            L2:
+            move r8 2 #a
+            L1:
+            "
+        }
+    );
+
+    Ok(())
+}
+
+#[test]
+fn test_if_else_if_statement() -> anyhow::Result<()> {
+    let compiled = compile! {
+        debug
+        "
+        let a = 0;
+        if (a == 1) {
+            a = 10;
+        } else if (a == 2) {
+            a = 20;
+        } else {
+            a = 30;
+        }
+        "
+    };
+
+    assert_eq!(
+        compiled,
+        indoc! {
+            "
+            j main
+            main:
+            move r8 0 #a
+            seq r1 r8 1
+            beq r1 0 L2
+            move r8 10 #a
+            j L1
+            L2:
+            seq r2 r8 2
+            beq r2 0 L4
+            move r8 20 #a
+            j L3
+            L4:
+            move r8 30 #a
+            L3:
+            L1:
+            "
+        }
+    );
+
+    Ok(())
+}
+
+#[test]
+fn test_spilled_variable_update_in_branch() -> anyhow::Result<()> {
+    let compiled = compile! {
+        debug
+        "
+        let a = 1;
+        let b = 2;
+        let c = 3;
+        let d = 4;
+        let e = 5;
+        let f = 6;
+        let g = 7;
+        let h = 8; // Spilled to stack (offset 0)
+
+        if (a == 1) {
+            h = 99;
+        }
+        "
+    };
+
+    assert_eq!(
+        compiled,
+        indoc! {
+            "
+            j main
+            main:
+            move r8 1 #a
+            move r9 2 #b
+            move r10 3 #c
+            move r11 4 #d
+            move r12 5 #e
+            move r13 6 #f
+            move r14 7 #g
+            push 8 #h
+            seq r1 r8 1
+            beq r1 0 L1
+            sub r0 sp 1
+            put db r0 99 #h
+            L1:
+            "
+        }
+    );
+
+    Ok(())
+}
+

libs/compiler/src/test/mod.rs

@@ -41,6 +41,7 @@ macro_rules! compile {
     }};
 }
 mod binary_expression;
+mod branching;
 mod declaration_function_invocation;
 mod declaration_literal;
 mod function_declaration;

libs/compiler/src/v1.rs

@@ -2,8 +2,9 @@ use crate::variable_manager::{self, LocationRequest, VariableLocation, VariableS
 use parser::{
     Parser as ASTParser,
     tree_node::{
-        BinaryExpression, BlockExpression, DeviceDeclarationExpression, Expression,
-        FunctionExpression, InvocationExpression, Literal, LogicalExpression,
+        AssignmentExpression, BinaryExpression, BlockExpression, DeviceDeclarationExpression,
+        Expression, FunctionExpression, IfExpression, InvocationExpression, Literal,
+        LogicalExpression,
     },
 };
 use quick_error::quick_error;
@@ -75,6 +76,7 @@ pub struct Compiler<'a, W: std::io::Write> {
     declared_main: bool,
     config: CompilerConfig,
     temp_counter: usize,
+    label_counter: usize,
 }
 
 impl<'a, W: std::io::Write> Compiler<'a, W> {
@@ -93,6 +95,7 @@ impl<'a, W: std::io::Write> Compiler<'a, W> {
             declared_main: false,
             config: config.unwrap_or_default(),
             temp_counter: 0,
+            label_counter: 0,
         }
     }
 
@@ -120,6 +123,11 @@ impl<'a, W: std::io::Write> Compiler<'a, W> {
         format!("__binary_temp_{}", self.temp_counter)
     }
 
+    fn next_label_name(&mut self) -> String {
+        self.label_counter += 1;
+        format!("L{}", self.label_counter)
+    }
+
     fn expression<'v>(
         &mut self,
         expr: Expression,
@@ -134,6 +142,10 @@ impl<'a, W: std::io::Write> Compiler<'a, W> {
                 self.expression_block(expr_block, scope)?;
                 Ok(None)
             }
+            Expression::If(expr_if) => {
+                self.expression_if(expr_if, scope)?;
+                Ok(None)
+            }
             Expression::DeviceDeclaration(expr_dev) => {
                 self.expression_device(expr_dev)?;
                 Ok(None)
@@ -145,6 +157,10 @@ impl<'a, W: std::io::Write> Compiler<'a, W> {
                     temp_name: None,
                 }))
             }
+            Expression::Assignment(assign_expr) => {
+                self.expression_assignment(assign_expr, scope)?;
+                Ok(None)
+            }
             Expression::Invocation(expr_invoke) => {
                 self.expression_function_invocation(expr_invoke, scope)?;
                 // Invocation returns result in r15 (RETURN_REGISTER).
@@ -353,6 +369,50 @@ impl<'a, W: std::io::Write> Compiler<'a, W> {
         Ok(Some(loc))
     }
 
+    fn expression_assignment<'v>(
+        &mut self,
+        expr: AssignmentExpression,
+        scope: &mut VariableScope<'v>,
+    ) -> Result<(), Error> {
+        let AssignmentExpression {
+            identifier,
+            expression,
+        } = expr;
+
+        let location = scope.get_location_of(&identifier)?;
+        let (val_str, cleanup) = self.compile_operand(*expression, scope)?;
+
+        let debug_tag = if self.config.debug {
+            format!(" #{}", identifier)
+        } else {
+            String::new()
+        };
+
+        match location {
+            VariableLocation::Temporary(reg) | VariableLocation::Persistant(reg) => {
+                self.write_output(format!("move r{reg} {val_str}{debug_tag}"))?;
+            }
+            VariableLocation::Stack(offset) => {
+                // Calculate address: sp - offset
+                self.write_output(format!(
+                    "sub r{0} sp {offset}",
+                    VariableScope::TEMP_STACK_REGISTER
+                ))?;
+                // Store value to stack/db at address
+                self.write_output(format!(
+                    "put db r{0} {val_str}{debug_tag}",
+                    VariableScope::TEMP_STACK_REGISTER
+                ))?;
+            }
+        }
+
+        if let Some(name) = cleanup {
+            scope.free_temp(name)?;
+        }
+
+        Ok(())
+    }
+
     fn expression_function_invocation(
         &mut self,
         invoke_expr: InvocationExpression,
@@ -467,6 +527,49 @@ impl<'a, W: std::io::Write> Compiler<'a, W> {
         Ok(())
     }
 
+    fn expression_if<'v>(
+        &mut self,
+        expr: IfExpression,
+        scope: &mut VariableScope<'v>,
+    ) -> Result<(), Error> {
+        let end_label = self.next_label_name();
+        let else_label = if expr.else_branch.is_some() {
+            self.next_label_name()
+        } else {
+            end_label.clone()
+        };
+
+        // Compile Condition
+        let (cond_str, cleanup) = self.compile_operand(*expr.condition, scope)?;
+
+        // If condition is FALSE (0), jump to else_label
+        self.write_output(format!("beq {cond_str} 0 {else_label}"))?;
+
+        if let Some(name) = cleanup {
+            scope.free_temp(name)?;
+        }
+
+        // Compile Body
+        // Scope variables in body are ephemeral to the block, handled by expression_block
+        self.expression_block(expr.body, scope)?;
+
+        // If we have an else branch, we need to jump over it after the 'if' body
+        if expr.else_branch.is_some() {
+            self.write_output(format!("j {end_label}"))?;
+            self.write_output(format!("{else_label}:"))?;
+
+            match *expr.else_branch.unwrap() {
+                Expression::Block(block) => self.expression_block(block, scope)?,
+                Expression::If(if_expr) => self.expression_if(if_expr, scope)?,
+                _ => unreachable!("Parser ensures else branch is Block or If"),
+            }
+        }
+
+        self.write_output(format!("{end_label}:"))?;
+
+        Ok(())
+    }
+
     /// Helper to resolve a location to a register string (e.g., "r0").
     /// Note: This does not handle Stack locations automatically, as they require
     /// instruction emission to load. Use `compile_operand` for general handling.