From c5c5e885bf3850afdbc1cf974fecbc8c5e787bf3 Mon Sep 17 00:00:00 2001
From: Selimaj Dev <klesti.selimay@icloud.com>
Date: Fri, 22 May 2026 03:25:05 +0200
Subject: [PATCH] Improved loop bodies

---
 codegen/src/expr.rs                   | 12 ++++--------
 codegen/src/statement.rs              | 10 +++++-----
 mist-test/src/app.mist                | 13 +++++++++++--
 parser/src/ast/expr.rs                | 10 ++++++++++
 parser/src/ast/statement.rs           | 25 ++++++++++++++++++++-----
 parser/src/parser/common/statement.rs |  2 +-
 6 files changed, 51 insertions(+), 21 deletions(-)

diff --git a/codegen/src/expr.rs b/codegen/src/expr.rs
index acf1e9c..2b51f26 100644
--- a/codegen/src/expr.rs
+++ b/codegen/src/expr.rs
@@ -100,14 +100,10 @@ impl GenRust for Expression {
             }
         }
 
-        if let Expression::Statement(stmt) = self {
-            if !stmt.is_block() && ensure_semicolon {
-                ctx.expr_ensure_semicolon = true;
-                cg.add(";");
-            }
-        } else {
-            if ensure_semicolon {
-                ctx.expr_ensure_semicolon = true;
+        if ensure_semicolon {
+            ctx.expr_ensure_semicolon = true;
+
+            if !self.is_block() {
                 cg.add(";");
             }
         }
diff --git a/codegen/src/statement.rs b/codegen/src/statement.rs
index d7c4f9d..8f16fcb 100644
--- a/codegen/src/statement.rs
+++ b/codegen/src/statement.rs
@@ -76,18 +76,18 @@ impl GenRust for Statement {
                 cg.add("if ");
                 initial.condition.gen_rust(ctx, cg);
                 cg.add(" ");
-                cg.ensure_brackets(ctx, &initial.body);
+                cg.ensure_brackets_expr(ctx, &initial.body);
 
                 for else_if_branch in else_if {
                     cg.add("else if");
                     else_if_branch.condition.gen_rust(ctx, cg);
                     cg.add(" ");
-                    cg.ensure_brackets(ctx, &else_if_branch.body);
+                    cg.ensure_brackets_expr(ctx, &else_if_branch.body);
                 }
 
                 if let Some(else_br) = else_branch {
                     cg.add(" else ");
-                    cg.ensure_brackets(ctx, else_br);
+                    cg.ensure_brackets_expr(ctx, else_br);
                 }
             }
 
@@ -95,7 +95,7 @@ impl GenRust for Statement {
                 cg.add("while ");
                 condition.gen_rust(ctx, cg);
                 cg.add(" ");
-                cg.ensure_brackets(ctx, body);
+                cg.ensure_brackets_expr(ctx, body);
             }
 
             Statement::Loop(body) => {
@@ -156,7 +156,7 @@ impl GenRust for Statement {
                 pattern.gen_rust(ctx, cg);
                 cg.add(" in ");
                 iterator.gen_rust(ctx, cg);
-                cg.ensure_brackets(ctx, body);
+                cg.ensure_brackets_expr(ctx, body);
             }
 
             Statement::Return(expr) => {
diff --git a/mist-test/src/app.mist b/mist-test/src/app.mist
index 644a5bf..e4c0e5d 100644
--- a/mist-test/src/app.mist
+++ b/mist-test/src/app.mist
@@ -92,7 +92,7 @@ pub class Application {
         }
     }
 
-    pub void loops_test() {
+    pub i32 loops_test() {
         for (x : 0..10) {
             println!("{x}");
         }
@@ -110,5 +110,14 @@ pub class Application {
         loop {
             println!("loop!");
         }
+
+        if (x == 0) bro();
+        else bro();
+
+        0
     }
-}
\ No newline at end of file
+}
+
+pub i32 bro() {
+        0
+    }
\ No newline at end of file
diff --git a/parser/src/ast/expr.rs b/parser/src/ast/expr.rs
index 77962cb..75583e0 100644
--- a/parser/src/ast/expr.rs
+++ b/parser/src/ast/expr.rs
@@ -81,3 +81,13 @@ pub enum Generic {
     Lifetime(Identifier),
     Type(TypeExpr),
 }
+
+impl Expression {
+    pub fn is_block(&self) -> bool {
+        if let Expression::Statement(stmt) = self {
+            stmt.is_block()
+        } else {
+            false
+        }
+    }
+}
diff --git a/parser/src/ast/statement.rs b/parser/src/ast/statement.rs
index fb8b7e0..f0682c5 100644
--- a/parser/src/ast/statement.rs
+++ b/parser/src/ast/statement.rs
@@ -14,7 +14,7 @@ pub enum Statement {
     If {
         initial: StatementBranch,
         else_if: Vec<StatementBranch>,
-        else_branch: Option<Box<Statement>>,
+        else_branch: Option<Box<Expression>>,
     },
     Loop(Expression),
     While(StatementBranch),
@@ -28,7 +28,7 @@ pub enum Statement {
         mutable: bool,
         pattern: Pattern,
         iterator: Expression,
-        body: Box<Statement>,
+        body: Box<Expression>,
     },
     Match(Expression, Vec<(Vec<Pattern>, Expression)>),
 
@@ -54,14 +54,29 @@ pub struct VarDeclStmt {
 #[derive(Debug, Clone, Serialize)]
 pub struct StatementBranch {
     pub condition: Expression,
-    pub body: Box<Statement>,
+    pub body: Box<Expression>,
 }
 
 impl Statement {
     pub fn is_block(&self) -> bool {
         match self {
-            Self::VarDecl(_) | Self::Return(_) | Self::Break | Self::Continue => false,
-            _ => true,
+            Self::Block(_) => true,
+            Self::If {
+                initial,
+                else_if,
+                else_branch,
+            } => {
+                else_branch
+                    .as_ref()
+                    .map(|v| v.is_block())
+                    .unwrap_or_default()
+                    || else_if
+                        .last()
+                        .map(|b| b.body.is_block())
+                        .unwrap_or_default()
+                    || initial.body.is_block()
+            }
+            _ => false,
         }
     }
 }
diff --git a/parser/src/parser/common/statement.rs b/parser/src/parser/common/statement.rs
index 7736e97..7ed8227 100644
--- a/parser/src/parser/common/statement.rs
+++ b/parser/src/parser/common/statement.rs
@@ -65,7 +65,7 @@ impl<'a> TryFrom<pest::iterators::Pair<'a, Rule>> for Statement {
                     else_if: collect_recovered(inner.next().unwrap().into_inner()),
                     else_branch: inner
                         .next()
-                        .map(Statement::try_from)
+                        .map(Expression::try_from)
                         .transpose()
                         .map(|v| v.map(Box::new))
                         .get_map(|v| { Some(Box::new(v)) }),