diff --git a/kclvm/ast/src/ast.rs b/kclvm/ast/src/ast.rs index f433bd75d..06c2a28e3 100644 --- a/kclvm/ast/src/ast.rs +++ b/kclvm/ast/src/ast.rs @@ -478,6 +478,18 @@ impl Program { pub fn get_module_ref(&self, module_path: &str) -> Option>> { self.modules.get(module_path).cloned() } + + pub fn get_modules_for_pkg(&self, pkg_name: &str) -> Vec>> { + 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. diff --git a/kclvm/evaluator/src/lib.rs b/kclvm/evaluator/src/lib.rs index f1693af7a..064fa0a2d 100644 --- a/kclvm/evaluator/src/lib.rs +++ b/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 = 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 = 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() } } diff --git a/kclvm/evaluator/src/module.rs b/kclvm/evaluator/src/module.rs index d79671521..108a1947f 100644 --- a/kclvm/evaluator/src/module.rs +++ b/kclvm/evaluator/src/module.rs @@ -1,5 +1,7 @@ // 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; @@ -7,7 +9,7 @@ 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>>) -> 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 diff --git a/kclvm/evaluator/src/node.rs b/kclvm/evaluator/src/node.rs index c076b3267..3490e5c62 100644 --- a/kclvm/evaluator/src/node.rs +++ b/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 = modules + let modules: Vec>> = 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);