feat: lsp compile unit cache. Add cache of lookup_compile_unit() and …

…watcher for config file(`kcl.yaml` and `kcl.mod`).

Signed-off-by: he1pa <18012015693@163.com>

kclvm/tools/src/LSP/src/find_refs.rs

@@ -3,7 +3,7 @@ use crate::goto_def::{find_def_with_gs, goto_definition_with_gs};
 use crate::to_lsp::lsp_location;
 use crate::util::{compile_with_params, Params};
 
-use crate::state::{KCLVfs, KCLWordIndexMap};
+use crate::state::{KCLCompileUnitCache, KCLVfs, KCLWordIndexMap};
 use anyhow::Result;
 use kclvm_ast::ast::Program;
 use kclvm_error::Position as KCLPos;
@@ -24,6 +24,7 @@ pub(crate) fn find_refs<F: Fn(String) -> Result<(), anyhow::Error>>(
     gs: &GlobalState,
     module_cache: Option<KCLModuleCache>,
     scope_cache: Option<KCLScopeCache>,
+    compile_unit_cache: Option<KCLCompileUnitCache>,
 ) -> Result<Vec<Location>, String> {
     let def = find_def_with_gs(kcl_pos, gs, true);
     match def {
@@ -42,6 +43,7 @@ pub(crate) fn find_refs<F: Fn(String) -> Result<(), anyhow::Error>>(
                         logger,
                         module_cache,
                         scope_cache,
+                        compile_unit_cache,
                     ))
                 } else {
                     Err(format!("Invalid file path: {0}", start.filename))
@@ -67,6 +69,7 @@ pub(crate) fn find_refs_from_def<F: Fn(String) -> Result<(), anyhow::Error>>(
     logger: F,
     module_cache: Option<KCLModuleCache>,
     scope_cache: Option<KCLScopeCache>,
+    compile_unit_cache: Option<KCLCompileUnitCache>,
 ) -> Vec<Location> {
     let mut ref_locations = vec![];
     for word_index in (*word_index_map.write()).values_mut() {
@@ -92,6 +95,7 @@ pub(crate) fn find_refs_from_def<F: Fn(String) -> Result<(), anyhow::Error>>(
                                 module_cache: module_cache.clone(),
                                 scope_cache: scope_cache.clone(),
                                 vfs: vfs.clone(),
+                                compile_unit_cache: compile_unit_cache.clone(),
                             }) {
                                 Ok((prog, _, gs)) => {
                                     let ref_pos = kcl_pos(&file_path, ref_loc.range.start);
@@ -216,6 +220,7 @@ mod tests {
                         logger,
                         None,
                         None,
+                        None,
                     ),
                 );
             }
@@ -273,6 +278,7 @@ mod tests {
                         logger,
                         None,
                         None,
+                        None,
                     ),
                 );
             }
@@ -330,6 +336,7 @@ mod tests {
                         logger,
                         None,
                         None,
+                        None,
                     ),
                 );
             }
@@ -380,6 +387,7 @@ mod tests {
                         logger,
                         None,
                         None,
+                        None,
                     ),
                 );
             }

kclvm/tools/src/LSP/src/notification.rs

@@ -1,3 +1,4 @@
+use kclvm_config::{modfile::KCL_MOD_FILE, settings::DEFAULT_SETTING_FILE};
 use lsp_types::notification::{
     Cancel, DidChangeTextDocument, DidChangeWatchedFiles, DidCloseTextDocument,
     DidOpenTextDocument, DidSaveTextDocument,
@@ -141,8 +142,14 @@ impl LanguageServerState {
     ) -> anyhow::Result<()> {
         for change in params.changes {
             let path = from_lsp::abs_path(&change.uri)?;
-            self.loader.handle.invalidate(path);
+            self.loader.handle.invalidate(path.clone());
+            if KCL_CONFIG_FILE.contains(&path.file_name().unwrap().to_str().unwrap()) {
+                self.compile_unit_cache.write().clear();
+            }
         }
+
         Ok(())
     }
 }
+
+const KCL_CONFIG_FILE: [&'static str; 2] = [DEFAULT_SETTING_FILE, KCL_MOD_FILE];

kclvm/tools/src/LSP/src/quick_fix.rs

@@ -155,6 +155,7 @@ mod tests {
             module_cache: None,
             scope_cache: None,
             vfs: Some(KCLVfs::default()),
+            compile_unit_cache: None,
         })
         .unwrap();
 

kclvm/tools/src/LSP/src/request.rs

@@ -267,7 +267,8 @@ pub(crate) fn handle_reference(
     let pos = kcl_pos(&file, params.text_document_position.position);
     let log = |msg: String| log_message(msg, &sender);
     let module_cache = snapshot.module_cache.clone();
-    let scope_cache = snapshot.scope_cache.clone();
+    let _scope_cache = snapshot.scope_cache.clone();
+    let compile_unit_cache = snapshot.compile_unit_cache.clone();
     match find_refs(
         &db.prog,
         &pos,
@@ -276,8 +277,9 @@ pub(crate) fn handle_reference(
         Some(snapshot.vfs.clone()),
         log,
         &db.gs,
-        module_cache,
-        scope_cache,
+        Some(module_cache),
+        None,
+        Some(compile_unit_cache),
     ) {
         core::result::Result::Ok(locations) => Ok(Some(locations)),
         Err(msg) => {
@@ -319,9 +321,10 @@ pub(crate) fn handle_completion(
 
                     match compile_with_params(Params {
                         file: file.clone(),
-                        module_cache: snapshot.module_cache,
+                        module_cache: Some(Arc::clone(&snapshot.module_cache)),
                         scope_cache: None,
                         vfs: Some(snapshot.vfs.clone()),
+                        compile_unit_cache: Some(Arc::clone(&snapshot.compile_unit_cache)),
                     }) {
                         Ok((prog, diags, gs)) => Arc::new(AnalysisDatabase {
                             prog,
@@ -416,8 +419,9 @@ pub(crate) fn handle_rename(
         Some(snapshot.vfs.clone()),
         log,
         &db.gs,
-        snapshot.module_cache.clone(),
-        snapshot.scope_cache.clone(),
+        Some(snapshot.module_cache),
+        Some(snapshot.scope_cache),
+        Some(snapshot.compile_unit_cache),
     );
     match references {
         Result::Ok(locations) => {

kclvm/tools/src/LSP/src/state.rs

@@ -7,7 +7,7 @@ use crate::util::{compile_with_params, get_file_name, to_json, Params};
 use crate::word_index::build_word_index;
 use anyhow::Result;
 use crossbeam_channel::{select, unbounded, Receiver, Sender};
-use kclvm_parser::KCLModuleCache;
+use kclvm_parser::{KCLModuleCache, LoadProgramOptions};
 use kclvm_sema::resolver::scope::{CachedScope, KCLScopeCache};
 use lsp_server::{ReqQueue, Request, Response};
 use lsp_types::Url;
@@ -48,6 +48,8 @@ pub(crate) struct Handle<H, C> {
 
 pub(crate) type KCLVfs = Arc<RwLock<Vfs>>;
 pub(crate) type KCLWordIndexMap = Arc<RwLock<HashMap<Url, HashMap<String, Vec<Location>>>>>;
+pub(crate) type KCLCompileUnitCache =
+    Arc<RwLock<HashMap<String, (Vec<String>, Option<LoadProgramOptions>)>>>;
 
 /// State for the language server
 pub(crate) struct LanguageServerState {
@@ -88,10 +90,13 @@ pub(crate) struct LanguageServerState {
     pub word_index_map: KCLWordIndexMap,
 
     /// KCL parse cache
-    pub module_cache: Option<KCLModuleCache>,
+    pub module_cache: KCLModuleCache,
 
     /// KCL resolver cache
-    pub scope_cache: Option<KCLScopeCache>,
+    pub scope_cache: KCLScopeCache,
+
+    /// KCL compile unit cache cache
+    pub compile_unit_cache: KCLCompileUnitCache,
 }
 
 /// A snapshot of the state of the language server
@@ -106,9 +111,11 @@ pub(crate) struct LanguageServerSnapshot {
     /// The word index map
     pub word_index_map: KCLWordIndexMap,
     /// KCL parse cache
-    pub module_cache: Option<KCLModuleCache>,
+    pub module_cache: KCLModuleCache,
     /// KCL resolver cache
-    pub scope_cache: Option<KCLScopeCache>,
+    pub scope_cache: KCLScopeCache,
+    /// KCL compile unit cache cache
+    pub compile_unit_cache: KCLCompileUnitCache,
 }
 
 #[allow(unused)]
@@ -141,8 +148,9 @@ impl LanguageServerState {
             opened_files: Arc::new(RwLock::new(HashMap::new())),
             word_index_map: Arc::new(RwLock::new(HashMap::new())),
             loader,
-            module_cache: Some(KCLModuleCache::default()),
-            scope_cache: Some(Arc::new(Mutex::new(CachedScope::default()))),
+            module_cache: KCLModuleCache::default(),
+            scope_cache: Arc::new(Mutex::new(CachedScope::default())),
+            compile_unit_cache: Arc::new(RwLock::new(HashMap::default())),
         };
 
         let word_index_map = state.word_index_map.clone();
@@ -228,18 +236,20 @@ impl LanguageServerState {
                         self.thread_pool.execute({
                             let mut snapshot = self.snapshot();
                             let sender = self.task_sender.clone();
-                            let module_cache = self.module_cache.clone();
-                            let scope_cache = self.scope_cache.clone();
+                            let module_cache = Arc::clone(&self.module_cache);
+                            let scope_cache = Arc::clone(&self.scope_cache);
+                            let compile_unit_cache = Arc::clone(&self.compile_unit_cache);
                             move || match url(&snapshot, file.file_id) {
                                 Ok(uri) => {
                                     let version =
                                         snapshot.opened_files.read().get(&file.file_id).cloned();
                                     let mut db = snapshot.db.write();
                                     match compile_with_params(Params {
                                         file: filename.clone(),
-                                        module_cache,
+                                        module_cache: Some(module_cache),
                                         scope_cache: None,
                                         vfs: Some(snapshot.vfs),
+                                        compile_unit_cache: Some(compile_unit_cache),
                                     }) {
                                         Ok((prog, diags, gs)) => {
                                             let current_version = snapshot
@@ -374,6 +384,7 @@ impl LanguageServerState {
             word_index_map: self.word_index_map.clone(),
             module_cache: self.module_cache.clone(),
             scope_cache: self.scope_cache.clone(),
+            compile_unit_cache: self.compile_unit_cache.clone(),
         }
     }
 

kclvm/tools/src/LSP/src/tests.rs

@@ -71,6 +71,7 @@ use crate::from_lsp::file_path_from_url;
 use crate::goto_def::goto_definition_with_gs;
 use crate::hover::hover;
 use crate::main_loop::main_loop;
+use crate::state::KCLCompileUnitCache;
 use crate::state::KCLVfs;
 use crate::to_lsp::kcl_diag_to_lsp_diags;
 use crate::util::to_json;
@@ -125,6 +126,7 @@ pub(crate) fn compile_test_file(
         module_cache: Some(KCLModuleCache::default()),
         scope_cache: Some(KCLScopeCache::default()),
         vfs: Some(KCLVfs::default()),
+        compile_unit_cache: Some(KCLCompileUnitCache::default()),
     })
     .unwrap();
     (file, program, diags, gs)
@@ -286,6 +288,7 @@ fn diagnostics_test() {
         module_cache: None,
         scope_cache: None,
         vfs: Some(KCLVfs::default()),
+        compile_unit_cache: Some(KCLCompileUnitCache::default()),
     })
     .unwrap();
 
@@ -472,6 +475,7 @@ fn complete_import_external_file_test() {
         module_cache: None,
         scope_cache: None,
         vfs: Some(KCLVfs::default()),
+        compile_unit_cache: Some(KCLCompileUnitCache::default()),
     })
     .unwrap();
 
@@ -529,6 +533,7 @@ fn goto_import_external_file_test() {
         module_cache: None,
         scope_cache: None,
         vfs: Some(KCLVfs::default()),
+        compile_unit_cache: Some(KCLCompileUnitCache::default()),
     })
     .unwrap();
 
@@ -1264,6 +1269,7 @@ fn konfig_goto_def_test_base() {
         module_cache: None,
         scope_cache: None,
         vfs: Some(KCLVfs::default()),
+        compile_unit_cache: Some(KCLCompileUnitCache::default()),
     })
     .unwrap();
 
@@ -1355,6 +1361,7 @@ fn konfig_goto_def_test_main() {
         module_cache: None,
         scope_cache: None,
         vfs: Some(KCLVfs::default()),
+        compile_unit_cache: Some(KCLCompileUnitCache::default()),
     })
     .unwrap();
 
@@ -1418,6 +1425,7 @@ fn konfig_completion_test_main() {
         module_cache: None,
         scope_cache: None,
         vfs: Some(KCLVfs::default()),
+        compile_unit_cache: Some(KCLCompileUnitCache::default()),
     })
     .unwrap();
 
@@ -1527,6 +1535,7 @@ fn konfig_hover_test_main() {
         module_cache: None,
         scope_cache: None,
         vfs: Some(KCLVfs::default()),
+        compile_unit_cache: Some(KCLCompileUnitCache::default()),
     })
     .unwrap();
 
@@ -1952,6 +1961,7 @@ fn compile_unit_test() {
         module_cache: None,
         scope_cache: None,
         vfs: Some(KCLVfs::default()),
+        compile_unit_cache: Some(KCLCompileUnitCache::default()),
     })
     .unwrap();
     // b.k is not contained in kcl.yaml but need to be contained in main pkg

kclvm/tools/src/LSP/src/util.rs

@@ -19,7 +19,7 @@ use kclvm_sema::resolver::resolve_program_with_opts;
 use kclvm_sema::resolver::scope::KCLScopeCache;
 
 use crate::from_lsp;
-use crate::state::KCLVfs;
+use crate::state::{KCLCompileUnitCache, KCLVfs};
 use lsp_types::Url;
 use parking_lot::RwLockReadGuard;
 use ra_ap_vfs::{FileId, Vfs};
@@ -64,13 +64,30 @@ pub(crate) struct Params {
     pub module_cache: Option<KCLModuleCache>,
     pub scope_cache: Option<KCLScopeCache>,
     pub vfs: Option<KCLVfs>,
+    pub compile_unit_cache: Option<KCLCompileUnitCache>,
 }
 
 pub(crate) fn compile_with_params(
     params: Params,
 ) -> anyhow::Result<(Program, IndexSet<Diagnostic>, GlobalState)> {
     // Lookup compile unit (kcl.mod or kcl.yaml) from the entry file.
-    let (mut files, opt) = lookup_compile_unit(&params.file, true);
+    let (mut files, opt) = match &params.compile_unit_cache {
+        Some(cache) => {
+            let map = cache.read();
+            match map.get(&params.file) {
+                Some(compile_unit) => compile_unit.clone(),
+                None => lookup_compile_unit(&params.file, true),
+            }
+        }
+        None => lookup_compile_unit(&params.file, true),
+    };
+
+    if let Some(cache) = params.compile_unit_cache {
+        cache
+            .write()
+            .insert(params.file.clone(), (files.clone(), opt.clone()));
+    }
+
     if !files.contains(&params.file) {
         files.push(params.file.clone());
     }