fix deep clone in evaluator

Signed-off-by: he1pa <[email protected]>

kclvm/ast/src/ast.rs

@@ -478,6 +478,18 @@ impl Program {
     pub fn get_module_ref(&self, module_path: &str) -> Option<Arc<RwLock<Module>>> {
         self.modules.get(module_path).cloned()
     }
+
+    pub fn get_modules_for_pkg(&self, pkg_name: &str) -> Vec<Arc<RwLock<Module>>> {
+        let mut result = Vec::new();
+        if let Some(module_names) = self.pkgs.get(pkg_name) {
+            for module_name in module_names {
+                if let Some(module) = self.get_module_ref(module_name) {
+                    result.push(module);
+                }
+            }
+        }
+        result
+    }
 }
 
 /// Module is an abstract syntax tree for a single KCL file.

kclvm/evaluator/src/lib.rs

@@ -38,7 +38,7 @@ use std::{cell::RefCell, panic::UnwindSafe};
 
 use crate::error as kcl_error;
 use anyhow::Result;
-use kclvm_ast::ast::{self, AstIndex, Module};
+use kclvm_ast::ast::{self, AstIndex};
 use kclvm_runtime::{Context, ValueRef};
 
 /// SCALAR_KEY denotes the temp scalar key for the global variable json plan process.
@@ -151,20 +151,9 @@ impl<'ctx> Evaluator<'ctx> {
 
     /// Evaluate the program and return the JSON and YAML result.
     pub fn run(self: &Evaluator<'ctx>) -> Result<(String, String)> {
-        if let Some(modules) = self.program.pkgs.get(kclvm_ast::MAIN_PKG) {
-            self.init_scope(kclvm_ast::MAIN_PKG);
-            let modules: Vec<Module> = modules
-                .iter()
-                .map(|m| {
-                    self.program
-                        .get_module(m)
-                        .expect("Failed to acquire module lock")
-                        .expect(&format!("module {:?} not found in program", m))
-                        .clone()
-                })
-                .collect();
-            self.compile_ast_modules(&modules);
-        }
+        let modules = self.program.get_modules_for_pkg(kclvm_ast::MAIN_PKG);
+        self.init_scope(kclvm_ast::MAIN_PKG);
+        self.compile_ast_modules(&modules);
         Ok(self.plan_globals_to_string())
     }
 
@@ -173,21 +162,12 @@ impl<'ctx> Evaluator<'ctx> {
     /// return the result of the function run, rather than a dictionary composed of each
     /// configuration attribute.
     pub fn run_as_function(self: &Evaluator<'ctx>) -> ValueRef {
-        if let Some(modules) = self.program.pkgs.get(kclvm_ast::MAIN_PKG) {
+        let modules = self.program.get_modules_for_pkg(kclvm_ast::MAIN_PKG);
+        if modules.is_empty() {
+            ValueRef::undefined()
+        } else {
             self.init_scope(kclvm_ast::MAIN_PKG);
-            let modules: Vec<Module> = modules
-                .iter()
-                .map(|m| {
-                    self.program
-                        .get_module(m)
-                        .expect("Failed to acquire module lock")
-                        .expect(&format!("module {:?} not found in program", m))
-                        .clone()
-                })
-                .collect();
             self.compile_ast_modules(&modules)
-        } else {
-            ValueRef::undefined()
         }
     }
 

kclvm/evaluator/src/module.rs

@@ -1,13 +1,15 @@
 // Copyright The KCL Authors. All rights reserved.
 
+use std::sync::{Arc, RwLock};
+
 use kclvm_ast::ast;
 use kclvm_ast::walker::TypedResultWalker;
 use kclvm_runtime::ValueRef;
 
 use super::Evaluator;
 use crate::error as kcl_error;
 
-impl<'ctx> Evaluator<'ctx> {
+impl<'ctx> Evaluator<'_> {
     pub fn compile_module_import_and_types(&self, module: &'ctx ast::Module) {
         for stmt in &module.body {
             match &stmt.node {
@@ -81,21 +83,24 @@ impl<'ctx> Evaluator<'ctx> {
     /// 1. scan all possible global variables and allocate undefined values to global pointers.
     /// 2. build all user-defined schema/rule types.
     /// 3. evaluate all codes for the third time.
-    pub(crate) fn compile_ast_modules(&self, modules: &'ctx [ast::Module]) -> ValueRef {
+    pub(crate) fn compile_ast_modules(&self, modules: &Vec<Arc<RwLock<ast::Module>>>) -> ValueRef {
         // Scan global variables
         for ast_module in modules {
+            let ast_module = ast_module.read().expect("Failed to acquire module lock");
             // Pre define global variables with undefined values
-            self.predefine_global_vars(ast_module);
+            self.predefine_global_vars(&ast_module);
         }
         // Scan global types
         for ast_module in modules {
-            self.compile_module_import_and_types(ast_module);
+            let ast_module = ast_module.read().expect("Failed to acquire module lock");
+            self.compile_module_import_and_types(&ast_module);
         }
         let mut result = ValueRef::undefined();
         // Compile the ast module in the pkgpath.
         for ast_module in modules {
+            let ast_module = ast_module.read().expect("Failed to acquire module lock");
             result = self
-                .walk_module(ast_module)
+                .walk_module(&ast_module)
                 .expect(kcl_error::RUNTIME_ERROR_MSG);
         }
         result

kclvm/evaluator/src/node.rs

@@ -2,7 +2,7 @@
 
 use std::cell::RefCell;
 use std::rc::Rc;
-use std::sync::Arc;
+use std::sync::{Arc, RwLock};
 
 use anyhow::Ok;
 use generational_arena::Index;
@@ -233,14 +233,14 @@ impl<'ctx> TypedResultWalker<'ctx> for Evaluator<'ctx> {
             if let Some(modules) = self.program.pkgs.get(&import_stmt.path.node) {
                 self.push_pkgpath(&pkgpath);
                 self.init_scope(&pkgpath);
-                let modules: Vec<Module> = modules
+                let modules: Vec<Arc<RwLock<Module>>> = modules
                     .iter()
                     .map(|m| {
-                        self.program
-                            .get_module(&m)
-                            .expect("Failed to acquire module lock")
-                            .expect(&format!("module {:?} not found in program", m))
-                            .clone()
+                        let m = self
+                            .program
+                            .get_module_ref(&m)
+                            .expect(&format!("module {:?} not found in program", m));
+                        m
                     })
                     .collect();
                 self.compile_ast_modules(&modules);