feat: program.pkg Module to Arc<Module>

kclvm/api/src/service/service_impl.rs

@@ -6,6 +6,7 @@ use std::string::String;
 use crate::gpyrpc::*;
 
 use kcl_language_server::rename;
+use kclvm_ast::ast::SerializeProgram;
 use kclvm_config::settings::build_settings_pathbuf;
 use kclvm_loader::option::list_options;
 use kclvm_loader::{load_packages_with_cache, LoadPackageOptions};
@@ -147,7 +148,8 @@ impl KclvmServiceImpl {
             }),
             Some(KCLModuleCache::default()),
         )?;
-        let ast_json = serde_json::to_string(&result.program)?;
+        let serialize_program: SerializeProgram = result.program.into();
+        let ast_json = serde_json::to_string(&serialize_program)?;
 
         Ok(ParseProgramResult {
             ast_json,
@@ -264,7 +266,8 @@ impl KclvmServiceImpl {
             // Thread local options
             kclvm_ast::ast::set_should_serialize_id(true);
         }
-        let program_json = serde_json::to_string(&packages.program)?;
+        let serialize_program: SerializeProgram = packages.program.into();
+        let program_json = serde_json::to_string(&serialize_program)?;
         let mut node_symbol_map = HashMap::new();
         let mut symbol_node_map = HashMap::new();
         let mut fully_qualified_name_map = HashMap::new();

kclvm/ast/src/ast.rs

@@ -39,6 +39,7 @@ use std::collections::HashMap;
 
 use compiler_base_span::{Loc, Span};
 use std::fmt::Debug;
+use std::sync::Arc;
 use uuid;
 
 use super::token;
@@ -374,11 +375,35 @@ pub struct SymbolSelectorSpec {
 
 /// Program is the AST collection of all files of the running KCL program.
 #[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq)]
-pub struct Program {
+pub struct SerializeProgram {
     pub root: String,
     pub pkgs: HashMap<String, Vec<Module>>,
 }
 
+impl Into<SerializeProgram> for Program {
+    fn into(self) -> SerializeProgram {
+        SerializeProgram {
+            root: self.root,
+            pkgs: self
+                .pkgs
+                .iter()
+                .map(|(name, modules)| {
+                    (
+                        name.clone(),
+                        modules.iter().map(|m| m.as_ref().clone()).collect(),
+                    )
+                })
+                .collect(),
+        }
+    }
+}
+
+#[derive(Debug, Clone, Default, PartialEq)]
+pub struct Program {
+    pub root: String,
+    pub pkgs: HashMap<String, Vec<Arc<Module>>>,
+}
+
 impl Program {
     /// Get main entry files.
     pub fn get_main_files(&self) -> Vec<String> {
@@ -388,9 +413,9 @@ impl Program {
         }
     }
     /// Get the first module in the main package.
-    pub fn get_main_package_first_module(&self) -> Option<&Module> {
+    pub fn get_main_package_first_module(&self) -> Option<Arc<Module>> {
         match self.pkgs.get(crate::MAIN_PKG) {
-            Some(modules) => modules.first(),
+            Some(modules) => modules.first().cloned(),
             None => None,
         }
     }

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};
+use kclvm_ast::ast::{self, AstIndex, Module};
 use kclvm_runtime::{Context, ValueRef};
 
 /// SCALAR_KEY denotes the temp scalar key for the global variable json plan process.
@@ -153,7 +153,11 @@ impl<'ctx> Evaluator<'ctx> {
     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);
-            self.compile_ast_modules(modules);
+            let modules: Vec<Module> = modules
+                .iter()
+                .map(|arc| arc.clone().as_ref().clone())
+                .collect();
+            self.compile_ast_modules(&modules);
         }
         Ok(self.plan_globals_to_string())
     }
@@ -165,7 +169,11 @@ impl<'ctx> Evaluator<'ctx> {
     pub fn run_as_function(self: &Evaluator<'ctx>) -> ValueRef {
         if let Some(modules) = self.program.pkgs.get(kclvm_ast::MAIN_PKG) {
             self.init_scope(kclvm_ast::MAIN_PKG);
-            self.compile_ast_modules(modules)
+            let modules: Vec<Module> = modules
+                .iter()
+                .map(|arc| arc.clone().as_ref().clone())
+                .collect();
+            self.compile_ast_modules(&modules)
         } else {
             ValueRef::undefined()
         }

kclvm/evaluator/src/node.rs

@@ -6,7 +6,7 @@ use std::sync::Arc;
 
 use anyhow::Ok;
 use generational_arena::Index;
-use kclvm_ast::ast::{self, CallExpr, ConfigEntry, NodeRef};
+use kclvm_ast::ast::{self, CallExpr, ConfigEntry, Module, NodeRef};
 use kclvm_ast::walker::TypedResultWalker;
 use kclvm_runtime::{
     schema_assert, schema_runtime_type, ConfigEntryOperationKind, DecoratorValue, RuntimeErrorType,
@@ -233,7 +233,11 @@ 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);
-                self.compile_ast_modules(modules);
+                let modules: Vec<Module> = modules
+                    .iter()
+                    .map(|arc| arc.clone().as_ref().clone())
+                    .collect();
+                self.compile_ast_modules(&modules);
                 self.pop_pkgpath();
             }
         }

kclvm/evaluator/src/scope.rs

@@ -4,7 +4,7 @@ use crate::{
     error as kcl_error, lazy::LazyEvalScope, rule::RuleEvalContextRef, schema::SchemaEvalContextRef,
 };
 use indexmap::{IndexMap, IndexSet};
-use kclvm_ast::ast;
+use kclvm_ast::ast::{self, Module};
 use kclvm_ast::walker::TypedResultWalker;
 use kclvm_runtime::{ValueRef, _kclvm_get_fn_ptr_by_name, MAIN_PKG_PATH};
 use kclvm_sema::{builtin, plugin};
@@ -39,20 +39,31 @@ impl<'ctx> Evaluator<'ctx> {
         }
         let msg = format!("pkgpath {} is not found", pkgpath);
         // Get the AST module list in the package path.
-        let module_list: &Vec<ast::Module> = if self.program.pkgs.contains_key(pkgpath) {
-            self.program.pkgs.get(pkgpath).expect(&msg)
+        let module_list: Vec<ast::Module> = if self.program.pkgs.contains_key(pkgpath) {
+            let modules = self.program.pkgs.get(pkgpath).expect(&msg);
+            let modules: Vec<Module> = modules
+                .iter()
+                .map(|arc| arc.clone().as_ref().clone())
+                .collect();
+            modules
         } else if pkgpath.starts_with(kclvm_runtime::PKG_PATH_PREFIX)
             && self.program.pkgs.contains_key(&pkgpath[1..])
         {
-            self.program
+            let modules = self
+                .program
                 .pkgs
                 .get(&pkgpath[1..])
-                .expect(kcl_error::INTERNAL_ERROR_MSG)
+                .expect(kcl_error::INTERNAL_ERROR_MSG);
+            let modules: Vec<Module> = modules
+                .iter()
+                .map(|arc| arc.clone().as_ref().clone())
+                .collect();
+            modules
         } else {
             panic!("pkgpath {} not found", pkgpath);
         };
         // Init all global types including schema and rule.
-        for module in module_list {
+        for module in &module_list {
             for stmt in &module.body {
                 let name = match &stmt.node {
                     ast::Stmt::Schema(schema_stmt) => schema_stmt.name.node.clone(),

kclvm/parser/src/lib.rs

@@ -34,6 +34,7 @@ use parser::Parser;
 use std::collections::{HashMap, HashSet, VecDeque};
 use std::path::PathBuf;
 use std::sync::{Arc, RwLock};
+use std::time::Instant;
 
 use kclvm_span::create_session_globals_then;
 
@@ -118,7 +119,7 @@ pub fn parse_single_file(filename: &str, code: Option<String>) -> Result<ParseFi
     );
     let result = loader.load_main()?;
     let module = match result.program.get_main_package_first_module() {
-        Some(module) => module.clone(),
+        Some(module) => module.as_ref().clone(),
         None => ast::Module::default(),
     };
     let file_graph = match loader.file_graph.read() {
@@ -665,7 +666,7 @@ pub fn parse_file(
     file: PkgFile,
     src: Option<String>,
     module_cache: KCLModuleCache,
-    pkgs: &mut HashMap<String, Vec<Module>>,
+    pkgs: &mut HashMap<String, Vec<Arc<Module>>>,
     pkgmap: &mut PkgMap,
     file_graph: FileGraphCache,
     opts: &LoadProgramOptions,
@@ -712,7 +713,7 @@ pub fn parse_file(
 pub fn get_deps(
     file: &PkgFile,
     m: &Module,
-    modules: &mut HashMap<String, Vec<Module>>,
+    modules: &mut HashMap<String, Vec<Arc<Module>>>,
     pkgmap: &mut PkgMap,
     opts: &LoadProgramOptions,
     sess: ParseSessionRef,
@@ -772,7 +773,7 @@ pub fn parse_pkg(
     sess: ParseSessionRef,
     files: Vec<(PkgFile, Option<String>)>,
     module_cache: KCLModuleCache,
-    pkgs: &mut HashMap<String, Vec<Module>>,
+    pkgs: &mut HashMap<String, Vec<Arc<Module>>>,
     pkgmap: &mut PkgMap,
     file_graph: FileGraphCache,
     opts: &LoadProgramOptions,
@@ -798,7 +799,7 @@ pub fn parse_entry(
     sess: ParseSessionRef,
     entry: &entry::Entry,
     module_cache: KCLModuleCache,
-    pkgs: &mut HashMap<String, Vec<Module>>,
+    pkgs: &mut HashMap<String, Vec<Arc<Module>>>,
     pkgmap: &mut PkgMap,
     file_graph: FileGraphCache,
     opts: &LoadProgramOptions,
@@ -915,9 +916,10 @@ pub fn parse_program(
 ) -> Result<LoadProgramResult> {
     let compile_entries = get_compile_entries_from_paths(&paths, &opts)?;
     let workdir = compile_entries.get_root_path().to_string();
-    let mut pkgs: HashMap<String, Vec<Module>> = HashMap::new();
+    let mut pkgs: HashMap<String, Vec<Arc<Module>>> = HashMap::new();
     let mut pkgmap = PkgMap::new();
     let mut new_files = HashSet::new();
+
     for entry in compile_entries.iter() {
         new_files.extend(parse_entry(
             sess.clone(),
@@ -962,43 +964,39 @@ pub fn parse_program(
         }
         Err(e) => return Err(anyhow::anyhow!("Parse program failed: {e}")),
     };
-
     for file in files.iter() {
-        let mut m = match module_cache.read() {
+        let mut m_ref = match module_cache.read() {
             Ok(module_cache) => module_cache
                 .ast_cache
                 .get(&file.canonicalize())
                 .expect(&format!(
                     "Module not found in module: {:?}",
                     file.canonicalize()
                 ))
-                .as_ref()
                 .clone(),
             Err(e) => return Err(anyhow::anyhow!("Parse program failed: {e}")),
         };
         if new_files.contains(file) {
             let pkg = pkgmap.get(file).expect("file not in pkgmap");
+            let mut m = Arc::make_mut(&mut m_ref);
             fix_rel_import_path_with_file(&pkg.pkg_root, &mut m, file, &pkgmap, opts, sess.clone());
-            let m = Arc::new(m.clone());
-            match &mut module_cache.write() {
-                Ok(module_cache) => module_cache.ast_cache.insert(file.canonicalize(), m),
-                Err(e) => return Err(anyhow::anyhow!("Parse program failed: {e}")),
-            };
         }
 
         match pkgs.get_mut(&file.pkg_path) {
             Some(modules) => {
-                modules.push(m);
+                modules.push(m_ref);
             }
             None => {
-                pkgs.insert(file.pkg_path.clone(), vec![m]);
+                pkgs.insert(file.pkg_path.clone(), vec![m_ref]);
             }
         }
     }
+
     let program = ast::Program {
         root: workdir,
         pkgs,
     };
+
     Ok(LoadProgramResult {
         program,
         errors: sess.1.read().diagnostics.clone(),

kclvm/query/src/override.rs

@@ -1,4 +1,5 @@
 use std::collections::HashSet;
+use std::sync::Arc;
 
 use anyhow::{anyhow, Result};
 
@@ -47,6 +48,7 @@ pub fn apply_overrides(
     for o in overrides {
         if let Some(modules) = prog.pkgs.get_mut(MAIN_PKG) {
             for m in modules.iter_mut() {
+                let m = Arc::make_mut(m);
                 if apply_override_on_module(m, o, import_paths)? && print_ast {
                     let code_str = print_ast_module(m);
                     std::fs::write(&m.filename, &code_str)?

kclvm/runner/src/lib.rs

@@ -1,4 +1,4 @@
-use std::{collections::HashMap, ffi::OsStr, path::Path};
+use std::{collections::HashMap, ffi::OsStr, path::Path, sync::Arc};
 
 use anyhow::{anyhow, bail, Result};
 use assembler::KclvmLibAssembler;
@@ -248,7 +248,7 @@ pub fn execute(
 /// **Note that it is not thread safe.**
 pub fn execute_module(m: Module) -> Result<ExecProgramResult> {
     let mut pkgs = HashMap::new();
-    pkgs.insert(MAIN_PKG.to_string(), vec![m]);
+    pkgs.insert(MAIN_PKG.to_string(), vec![Arc::new(m)]);
 
     let prog = Program {
         root: MAIN_PKG.to_string(),

kclvm/runner/src/tests.rs

@@ -179,7 +179,7 @@ fn parse_program(test_kcl_case_path: &str) -> Program {
 ///     Program.root = "__main__"
 fn construct_program(module: Module) -> Program {
     let mut pkgs_ast = HashMap::new();
-    pkgs_ast.insert(MAIN_PKG_NAME.to_string(), vec![module]);
+    pkgs_ast.insert(MAIN_PKG_NAME.to_string(), vec![Arc::new(module)]);
     Program {
         root: MAIN_PKG_NAME.to_string(),
         pkgs: pkgs_ast,