From 43ce27c12e542bd569ca9a2306a405662537fb35 Mon Sep 17 00:00:00 2001
From: Klesti Selimaj <klesti.selimay@icloud.com>
Date: Tue, 14 Apr 2026 06:13:55 +0200
Subject: [PATCH] Fixing for loop

---
 src/parser/ast.rs       | 17 +++++--------
 src/parser/grammar.pest | 12 ++++++---
 src/parser/mod.rs       | 54 ++++++++++++++++++++++++-----------------
 test/src/main.ms        |  2 ++
 4 files changed, 48 insertions(+), 37 deletions(-)

diff --git a/src/parser/ast.rs b/src/parser/ast.rs
index 7fd5c76..08e305d 100644
--- a/src/parser/ast.rs
+++ b/src/parser/ast.rs
@@ -70,6 +70,11 @@ pub enum Statement {
         init: Option<Expression>,
     },
 
+    VarAssign {
+        target: Expression,
+        value: Expression,
+    },
+
     // if (...) stmt else stmt
     If {
         condition: Expression,
@@ -85,7 +90,7 @@ pub enum Statement {
 
     // for (...) stmt
     For {
-        init: Option<ForInit>,
+        init: (VarKind, String, Option<Expression>),
         condition: Option<Expression>,
         update: Option<Expression>,
         body: Box<Statement>,
@@ -118,13 +123,3 @@ pub enum VarKind {
     Const,
     Var,
 }
-
-#[derive(Debug, Clone, Serialize)]
-pub enum ForInit {
-    VarDecl {
-        kind: VarKind,
-        name: String,
-        init: Option<Expression>,
-    },
-    Expr(Expression),
-}
diff --git a/src/parser/grammar.pest b/src/parser/grammar.pest
index a20eb82..ef15bad 100644
--- a/src/parser/grammar.pest
+++ b/src/parser/grammar.pest
@@ -39,6 +39,7 @@ statement_list = { statement* }
 statement = _{
     expr_stmt
   | var_decl
+  | var_assign
   | if_stmt
   | while_stmt
   | for_stmt
@@ -66,6 +67,10 @@ var_decl = {
     var_decl_kind ~ identifier ~ ("=" ~ expr)? ~ ";"
 }
 
+var_assign = {
+    expr ~ "=" ~ expr ~ ";"
+}
+
 // --- control flow ---
 
 if_stmt = {
@@ -77,12 +82,11 @@ while_stmt = {
 }
 
 for_stmt = {
-    "for" ~ "(" ~ (var_decl_no_semicolon | expr_stmt | ";") ~ expr? ~ ";" ~ expr? ~ ")" ~ statement
+    "for" ~ "(" ~ var_decl ~ expr ~ ";" ~ var_assign_no_semicolon ~ ")" ~ statement
 }
 
-// helper (same as var_decl but no trailing ;)
-var_decl_no_semicolon = {
-    var_decl_kind ~ identifier ~ ("=" ~ expr)?
+var_assign_no_semicolon = {
+    expr ~ "=" ~ expr
 }
 
 // ================= EXPRESSIONS =================
diff --git a/src/parser/mod.rs b/src/parser/mod.rs
index d98087d..4fb42c3 100644
--- a/src/parser/mod.rs
+++ b/src/parser/mod.rs
@@ -229,30 +229,32 @@ impl Statement {
             Rule::for_stmt => {
                 let mut inner = pair.into_inner();
 
-                let init = inner.next().map(|p| match p.as_rule() {
-                    Rule::var_decl_no_semicolon => {
-                        let mut it = p.into_inner();
+                let init = inner
+                    .next()
+                    .map(|p| match p.as_rule() {
+                        Rule::var_decl => {
+                            let mut it = p.into_inner();
 
-                        let kind = match it.next().unwrap().as_str() {
-                            "let" => VarKind::Let,
-                            "const" => VarKind::Const,
-                            "var" => VarKind::Var,
-                            _ => unreachable!(),
-                        };
+                            let kind = match it.next().unwrap().as_str() {
+                                "let" => VarKind::Let,
+                                "const" => VarKind::Const,
+                                "var" => VarKind::Var,
+                                _ => unreachable!(),
+                            };
 
-                        let name = it.next().unwrap().as_str().to_string();
-                        let init_expr = it
-                            .next()
-                            .map(|e| Expression::from_pair(e.into_inner().next().unwrap()));
+                            let name = it.next().unwrap().as_str().to_string();
+                            let init_expr = it
+                                .next()
+                                .map(|e| Expression::from_pair(e.into_inner().next().unwrap()));
 
-                        ForInit::VarDecl {
-                            kind,
-                            name,
-                            init: init_expr,
+                            (kind, name, init_expr)
                         }
-                    }
-                    _ => ForInit::Expr(Expression::from_pair(p)),
-                });
+                        _ => unimplemented!(
+                            "For loop init parsing not implemented yet: {:?}",
+                            p.as_rule()
+                        ),
+                    })
+                    .unwrap();
 
                 let condition = inner.next().map(Expression::from_pair);
                 let update = inner.next().map(Expression::from_pair);
@@ -260,12 +262,20 @@ impl Statement {
 
                 Statement::For {
                     init,
-                    condition,
-                    update,
+                    condition: None,
+                    update: None,
                     body: Box::new(body),
                 }
             }
 
+            Rule::var_assign => {
+                let mut inner = pair.into_inner();
+                let target = Expression::from_pair(inner.next().unwrap());
+                let value = Expression::from_pair(inner.next().unwrap());
+
+                Statement::VarAssign { target, value }
+            }
+
             _ => unimplemented!(
                 "Statement parsing not implemented yet: {:?}",
                 pair.as_rule()
diff --git a/test/src/main.ms b/test/src/main.ms
index 977dd57..755f17d 100644
--- a/test/src/main.ms
+++ b/test/src/main.ms
@@ -1,11 +1,13 @@
 function add(a: int, b: int): int {
     let sum = 0;
 
+    for (let i = 0; i < 10; i = i + 1) {
         if (i == 5) {
             continue;
         }
 
         sum = sum + i;
+    }
 
     if (sum > 20) {
         return sum;