From c4017715d24c820a165cba6ca207c8abae03ca5a Mon Sep 17 00:00:00 2001
From: sofia <sofia@teascade.net>
Date: Wed, 23 Jul 2025 21:07:12 +0300
Subject: [PATCH] Codegen for/while loops

---
 reid/src/ast/parse.rs     |  2 +-
 reid/src/codegen.rs       | 50 ++++++++++++++++++++++++++++++++++++++-
 reid/src/mir/typecheck.rs | 14 +++++++++--
 reid/tests/e2e.rs         |  4 ++++
 4 files changed, 66 insertions(+), 4 deletions(-)

diff --git a/reid/src/ast/parse.rs b/reid/src/ast/parse.rs
index dfc519e..757cabe 100644
--- a/reid/src/ast/parse.rs
+++ b/reid/src/ast/parse.rs
@@ -696,7 +696,7 @@ impl Parse for BlockLevelStatement {
 pub struct ForStatement(String, TokenRange, Expression, Expression, Block);
 
 #[derive(Debug)]
-pub struct WhileStatement(Expression, Block);
+pub struct WhileStatement(pub Expression, pub Block);
 
 impl Parse for ForStatement {
     fn parse(mut stream: TokenStream) -> Result<Self, Error> {
diff --git a/reid/src/codegen.rs b/reid/src/codegen.rs
index ddb69c6..421f2ab 100644
--- a/reid/src/codegen.rs
+++ b/reid/src/codegen.rs
@@ -19,6 +19,7 @@ use crate::{
     mir::{
         self, implement::TypeCategory, CustomTypeKey, Metadata, NamedVariableRef, SourceModuleId,
         StructField, StructType, TypeDefinition, TypeDefinitionKind, TypeKind, VagueLiteral,
+        WhileStatement,
     },
     util::try_all,
 };
@@ -598,7 +599,54 @@ impl mir::Statement {
             }
             mir::StmtKind::Import(_) => todo!(),
             mir::StmtKind::Expression(expression) => expression.codegen(scope, state),
-            mir::StmtKind::While(_) => todo!(),
+            mir::StmtKind::While(WhileStatement {
+                condition, block, ..
+            }) => {
+                let condition_block = scope.function.ir.block("condition_block");
+                let condition_true_block = scope.function.ir.block("condition_true");
+                let condition_failed_block = scope.function.ir.block("condition_failed");
+
+                scope
+                    .block
+                    .terminate(Term::Br(condition_block.value()))
+                    .unwrap();
+                let mut condition_scope = scope.with_block(condition_block);
+                let condition_res = condition.codegen(&mut condition_scope, state)?.unwrap();
+                let true_instr = condition_scope
+                    .block
+                    .build(Instr::Constant(ConstValue::Bool(true)))
+                    .unwrap();
+                let check = condition_scope
+                    .block
+                    .build(Instr::ICmp(
+                        CmpPredicate::EQ,
+                        condition_res.instr(),
+                        true_instr,
+                    ))
+                    .unwrap();
+
+                condition_scope
+                    .block
+                    .terminate(Term::CondBr(
+                        check,
+                        condition_true_block.value(),
+                        condition_failed_block.value(),
+                    ))
+                    .unwrap();
+
+                let mut condition_true_scope = scope.with_block(condition_true_block);
+                block.codegen(&mut condition_true_scope, state)?;
+
+                condition_true_scope
+                    .block
+                    .terminate(Term::Br(condition_scope.block.value()))
+                    // Can hard return inside the condition_true_scope
+                    .ok();
+
+                scope.swap_block(condition_failed_block);
+
+                Ok(None)
+            }
         }
     }
 }
diff --git a/reid/src/mir/typecheck.rs b/reid/src/mir/typecheck.rs
index 653b256..8e8f297 100644
--- a/reid/src/mir/typecheck.rs
+++ b/reid/src/mir/typecheck.rs
@@ -336,9 +336,19 @@ impl Block {
                     }
                 }
                 StmtKind::While(WhileStatement {
-                    condition, block, ..
+                    condition,
+                    block,
+                    meta,
                 }) => {
-                    condition.typecheck(&mut state, typerefs, Some(&TypeKind::Bool))?;
+                    let condition_ty =
+                        condition.typecheck(&mut state, typerefs, Some(&TypeKind::Bool))?;
+                    if condition_ty.assert_known(typerefs, &state)? != TypeKind::Bool {
+                        state.note_errors(
+                            &vec![ErrorKind::TypesIncompatible(condition_ty, TypeKind::Bool)],
+                            *meta,
+                        );
+                    }
+
                     block.typecheck(&mut state, typerefs, None)?;
 
                     None
diff --git a/reid/tests/e2e.rs b/reid/tests/e2e.rs
index abf94ef..5685490 100644
--- a/reid/tests/e2e.rs
+++ b/reid/tests/e2e.rs
@@ -120,3 +120,7 @@ fn strings_compiles_well() {
 fn struct_compiles_well() {
     test(include_str!("../../examples/struct.reid"), "test", 17);
 }
+#[test]
+fn loops_compiles_well() {
+    test(include_str!("../../examples/loops.reid"), "test", 10);
+}