feat: add while loops, break and continue statements

src/middle/builder.rs

@@ -85,6 +85,9 @@ struct FuncBuilder {
 
     /// Scoped mapping from user-defined variable names to their corresponding `LocalId`.
     scopes: Vec<HashMap<String, LocalId>>,
+
+    /// Stack of `(continue_target, break_target)` for nested loops
+    loop_stack: Vec<(BlockId, BlockId)>,
 }
 
 impl FuncBuilder {
@@ -100,6 +103,7 @@ impl FuncBuilder {
             current_statements: Vec::new(),
             next_block_id: 0,
             scopes: vec![HashMap::new()],
+            loop_stack: Vec::new(),
         }
     }
 
@@ -249,6 +253,57 @@ impl FuncBuilder {
 
                 self.switch_to_block(merge_block);
             }
+            TypedStmtKind::While { condition, body } => {
+                let cond_block = self.new_block();
+                let body_block = self.new_block();
+                let merge_block = self.new_block();
+
+                self.terminate(Terminator {
+                    kind: TerminatorKind::Goto { target: cond_block },
+                    span: stmt.span,
+                });
+
+                self.switch_to_block(cond_block);
+                let cond_op = self.lower_expr(condition);
+                self.terminate(Terminator {
+                    kind: TerminatorKind::CondBranch {
+                        cond: cond_op,
+                        target_true: body_block,
+                        target_false: merge_block,
+                    },
+                    span: condition.span,
+                });
+
+                self.switch_to_block(body_block);
+                self.loop_stack.push((cond_block, merge_block));
+                self.lower_stmt(body);
+                self.loop_stack.pop();
+
+                self.terminate(Terminator {
+                    kind: TerminatorKind::Goto { target: cond_block },
+                    span: body.span,
+                });
+
+                self.switch_to_block(merge_block);
+            }
+            TypedStmtKind::Break => {
+                if let Some(&(_, merge_block)) = self.loop_stack.last() {
+                    self.terminate(Terminator {
+                        kind: TerminatorKind::Goto {
+                            target: merge_block,
+                        },
+                        span: stmt.span,
+                    });
+                }
+            }
+            TypedStmtKind::Continue => {
+                if let Some(&(cond_block, _)) = self.loop_stack.last() {
+                    self.terminate(Terminator {
+                        kind: TerminatorKind::Goto { target: cond_block },
+                        span: stmt.span,
+                    });
+                }
+            }
             TypedStmtKind::Return { value } => {
                 let val_op = value.as_ref().map(|v| self.lower_expr(v));
                 self.terminate(Terminator {
@@ -330,3 +385,94 @@ impl FuncBuilder {
         }
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::frontend::parser::Parser;
+    use crate::frontend::sema::Sema;
+
+    /// Helper function to parse, analyze, and build a MIR module from source code.
+    fn build_mir(source: &str) -> MirModule {
+        let mut parser = Parser::new(source);
+        let module = parser.parse_module();
+        if let Some(errors) = parser.errors() {
+            panic!("Parse errors: {:?}", errors);
+        }
+
+        let mut sema = Sema::new();
+        let typed_module = sema.analyze_module(&module);
+        if let Some(errors) = sema.errors() {
+            panic!("Semantic errors: {:?}", errors);
+        }
+
+        MirBuilder::build(&typed_module)
+    }
+
+    #[test]
+    fn test_lower_while_loop() {
+        let mir = build_mir("fn main() { while true { } }");
+        let func = &mir.functions[0];
+
+        // Ensure exactly 4 basic blocks are generated
+        assert_eq!(func.blocks.len(), 4);
+
+        // Block 0: Entry -> Goto(1)
+        assert!(matches!(
+            func.blocks[0].terminator.kind,
+            TerminatorKind::Goto { target: BlockId(1) }
+        ));
+        // Block 1: Cond -> CondBranch(true -> 2, false -> 3)
+        assert!(matches!(
+            func.blocks[1].terminator.kind,
+            TerminatorKind::CondBranch {
+                target_true: BlockId(2),
+                target_false: BlockId(3),
+                ..
+            }
+        ));
+        // Block 2: Body -> Goto(1)
+        assert!(matches!(
+            func.blocks[2].terminator.kind,
+            TerminatorKind::Goto { target: BlockId(1) }
+        ));
+        // Block 3: Merge -> Return
+        assert!(matches!(
+            func.blocks[3].terminator.kind,
+            TerminatorKind::Return { value: None }
+        ));
+    }
+
+    #[test]
+    fn test_lower_break_and_continue() {
+        let mir = build_mir("fn main() { while true { continue; break; } }");
+        let func = &mir.functions[0];
+
+        // The body block (BlockId(2)) hits `continue` first, meaning it terminates immediately with a Goto back to the condition block.
+        // The trailing `break` is correctly skipped by the builder as automatic dead code!
+        assert!(matches!(
+            func.blocks[2].terminator.kind,
+            TerminatorKind::Goto { target: BlockId(1) }
+        ));
+    }
+
+    #[test]
+    fn test_lower_let_and_assign() {
+        let mir = build_mir("fn main() { let a = 5; a = 10; }");
+        let func = &mir.functions[0];
+
+        assert_eq!(func.locals.len(), 1); // Only 1 variable 'a' (No temporaries generated for literals)
+        assert_eq!(func.blocks.len(), 1); // No branches, so 1 block
+
+        let block = &func.blocks[0];
+        assert_eq!(block.statements.len(), 2); // Includes both the let initialization and the later assignment
+        assert!(matches!(
+            block.statements[0].kind,
+            StatementKind::Assign(LocalId(0), _)
+        ));
+        assert!(matches!(
+            block.statements[1].kind,
+            StatementKind::Assign(LocalId(0), _)
+        ));
+    }
+}