loops

libs/compiler/src/test/loops.rs

@@ -0,0 +1,108 @@
+use crate::compile;
+use indoc::indoc;
+use pretty_assertions::assert_eq;
+
+#[test]
+fn test_infinite_loop() -> anyhow::Result<()> {
+    let compiled = compile! {
+        debug
+        "
+        let a = 0;
+        loop {
+            a = a + 1;
+        }
+        "
+    };
+
+    // Labels: L1 (start), L2 (end)
+    assert_eq!(
+        compiled,
+        indoc! {
+            "
+            j main
+            main:
+            move r8 0 #a
+            L1:
+            add r1 r8 1
+            move r8 r1 #a
+            j L1
+            L2:
+            "
+        }
+    );
+
+    Ok(())
+}
+
+#[test]
+fn test_loop_break() -> anyhow::Result<()> {
+    let compiled = compile! {
+        debug
+        "
+        let a = 0;
+        loop {
+            a = a + 1;
+            if (a > 10) {
+                break;
+            }
+        }
+        "
+    };
+
+    // Labels: L1 (start), L2 (end), L3 (if end - implicit else label)
+    assert_eq!(
+        compiled,
+        indoc! {
+            "
+            j main
+            main:
+            move r8 0 #a
+            L1:
+            add r1 r8 1
+            move r8 r1 #a
+            sgt r2 r8 10
+            beq r2 0 L3
+            j L2
+            L3:
+            j L1
+            L2:
+            "
+        }
+    );
+
+    Ok(())
+}
+
+#[test]
+fn test_while_loop() -> anyhow::Result<()> {
+    let compiled = compile! {
+        debug
+        "
+        let a = 0;
+        while (a < 10) {
+            a = a + 1;
+        }
+        "
+    };
+
+    // Labels: L1 (start), L2 (end)
+    assert_eq!(
+        compiled,
+        indoc! {
+            "
+            j main
+            main:
+            move r8 0 #a
+            L1:
+            slt r1 r8 10
+            beq r1 0 L2
+            add r2 r8 1
+            move r8 r2 #a
+            j L1
+            L2:
+            "
+        }
+    );
+
+    Ok(())
+}

libs/compiler/src/test/mod.rs

@@ -46,3 +46,4 @@ mod declaration_function_invocation;
 mod declaration_literal;
 mod function_declaration;
 mod logic_expression;
+mod loops;

libs/compiler/src/v1.rs

@@ -4,7 +4,7 @@ use parser::{
     tree_node::{
         AssignmentExpression, BinaryExpression, BlockExpression, DeviceDeclarationExpression,
         Expression, FunctionExpression, IfExpression, InvocationExpression, Literal,
-        LogicalExpression,
+        LogicalExpression, LoopExpression, WhileExpression,
     },
 };
 use quick_error::quick_error;
@@ -77,6 +77,7 @@ pub struct Compiler<'a, W: std::io::Write> {
     config: CompilerConfig,
     temp_counter: usize,
     label_counter: usize,
+    loop_stack: Vec<String>, // Stores the 'end' label of the current loops
 }
 
 impl<'a, W: std::io::Write> Compiler<'a, W> {
@@ -96,6 +97,7 @@ impl<'a, W: std::io::Write> Compiler<'a, W> {
             config: config.unwrap_or_default(),
             temp_counter: 0,
             label_counter: 0,
+            loop_stack: Vec::new(),
         }
     }
 
@@ -146,6 +148,18 @@ impl<'a, W: std::io::Write> Compiler<'a, W> {
                 self.expression_if(expr_if, scope)?;
                 Ok(None)
             }
+            Expression::Loop(expr_loop) => {
+                self.expression_loop(expr_loop, scope)?;
+                Ok(None)
+            }
+            Expression::While(expr_while) => {
+                self.expression_while(expr_while, scope)?;
+                Ok(None)
+            }
+            Expression::Break => {
+                self.expression_break()?;
+                Ok(None)
+            }
             Expression::DeviceDeclaration(expr_dev) => {
                 self.expression_device(expr_dev)?;
                 Ok(None)
@@ -570,6 +584,77 @@ impl<'a, W: std::io::Write> Compiler<'a, W> {
         Ok(())
     }
 
+    fn expression_loop<'v>(
+        &mut self,
+        expr: LoopExpression,
+        scope: &mut VariableScope<'v>,
+    ) -> Result<(), Error> {
+        let start_label = self.next_label_name();
+        let end_label = self.next_label_name();
+
+        // Push end label to stack for 'break'
+        self.loop_stack.push(end_label.clone());
+
+        self.write_output(format!("{start_label}:"))?;
+
+        // Compile Body
+        self.expression_block(expr.body, scope)?;
+
+        // Jump back to start
+        self.write_output(format!("j {start_label}"))?;
+        self.write_output(format!("{end_label}:"))?;
+
+        self.loop_stack.pop();
+
+        Ok(())
+    }
+
+    fn expression_while<'v>(
+        &mut self,
+        expr: WhileExpression,
+        scope: &mut VariableScope<'v>,
+    ) -> Result<(), Error> {
+        let start_label = self.next_label_name();
+        let end_label = self.next_label_name();
+
+        // Push end label to stack for 'break'
+        self.loop_stack.push(end_label.clone());
+
+        self.write_output(format!("{start_label}:"))?;
+
+        // Compile Condition
+        let (cond_str, cleanup) = self.compile_operand(*expr.condition, scope)?;
+
+        // If condition is FALSE, jump to end
+        self.write_output(format!("beq {cond_str} 0 {end_label}"))?;
+
+        if let Some(name) = cleanup {
+            scope.free_temp(name)?;
+        }
+
+        // Compile Body
+        self.expression_block(expr.body, scope)?;
+
+        // Jump back to start
+        self.write_output(format!("j {start_label}"))?;
+        self.write_output(format!("{end_label}:"))?;
+
+        self.loop_stack.pop();
+
+        Ok(())
+    }
+
+    fn expression_break(&mut self) -> Result<(), Error> {
+        if let Some(label) = self.loop_stack.last() {
+            self.write_output(format!("j {label}"))?;
+            Ok(())
+        } else {
+            // This is a semantic error, but for now we can return a generic error
+            // Ideally we'd have a specific error type for this
+            Err(Error::Unknown("Break statement outside of loop".into()))
+        }
+    }
+
     /// 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,7 +204,7 @@ impl Parser {
 
         let expr = match current_token.token_type {
             // match unsupported keywords
-            TokenType::Keyword(e) if matches_keyword!(e, Keyword::Enum, Keyword::While) => {
+            TokenType::Keyword(e) if matches_keyword!(e, Keyword::Enum) => {
                 return Err(Error::UnsupportedKeyword(current_token.clone()));
             }
 
@@ -219,6 +219,22 @@ impl Parser {
             // match if statements
             TokenType::Keyword(Keyword::If) => Expression::If(self.if_expression()?),
 
+            // match loop statements
+            TokenType::Keyword(Keyword::Loop) => Expression::Loop(self.loop_expression()?),
+
+            // match while statements
+            TokenType::Keyword(Keyword::While) => Expression::While(self.while_expression()?),
+
+            // match break statements
+            TokenType::Keyword(Keyword::Break) => {
+                // make sure the next token is a semi-colon
+                let next = token_from_option!(self.get_next()?);
+                if !token_matches!(next, TokenType::Symbol(Symbol::Semicolon)) {
+                    return Err(Error::UnexpectedToken(next.clone()));
+                }
+                Expression::Break
+            }
+
             // match syscalls with a `syscall` keyword
             TokenType::Identifier(ref id) if SysCall::is_syscall(id) => {
                 Expression::Syscall(self.syscall()?)
@@ -832,6 +848,61 @@ impl Parser {
         })
     }
 
+    fn loop_expression(&mut self) -> Result<LoopExpression, Error> {
+        let current_token = token_from_option!(self.current_token);
+        if !self_matches_current!(self, TokenType::Keyword(Keyword::Loop)) {
+            return Err(Error::UnexpectedToken(current_token.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()?;
+
+        Ok(LoopExpression { body })
+    }
+
+    fn while_expression(&mut self) -> Result<WhileExpression, Error> {
+        let current_token = token_from_option!(self.current_token);
+        if !self_matches_current!(self, TokenType::Keyword(Keyword::While)) {
+            return Err(Error::UnexpectedToken(current_token.clone()));
+        }
+
+        // consume 'while'
+        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()?;
+
+        Ok(WhileExpression {
+            condition: boxed!(condition),
+            body,
+        })
+    }
+
     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
@@ -1121,4 +1192,3 @@ impl Parser {
         }
     }
 }
-

libs/parser/src/tree_node.rs

@@ -185,11 +185,35 @@ impl std::fmt::Display for IfExpression {
     }
 }
 
+#[derive(Debug, PartialEq, Eq)]
+pub struct LoopExpression {
+    pub body: BlockExpression,
+}
+
+impl std::fmt::Display for LoopExpression {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "(loop {})", self.body)
+    }
+}
+
+#[derive(Debug, PartialEq, Eq)]
+pub struct WhileExpression {
+    pub condition: Box<Expression>,
+    pub body: BlockExpression,
+}
+
+impl std::fmt::Display for WhileExpression {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "(while {} {})", self.condition, self.body)
+    }
+}
+
 #[derive(Debug, PartialEq, Eq)]
 pub enum Expression {
     Assignment(AssignmentExpression),
     Binary(BinaryExpression),
     Block(BlockExpression),
+    Break,
     Declaration(String, Box<Expression>),
     DeviceDeclaration(DeviceDeclarationExpression),
     Function(FunctionExpression),
@@ -197,11 +221,13 @@ pub enum Expression {
     Invocation(InvocationExpression),
     Literal(Literal),
     Logical(LogicalExpression),
+    Loop(LoopExpression),
     Negation(Box<Expression>),
     Priority(Box<Expression>),
     Return(Box<Expression>),
     Syscall(SysCall),
     Variable(String),
+    While(WhileExpression),
 }
 
 impl std::fmt::Display for Expression {
@@ -209,19 +235,23 @@ impl std::fmt::Display for Expression {
         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::Logical(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::Break => write!(f, "break"),
+            Expression::Declaration(id, e) => write!(f, "(let {} = {})", id, e),
+            Expression::DeviceDeclaration(e) => write!(f, "{}", e),
+            Expression::Function(e) => write!(f, "{}", e),
             Expression::If(e) => write!(f, "{}", e),
             Expression::Invocation(e) => write!(f, "{}", e),
-            Expression::Variable(id) => write!(f, "{}", id),
+            Expression::Literal(l) => write!(f, "{}", l),
+            Expression::Logical(e) => write!(f, "{}", e),
+            Expression::Loop(e) => write!(f, "{}", e),
+            Expression::Negation(e) => write!(f, "(-{})", e),
             Expression::Priority(e) => write!(f, "({})", e),
             Expression::Return(e) => write!(f, "(return {})", e),
-            Expression::DeviceDeclaration(e) => write!(f, "{}", e),
             Expression::Syscall(e) => write!(f, "{}", e),
+            Expression::Variable(id) => write!(f, "{}", id),
+            Expression::While(e) => write!(f, "{}", e),
         }
     }
 }
+