diff --git a/kclvm/tools/src/LSP/src/find_refs.rs b/kclvm/tools/src/LSP/src/find_refs.rs index c3e190053..52fddf571 100644 --- a/kclvm/tools/src/LSP/src/find_refs.rs +++ b/kclvm/tools/src/LSP/src/find_refs.rs @@ -1,11 +1,12 @@ 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 Result<(), anyhow::Error>>( + word_index_map: HashMap>>, def_loc: Location, name: String, cursor_path: String, @@ -13,65 +14,59 @@ pub(crate) fn find_refs Result<(), anyhow::Error>>( ) -> anyhow::Result>> { // 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 = 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>> { + 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"), diff --git a/kclvm/tools/src/LSP/src/main.rs b/kclvm/tools/src/LSP/src/main.rs index 6c87df6f4..62e576694 100644 --- a/kclvm/tools/src/LSP/src/main.rs +++ b/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(()) } diff --git a/kclvm/tools/src/LSP/src/main_loop.rs b/kclvm/tools/src/LSP/src/main_loop.rs index 4e2efddcb..d3ea01386 100644 --- a/kclvm/tools/src/LSP/src/main_loop.rs +++ b/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)] diff --git a/kclvm/tools/src/LSP/src/request.rs b/kclvm/tools/src/LSP/src/request.rs index 1d4a3112d..6f386f973 100644 --- a/kclvm/tools/src/LSP/src/request.rs +++ b/kclvm/tools/src/LSP/src/request.rs @@ -146,6 +146,7 @@ pub(crate) fn handle_reference( let path = from_lsp::abs_path(¶ms.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), diff --git a/kclvm/tools/src/LSP/src/state.rs b/kclvm/tools/src/LSP/src/state.rs index e64a485dd..c51365826 100644 --- a/kclvm/tools/src/LSP/src/state.rs +++ b/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, /// The word index map - pub word_index: HashMap>, + pub word_index_map: HashMap>>, } /// 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, /// The word index map - pub word_index: HashMap>, + pub word_index_map: HashMap>>, } #[allow(unused)] impl LanguageServerState { - pub fn new(sender: Sender, config: Config) -> Self { + pub fn new(sender: Sender, config: Config, initialize_params: InitializeParams) -> Self { let (task_sender, task_receiver) = unbounded::(); let (vfs_sender, receiver) = unbounded::(); 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; - + + // 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(), } } diff --git a/kclvm/tools/src/LSP/src/tests.rs b/kclvm/tools/src/LSP/src/tests.rs index 8dccd981d..d8cc6b53c 100644 --- a/kclvm/tools/src/LSP/src/tests.rs +++ b/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::( @@ -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::( @@ -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::( @@ -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::( @@ -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::( @@ -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