feat: program.pkg Module to Arc<Module> (#1715)

* feat: program.pkg HashMap<String, Vec<Module>> -> HashMap<String, Vec<Arc<Module>>> . Reduce the performance loss of module ast clone Signed-off-by: he1pa <[email protected]> * feat: use Arc<RwLock<Module>> in program to reduce deep copy when compiling Signed-off-by: he1pa <[email protected]> * fix deep clone in evaluator Signed-off-by: he1pa <[email protected]> * s/get_mut_module/get_module_mut Signed-off-by: he1pa <[email protected]> * Refactor some duplicate code Signed-off-by: he1pa <[email protected]> * replace std::Result with anyhow::Result Signed-off-by: he1pa <[email protected]> --------- Signed-off-by: he1pa <[email protected]>
kcl-lang · Nov 1, 2024 · a6381df · a6381df
1 parent caa8d47
commit a6381df
Show file tree

Hide file tree

Showing 40 changed files with 487 additions and 262 deletions.
diff --git a/kclvm/Cargo.lock b/kclvm/Cargo.lock
diff --git a/kclvm/api/src/service/service_impl.rs b/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();

diff --git a/kclvm/ast/Cargo.toml b/kclvm/ast/Cargo.toml
@@ -15,6 +15,7 @@ kclvm-span = { path = "../span" }
 kclvm-error = { path = "../error" }
 thread_local = "1.1.7"
 kclvm-utils = {path = "../utils"}
+anyhow = "1.0"
 
 [dev-dependencies]
 kclvm-parser = { path = "../parser" }
diff --git a/kclvm/ast/src/ast.rs b/kclvm/ast/src/ast.rs
@@ -35,10 +35,14 @@
 
 use kclvm_utils::path::PathPrefix;
 use serde::{ser::SerializeStruct, Deserialize, Serialize, Serializer};
-use std::collections::HashMap;
+use std::{
+    collections::HashMap,
+    sync::{RwLock, RwLockReadGuard, RwLockWriteGuard},
+};
 
 use compiler_base_span::{Loc, Span};
 use std::fmt::Debug;
+use std::sync::Arc;
 use uuid;
 
 use super::token;
@@ -374,47 +378,117 @@ 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.clone(),
+            pkgs: self
+                .pkgs
+                .iter()
+                .map(|(name, modules)| {
+                    (
+                        name.clone(),
+                        modules
+                            .iter()
+                            .map(|m| {
+                                self.get_module(m)
+                                    .expect("Failed to acquire module lock")
+                                    .expect(&format!("module {:?} not found in program", m))
+                                    .clone()
+                            })
+                            .collect(),
+                    )
+                })
+                .collect(),
+        }
+    }
+}
+
+#[derive(Debug, Clone, Default)]
+pub struct Program {
+    pub root: String,
+    pub pkgs: HashMap<String, Vec<String>>,
+    pub modules: HashMap<String, Arc<RwLock<Module>>>,
+}
+
 impl Program {
     /// Get main entry files.
     pub fn get_main_files(&self) -> Vec<String> {
         match self.pkgs.get(crate::MAIN_PKG) {
-            Some(modules) => modules.iter().map(|m| m.filename.clone()).collect(),
+            Some(modules) => modules.clone(),
             None => vec![],
         }
     }
     /// 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<RwLockReadGuard<'_, Module>> {
         match self.pkgs.get(crate::MAIN_PKG) {
-            Some(modules) => modules.first(),
+            Some(modules) => match modules.first() {
+                Some(first_module_path) => self.get_module(&first_module_path).unwrap_or(None),
+                None => None,
+            },
             None => None,
         }
     }
     /// Get stmt on position
     pub fn pos_to_stmt(&self, pos: &Position) -> Option<Node<Stmt>> {
         for (_, v) in &self.pkgs {
             for m in v {
-                if m.filename == pos.filename {
-                    return m.pos_to_stmt(pos);
+                if let Ok(m) = self.get_module(m) {
+                    let m = m?;
+                    if m.filename == pos.filename {
+                        return m.pos_to_stmt(pos);
+                    }
                 }
             }
         }
         None
     }
 
-    pub fn get_module(&self, module_name: &str) -> Option<&Module> {
-        for (_, modules) in &self.pkgs {
-            for module in modules {
-                if module.filename == module_name {
-                    return Some(module);
+    pub fn get_module(
+        &self,
+        module_path: &str,
+    ) -> anyhow::Result<Option<RwLockReadGuard<'_, Module>>> {
+        match self.modules.get(module_path) {
+            Some(module_ref) => match module_ref.read() {
+                Ok(m) => Ok(Some(m)),
+                Err(_) => Err(anyhow::anyhow!("Failed to acquire module lock")),
+            },
+            None => Ok(None),
+        }
+    }
+
+    pub fn get_module_mut(
+        &self,
+        module_path: &str,
+    ) -> anyhow::Result<Option<RwLockWriteGuard<'_, Module>>> {
+        match self.modules.get(module_path) {
+            Some(module_ref) => match module_ref.write() {
+                Ok(m) => Ok(Some(m)),
+                Err(_) => Err(anyhow::anyhow!("Failed to acquire module lock")),
+            },
+            None => Ok(None),
+        }
+    }
+
+    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);
                 }
             }
         }
-        None
+        result
     }
 }
 

diff --git a/kclvm/evaluator/src/lib.rs b/kclvm/evaluator/src/lib.rs
@@ -151,10 +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);
-            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())
     }
 
@@ -163,11 +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) {
-            self.init_scope(kclvm_ast::MAIN_PKG);
-            self.compile_ast_modules(modules)
-        } else {
+        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);
+            self.compile_ast_modules(&modules)
         }
     }
 

diff --git a/kclvm/evaluator/src/module.rs b/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: &[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

diff --git a/kclvm/evaluator/src/node.rs b/kclvm/evaluator/src/node.rs
@@ -2,11 +2,11 @@
 
 use std::cell::RefCell;
 use std::rc::Rc;
-use std::sync::Arc;
+use std::sync::{Arc, RwLock};
 
 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,17 @@ 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<Arc<RwLock<Module>>> = modules
+                    .iter()
+                    .map(|m| {
+                        let m = self
+                            .program
+                            .get_module_ref(&m)
+                            .expect(&format!("module {:?} not found in program", m));
+                        m
+                    })
+                    .collect();
+                self.compile_ast_modules(&modules);
                 self.pop_pkgpath();
             }
         }
@@ -1154,6 +1164,11 @@ impl<'ctx> Evaluator<'ctx> {
                     .get(&pkgpath_without_prefix!(frame.pkgpath))
                 {
                     if let Some(module) = module_list.get(*index) {
+                        let module = self
+                            .program
+                            .get_module(module)
+                            .expect("Failed to acquire module lock")
+                            .expect(&format!("module {:?} not found in program", module));
                         if let Some(stmt) = module.body.get(setter.stmt) {
                             self.push_backtrack_meta(setter);
                             self.walk_stmt(stmt).expect(INTERNAL_ERROR_MSG);

diff --git a/kclvm/evaluator/src/scope.rs b/kclvm/evaluator/src/scope.rs
@@ -1,3 +1,5 @@
+use std::sync::{Arc, RwLock};
+
 use crate::func::ClosureMap;
 use crate::lazy::merge_setters;
 use crate::{
@@ -37,22 +39,20 @@ impl<'ctx> Evaluator<'ctx> {
             let scopes = vec![Scope::default()];
             pkg_scopes.insert(String::from(pkgpath), scopes);
         }
-        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<Arc<RwLock<ast::Module>>> = if self.program.pkgs.contains_key(pkgpath)
+        {
+            self.program.get_modules_for_pkg(pkgpath)
         } else if pkgpath.starts_with(kclvm_runtime::PKG_PATH_PREFIX)
             && self.program.pkgs.contains_key(&pkgpath[1..])
         {
-            self.program
-                .pkgs
-                .get(&pkgpath[1..])
-                .expect(kcl_error::INTERNAL_ERROR_MSG)
+            self.program.get_modules_for_pkg(&pkgpath[1..])
         } else {
             panic!("pkgpath {} not found", pkgpath);
         };
         // Init all global types including schema and rule.
-        for module in module_list {
+        for module in &module_list {
+            let module = module.read().expect("Failed to acquire module lock");
             for stmt in &module.body {
                 let name = match &stmt.node {
                     ast::Stmt::Schema(schema_stmt) => schema_stmt.name.node.clone(),
@@ -77,6 +77,7 @@ impl<'ctx> Evaluator<'ctx> {
             let mut setters = IndexMap::new();
             for (index, module) in module_list.iter().enumerate() {
                 let index = self.add_global_body(index);
+                let module = module.read().expect("Failed to acquire module lock");
                 merge_setters(&mut setters, &self.emit_setters(&module.body, Some(index)))
             }
             if !lazy_scopes.contains_key(pkgpath) {

diff --git a/kclvm/loader/src/option.rs b/kclvm/loader/src/option.rs
@@ -55,7 +55,12 @@ pub fn list_options(opts: &LoadPackageOptions) -> Result<Vec<OptionHelp>> {
     for (pkgpath, modules) in &packages.program.pkgs {
         extractor.pkgpath = pkgpath.clone();
         for module in modules {
-            extractor.walk_module(module)
+            let module = packages
+                .program
+                .get_module(module)
+                .expect("Failed to acquire module lock")
+                .expect(&format!("module {:?} not found in program", module));
+            extractor.walk_module(&module)
         }
     }
     Ok(extractor.options)