From d5fdd9c5f419b77f3334e2eb5a6ed48f4825f445 Mon Sep 17 00:00:00 2001
From: peefy <xpf6677@163.com>
Date: Fri, 28 Jun 2024 20:02:19 +0800
Subject: [PATCH] fix: local var scope in the evaluator

Signed-off-by: peefy <xpf6677@163.com>
---
 kclvm/evaluator/src/context.rs                | 15 ++++++++++-
 kclvm/evaluator/src/node.rs                   | 11 ++++++--
 .../schema/index_signature/normal_11/main.k   | 25 +++++++++++++++++++
 .../index_signature/normal_11/stdout.golden   |  9 +++++++
 test/grammar/schema/init/init_schema_6/main.k | 11 ++++++++
 .../schema/init/init_schema_6/stdout.golden   | 10 ++++++++
 6 files changed, 78 insertions(+), 3 deletions(-)
 create mode 100644 test/grammar/schema/index_signature/normal_11/main.k
 create mode 100644 test/grammar/schema/index_signature/normal_11/stdout.golden
 create mode 100644 test/grammar/schema/init/init_schema_6/main.k
 create mode 100644 test/grammar/schema/init/init_schema_6/stdout.golden

diff --git a/kclvm/evaluator/src/context.rs b/kclvm/evaluator/src/context.rs
index 6fef0d73c..98730ff50 100644
--- a/kclvm/evaluator/src/context.rs
+++ b/kclvm/evaluator/src/context.rs
@@ -1,4 +1,4 @@
-use std::rc::Rc;
+use std::{collections::HashSet, rc::Rc};
 
 use generational_arena::Index;
 use kclvm_ast::ast;
@@ -146,6 +146,19 @@ impl<'ctx> Evaluator<'ctx> {
         self.local_vars.borrow_mut().clear();
     }
 
+    #[inline]
+    pub(crate) fn clean_and_cloned_local_vars(&self) -> HashSet<String> {
+        let mut local_vars = self.local_vars.borrow_mut();
+        let r = local_vars.clone();
+        local_vars.clear();
+        r
+    }
+
+    #[inline]
+    pub(crate) fn set_local_vars(&self, vars: HashSet<String>) {
+        self.local_vars.borrow_mut().extend(vars);
+    }
+
     #[inline]
     pub(crate) fn add_target_var(&self, name: &str) {
         self.target_vars.borrow_mut().push(name.to_string());
diff --git a/kclvm/evaluator/src/node.rs b/kclvm/evaluator/src/node.rs
index 3cda66e8b..121b5e672 100644
--- a/kclvm/evaluator/src/node.rs
+++ b/kclvm/evaluator/src/node.rs
@@ -440,12 +440,14 @@ impl<'ctx> TypedResultWalker<'ctx> for Evaluator<'ctx> {
                 ast::QuantOperation::All => {
                     if !is_truth {
                         self.leave_scope();
+                        self.clear_local_vars();
                         return Ok(self.bool_value(false));
                     }
                 }
                 ast::QuantOperation::Any => {
                     if is_truth {
                         self.leave_scope();
+                        self.clear_local_vars();
                         return Ok(self.bool_value(true));
                     }
                 }
@@ -466,11 +468,13 @@ impl<'ctx> TypedResultWalker<'ctx> for Evaluator<'ctx> {
             }
         }
         self.leave_scope();
+        self.clear_local_vars();
         // End for block.
         Ok(result)
     }
 
     fn walk_schema_attr(&self, schema_attr: &'ctx ast::SchemaAttr) -> Self::Result {
+        self.clear_local_vars();
         let name = schema_attr.name.node.as_str();
         self.add_target_var(name);
         for decorator in &schema_attr.decorators {
@@ -687,7 +691,8 @@ impl<'ctx> TypedResultWalker<'ctx> for Evaluator<'ctx> {
             };
             self.dict_insert_value(&mut dict_value, name.node.as_str(), &value);
         }
-        if let Some(proxy) = func.try_get_proxy() {
+        let vars = self.clean_and_cloned_local_vars();
+        let result = if let Some(proxy) = func.try_get_proxy() {
             // Invoke user defined functions, schemas or rules.
             Ok(self.invoke_proxy_function(proxy, &list_value, &dict_value))
         } else {
@@ -698,7 +703,9 @@ impl<'ctx> TypedResultWalker<'ctx> for Evaluator<'ctx> {
                 &dict_value,
                 &mut self.runtime_ctx.borrow_mut(),
             ))
-        }
+        };
+        self.set_local_vars(vars);
+        result
     }
 
     fn walk_subscript(&self, subscript: &'ctx ast::Subscript) -> Self::Result {
diff --git a/test/grammar/schema/index_signature/normal_11/main.k b/test/grammar/schema/index_signature/normal_11/main.k
new file mode 100644
index 000000000..1966d0b5e
--- /dev/null
+++ b/test/grammar/schema/index_signature/normal_11/main.k
@@ -0,0 +1,25 @@
+schema User:
+    id: str
+
+schema Users:
+    [str]: User
+
+schema DB:
+    users: Users = {}
+
+    check:
+        all user in users {
+            user == users[user].id
+        }
+
+schema DBs:
+    [str]: DB
+
+dbs_user: DBs = {
+    user = DB {
+        users: {
+            app = User {id = "app"}
+        }
+    }
+}
+db_user = dbs_user.user
diff --git a/test/grammar/schema/index_signature/normal_11/stdout.golden b/test/grammar/schema/index_signature/normal_11/stdout.golden
new file mode 100644
index 000000000..19548ed0d
--- /dev/null
+++ b/test/grammar/schema/index_signature/normal_11/stdout.golden
@@ -0,0 +1,9 @@
+dbs_user:
+  user:
+    users:
+      app:
+        id: app
+db_user:
+  users:
+    app:
+      id: app
diff --git a/test/grammar/schema/init/init_schema_6/main.k b/test/grammar/schema/init/init_schema_6/main.k
new file mode 100644
index 000000000..55dda4b66
--- /dev/null
+++ b/test/grammar/schema/init/init_schema_6/main.k
@@ -0,0 +1,11 @@
+schema Config:
+    name?: str
+
+makeCopy = lambda p: Config -> Config {
+    Config {name = p.name + "-copy"}
+}
+configs = {
+    "foo": Config {name = "foo"}
+    "bar": Config {name = "bar"}
+}
+copies = {"${name}-copy": makeCopy(config) for name, config in configs}
diff --git a/test/grammar/schema/init/init_schema_6/stdout.golden b/test/grammar/schema/init/init_schema_6/stdout.golden
new file mode 100644
index 000000000..a14973ede
--- /dev/null
+++ b/test/grammar/schema/init/init_schema_6/stdout.golden
@@ -0,0 +1,10 @@
+configs:
+  foo:
+    name: foo
+  bar:
+    name: bar
+copies:
+  foo-copy:
+    name: foo-copy
+  bar-copy:
+    name: bar-copy