fix: Fix generator variable allocation bug

The code changes in `src/ast/node/control.rs` fix a bug related to variable allocation in generator functions. Previously, when a variable was used in a generator, it needed to be stored in the context. However, there was an issue where the variable was not properly allocated in the context, leading to incorrect behavior.

This commit addresses the bug by correctly allocating the variable in the context when inside a generator. It also adds a check to ensure that the allocation is only performed when necessary.

I also added some code in iter test, to digging deeper into possibility of implement effect system based on generator.

Author: Chronostasys

src/ast/node/control.rs

@@ -434,6 +434,19 @@ impl MatchArmCondition {
             }
             MatchArmCondition::Var(a) => {
                 ctx.push_semantic_token(a.range, SemanticTokenType::VARIABLE, 0);
+                let v = if ctx.generator_data.is_some() {
+                    // special case: in generator, every variable needs to be stored in the context
+                    let alloca = builder.alloc(
+                        a.name.as_str(),
+                        &PLType::new_i8_ptr(),
+                        ctx,
+                        Some(range.start),
+                    );
+                    builder.build_store(alloca, v);
+                    alloca
+                } else {
+                    v
+                };
                 _ = ctx.add_symbol(a.name, v, ty.clone(), a.range, false, false);
                 ctx.push_type_hints(a.range, ty);
                 let i = builder.int_value(&PriType::BOOL, 1, false);

src/ast/node/node_result.rs

@@ -18,7 +18,7 @@ impl TerminatorEnum {
         self == TerminatorEnum::None || self == TerminatorEnum::YieldReturn
     }
     pub fn is_return(self) -> bool {
-        self == TerminatorEnum::Return || self == TerminatorEnum::YieldReturn
+        self == TerminatorEnum::Return
     }
 }
 

test/main.pi

@@ -71,6 +71,9 @@ pub fn main() i64 {
     _io::test_io();
     future_test::test_future();
     std_test::test_std();
+
+
 
     return 0;
 }
+

test/test/iter.pi

@@ -46,18 +46,54 @@ gen fn generator(b: i64, c: i64) Iterator<i64> {
 
 var effect1 = 1;
 var effect2 = 2;
+var eff_ret = 10086;
+
+gen fn generator1() Iterator<(i64, |*||=>void,*||=>void| =>void)> {
+    let retgen2 = 0;
+    let generator = generator2();
+    while true {
+        let nx = generator.next();
+        if nx is None {
+            break;
+        }
+        let nx_v = nx as (i64, |*||=>void,*||=>void| =>void)!;
+        match nx_v {
+            (10086, f) => {
+                let wrapped_f = |iter:*||=>void,handler:*||=>void| => {
+                    let real_handler = unsafe_cast<|||=>void|=>void>(handler);
+                    let ret_handler = |resolve:||=>void, ret|=>{
+                        retgen2 = ret;
+                        (*real_handler)(resolve);
+                        return;
+                    };
+                    let clp = &ret_handler;
+                    let arg2 = unsafe_cast<||=>void>(clp);
+                    f(iter,arg2);
+
+                    return;
+                };
+                let yr= (eff_ret,wrapped_f);
+                yield return yr;
+                break;
+            }
+            yr => {
+                yield return yr;
+            }
+        }
+    }
+
+    println!("got retgen2 ", retgen2);
+
 
-gen fn generator1() Iterator<(i64, |*Iterator<()>,*||=>void| =>void)> {
     // executeGenerator(generator2());
-    let a = 1;
+    let a = retgen2;
     let re = 0;
 
-    let f = |iter:*Iterator<()>,handler:*||=>void| => {
+    let f = |next:*||=>void,handler:*||=>void| => {
         let real_handler = unsafe_cast<||i64|=>void,i64|=>void>(handler);
         let resolve = |ret|=> {
             re = ret;
-            println!(ret);
-            (*iter).next();
+            (*next)();
             return;
         };
         (*real_handler)(resolve,a);
@@ -70,26 +106,82 @@ gen fn generator1() Iterator<(i64, |*Iterator<()>,*||=>void| =>void)> {
 }
 
 
-fn executeGenerator(g:Iterator<(i64, |*Iterator<()>,*||=>void| =>void)>) void {
+gen fn generator2() Iterator<(i64, |*||=>void,*||=>void| =>void)> {
+    // executeGenerator(generator2());
+    let a = 1;
+    let re = 0;
+
+    let f = |next:*||=>void,handler:*||=>void| => {
+        let real_handler = unsafe_cast<||i64|=>void,i64|=>void>(handler);
+        let resolve = |ret|=> {
+            re = ret;
+            (*next)();
+            return;
+        };
+        (*real_handler)(resolve,a);
+
+        return;
+    };
+    let yr= (effect1,f);
+    yield return yr;
+    println!("g2 got re ", re);
+
+    f = |next:*||=>void,handler:*||=>void| => {
+        let real_handler = unsafe_cast<|||=>void,i64|=>void>(handler);
+        let resolve = ||=> {
+            (*next)();
+            return;
+        };
+        (*real_handler)(resolve,100);
+
+        return;
+    };
+    // return
+    yr= (eff_ret,f);
+    yield return yr;
+}
+
+
+fn executeGenerator(g:Iterator<(i64, |*||=>void,*||=>void| =>void)>) void {
     let eff = g.next();
     if eff is None {
+        println!("ret none");
         return;
     }
-    match eff as (i64, |*Iterator<()>,*||=>void| =>void)! {
+    match eff as (i64, |*||=>void,*||=>void| =>void)! {
         (1,r) => {
+            println!("handle eff1");
             let cl =|resolve, a| => {
                 eff1handler(resolve, a);
                 return;
             };
-            let it = &g;
+            let next = & ||=>{
+                executeGenerator(g);
+                return;
+            };
             let clp = &cl;
-            let arg1 = unsafe_cast<Iterator<()>>(it);
             let arg2 = unsafe_cast<||=>void>(clp);
-            r(arg1,arg2);
+            r(next,arg2);
         }
         (2,r) => {
+            println!("handle eff2");
+        }
+        (10086,r) => {
+            println!("handle effret");
+            let cl =|resolve:||=>void| => {
+                resolve();
+                return;
+            };
+            let next = & ||=>{
+                executeGenerator(g);
+                return;
+            };
+            let clp = &cl;
+            let arg2 = unsafe_cast<||=>void>(clp);
+            r(next,arg2);
         }
-        _ => {
+        (n, f) => {
+            println!("handle eff", n);
         }
     }
     return;
@@ -109,7 +201,3 @@ fn eff2handler(resolve: ||=>void) void {
     return;
 }
 
-
-gen fn generator2() Iterator<i64> {
-    yield return effect2;
-}