From c9d253f5f07bd3232b5c8ff65fcc53df85cfb9fc Mon Sep 17 00:00:00 2001
From: kata <kata.choi@gmail.com>
Date: Thu, 24 Oct 2024 17:05:03 +0800
Subject: [PATCH] allow generic method in forloop

---
 examples/generic_for_loop.no | 19 ++++++++
 src/negative_tests.rs        | 88 ++++++++++++++++++++++++++++++++++--
 src/type_checker/checker.rs  | 17 ++++++-
 3 files changed, 119 insertions(+), 5 deletions(-)

diff --git a/examples/generic_for_loop.no b/examples/generic_for_loop.no
index 1fda688d0..2c01ddb29 100644
--- a/examples/generic_for_loop.no
+++ b/examples/generic_for_loop.no
@@ -1,3 +1,15 @@
+struct Thing {
+    xx: Field,
+}
+
+fn Thing.clone(self, const LEN: Field) -> [Field; LEN] {
+    let mut arr = [0; LEN];
+    for idx in 0..LEN {
+        arr[idx] = self.xx;
+    }
+    return arr;
+}
+
 fn join(const LEN: Field, arr1: [Field; LLEN], arr2: [Field; RLEN]) -> [Field; LEN] {
     let mut arr = [0; LEN];
     for ii in 0..LLEN {
@@ -47,6 +59,13 @@ fn main(pub xx: Field) {
         arr3[idx] = cloned;
     }
 
+    // test that the generic method is callable within a for loop
+    let thing = Thing { xx: 5 };
+    for idx in 0..2 {
+        let cloned = thing.clone(2);
+        arr3[idx] = cloned;
+    }
+
     let init_val = 1;
     let res = accumulate_mut(init_val);
     assert_eq(res, 8);
diff --git a/src/negative_tests.rs b/src/negative_tests.rs
index 7d148a489..183228832 100644
--- a/src/negative_tests.rs
+++ b/src/negative_tests.rs
@@ -162,6 +162,90 @@ fn test_generic_array_for_loop() {
     ));
 }
 
+#[test]
+fn test_generic_method_cst_for_loop() {
+    let code = r#"
+        struct Thing {
+            xx: Field,
+        }
+
+        // generic on const argument
+        fn Thing.gen(const LEN: Field) -> [Field; LEN] {
+            return [0; LEN];
+        }
+
+        fn loop() {
+            let thing = Thing { xx: 3 };
+            for ii in 0..3 {
+                thing.gen(ii);
+            }
+        }
+        "#;
+
+    let res = tast_pass(code).0;
+
+    assert!(matches!(
+        res.unwrap_err().kind,
+        ErrorKind::VarAccessForbiddenInForLoop(..)
+    ));
+}
+
+#[test]
+fn test_generic_method_array_for_loop() {
+    let code = r#"
+        struct Thing {
+            xx: Field,
+        }
+
+        // generic on array argument
+        fn Thing.gen(arr: [Field; LEN]) -> [Field; LEN] {
+            return arr;
+        }
+
+        fn loop() {
+            let thing = Thing { xx: 3 };
+            for ii in 0..3 {
+                thing.gen([0; ii]);
+            }
+        }
+        "#;
+
+    let res = tast_pass(code).0;
+    assert!(matches!(
+        res.unwrap_err().kind,
+        ErrorKind::VarAccessForbiddenInForLoop(..)
+    ));
+}
+
+#[test]
+fn test_generic_method_nested_for_loop() {
+    let code = r#"
+        struct Thing {
+            xx: Field,
+        }
+
+        // generic on array argument
+        fn Thing.gen(arr: [Field; LEN]) -> [Field; LEN] {
+            return arr;
+        }
+
+        fn loop() {
+            let thing = Thing { xx: 3 };
+            for ii in 0..3 {
+                for jj in 0..3 {
+                    thing.gen([0; ii]);
+                }
+            }
+        }
+        "#;
+
+    let res = tast_pass(code).0;
+    assert!(matches!(
+        res.unwrap_err().kind,
+        ErrorKind::VarAccessForbiddenInForLoop(..)
+    ));
+}
+
 #[test]
 fn test_generic_missing_parameter_arg() {
     let code = r#"
@@ -189,10 +273,6 @@ fn test_generic_symbolic_size_mismatched() {
         }
         "#;
 
-    // in theory, this can be caught by the tast phase as it can be checked symbolically.
-    // but we can't archive this until
-    // - both Array and GenericArray are abstracted into one type with symbolic size.
-    // - the `match_expected` and `same_as` functions are replaced by checking rules for different contexts.
     let res = mast_pass(code).err();
 
     assert!(matches!(
diff --git a/src/type_checker/checker.rs b/src/type_checker/checker.rs
index 5ffdbd40f..06216332b 100644
--- a/src/type_checker/checker.rs
+++ b/src/type_checker/checker.rs
@@ -199,7 +199,22 @@ impl<B: Backend> TypeChecker<B> {
 
                 // check if generic is allowed
                 if method_type.sig.require_monomorphization() && typed_fn_env.is_in_forloop() {
-                    return Err(self.error(ErrorKind::GenericInForLoop, expr.span));
+                    for (observed_arg, expected_arg) in
+                        args.iter().zip(method_type.sig.arguments.iter())
+                    {
+                        // check if the arg involves generic vars
+                        if !expected_arg.extract_generic_names().is_empty() {
+                            let mut forbidden_env = typed_fn_env.clone();
+                            forbidden_env.forbid_forloop_scope();
+
+                            // rewalk the observed arg expression
+                            // it should throw an error if the arg contains generic vars relating to the variables in the forloop scope
+                            self.compute_type(observed_arg, &mut forbidden_env)?;
+
+                            // release the forbidden flag
+                            forbidden_env.allow_forloop_scope();
+                        }
+                    }
                 }
 
                 // type check the method call