lsp: cache word index in lsp state

Signed-off-by: xiarui.xr <[email protected]>

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

@@ -1,77 +1,72 @@
 use crate::from_lsp::kcl_pos;
 use crate::goto_def::goto_definition;
-use crate::util::{build_word_index, parse_param_and_compile, Param};
+use crate::util::{parse_param_and_compile, Param};
 use anyhow;
-use kclvm_config::modfile::get_pkg_root;
 use lsp_types::Location;
+use std::collections::HashMap;
 
 pub(crate) fn find_refs<F: Fn(String) -> Result<(), anyhow::Error>>(
+    word_index_map: HashMap<String, HashMap<String, Vec<Location>>>,
     def_loc: Location,
     name: String,
     cursor_path: String,
     logger: F,
 ) -> anyhow::Result<Option<Vec<Location>>> {
     // todo: decide the scope by the workspace root and the kcl.mod both, use the narrower scope
     // todo: should use the current file path
-    if let Some(root) = get_pkg_root(def_loc.uri.path()) {
-        match build_word_index(root) {
-            std::result::Result::Ok(word_index) => {
-                if let Some(locs) = word_index.get(name.as_str()).cloned() {
-                    return anyhow::Ok(Some(
-                        locs.into_iter()
-                            .filter(|ref_loc| {
-                                // from location to real def
-                                // return if the real def location matches the def_loc
-                                let file_path = ref_loc.uri.path().to_string();
-                                match parse_param_and_compile(
-                                    Param {
-                                        file: file_path.clone(),
-                                    },
-                                    None,
-                                ) {
-                                    Ok((prog, scope, _)) => {
-                                        let ref_pos = kcl_pos(&file_path, ref_loc.range.start);
-                                        // find def from the ref_pos
-                                        if let Some(real_def) =
-                                            goto_definition(&prog, &ref_pos, &scope)
-                                        {
-                                            match real_def {
-                                                lsp_types::GotoDefinitionResponse::Scalar(
-                                                    real_def_loc,
-                                                ) => real_def_loc == def_loc,
-                                                _ => false,
-                                            }
-                                        } else {
-                                            false
-                                        }
-                                    }
-                                    Err(_) => {
-                                        let _ = logger(format!("{cursor_path} compilation failed"));
-                                        return false;
-                                    }
+
+    let mut ref_locations = vec![];
+
+    for (_, word_index) in word_index_map {
+        if let Some(locs) = word_index.get(name.as_str()).cloned() {
+            let matched_locs: Vec<Location> = locs.into_iter()
+                .filter(|ref_loc| {
+                    // from location to real def
+                    // return if the real def location matches the def_loc
+                    let file_path = ref_loc.uri.path().to_string();
+                    match parse_param_and_compile(
+                        Param {
+                            file: file_path.clone(),
+                        },
+                        None,
+                    ) {
+                        Ok((prog, scope, _)) => {
+                            let ref_pos = kcl_pos(&file_path, ref_loc.range.start);
+                            // find def from the ref_pos
+                            if let Some(real_def) =
+                                goto_definition(&prog, &ref_pos, &scope)
+                            {
+                                match real_def {
+                                    lsp_types::GotoDefinitionResponse::Scalar(
+                                        real_def_loc,
+                                    ) => real_def_loc == def_loc,
+                                    _ => false,
                                 }
-                            })
-                            .collect(),
-                    ));
-                } else {
-                    return Ok(None);
-                }
-            }
-            Err(_) => {
-                logger("build word index failed".to_string())?;
-                return Ok(None);
-            }
+                            } else {
+                                false
+                            }
+                        }
+                        Err(_) => {
+                            let _ = logger(format!("{cursor_path} compilation failed"));
+                            return false;
+                        }
+                    }
+                })
+                .collect();
+            ref_locations.extend(matched_locs);
         }
-    } else {
-        return Ok(None);
     }
+    anyhow::Ok(Some(ref_locations))
 }
 
 #[cfg(test)]
 mod tests {
     use super::find_refs;
     use lsp_types::{Location, Position, Range};
+    use std::ops::Index;
     use std::path::PathBuf;
+    use std::collections::HashMap;
+    use crate::util::build_word_index;
 
     fn logger(msg: String) -> Result<(), anyhow::Error> {
         println!("{}", msg);
@@ -91,6 +86,10 @@ mod tests {
         }
     }
 
+    fn setup_word_index_map(root: &str) -> HashMap<String, HashMap<String, Vec<Location>>> {
+        HashMap::from([("default".to_string(), build_word_index(root.to_string()).unwrap())])
+    }
+
     #[test]
     fn find_refs_from_variable_test() {
         let root = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
@@ -138,7 +137,7 @@ mod tests {
                 ];
                 check_locations_match(
                     expect,
-                    find_refs(def_loc, "a".to_string(), path.to_string(), logger),
+                    find_refs(setup_word_index_map(path), def_loc, "a".to_string(), path.to_string(), logger),
                 );
             }
             Err(_) => assert!(false, "file not found"),
@@ -185,7 +184,7 @@ mod tests {
                 ];
                 check_locations_match(
                     expect,
-                    find_refs(def_loc, "Name".to_string(), path.to_string(), logger),
+                    find_refs(setup_word_index_map(path), def_loc, "Name".to_string(), path.to_string(), logger),
                 );
             }
             Err(_) => assert!(false, "file not found"),

kclvm/tools/src/LSP/src/main.rs

@@ -1,6 +1,8 @@
 use crate::main_loop::main_loop;
 use config::Config;
 use main_loop::app;
+use std::collections::HashMap;
+
 mod analysis;
 mod capabilities;
 mod completion;
@@ -78,9 +80,8 @@ fn run_server() -> anyhow::Result<()> {
         .map_err(|_| anyhow::anyhow!("Initialize result error"))?;
 
     connection.initialize_finish(initialize_id, initialize_result)?;
-
     let config = Config::default();
-    main_loop(connection, config)?;
+    main_loop(connection, config, initialize_params)?;
     io_threads.join()?;
     Ok(())
 }

kclvm/tools/src/LSP/src/main_loop.rs

@@ -2,11 +2,12 @@ use crate::config::Config;
 use crate::state::LanguageServerState;
 use clap::Command;
 use lsp_server::Connection;
+use lsp_types::InitializeParams;
 
 #[allow(dead_code)]
 /// Runs the main loop of the language server. This will receive requests and handle them.
-pub(crate) fn main_loop(connection: Connection, config: Config) -> anyhow::Result<()> {
-    LanguageServerState::new(connection.sender, config).run(connection.receiver)
+pub(crate) fn main_loop(connection: Connection, config: Config, initialize_params: InitializeParams) -> anyhow::Result<()> {
+    LanguageServerState::new(connection.sender, config, initialize_params).run(connection.receiver)
 }
 
 #[allow(dead_code)]

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

@@ -146,6 +146,7 @@ pub(crate) fn handle_reference(
     let path = from_lsp::abs_path(&params.text_document_position.text_document.uri)?;
     let db = snapshot.get_db(&path.clone().into())?;
     let pos = kcl_pos(&file, params.text_document_position.position);
+    let word_index_map = snapshot.word_index_map.clone();
 
     let log = |msg: String| log_message(msg, &sender);
 
@@ -163,7 +164,7 @@ pub(crate) fn handle_reference(
                     },
                     None => None,
                 } {
-                    return find_refs(def_loc, def_name, file, log);
+                    return find_refs(word_index_map, def_loc, def_name, file, log);
                 }
             }
             _ => return Ok(None),

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

@@ -2,13 +2,14 @@ use crate::analysis::Analysis;
 use crate::config::Config;
 use crate::db::AnalysisDatabase;
 use crate::to_lsp::{kcl_diag_to_lsp_diags, url};
-use crate::util::{self, get_file_name, parse_param_and_compile, to_json, Param};
+use crate::util::{get_file_name, parse_param_and_compile, to_json, Param, build_word_index};
 use crossbeam_channel::{select, unbounded, Receiver, Sender};
 use indexmap::IndexSet;
 use lsp_server::{ReqQueue, Response};
 use lsp_types::{
     notification::{Notification, PublishDiagnostics},
     Diagnostic, Location, PublishDiagnosticsParams,
+    InitializeParams,
 };
 use parking_lot::RwLock;
 use ra_ap_vfs::{FileId, Vfs};
@@ -69,7 +70,7 @@ pub(crate) struct LanguageServerState {
     pub vfs_handle: Box<dyn ra_ap_vfs::loader::Handle>,
 
     /// The word index map
-    pub word_index: HashMap<String, Vec<Location>>,
+    pub word_index_map: HashMap<String, HashMap<String, Vec<Location>>>,
 }
 
 /// A snapshot of the state of the language server
@@ -82,19 +83,35 @@ pub(crate) struct LanguageServerSnapshot {
     /// Documents that are currently kept in memory from the client
     pub opened_files: IndexSet<FileId>,
     /// The word index map
-    pub word_index: HashMap<String, Vec<Location>>,
+    pub word_index_map: HashMap<String, HashMap<String, Vec<Location>>>,
 }
 
 #[allow(unused)]
 impl LanguageServerState {
-    pub fn new(sender: Sender<lsp_server::Message>, config: Config) -> Self {
+    pub fn new(sender: Sender<lsp_server::Message>, config: Config, initialize_params: InitializeParams) -> Self {
         let (task_sender, task_receiver) = unbounded::<Task>();
 
         let (vfs_sender, receiver) = unbounded::<ra_ap_vfs::loader::Message>();
         let handle: NotifyHandle =
             ra_ap_vfs::loader::Handle::spawn(Box::new(move |msg| vfs_sender.send(msg).unwrap()));
         let handle = Box::new(handle) as Box<dyn ra_ap_vfs::loader::Handle>;
-
+
+        // build word index for all the workspace folders
+        // todo: async
+        let mut word_index_map = HashMap::new();
+        if let Some(workspace_folders) = initialize_params.workspace_folders {
+            for folder in workspace_folders {
+                let path = folder.uri.path();
+                if let Ok(word_index) = build_word_index(path.to_string()) {
+                    word_index_map.insert(folder.name, word_index);
+                }
+            }
+        } else if let Some(root_uri) = initialize_params.root_uri {
+            let path = root_uri.path();
+            if let Ok(word_index) = build_word_index(path.to_string()) {
+                word_index_map.insert("default".to_string(), word_index);
+            }
+        }
         LanguageServerState {
             sender,
             request_queue: ReqQueue::default(),
@@ -107,7 +124,7 @@ impl LanguageServerState {
             analysis: Analysis::default(),
             opened_files: IndexSet::new(),
             vfs_handle: handle,
-            word_index: HashMap::new(),
+            word_index_map: word_index_map,
         }
     }
 
@@ -251,7 +268,7 @@ impl LanguageServerState {
             vfs: self.vfs.clone(),
             db: self.analysis.db.clone(),
             opened_files: self.opened_files.clone(),
-            word_index: self.word_index.clone(),
+            word_index_map: self.word_index_map.clone(),
         }
     }
 

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

@@ -23,6 +23,10 @@ use lsp_types::TextDocumentIdentifier;
 use lsp_types::TextDocumentItem;
 use lsp_types::TextDocumentPositionParams;
 use lsp_types::TextEdit;
+use lsp_types::InitializeParams;
+use lsp_types::WorkspaceFolder;
+use lsp_types::Url;
+
 use serde::Serialize;
 use std::cell::Cell;
 use std::cell::RefCell;
@@ -43,7 +47,6 @@ use lsp_types::DiagnosticRelatedInformation;
 use lsp_types::DiagnosticSeverity;
 use lsp_types::Location;
 use lsp_types::NumberOrString;
-use lsp_types::Url;
 use lsp_types::{Position, Range, TextDocumentContentChangeEvent};
 use parking_lot::RwLock;
 use proc_macro_crate::bench_test;
@@ -415,9 +418,9 @@ pub struct Project {}
 
 impl Project {
     /// Instantiates a language server for this project.
-    pub fn server(self) -> Server {
+    pub fn server(self, initialize_params: InitializeParams) -> Server {
         let config = Config::default();
-        Server::new(config)
+        Server::new(config, initialize_params)
     }
 }
 
@@ -432,11 +435,11 @@ pub struct Server {
 
 impl Server {
     /// Constructs and initializes a new `Server`
-    pub fn new(config: Config) -> Self {
+    pub fn new(config: Config, initialize_params: InitializeParams) -> Self {
         let (connection, client) = Connection::memory();
 
         let worker = std::thread::spawn(move || {
-            main_loop(connection, config).unwrap();
+            main_loop(connection, config, initialize_params).unwrap();
         });
 
         Self {
@@ -539,7 +542,7 @@ fn notification_test() {
 
     let path = path.to_str().unwrap();
     let src = std::fs::read_to_string(path.clone()).unwrap();
-    let server = Project {}.server();
+    let server = Project {}.server(InitializeParams::default());
 
     // Mock open file
     server.notification::<lsp_types::notification::DidOpenTextDocument>(
@@ -588,7 +591,7 @@ fn goto_def_test() {
 
     let path = path.to_str().unwrap();
     let src = std::fs::read_to_string(path.clone()).unwrap();
-    let server = Project {}.server();
+    let server = Project {}.server(InitializeParams::default());
 
     // Mock open file
     server.notification::<lsp_types::notification::DidOpenTextDocument>(
@@ -645,7 +648,7 @@ fn complete_test() {
 
     let path = path.to_str().unwrap();
     let src = std::fs::read_to_string(path.clone()).unwrap();
-    let server = Project {}.server();
+    let server = Project {}.server(InitializeParams::default());
 
     // Mock open file
     server.notification::<lsp_types::notification::DidOpenTextDocument>(
@@ -709,7 +712,7 @@ fn hover_test() {
 
     let path = path.to_str().unwrap();
     let src = std::fs::read_to_string(path.clone()).unwrap();
-    let server = Project {}.server();
+    let server = Project {}.server(InitializeParams::default());
 
     // Mock open file
     server.notification::<lsp_types::notification::DidOpenTextDocument>(
@@ -821,7 +824,7 @@ fn formatting_test() {
 
     let path = path.to_str().unwrap();
     let src = std::fs::read_to_string(path.clone()).unwrap();
-    let server = Project {}.server();
+    let server = Project {}.server(InitializeParams::default());
 
     // Mock open file
     server.notification::<lsp_types::notification::DidOpenTextDocument>(
@@ -1314,7 +1317,12 @@ fn test_find_refs() {
 
     let path = path.to_str().unwrap();
     let src = std::fs::read_to_string(path.clone()).unwrap();
-    let server = Project {}.server();
+    let mut initialize_params = InitializeParams::default();
+    initialize_params.workspace_folders = Some(vec![WorkspaceFolder{
+        uri: Url::from_file_path(root.clone()).unwrap(),
+        name: "test".to_string(),
+    }]);
+    let server = Project {}.server(initialize_params);
     let url = Url::from_file_path(path).unwrap();
 
     // Mock open file