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,
+        AssignmentExpression, BinaryExpression, BlockExpression, DeviceDeclarationExpression,
-        FunctionExpression, InvocationExpression, Literal, LogicalExpression,
+        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.

libs/parser/src/lib.rs

@@ -204,9 +204,7 @@ impl Parser {
 
         let expr = match current_token.token_type {
             // match unsupported keywords
-            TokenType::Keyword(e)
+            TokenType::Keyword(e) if matches_keyword!(e, Keyword::Enum, Keyword::While) => {
-                if matches_keyword!(e, Keyword::Enum, Keyword::If, Keyword::Else) =>
-            {
                 return Err(Error::UnsupportedKeyword(current_token.clone()));
             }
 
@@ -218,6 +216,9 @@ impl Parser {
             // match functions with a `fn` keyword
             TokenType::Keyword(Keyword::Fn) => Expression::Function(self.function()?),
 
+            // match if statements
+            TokenType::Keyword(Keyword::If) => Expression::If(self.if_expression()?),
+
             // match syscalls with a `syscall` keyword
             TokenType::Identifier(ref id) if SysCall::is_syscall(id) => {
                 Expression::Syscall(self.syscall()?)
@@ -711,10 +712,19 @@ impl Parser {
             let expression = self.expression()?.ok_or(Error::UnexpectedEOF)?;
             let return_expr = Expression::Return(boxed!(expression));
             expressions.push(return_expr);
-            self.assign_next()?;
-        }
 
-        self.assign_next()?;
+            // check for semicolon
+            let next = token_from_option!(self.get_next()?);
+            if !token_matches!(next, TokenType::Symbol(Symbol::Semicolon)) {
+                return Err(Error::UnexpectedToken(next.clone()));
+            }
+
+            // check for right brace
+            let next = token_from_option!(self.get_next()?);
+            if !token_matches!(next, TokenType::Symbol(Symbol::RBrace)) {
+                return Err(Error::UnexpectedToken(next.clone()));
+            }
+        }
 
         Ok(BlockExpression(expressions))
     }
@@ -763,6 +773,65 @@ impl Parser {
         Ok(literal)
     }
 
+    fn if_expression(&mut self) -> Result<IfExpression, Error> {
+        let current_token = token_from_option!(self.current_token);
+        if !self_matches_current!(self, TokenType::Keyword(Keyword::If)) {
+            return Err(Error::UnexpectedToken(current_token.clone()));
+        }
+
+        // consume 'if'
+        let next = token_from_option!(self.get_next()?);
+        if !token_matches!(next, TokenType::Symbol(Symbol::LParen)) {
+            return Err(Error::UnexpectedToken(next.clone()));
+        }
+        self.assign_next()?;
+
+        // parse condition
+        let condition = self.expression()?.ok_or(Error::UnexpectedEOF)?;
+
+        // check for ')'
+        let next = token_from_option!(self.get_next()?);
+        if !token_matches!(next, TokenType::Symbol(Symbol::RParen)) {
+            return Err(Error::UnexpectedToken(next.clone()));
+        }
+
+        // check for '{'
+        let next = token_from_option!(self.get_next()?);
+        if !token_matches!(next, TokenType::Symbol(Symbol::LBrace)) {
+            return Err(Error::UnexpectedToken(next.clone()));
+        }
+
+        // parse body
+        let body = self.block()?;
+
+        // check for 'else'
+        let else_branch = if self_matches_peek!(self, TokenType::Keyword(Keyword::Else)) {
+            self.assign_next()?; // consume 'else'
+
+            if self_matches_peek!(self, TokenType::Keyword(Keyword::If)) {
+                // else if ...
+                self.assign_next()?;
+                Some(boxed!(Expression::If(self.if_expression()?)))
+            } else if self_matches_peek!(self, TokenType::Symbol(Symbol::LBrace)) {
+                // else { ... }
+                self.assign_next()?;
+                Some(boxed!(Expression::Block(self.block()?)))
+            } else {
+                return Err(Error::UnexpectedToken(
+                    token_from_option!(self.get_next()?).clone(),
+                ));
+            }
+        } else {
+            None
+        };
+
+        Ok(IfExpression {
+            condition: boxed!(condition),
+            body,
+            else_branch,
+        })
+    }
+
     fn function(&mut self) -> Result<FunctionExpression, Error> {
         let current_token = token_from_option!(self.current_token);
         // Sanify check that the current token is a `fn` keyword

libs/parser/src/tree_node.rs

@@ -168,35 +168,54 @@ impl std::fmt::Display for DeviceDeclarationExpression {
     }
 }
 
+#[derive(Debug, PartialEq, Eq)]
+pub struct IfExpression {
+    pub condition: Box<Expression>,
+    pub body: BlockExpression,
+    pub else_branch: Option<Box<Expression>>,
+}
+
+impl std::fmt::Display for IfExpression {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "(if ({}) {}", self.condition, self.body)?;
+        if let Some(else_branch) = &self.else_branch {
+            write!(f, " else {}", else_branch)?;
+        }
+        write!(f, ")")
+    }
+}
+
 #[derive(Debug, PartialEq, Eq)]
 pub enum Expression {
     Assignment(AssignmentExpression),
     Binary(BinaryExpression),
     Block(BlockExpression),
     Declaration(String, Box<Expression>),
+    DeviceDeclaration(DeviceDeclarationExpression),
     Function(FunctionExpression),
+    If(IfExpression),
     Invocation(InvocationExpression),
     Literal(Literal),
     Logical(LogicalExpression),
     Negation(Box<Expression>),
     Priority(Box<Expression>),
     Return(Box<Expression>),
-    Variable(String),
-    DeviceDeclaration(DeviceDeclarationExpression),
     Syscall(SysCall),
+    Variable(String),
 }
 
 impl std::fmt::Display for Expression {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
+            Expression::Assignment(e) => write!(f, "{}", e),
+            Expression::Binary(e) => write!(f, "{}", e),
             Expression::Literal(l) => write!(f, "{}", l),
             Expression::Negation(e) => write!(f, "(-{})", e),
-            Expression::Binary(e) => write!(f, "{}", e),
             Expression::Logical(e) => write!(f, "{}", e),
-            Expression::Assignment(e) => write!(f, "{}", e),
             Expression::Declaration(id, e) => write!(f, "(let {} = {})", id, e),
             Expression::Function(e) => write!(f, "{}", e),
             Expression::Block(e) => write!(f, "{}", e),
+            Expression::If(e) => write!(f, "{}", e),
             Expression::Invocation(e) => write!(f, "{}", e),
             Expression::Variable(id) => write!(f, "{}", id),
             Expression::Priority(e) => write!(f, "({})", e),